Add draft for basic extension type support

Author: tschwarzinger

Cargo.lock

@@ -81,6 +81,7 @@ paste = { workspace = true }
 recursive = { workspace = true, optional = true }
 sqlparser = { workspace = true, optional = true }
 tokio = { workspace = true }
+uuid = { workspace = true, features = ["v4"] }
 
 [target.'cfg(target_family = "wasm")'.dependencies]
 web-time = "1.1.0"

datafusion/common/Cargo.toml

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

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

@@ -0,0 +1,94 @@
+use crate::error::_internal_err;
+use crate::types::extension::DFExtensionType;
+use arrow::array::{Array, FixedSizeBinaryArray};
+use arrow::datatypes::DataType;
+use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult};
+use std::fmt::Write;
+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 {
+    fn create_array_formatter<'fmt>(
+        &self,
+        array: &'fmt dyn Array,
+        options: &FormatOptions<'fmt>,
+    ) -> crate::Result<Option<ArrayFormatter<'fmt>>> {
+        if array.data_type() != &DataType::FixedSizeBinary(16) {
+            return _internal_err!("Wrong array type for Uuid");
+        }
+
+        let display_index = UuidValueDisplayIndex {
+            array: array.as_any().downcast_ref().unwrap(),
+            null_str: options.null(),
+        };
+        Ok(Some(ArrayFormatter::new(
+            Box::new(display_index),
+            options.safe(),
+        )))
+    }
+}
+
+/// Pretty printer for binary UUID values.
+#[derive(Debug, Clone, Copy)]
+struct UuidValueDisplayIndex<'a> {
+    array: &'a FixedSizeBinaryArray,
+    null_str: &'a str,
+}
+
+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)?;
+            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})")?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ScalarValue;
+
+    #[test]
+    pub fn test_pretty_print_uuid() {
+        let my_uuid = Uuid::nil();
+        let uuid = ScalarValue::FixedSizeBinary(16, Some(my_uuid.as_bytes().to_vec()))
+            .to_array_of_size(1)
+            .unwrap();
+
+        let extension_type = UuidDFExtensionType::new();
+        let formatter = extension_type
+            .create_array_formatter(uuid.as_ref(), &FormatOptions::default())
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(
+            formatter.value(0).to_string(),
+            "arrow.uuid(00000000-0000-0000-0000-000000000000)"
+        );
+    }
+}

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

@@ -0,0 +1,53 @@
+use crate::error::Result;
+use arrow::array::Array;
+use arrow::datatypes::DataType;
+use arrow::util::display::{ArrayFormatter, FormatOptions};
+use std::fmt::Debug;
+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:
+/// - 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.
+///
+/// Furthermore, Arrow's current trait is not dyn-compatible which we need for implementing
+/// extension type registries. In the future, the two implementations may increasingly converge.
+///
+/// # Example
+///
+///
+pub trait DFExtensionType: Debug + Send + Sync {
+    /// Returns an [`ArrayFormatter`] that can format values of this type.
+    ///
+    /// If `Ok(None)` is returned, the default implementation will be used.
+    /// If an error is returned, there was an error creating the formatter.
+    fn create_array_formatter<'fmt>(
+        &self,
+        _array: &'fmt dyn Array,
+        _options: &FormatOptions<'fmt>,
+    ) -> Result<Option<ArrayFormatter<'fmt>>> {
+        Ok(None)
+    }
+}

datafusion/common/src/types/extension.rs

@@ -16,11 +16,15 @@
 // under the License.
 
 mod builtin;
+mod canonical_extensions;
+mod extension;
 mod field;
 mod logical;
 mod native;
 
 pub use builtin::*;
+pub use canonical_extensions::*;
+pub use extension::*;
 pub use field::*;
 pub use logical::*;
 pub use native::*;

datafusion/common/src/types/mod.rs

@@ -69,6 +69,7 @@ use datafusion_functions_aggregate::expr_fn::{
     avg, count, max, median, min, stddev, sum,
 };
 
+use crate::extension_types::DFArrayFormatterFactory;
 use async_trait::async_trait;
 use datafusion_catalog::Session;
 
@@ -1516,6 +1517,11 @@ impl DataFrame {
         let options = self.session_state.config().options().format.clone();
         let arrow_options: arrow::util::display::FormatOptions = (&options).try_into()?;
 
+        let registry = self.session_state.extension_type_registry();
+        let formatter_factory = DFArrayFormatterFactory::new(Arc::clone(registry));
+        let arrow_options =
+            arrow_options.with_formatter_factory(Some(&formatter_factory));
+
         let results = self.collect().await?;
         Ok(
             pretty::pretty_format_batches_with_options(&results, &arrow_options)?

datafusion/core/src/dataframe/mod.rs

@@ -30,6 +30,7 @@ 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::{
@@ -56,7 +57,7 @@ use datafusion_expr::expr_rewriter::FunctionRewrite;
 use datafusion_expr::planner::ExprPlanner;
 #[cfg(feature = "sql")]
 use datafusion_expr::planner::{RelationPlanner, TypePlanner};
-use datafusion_expr::registry::{FunctionRegistry, SerializerRegistry};
+use datafusion_expr::registry::{ExtensionTypeRegistration, ExtensionTypeRegistrationRef, ExtensionTypeRegistry, ExtensionTypeRegistryRef, FunctionRegistry, MemoryExtensionTypeRegistry, SerializerRegistry, SimpleExtensionTypeRegistration};
 use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::{AggregateUDF, Explain, Expr, LogicalPlan, ScalarUDF, WindowUDF};
 use datafusion_optimizer::simplify_expressions::ExprSimplifier;
@@ -77,6 +78,7 @@ 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;
@@ -158,6 +160,8 @@ pub struct SessionState {
     aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
     /// Window functions registered in the context
     window_functions: HashMap<String, Arc<WindowUDF>>,
+    /// Extension types registry for extensions.
+    extension_types: ExtensionTypeRegistryRef,
     /// Deserializer registry for extensions.
     serializer_registry: Arc<dyn SerializerRegistry>,
     /// Holds registered external FileFormat implementations
@@ -266,6 +270,10 @@ impl Session for SessionState {
         &self.window_functions
     }
 
+    fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef {
+        &self.extension_types
+    }
+
     fn runtime_env(&self) -> &Arc<RuntimeEnv> {
         self.runtime_env()
     }
@@ -986,6 +994,7 @@ pub struct SessionStateBuilder {
     scalar_functions: Option<Vec<Arc<ScalarUDF>>>,
     aggregate_functions: Option<Vec<Arc<AggregateUDF>>>,
     window_functions: Option<Vec<Arc<WindowUDF>>>,
+    extension_types: Option<ExtensionTypeRegistryRef>,
     serializer_registry: Option<Arc<dyn SerializerRegistry>>,
     file_formats: Option<Vec<Arc<dyn FileFormatFactory>>>,
     config: Option<SessionConfig>,
@@ -1026,6 +1035,7 @@ impl SessionStateBuilder {
             scalar_functions: None,
             aggregate_functions: None,
             window_functions: None,
+            extension_types: None,
             serializer_registry: None,
             file_formats: None,
             table_options: None,
@@ -1081,6 +1091,7 @@ impl SessionStateBuilder {
                 existing.aggregate_functions.into_values().collect_vec(),
             ),
             window_functions: Some(existing.window_functions.into_values().collect_vec()),
+            extension_types: Some(existing.extension_types),
             serializer_registry: Some(existing.serializer_registry),
             file_formats: Some(existing.file_formats.into_values().collect_vec()),
             config: Some(new_config),
@@ -1126,6 +1137,11 @@ impl SessionStateBuilder {
             .get_or_insert_with(Vec::new)
             .extend(SessionStateDefaults::default_window_functions());
 
+        self.extension_types
+            .get_or_insert_with(|| Arc::new(MemoryExtensionTypeRegistry::new()))
+            .extend(&SessionStateDefaults::default_extension_types())
+            .expect("MemoryExtensionTypeRegistry is not read-only.");
+
         self.table_functions
             .get_or_insert_with(HashMap::new)
             .extend(
@@ -1316,6 +1332,44 @@ impl SessionStateBuilder {
         self
     }
 
+    /// Set the map of [`ExtensionTypeRegistration`]s
+    pub fn with_extension_type(
+        mut self,
+        registry: ExtensionTypeRegistryRef,
+    ) -> Self {
+        self.extension_types = Some(registry);
+        self
+    }
+
+    /// Registers [canonical extension types](https://arrow.apache.org/docs/format/CanonicalExtensions.html)
+    /// in DataFusion's extension type registry. For more information see [`ExtensionTypeRegistry`].
+    ///
+    /// # Errors
+    ///
+    /// 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()),
+        )];
+
+        match &self.extension_types {
+            None => {
+                let registry = Arc::new(MemoryExtensionTypeRegistry::new());
+                registry
+                    .extend(&canonical_extension_types)
+                    .expect("Adding valid extension types to MemoryExtensionTypeRegistry always succeeds.");
+                self.extension_types = Some(registry);
+            }
+            Some(registry) => {
+                registry.extend(&canonical_extension_types)?;
+            }
+        }
+
+        Ok(self)
+    }
+
     /// Set the [`SerializerRegistry`]
     pub fn with_serializer_registry(
         mut self,
@@ -1454,6 +1508,7 @@ impl SessionStateBuilder {
             scalar_functions,
             aggregate_functions,
             window_functions,
+            extension_types,
             serializer_registry,
             file_formats,
             table_options,
@@ -1490,6 +1545,7 @@ impl SessionStateBuilder {
             scalar_functions: HashMap::new(),
             aggregate_functions: HashMap::new(),
             window_functions: HashMap::new(),
+            extension_types: Arc::new(MemoryExtensionTypeRegistry::default()),
             serializer_registry: serializer_registry
                 .unwrap_or_else(|| Arc::new(EmptySerializerRegistry)),
             file_formats: HashMap::new(),
@@ -1559,6 +1615,10 @@ impl SessionStateBuilder {
             });
         }
 
+        if let Some(extension_types) = extension_types {
+            state.extension_types = extension_types;
+        }
+
         if state.config.create_default_catalog_and_schema() {
             let default_catalog = SessionStateDefaults::default_catalog(
                 &state.config,
@@ -2071,6 +2131,35 @@ impl datafusion_execution::TaskContextProvider for SessionState {
     }
 }
 
+impl ExtensionTypeRegistry for SessionState {
+    fn extension_type_registration(
+        &self,
+        name: &str,
+    ) -> datafusion_common::Result<ExtensionTypeRegistrationRef> {
+        self.extension_types.extension_type_registration(name)
+    }
+
+    fn extension_type_registrations(&self) -> Vec<Arc<dyn ExtensionTypeRegistration>> {
+        self.extension_types.extension_type_registrations()
+    }
+
+    fn add_extension_type_registration(
+        &self,
+        extension_type: ExtensionTypeRegistrationRef,
+    ) -> datafusion_common::Result<Option<ExtensionTypeRegistrationRef>> {
+        self.extension_types
+            .add_extension_type_registration(extension_type)
+    }
+
+    fn remove_extension_type_registration(
+        &self,
+        name: &str,
+    ) -> datafusion_common::Result<Option<ExtensionTypeRegistrationRef>> {
+        self.extension_types
+            .remove_extension_type_registration(name)
+    }
+}
+
 impl OptimizerConfig for SessionState {
     fn query_execution_start_time(&self) -> Option<DateTime<Utc>> {
         self.execution_props.query_execution_start_time