Skip to main content

arrow_schema/extension/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Extension types.
6
7#[cfg(feature = "canonical_extension_types")]
8mod canonical;
9#[cfg(feature = "canonical_extension_types")]
10pub use canonical::*;
11
12use crate::{ArrowError, DataType};
13
14/// The metadata key for the string name identifying an [`ExtensionType`].
15pub const EXTENSION_TYPE_NAME_KEY: &str = "ARROW:extension:name";
16
17/// The metadata key for a serialized representation of the [`ExtensionType`]
18/// necessary to reconstruct the custom type.
19pub const EXTENSION_TYPE_METADATA_KEY: &str = "ARROW:extension:metadata";
20
21/// Extension types.
22///
23/// User-defined “extension” types can be defined setting certain key value
24/// pairs in the [`Field`] metadata structure. These extension keys are:
25/// - [`EXTENSION_TYPE_NAME_KEY`]
26/// - [`EXTENSION_TYPE_METADATA_KEY`]
27///
28/// Canonical extension types support in this crate requires the
29/// `canonical_extension_types` feature.
30///
31/// Extension types may or may not use the [`EXTENSION_TYPE_METADATA_KEY`]
32/// field.
33///
34/// # Example
35///
36/// The example below demonstrates how to implement this trait for a `Uuid`
37/// type. Note this is not the canonical extension type for `Uuid`, which does
38/// not include information about the `Uuid` version.
39///
40/// ```
41/// # use arrow_schema::ArrowError;
42/// # fn main() -> Result<(), ArrowError> {
43/// use arrow_schema::{DataType, extension::ExtensionType, Field};
44/// use std::{fmt, str::FromStr};
45///
46/// /// The different Uuid versions.
47/// #[derive(Clone, Copy, Debug, PartialEq)]
48/// enum UuidVersion {
49///     V1,
50///     V2,
51///     V3,
52///     V4,
53///     V5,
54///     V6,
55///     V7,
56///     V8,
57/// }
58///
59/// // We'll use `Display` to serialize.
60/// impl fmt::Display for UuidVersion {
61///     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62///         write!(
63///             f,
64///             "{}",
65///             match self {
66///                 Self::V1 => "V1",
67///                 Self::V2 => "V2",
68///                 Self::V3 => "V3",
69///                 Self::V4 => "V4",
70///                 Self::V5 => "V5",
71///                 Self::V6 => "V6",
72///                 Self::V7 => "V7",
73///                 Self::V8 => "V8",
74///             }
75///         )
76///     }
77/// }
78///
79/// // And `FromStr` to deserialize.
80/// impl FromStr for UuidVersion {
81///     type Err = ArrowError;
82///
83///     fn from_str(s: &str) -> Result<Self, Self::Err> {
84///         match s {
85///             "V1" => Ok(Self::V1),
86///             "V2" => Ok(Self::V2),
87///             "V3" => Ok(Self::V3),
88///             "V4" => Ok(Self::V4),
89///             "V5" => Ok(Self::V5),
90///             "V6" => Ok(Self::V6),
91///             "V7" => Ok(Self::V7),
92///             "V8" => Ok(Self::V8),
93///             _ => Err(ArrowError::ParseError("Invalid UuidVersion".to_owned())),
94///         }
95///     }
96/// }
97///
98/// /// This is the extension type, not the container for Uuid values. It
99/// /// stores the Uuid version (this is the metadata of this extension type).
100/// #[derive(Clone, Copy, Debug, PartialEq)]
101/// struct Uuid(UuidVersion);
102///
103/// impl ExtensionType for Uuid {
104///     // We use a namespace as suggested by the specification.
105///     const NAME: &'static str = "myorg.example.uuid";
106///
107///     // The metadata type is the Uuid version.
108///     type Metadata = UuidVersion;
109///
110///     // We just return a reference to the Uuid version.
111///     fn metadata(&self) -> &Self::Metadata {
112///         &self.0
113///     }
114///
115///     // We use the `Display` implementation to serialize the Uuid
116///     // version.
117///     fn serialize_metadata(&self) -> Option<String> {
118///         Some(self.0.to_string())
119///     }
120///
121///     // We use the `FromStr` implementation to deserialize the Uuid
122///     // version.
123///     fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
124///         metadata.map_or_else(
125///             || {
126///                 Err(ArrowError::InvalidArgumentError(
127///                     "Uuid extension type metadata missing".to_owned(),
128///                 ))
129///             },
130///             str::parse,
131///         )
132///     }
133///
134///     // The only supported data type is `FixedSizeBinary(16)`.
135///     fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
136///         match data_type {
137///             DataType::FixedSizeBinary(16) => Ok(()),
138///             data_type => Err(ArrowError::InvalidArgumentError(format!(
139///                 "Uuid data type mismatch, expected FixedSizeBinary(16), found {data_type}"
140///             ))),
141///         }
142///     }
143///
144///     // We should always check if the data type is supported before
145///     // constructing the extension type.
146///     fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError> {
147///         let uuid = Self(metadata);
148///         uuid.supports_data_type(data_type)?;
149///         Ok(uuid)
150///     }
151/// }
152///
153/// // We can now construct the extension type.
154/// let uuid_v1 = Uuid(UuidVersion::V1);
155///
156/// // And add it to a field.
157/// let mut field =
158///     Field::new("", DataType::FixedSizeBinary(16), false).with_extension_type(uuid_v1);
159///
160/// // And extract it from this field.
161/// assert_eq!(field.try_extension_type::<Uuid>()?, uuid_v1);
162///
163/// // When we try to add this to a field with an unsupported data type we
164/// // get an error.
165/// let result = Field::new("", DataType::Null, false).try_with_extension_type(uuid_v1);
166/// assert!(result.is_err());
167/// # Ok(()) }
168/// ```
169///
170/// <https://arrow.apache.org/docs/format/Columnar.html#extension-types>
171///
172/// [`Field`]: crate::Field
173pub trait ExtensionType: Sized {
174    /// The name identifying this extension type.
175    ///
176    /// This is the string value that is used for the
177    /// [`EXTENSION_TYPE_NAME_KEY`] in the [`Field::metadata`] of a [`Field`]
178    /// to identify this extension type.
179    ///
180    /// We recommend that you use a “namespace”-style prefix for extension
181    /// type names to minimize the possibility of conflicts with multiple Arrow
182    /// readers and writers in the same application. For example, use
183    /// `myorg.name_of_type` instead of simply `name_of_type`.
184    ///
185    /// Extension names beginning with `arrow.` are reserved for canonical
186    /// extension types, they should not be used for third-party extension
187    /// types.
188    ///
189    /// Extension names are case-sensitive.
190    ///
191    /// [`Field`]: crate::Field
192    /// [`Field::metadata`]: crate::Field::metadata
193    const NAME: &'static str;
194
195    /// The metadata type of this extension type.
196    ///
197    /// Implementations can use strongly or loosely typed data structures here
198    /// depending on the complexity of the metadata.
199    ///
200    /// Implementations can also use `Self` here if the extension type can be
201    /// constructed directly from its metadata.
202    ///
203    /// If an extension type defines no metadata it should use `()` to indicate
204    /// this.
205    type Metadata;
206
207    /// Returns a reference to the metadata of this extension type, or `&()` if
208    /// if this extension type defines no metadata (`Self::Metadata=()`).
209    fn metadata(&self) -> &Self::Metadata;
210
211    /// Returns the serialized representation of the metadata of this extension
212    /// type, or `None` if this extension type defines no metadata
213    /// (`Self::Metadata=()`).
214    ///
215    /// This is string value that is used for the
216    /// [`EXTENSION_TYPE_METADATA_KEY`] in the [`Field::metadata`] of a
217    /// [`Field`].
218    ///
219    /// [`Field`]: crate::Field
220    /// [`Field::metadata`]: crate::Field::metadata
221    fn serialize_metadata(&self) -> Option<String>;
222
223    /// Deserialize the metadata of this extension type from the serialized
224    /// representation of the metadata. An extension type that defines no
225    /// metadata should expect `None` for the serialized metadata and return
226    /// `Ok(())`.
227    ///
228    /// This function should return an error when
229    /// - expected metadata is missing (for extensions types with non-optional
230    ///   metadata)
231    /// - unexpected metadata is set (for extension types without metadata)
232    /// - deserialization of metadata fails
233    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError>;
234
235    /// Returns `Ok(())` iff the given data type is supported by this extension
236    /// type.
237    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError>;
238
239    /// Construct this extension type for a field with the given data type and
240    /// metadata.
241    ///
242    /// This should return an error if the given data type is not supported by
243    /// this extension type.
244    fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError>;
245
246    /// Validate this extension type for a field with the given data type and
247    /// metadata.
248    ///
249    /// The default implementation delegates to [`Self::try_new`]. Extension
250    /// types may override this to validate without constructing `Self`.
251    fn validate(data_type: &DataType, metadata: Self::Metadata) -> Result<(), ArrowError> {
252        Self::try_new(data_type, metadata).map(|_| ())
253    }
254
255    /// Construct this extension type from field metadata and data type.
256    ///
257    /// This is a provided method that extracts extension type information from
258    /// metadata (using [`EXTENSION_TYPE_NAME_KEY`] and
259    /// [`EXTENSION_TYPE_METADATA_KEY`]) and delegates to [`Self::try_new`].
260    ///
261    /// Returns an error if:
262    /// - The extension type name is missing or doesn't match [`Self::NAME`]
263    /// - Metadata deserialization fails
264    /// - The data type is not supported
265    ///
266    /// This method enables extension type checking without requiring a full
267    /// [`Field`] instance, useful when only metadata and data type are available.
268    ///
269    /// [`Field`]: crate::Field
270    fn try_new_from_field_metadata(
271        data_type: &DataType,
272        metadata: &crate::Metadata,
273    ) -> Result<Self, ArrowError> {
274        // Check the extension name in the metadata
275        match metadata.get(EXTENSION_TYPE_NAME_KEY).map(|s| s.as_str()) {
276            // It should match the name of the given extension type
277            Some(name) if name == Self::NAME => {
278                // Deserialize the metadata and try to construct the extension type
279                let ext_metadata = metadata
280                    .get(EXTENSION_TYPE_METADATA_KEY)
281                    .map(|s| s.as_str());
282                let parsed = Self::deserialize_metadata(ext_metadata)?;
283                Self::try_new(data_type, parsed)
284            }
285            // Name mismatch
286            Some(name) => Err(ArrowError::InvalidArgumentError(format!(
287                "Extension type name mismatch: expected {}, got {name}",
288                Self::NAME
289            ))),
290            // Name missing
291            None => Err(ArrowError::InvalidArgumentError(
292                "Extension type name missing".to_string(),
293            )),
294        }
295    }
296}