Remove requirement to arrow-rs ExtensionType

tschwarzinger · tschwarzinger · commit 28578aff3c84 · 2026-04-02T19:46:29.000+02:00
diff --git a/datafusion-examples/examples/extension_types/temperature.rs b/datafusion-examples/examples/extension_types/temperature.rs
@@ -21,7 +21,7 @@ use arrow::array::{
 };
 use arrow::datatypes::{Float32Type, Float64Type};
 use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult};
-use arrow_schema::extension::ExtensionType;
+use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
 use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef};
 use datafusion::dataframe::DataFrame;
 use datafusion::error::Result;
@@ -32,6 +32,7 @@ use datafusion_common::types::DFExtensionType;
 use datafusion_expr::registry::{
     DefaultExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry,
 };
+use std::collections::HashMap;
 use std::fmt::{Display, Write};
 use std::sync::Arc;
 
@@ -50,15 +51,15 @@ fn create_session_context() -> Result<SessionContext> {
     let registry = MemoryExtensionTypeRegistry::new_empty();
 
     // The registration creates a new instance of the extension type with the deserialized metadata.
-    let temp_registration =
-        DefaultExtensionTypeRegistration::<TemperatureExtensionType>::new_arc(
-            |storage_type, metadata| {
-                Ok(Arc::new(TemperatureExtensionType::try_new(
-                    storage_type,
-                    metadata,
-                )?))
-            },
-        );
+    let temp_registration = DefaultExtensionTypeRegistration::new_arc(
+        TemperatureExtensionType::NAME,
+        |storage_type, metadata| {
+            Ok(Arc::new(TemperatureExtensionType::try_new(
+                storage_type,
+                TemperatureUnit::try_from(metadata)?,
+            )?))
+        },
+    );
     registry.add_extension_type_registration(temp_registration)?;
 
     let state = SessionStateBuilder::default()
@@ -98,38 +99,15 @@ async fn register_temperature_table(ctx: &SessionContext) -> Result<DataFrame> {
 fn example_schema() -> SchemaRef {
     Arc::new(Schema::new(vec![
         Field::new("city", DataType::Utf8, false),
-        Field::new("celsius", DataType::Float64, false).with_extension_type(
-            TemperatureExtensionType::try_new(
-                &DataType::Float64,
-                TemperatureUnit::Celsius,
-            )
-            .expect("Valid Type"),
-        ),
-        Field::new("fahrenheit", DataType::Float64, false).with_extension_type(
-            TemperatureExtensionType::try_new(
-                &DataType::Float64,
-                TemperatureUnit::Fahrenheit,
-            )
-            .expect("Valid Type"),
-        ),
-        Field::new("kelvin", DataType::Float32, false).with_extension_type(
-            TemperatureExtensionType::try_new(
-                &DataType::Float32,
-                TemperatureUnit::Kelvin,
-            )
-            .expect("Valid Type"),
-        ),
+        Field::new("celsius", DataType::Float64, false)
+            .with_metadata(create_metadata(TemperatureUnit::Celsius)),
+        Field::new("fahrenheit", DataType::Float64, false)
+            .with_metadata(create_metadata(TemperatureUnit::Fahrenheit)),
+        Field::new("kelvin", DataType::Float32, false)
+            .with_metadata(create_metadata(TemperatureUnit::Kelvin)),
     ]))
 }
 
-/// Represents the unit of a temperature reading.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum TemperatureUnit {
-    Celsius,
-    Fahrenheit,
-    Kelvin,
-}
-
 /// Represents a float that semantically represents a temperature. The temperature can be one of
 /// the supported [`TemperatureUnit`]s.
 ///
@@ -157,51 +135,61 @@ pub struct TemperatureExtensionType {
 }
 
 impl TemperatureExtensionType {
+    /// The name of the extension type.
+    pub const NAME: &'static str = "custom.temperature";
+
     /// Creates a new [`TemperatureExtensionType`].
     pub fn try_new(
         storage_type: &DataType,
         temperature_unit: TemperatureUnit,
-    ) -> Result<Self> {
+    ) -> Result<Self, ArrowError> {
+        match storage_type {
+            DataType::Float32 | DataType::Float64 => {}
+            _ => {
+                return Err(ArrowError::InvalidArgumentError(format!(
+                    "Invalid data type: {storage_type} for temperature type, expected Float32 or Float64",
+                )));
+            }
+        }
+
         let result = Self {
             storage_type: storage_type.clone(),
             temperature_unit,
         };
-        result.supports_data_type(storage_type)?; // Validate the storage type
         Ok(result)
     }
 }
 
-/// Implementation of [`ExtensionType`] for [`TemperatureExtensionType`].
-///
-/// This implements the arrow-rs trait for reading, writing, and validating extension types.
-impl ExtensionType for TemperatureExtensionType {
-    /// Arrow extension type name that is stored in the `ARROW:extension:name` field.
-    const NAME: &'static str = "custom.temperature";
-    type Metadata = TemperatureUnit;
-
-    fn metadata(&self) -> &Self::Metadata {
-        &self.temperature_unit
-    }
+/// Represents the unit of a temperature reading.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum TemperatureUnit {
+    Celsius,
+    Fahrenheit,
+    Kelvin,
+}
 
+impl TemperatureUnit {
     /// Arrow extension type metadata is encoded as a string and stored using the
     /// `ARROW:extension:metadata` key. As we only store the name of the unit, a simple string
     /// suffices. Extension types can store more complex metadata using serialization formats like
     /// JSON.
-    fn serialize_metadata(&self) -> Option<String> {
-        let s = match self.temperature_unit {
+    pub fn serialize(self) -> String {
+        let result = match self {
             TemperatureUnit::Celsius => "celsius",
             TemperatureUnit::Fahrenheit => "fahrenheit",
             TemperatureUnit::Kelvin => "kelvin",
         };
-        Some(s.to_string())
+        result.to_owned()
     }
+}
 
-    /// Inverse operation of [`Self::serialize_metadata`]. This creates the [`TemperatureUnit`]
-    /// value from the serialized string.
-    fn deserialize_metadata(
-        metadata: Option<&str>,
-    ) -> std::result::Result<Self::Metadata, ArrowError> {
-        match metadata {
+/// Inverse operation of [`TemperatureUnit::serialize`]. This creates the [`TemperatureUnit`]
+/// value from the serialized string.
+impl TryFrom<Option<&str>> for TemperatureUnit {
+    type Error = ArrowError;
+
+    fn try_from(value: Option<&str>) -> std::result::Result<Self, Self::Error> {
+        match value {
             Some("celsius") => Ok(TemperatureUnit::Celsius),
             Some("fahrenheit") => Ok(TemperatureUnit::Fahrenheit),
             Some("kelvin") => Ok(TemperatureUnit::Kelvin),
@@ -213,28 +201,18 @@ impl ExtensionType for TemperatureExtensionType {
             )),
         }
     }
+}
 
-    /// Checks that the extension type supports a given [`DataType`].
-    fn supports_data_type(
-        &self,
-        data_type: &DataType,
-    ) -> std::result::Result<(), ArrowError> {
-        match data_type {
-            DataType::Float32 | DataType::Float64 => Ok(()),
-            _ => Err(ArrowError::InvalidArgumentError(format!(
-                "Invalid data type: {data_type} for temperature type, expected Float32 or Float64",
-            ))),
-        }
-    }
-
-    fn try_new(
-        data_type: &DataType,
-        metadata: Self::Metadata,
-    ) -> std::result::Result<Self, ArrowError> {
-        let instance = Self::try_new(data_type, metadata)?;
-        instance.supports_data_type(data_type)?;
-        Ok(instance)
-    }
+/// This creates a metadata map for the temperature type. Another way of writing the metadata can be
+/// implemented using arrow-rs' [`ExtensionType`](arrow_schema::extension::ExtensionType) trait.
+fn create_metadata(unit: TemperatureUnit) -> HashMap<String, String> {
+    HashMap::from([
+        (
+            EXTENSION_TYPE_NAME_KEY.to_owned(),
+            TemperatureExtensionType::NAME.to_owned(),
+        ),
+        (EXTENSION_TYPE_METADATA_KEY.to_owned(), unit.serialize()),
+    ])
 }
 
 /// Implementation of [`DFExtensionType`] for [`TemperatureExtensionType`].
@@ -246,7 +224,7 @@ impl DFExtensionType for TemperatureExtensionType {
     }
 
     fn serialize_metadata(&self) -> Option<String> {
-        ExtensionType::serialize_metadata(self)
+        Some(self.temperature_unit.serialize())
     }
 
     fn create_array_formatter<'fmt>(
diff --git a/datafusion/expr/src/registry.rs b/datafusion/expr/src/registry.rs
@@ -336,62 +336,54 @@ pub trait ExtensionTypeRegistry: Debug + Send + Sync {
 
 /// A factory that creates instances of extension types from a storage [`DataType`] and the
 /// metadata.
-pub type ExtensionTypeFactory<TExtensionType> = dyn Fn(
-        &DataType,
-        <TExtensionType as ExtensionType>::Metadata,
-    ) -> Result<DFExtensionTypeRef>
-    + Send
-    + Sync;
+pub type ExtensionTypeFactory =
+    dyn Fn(&DataType, Option<&str>) -> Result<DFExtensionTypeRef> + Send + Sync;
 
 /// A default implementation of [ExtensionTypeRegistration] that parses the metadata from the
-/// given extension type and passes it to a constructor function. The [`ExtensionType::NAME`] is
-/// used for registering the extension type.
-pub struct DefaultExtensionTypeRegistration<TExtensionType: ExtensionType + 'static> {
+/// given extension type and passes it to a constructor function.
+pub struct DefaultExtensionTypeRegistration {
+    /// The name of the extension type.
+    name: String,
     /// A function that creates an instance of [`DFExtensionTypeRef`] from the storage type and the
     /// metadata.
-    factory: Box<ExtensionTypeFactory<TExtensionType>>,
+    factory: Box<ExtensionTypeFactory>,
 }
 
-impl<TExtensionType: ExtensionType + 'static>
-    DefaultExtensionTypeRegistration<TExtensionType>
-{
+impl DefaultExtensionTypeRegistration {
     /// Creates a new registration for an extension type. The factory is required to validate that
     /// the storage [`DataType`] is compatible with the extension type.
     pub fn new_arc(
-        factory: impl Fn(&DataType, TExtensionType::Metadata) -> Result<DFExtensionTypeRef>
+        name: impl Into<String>,
+        factory: impl Fn(&DataType, Option<&str>) -> Result<DFExtensionTypeRef>
         + Send
         + Sync
         + 'static,
     ) -> ExtensionTypeRegistrationRef {
         Arc::new(Self {
+            name: name.into(),
             factory: Box::new(factory),
         })
     }
 }
 
-impl<TExtensionType: ExtensionType> ExtensionTypeRegistration
-    for DefaultExtensionTypeRegistration<TExtensionType>
-{
+impl ExtensionTypeRegistration for DefaultExtensionTypeRegistration {
     fn type_name(&self) -> &str {
-        TExtensionType::NAME
+        &self.name
     }
 
     fn create_df_extension_type(
         &self,
         storage_type: &DataType,
         metadata: Option<&str>,
     ) -> Result<DFExtensionTypeRef> {
-        let metadata = TExtensionType::deserialize_metadata(metadata)?;
         self.factory.as_ref()(storage_type, metadata)
     }
 }
 
-impl<TExtensionType: ExtensionType> Debug
-    for DefaultExtensionTypeRegistration<TExtensionType>
-{
+impl Debug for DefaultExtensionTypeRegistration {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("DefaultExtensionTypeRegistration")
-            .field("type_name", &TExtensionType::NAME)
+            .field("type_name", &self.name)
             .finish()
     }
 }
@@ -421,47 +413,66 @@ impl MemoryExtensionTypeRegistry {
     /// in the extension type registry.
     pub fn new_with_canonical_extension_types() -> Self {
         let mapping = [
-            DefaultExtensionTypeRegistration::<FixedShapeTensor>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                FixedShapeTensor::NAME,
                 |storage_type, metadata| {
                     Ok(Arc::new(DFFixedShapeTensor::try_new(
                         storage_type,
-                        metadata,
+                        FixedShapeTensor::deserialize_metadata(metadata)?,
                     )?))
                 },
             ),
-            DefaultExtensionTypeRegistration::<VariableShapeTensor>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                VariableShapeTensor::NAME,
                 |storage_type, metadata| {
                     Ok(Arc::new(DFVariableShapeTensor::try_new(
                         storage_type,
-                        metadata,
+                        VariableShapeTensor::deserialize_metadata(metadata)?,
                     )?))
                 },
             ),
-            DefaultExtensionTypeRegistration::<Json>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                Json::NAME,
                 |storage_type, metadata| {
-                    Ok(Arc::new(DFJson::try_new(storage_type, metadata)?))
+                    Ok(Arc::new(DFJson::try_new(
+                        storage_type,
+                        Json::deserialize_metadata(metadata)?,
+                    )?))
                 },
             ),
-            DefaultExtensionTypeRegistration::<Uuid>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                Uuid::NAME,
                 |storage_type, metadata| {
-                    Ok(Arc::new(DFUuid::try_new(storage_type, metadata)?))
+                    Ok(Arc::new(DFUuid::try_new(
+                        storage_type,
+                        Uuid::deserialize_metadata(metadata)?,
+                    )?))
                 },
             ),
-            DefaultExtensionTypeRegistration::<Opaque>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                Opaque::NAME,
                 |storage_type, metadata| {
-                    Ok(Arc::new(DFOpaque::try_new(storage_type, metadata)?))
+                    Ok(Arc::new(DFOpaque::try_new(
+                        storage_type,
+                        Opaque::deserialize_metadata(metadata)?,
+                    )?))
                 },
             ),
-            DefaultExtensionTypeRegistration::<Bool8>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                Bool8::NAME,
                 |storage_type, metadata| {
-                    Ok(Arc::new(DFBool8::try_new(storage_type, metadata)?))
+                    Ok(Arc::new(DFBool8::try_new(
+                        storage_type,
+                        Bool8::deserialize_metadata(metadata)?,
+                    )?))
                 },
             ),
-            DefaultExtensionTypeRegistration::<TimestampWithOffset>::new_arc(
+            DefaultExtensionTypeRegistration::new_arc(
+                TimestampWithOffset::NAME,
                 |storage_type, metadata| {
                     Ok(Arc::new(DFTimestampWithOffset::try_new(
                         storage_type,
-                        metadata,
+                        TimestampWithOffset::deserialize_metadata(metadata)?,
                     )?))
                 },
             ),
