Further improvements of the extension type API proposal

Author: tschwarzinger

Cargo.lock

@@ -133,7 +133,7 @@ cargo run --example dataframe -- dataframe
 
 | Subcommand | File Path | Description |
 | --- | --- | --- |
-| my_id | [`extension_types/event_id.rs`](examples/extension_types/event_id.rs) | A custom wrapper around integers that represent event ids |
+| event_id | [`extension_types/event_id.rs`](examples/extension_types/event_id.rs) | A custom wrapper around integers that represent event ids |
 
 ## External Dependency Examples
 

datafusion-examples/README.md

@@ -24,9 +24,9 @@ use datafusion::error::Result;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion_common::internal_err;
-use datafusion_common::types::{DFExtensionType, DFExtensionTypeRef};
+use datafusion_common::types::DFExtensionType;
 use datafusion_expr::registry::{
-    ExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry,
+    DefaultExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry,
 };
 use std::fmt::Write;
 use std::sync::Arc;
@@ -48,7 +48,9 @@ pub async fn event_id_example() -> Result<()> {
 fn create_session_context() -> Result<SessionContext> {
     // Create a registry with a reference to the custom extension type implementation.
     let registry = MemoryExtensionTypeRegistry::new();
-    let event_id_registration = Arc::new(EventIdExtensionTypeRegistration {});
+    let event_id_registration = DefaultExtensionTypeRegistration::new_arc(|metadata| {
+        Ok(EventIdExtensionType(metadata))
+    });
     registry.add_extension_type_registration(event_id_registration)?;
 
     // Set the extension type registry in the session state so that DataFusion can use it.
@@ -104,8 +106,8 @@ fn example_schema() -> SchemaRef {
     ]))
 }
 
-/// Represents a 32-bit custom identifier that represents a single event. Using this format is not
-/// a good idea in practice, but it is useful for demonstrating the API usage.
+/// Represents a 32-bit custom identifier that represents a single event. Using this format is
+/// probably not a good idea in practice, but it is useful for demonstrating the API usage.
 ///
 /// An event is constructed of three parts:
 /// - The year
@@ -274,30 +276,6 @@ impl DisplayIndex for EventIdDisplayIndex<'_> {
     }
 }
 
-/// The registration is the last piece missing for the extension type implementation. It contains
-/// the logic for deserializing the metadata from the arrow [`Field`]s and creating the extension
-/// type instance. We cannot use the trait from arrow-rs as it's not dyn-compatible (the Metadata
-/// type must be known at compile time).
-///
-/// If an extension type does not have any parameters, the [`SimpleExtensionTypeRegistration`]
-/// provides an easier way of registering it.
-#[derive(Debug)]
-pub struct EventIdExtensionTypeRegistration();
-
-impl ExtensionTypeRegistration for EventIdExtensionTypeRegistration {
-    fn type_name(&self) -> &str {
-        EventIdExtensionType::NAME
-    }
-
-    fn create_df_extension_type(
-        &self,
-        metadata: Option<&str>,
-    ) -> Result<DFExtensionTypeRef> {
-        let metadata = EventIdExtensionType::deserialize_metadata(metadata)?;
-        Ok(Arc::new(EventIdExtensionType(metadata)))
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion-examples/examples/extension_types/event_id.rs

@@ -21,13 +21,13 @@
 //!
 //! ## Usage
 //! ```bash
-//! cargo run --example dataframe -- [all|my_id]
+//! cargo run --example extension_types -- [all|event_id]
 //! ```
 //!
 //! Each subcommand runs a corresponding example:
 //! - `all` — run all examples included in this module
 //!
-//! - `my_id`
+//! - `event_id`
 //!   (file: event_id.rs, desc: A custom wrapper around integers that represent event ids)
 
 mod event_id;

datafusion-examples/examples/extension_types/main.rs

@@ -66,6 +66,7 @@ apache-avro = { workspace = true, features = [
     "zstandard",
 ], optional = true }
 arrow = { workspace = true }
+arrow-schema = { workspace = true, features = ["canonical_extension_types"] }
 arrow-ipc = { workspace = true }
 chrono = { workspace = true }
 half = { workspace = true }

datafusion/common/Cargo.toml

@@ -1,3 +1 @@
 mod uuid;
-
-pub use uuid::*;

datafusion/common/src/types/canonical_extensions/mod.rs

@@ -9,23 +9,7 @@ use uuid::{Bytes, Uuid};
 /// Defines the extension type logic for the canonical `arrow.uuid` extension type.
 ///
 /// See [`DFExtensionType`] for information on DataFusion's extension type mechanism.
-#[derive(Debug)]
-pub struct UuidDFExtensionType();
-
-impl UuidDFExtensionType {
-    /// Create a new instance of [`UuidDFExtensionType`].
-    pub fn new() -> Self {
-        Self {}
-    }
-}
-
-impl Default for UuidDFExtensionType {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl DFExtensionType for UuidDFExtensionType {
+impl DFExtensionType for arrow_schema::extension::Uuid {
     fn create_array_formatter<'fmt>(
         &self,
         array: &'fmt dyn Array,
@@ -56,14 +40,14 @@ struct UuidValueDisplayIndex<'a> {
 impl DisplayIndex for UuidValueDisplayIndex<'_> {
     fn write(&self, idx: usize, f: &mut dyn Write) -> FormatResult {
         if self.array.is_null(idx) {
-            write!(f, "arrow.uuid({})", self.null_str)?;
+            write!(f, "{}", self.null_str)?;
             return Ok(());
         }
 
         let bytes = Bytes::try_from(self.array.value(idx))
             .expect("FixedSizeBinaryArray length checked in create_array_formatter");
         let uuid = Uuid::from_bytes(bytes);
-        write!(f, "arrow.uuid({uuid})")?;
+        write!(f, "{uuid}")?;
         Ok(())
     }
 }
@@ -88,7 +72,7 @@ mod tests {
 
         assert_eq!(
             formatter.value(0).to_string(),
-            "arrow.uuid(00000000-0000-0000-0000-000000000000)"
+            "00000000-0000-0000-0000-000000000000"
         );
     }
 }

datafusion/common/src/types/canonical_extensions/uuid.rs

@@ -7,36 +7,38 @@ use std::sync::Arc;
 /// A cheaply cloneable pointer to a [`DFExtensionType`].
 pub type DFExtensionTypeRef = Arc<dyn DFExtensionType>;
 
-/// Represents an implementation of a DataFusion extension type, allowing users to customize the
-/// behavior of DataFusion for custom extension types.
-///
-/// Extension types may change the semantics of a column. For example, adding two values of
-/// [`DataType::Int64`] is a sensible thing to do. However, if the same data type is annotated with
-/// an extension type like `custom.id`, the correct interpretation of a column changes. For example,
-/// adding together two `custom.id` values (represented as a 64-bit integer) may no longer make
-/// sense.
-///
-/// Note that while helping users to navigate the semantic gap between the data type and extension
-/// types is a goal of this trait, DataFusion's extension type support is still evolving and does
-/// not cover all use cases. Currently, the following capabilities can be customized:
+/// Represents an implementation of a DataFusion extension type.
+///
+/// This allows users to customize the behavior of DataFusion for certain types. Having this ability
+/// is necessary because extension types affect how columns should be treated by the query engine.
+/// This effect includes which operations are possible on a column and what are the expected results
+/// from these operations. The extension type mechanism allows users to define how these operations
+/// apply to a particular extension type.
+///
+/// For example, adding two values of [`DataType::Int64`] is a sensible thing to do. However, if the
+/// same column is annotated with an extension type like `custom.id`, the correct interpretation of
+/// a column changes. Adding together two `custom.id` values, even though they are stored as
+/// integers, may no longer make sense.
+///
+/// Note that DataFusion's extension type support is still young and therefore might not cover all
+/// relevant use cases. Currently, the following operations can be customized:
 /// - Pretty-printing values in record batches
 ///
 /// # Relation to Arrow's `ExtensionType`
 ///
-/// The purpose of Arrow's `ExtensionType` trait, for the time being, is to provide a way to handle
-/// metadata of an extension type in a type-safe manner. The trait does not provide any
-/// customization options such that users can customize the behavior of any kernels (e.g.,
-/// [`DFExtensionType::create_array_formatter`] for formatting record batches). Therefore,
-/// downstream users (such as DataFusion) have the flexibility to implement the extension type
-/// mechanism according to their needs. [`DFExtensionType`] is DataFusion's implementation of this
-/// extension type mechanism.
+/// The purpose of Arrow's `ExtensionType` trait, for the time being, is to allow reading and
+/// writing extension type metadata in a type-safe manner. The trait does not provide any
+/// customization options. Therefore, downstream users (such as DataFusion) have the flexibility to
+/// implement the extension type mechanism according to their needs. [`DFExtensionType`] is
+/// DataFusion's implementation of this extension type mechanism.
 ///
-/// Furthermore, Arrow's current trait is not dyn-compatible which we need for implementing
+/// Furthermore, the current trait in arrow-rs is not dyn-compatible, which we need for implementing
 /// extension type registries. In the future, the two implementations may increasingly converge.
 ///
-/// # Example
-///
+/// # Examples
 ///
+/// Examples for using the extension type machinery can be found in the DataFusion examples
+/// directory.
 pub trait DFExtensionType: Debug + Send + Sync {
     /// Returns an [`ArrayFormatter`] that can format values of this type.
     ///

datafusion/common/src/types/extension.rs

@@ -23,7 +23,6 @@ mod logical;
 mod native;
 
 pub use builtin::*;
-pub use canonical_extensions::*;
 pub use extension::*;
 pub use field::*;
 pub use logical::*;

datafusion/common/src/types/mod.rs

@@ -30,7 +30,6 @@ use crate::datasource::provider_as_source;
 use crate::execution::SessionStateDefaults;
 use crate::execution::context::{EmptySerializerRegistry, FunctionFactory, QueryPlanner};
 use crate::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
-use arrow_schema::extension::ExtensionType;
 use arrow_schema::{DataType, FieldRef};
 use datafusion_catalog::MemoryCatalogProviderList;
 use datafusion_catalog::information_schema::{
@@ -58,9 +57,9 @@ use datafusion_expr::planner::ExprPlanner;
 #[cfg(feature = "sql")]
 use datafusion_expr::planner::{RelationPlanner, TypePlanner};
 use datafusion_expr::registry::{
-    ExtensionTypeRegistration, ExtensionTypeRegistrationRef, ExtensionTypeRegistry,
-    ExtensionTypeRegistryRef, FunctionRegistry, MemoryExtensionTypeRegistry,
-    SerializerRegistry, SimpleExtensionTypeRegistration,
+    DefaultExtensionTypeRegistration, ExtensionTypeRegistration,
+    ExtensionTypeRegistrationRef, ExtensionTypeRegistry, ExtensionTypeRegistryRef,
+    FunctionRegistry, MemoryExtensionTypeRegistry, SerializerRegistry,
 };
 use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::{AggregateUDF, Explain, Expr, LogicalPlan, ScalarUDF, WindowUDF};
@@ -82,7 +81,6 @@ use datafusion_sql::{
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
-use datafusion_common::types::UuidDFExtensionType;
 use itertools::Itertools;
 use log::{debug, info};
 use object_store::ObjectStore;
@@ -1353,10 +1351,10 @@ impl SessionStateBuilder {
     /// May fail if an already registered [`ExtensionTypeRegistry`] raises an error while
     /// registering the canonical extension types.
     pub fn with_canonical_extension_types(mut self) -> datafusion_common::Result<Self> {
-        let canonical_extension_types = vec![SimpleExtensionTypeRegistration::new_arc(
-            arrow_schema::extension::Uuid::NAME,
-            Arc::new(UuidDFExtensionType::new()),
-        )];
+        let uuid = DefaultExtensionTypeRegistration::new_arc(|_| {
+            Ok(arrow_schema::extension::Uuid {})
+        });
+        let canonical_extension_types = vec![uuid];
 
         match &self.extension_types {
             None => {