Add ExtensionType trait and CanonicalExtensionType enum

arrow-schema/src/datatype.rs

@@ -672,6 +672,82 @@ impl DataType {
     }
 }
 
+/// Canonical extension types.
+///
+/// The Arrow columnar format allows defining extension types so as to extend
+/// standard Arrow data types with custom semantics. Often these semantics will
+/// be specific to a system or application. However, it is beneficial to share
+/// the definitions of well-known extension types so as to improve
+/// interoperability between different systems integrating Arrow columnar data.
+///
+/// <https://arrow.apache.org/docs/format/CanonicalExtensions.html>
+#[non_exhaustive]
+#[derive(Debug, Clone)]
+pub enum ExtensionType {
+    /// Extension name: `arrow.uuid`.
+    ///
+    /// The storage type of the extension is `FixedSizeBinary` with a length of
+    /// 16 bytes.
+    ///
+    /// Note:
+    /// A specific UUID version is not required or guaranteed. This extension
+    /// represents UUIDs as FixedSizeBinary(16) with big-endian notation and
+    /// does not interpret the bytes in any way.
+    Uuid,
+}
+
+impl fmt::Display for ExtensionType {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}
+
+impl ExtensionType {
+    /// The metadata key for the string name identifying the custom data type.
+    pub const NAME_KEY: &'static str = "ARROW:extension:name";
+
+    /// The metadata key for a serialized representation of the ExtensionType
+    /// necessary to reconstruct the custom type.
+    pub const METADATA_KEY: &'static str = "ARROW:extension:metadata";
+
+    /// Returns the name of this extension type.
+    pub fn name(&self) -> &'static str {
+        match self {
+            ExtensionType::Uuid => "arrow.uuid",
+        }
+    }
+
+    /// Returns the metadata of this extension type.
+    pub fn metadata(&self) -> Option<String> {
+        match self {
+            ExtensionType::Uuid => None,
+        }
+    }
+
+    /// Returns `true` iff the given [`DataType`] can be used as storage type
+    /// for this extension type.
+    pub(crate) fn supports_storage_type(&self, data_type: &DataType) -> bool {
+        match self {
+            ExtensionType::Uuid => matches!(data_type, DataType::FixedSizeBinary(16)),
+        }
+    }
+
+    /// Extract an [`ExtensionType`] from the given [`Field`].
+    ///
+    /// This function returns `None` if the extension type is not supported or
+    /// recognized.
+    pub(crate) fn try_from_field(field: &Field) -> Option<Self> {
+        let metadata = field.metadata().get(ExtensionType::METADATA_KEY);
+        field
+            .metadata()
+            .get(ExtensionType::NAME_KEY)
+            .and_then(|name| match name.as_str() {
+                "arrow.uuid" if metadata.is_none() => Some(ExtensionType::Uuid),
+                _ => None,
+            })
+    }
+}
+
 /// The maximum precision for [DataType::Decimal128] values
 pub const DECIMAL128_MAX_PRECISION: u8 = 38;
 

arrow-schema/src/field.rs

@@ -23,7 +23,7 @@ use std::sync::Arc;
 
 use crate::datatype::DataType;
 use crate::schema::SchemaBuilder;
-use crate::{Fields, UnionFields, UnionMode};
+use crate::{ExtensionType, Fields, UnionFields, UnionMode};
 
 /// A reference counted [`Field`]
 pub type FieldRef = Arc<Field>;
@@ -337,6 +337,37 @@ impl Field {
         self
     }
 
+    /// Returns the canonical [`ExtensionType`] of this [`Field`], if set.
+    pub fn extension_type(&self) -> Option<ExtensionType> {
+        ExtensionType::try_from_field(self)
+    }
+
+    /// Updates the metadata of this [`Field`] with the [`ExtensionType::name`]
+    /// and [`ExtensionType::metadata`] of the given [`ExtensionType`].
+    ///
+    /// # Panics
+    ///
+    /// This function panics when the datatype of this field is not a valid
+    /// storage type for the given extension type.
+    pub fn with_extension_type(mut self, extension_type: ExtensionType) -> Self {
+        if extension_type.supports_storage_type(&self.data_type) {
+            self.metadata.insert(
+                ExtensionType::NAME_KEY.to_owned(),
+                extension_type.name().to_owned(),
+            );
+            if let Some(metadata) = extension_type.metadata() {
+                self.metadata
+                    .insert(ExtensionType::METADATA_KEY.to_owned(), metadata);
+            }
+            self
+        } else {
+            panic!(
+                "{extension_type} does not support {} as storage type",
+                self.data_type
+            );
+        }
+    }
+
     /// Indicates whether this [`Field`] supports null values.
     #[inline]
     pub const fn is_nullable(&self) -> bool {

parquet/src/arrow/schema/mod.rs

@@ -29,7 +29,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use arrow_ipc::writer;
-use arrow_schema::{DataType, Field, Fields, Schema, TimeUnit};
+use arrow_schema::{DataType, ExtensionType, Field, Fields, Schema, TimeUnit};
 
 use crate::basic::{
     ConvertedType, LogicalType, Repetition, TimeUnit as ParquetTimeUnit, Type as PhysicalType,
@@ -468,6 +468,10 @@ fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
                 .with_repetition(repetition)
                 .with_id(id)
                 .with_length(*length)
+                .with_logical_type(match field.extension_type() {
+                    Some(ExtensionType::Uuid) => Some(LogicalType::Uuid),
+                    _ => None,
+                })
                 .build()
         }
         DataType::BinaryView => Type::primitive_type_builder(name, PhysicalType::BYTE_ARRAY)