Add simple registration.

Pretty-printing record batches is working.

Author: tschwarzinger

datafusion/core/src/execution/session_state.rs

@@ -1125,6 +1125,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(

datafusion/core/src/execution/session_state_defaults.rs

@@ -30,12 +30,16 @@ use crate::execution::context::SessionState;
 #[cfg(feature = "nested_expressions")]
 use crate::functions_nested;
 use crate::{functions, functions_aggregate, functions_table, functions_window};
+use arrow_schema::extension::{ExtensionType, Uuid};
 use datafusion_catalog::TableFunction;
 use datafusion_catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion_execution::config::SessionConfig;
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_expr::planner::ExprPlanner;
+use datafusion_expr::registry::{
+    ExtensionTypeRegistrationRef, SimpleExtensionTypeRegistration,
+};
 use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -122,6 +126,14 @@ impl SessionStateDefaults {
         functions_window::all_default_window_functions()
     }
 
+    /// returns the list of default [`WindowUDF`]s
+    pub fn default_extension_types() -> Vec<ExtensionTypeRegistrationRef> {
+        vec![SimpleExtensionTypeRegistration::new_arc(
+            Uuid::NAME,
+            Arc::new(Uuid),
+        )]
+    }
+
     /// returns the list of default [`TableFunction`]s
     pub fn default_table_functions() -> Vec<Arc<TableFunction>> {
         functions_table::all_default_table_functions()

datafusion/core/tests/extension_types/pretty_print_data_frames.rs

@@ -42,11 +42,12 @@ async fn test_pretty_print_extension_types() -> Result<()> {
     assert_snapshot!(
         result,
         @r"
-    +---------------+
-    | ascii(test.a) |
-    +---------------+
-    | 97            |
-    +---------------+
+    +--------------------------------------------------+
+    | my_uuids                                         |
+    +--------------------------------------------------+
+    | arrow.uuid(00000000-0000-0000-0000-000000000000) |
+    | arrow.uuid(00010203-0405-0607-0809-000102030506) |
+    +--------------------------------------------------+
     "
     );
 

datafusion/expr/src/registry.rs

@@ -22,7 +22,7 @@ use crate::planner::ExprPlanner;
 use crate::{AggregateUDF, ScalarUDF, UserDefinedLogicalNode, WindowUDF};
 use arrow_schema::Field;
 use datafusion_common::types::{LogicalTypeRef, NativeType};
-use datafusion_common::{HashMap, Result, not_impl_err, plan_datafusion_err};
+use datafusion_common::{not_impl_err, plan_datafusion_err, HashMap, Result};
 use std::collections::HashSet;
 use std::fmt::Debug;
 use std::sync::{Arc, RwLock};
@@ -272,6 +272,17 @@ pub trait ExtensionTypeRegistry: Debug + Send + Sync {
         extension_type: ExtensionTypeRegistrationRef,
     ) -> Result<Option<ExtensionTypeRegistrationRef>>;
 
+    /// Extends the registry with the provided extension types.
+    ///
+    /// Returns an error if the type cannot be registered, for example, if the registry is
+    /// read-only.
+    fn extend(&self, extension_types: &[ExtensionTypeRegistrationRef]) -> Result<()> {
+        for extension_type in extension_types {
+            self.register_extension_type(extension_type.clone())?;
+        }
+        Ok(())
+    }
+
     /// Deregisters an extension type registration with the name `name`, returning the
     /// implementation that was deregistered.
     ///
@@ -283,6 +294,41 @@ pub trait ExtensionTypeRegistry: Debug + Send + Sync {
     ) -> Result<Option<ExtensionTypeRegistrationRef>>;
 }
 
+/// A simple implementation of [ExtensionTypeRegistration] where the logical type instance can be
+/// obtained by cloning the inner [`LogicalTypeRef`].
+///
+/// As a result, the logical type instance cannot depend on any parameters. For example, if the
+/// logical type is `Opaque(Uuid)`, and the `create_logical_type` method should return an instance
+/// of `Opaque` where `Uuid` is a parameter, then this implementation cannot be used, as the inner
+/// type (`Uuid`) needs to be extracted from the field and passed into the `Opaque` constructor.
+#[derive(Debug)]
+pub struct SimpleExtensionTypeRegistration {
+    /// The name of the extension type.
+    name: String,
+    /// The logical type instance.
+    logical_type: LogicalTypeRef,
+}
+
+impl SimpleExtensionTypeRegistration {
+    /// Creates a new registration for the given `name` and `logical_type`.
+    pub fn new_arc(name: &str, logical_type: LogicalTypeRef) -> Arc<Self> {
+        Arc::new(Self {
+            name: name.to_string(),
+            logical_type,
+        })
+    }
+}
+
+impl ExtensionTypeRegistration for SimpleExtensionTypeRegistration {
+    fn type_name(&self) -> &str {
+        &self.name
+    }
+
+    fn create_logical_type(&self, _field: &Field) -> Result<LogicalTypeRef> {
+        Ok(Arc::clone(&self.logical_type))
+    }
+}
+
 /// An [`ExtensionTypeRegistry`] that uses in memory [`HashMap`]s.
 #[derive(Clone, Debug)]
 pub struct MemoryExtensionTypeRegistry {