Add first draft of using type registry for uuids

Author: tschwarzinger

datafusion/catalog-listing/src/helpers.rs

@@ -541,6 +541,7 @@ mod tests {
     use std::ops::Not;
 
     use super::*;
+    use datafusion_expr::registry::MemoryExtensionTypeRegistry;
     use datafusion_expr::{
         case, col, lit, AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF,
     };
@@ -1059,6 +1060,10 @@ mod tests {
             unimplemented!()
         }
 
+        fn extension_types(&self) -> &MemoryExtensionTypeRegistry {
+            unimplemented!()
+        }
+
         fn runtime_env(&self) -> &Arc<RuntimeEnv> {
             unimplemented!()
         }

datafusion/common/src/types/canonical.rs

@@ -0,0 +1,41 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::types::{LogicalType, NativeType, TypeSignature};
+
+/// Represents the canonical [UUID extension type](https://arrow.apache.org/docs/format/CanonicalExtensions.html#uuid).
+pub struct UuidType;
+
+impl UuidType {
+    /// Creates a new [UuidType].
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl LogicalType for UuidType {
+    fn native(&self) -> &NativeType {
+        &NativeType::FixedSizeBinary(16)
+    }
+
+    fn signature(&self) -> TypeSignature<'_> {
+        TypeSignature::Extension {
+            name: "arrow.uuid",
+            parameters: &[],
+        }
+    }
+}

datafusion/common/src/types/mod.rs

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

datafusion/core/src/execution/session_state.rs

@@ -932,6 +932,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<Vec<LogicalTypeRef>>,
     serializer_registry: Option<Arc<dyn SerializerRegistry>>,
     file_formats: Option<Vec<Arc<dyn FileFormatFactory>>>,
     config: Option<SessionConfig>,
@@ -969,6 +970,7 @@ impl SessionStateBuilder {
             scalar_functions: None,
             aggregate_functions: None,
             window_functions: None,
+            extension_types: None,
             serializer_registry: None,
             file_formats: None,
             table_options: None,
@@ -1021,6 +1023,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.all_types()),
             serializer_registry: Some(existing.serializer_registry),
             file_formats: Some(existing.file_formats.into_values().collect_vec()),
             config: Some(new_config),
@@ -1407,6 +1410,7 @@ impl SessionStateBuilder {
             scalar_functions: HashMap::new(),
             aggregate_functions: HashMap::new(),
             window_functions: HashMap::new(),
+            extension_types: MemoryExtensionTypeRegistry::new(),
             serializer_registry: serializer_registry
                 .unwrap_or_else(|| Arc::new(EmptySerializerRegistry)),
             file_formats: HashMap::new(),
@@ -1456,6 +1460,15 @@ impl SessionStateBuilder {
             });
         }
 
+        if let Some(extension_types) = extension_types {
+            extension_types.into_iter().for_each(|ext_type| {
+                let existing_type = state.register_extension_type(ext_type);
+                if let Ok(Some(existing_type)) = existing_type {
+                    debug!("Overwrote an existing UDF: {}", existing_type);
+                }
+            });
+        }
+
         if state.config.create_default_catalog_and_schema() {
             let default_catalog = SessionStateDefaults::default_catalog(
                 &state.config,

datafusion/datasource/src/url.rs

@@ -411,6 +411,7 @@ mod tests {
     use datafusion_execution::runtime_env::RuntimeEnv;
     use datafusion_execution::TaskContext;
     use datafusion_expr::execution_props::ExecutionProps;
+    use datafusion_expr::registry::MemoryExtensionTypeRegistry;
     use datafusion_expr::{AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF};
     use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
     use datafusion_physical_plan::ExecutionPlan;
@@ -767,6 +768,10 @@ mod tests {
             unimplemented!()
         }
 
+        fn extension_types(&self) -> &MemoryExtensionTypeRegistry {
+            unimplemented!()
+        }
+
         fn runtime_env(&self) -> &Arc<RuntimeEnv> {
             &self.runtime_env
         }

datafusion/expr/src/registry.rs

@@ -245,14 +245,16 @@ pub trait ExtensionTypeRegistry {
 #[derive(Clone, Debug)]
 pub struct MemoryExtensionTypeRegistry {
     /// Holds a mapping between the name of an extension type and its logical type.
+    ///
+    /// TODO: Use TypeSignature to support arguments
     extension_types: HashMap<String, LogicalTypeRef>,
 }
 
 impl Default for MemoryExtensionTypeRegistry {
     fn default() -> Self {
-        let mut registry = MemoryExtensionTypeRegistry::new();
-        // TODO add canonical types
-        registry
+        MemoryExtensionTypeRegistry {
+            extension_types: HashMap::new(),
+        }
     }
 }
 
@@ -263,6 +265,29 @@ impl MemoryExtensionTypeRegistry {
             extension_types: HashMap::new(),
         }
     }
+
+    /// Creates a new [MemoryExtensionTypeRegistry] with the provided `types`.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if one of the `types` is a native type.
+    pub fn new_with_types(
+        types: impl IntoIterator<Item = LogicalTypeRef>,
+    ) -> Result<Self> {
+        let extension_types = types
+            .into_iter()
+            .map(|t| match t.signature() {
+                TypeSignature::Native(_) => todo!("TODO"),
+                TypeSignature::Extension { name, .. } => (name.to_owned(), t),
+            })
+            .collect();
+        Ok(Self { extension_types })
+    }
+
+    /// Returns a list of all registered types.
+    pub fn all_types(&self) -> Vec<LogicalTypeRef> {
+        self.extension_types.values().cloned().collect()
+    }
 }
 
 impl ExtensionTypeRegistry for MemoryExtensionTypeRegistry {

datafusion/functions/src/string/uuid.rs

@@ -18,14 +18,15 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::array::GenericStringBuilder;
-use arrow::datatypes::DataType;
+use arrow::array::FixedSizeBinaryBuilder;
 use arrow::datatypes::DataType::Utf8;
+use arrow::datatypes::{DataType, Field, FieldRef};
 use rand::Rng;
 use uuid::Uuid;
 
 use datafusion_common::{internal_err, Result};
-use datafusion_expr::{ColumnarValue, Documentation, Volatility};
+use datafusion_common::types::UuidType;
+use datafusion_expr::{ColumnarValue, Documentation, ReturnFieldArgs, Volatility};
 use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
 use datafusion_macros::user_doc;
 
@@ -75,7 +76,12 @@ impl ScalarUDFImpl for UuidFunc {
     }
 
     fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(Utf8)
+        unreachable!("return_field_from_args is overwritten")
+    }
+
+    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
+        // TODO: pass-in the registry
+        Ok(Arc::new(Field::new("output", Utf8, false).with_extension_type(UuidType::new())))
     }
 
     /// Prints random (v4) uuid values per row
@@ -90,18 +96,17 @@ impl ScalarUDFImpl for UuidFunc {
         let mut randoms = vec![0u128; args.number_rows];
         rng.fill(&mut randoms[..]);
 
-        let mut builder = GenericStringBuilder::<i32>::with_capacity(
-            args.number_rows,
-            args.number_rows * 36,
-        );
+        let mut builder = FixedSizeBinaryBuilder::with_capacity(args.number_rows, 16);
 
         let mut buffer = [0u8; 36];
         for x in &mut randoms {
             // From Uuid::new_v4(): Mask out the version and variant bits
             *x = *x & 0xFFFFFFFFFFFF4FFFBFFFFFFFFFFFFFFF | 0x40008000000000000000;
             let uuid = Uuid::from_u128(*x);
             let fmt = uuid::fmt::Hyphenated::from_uuid(uuid);
-            builder.append_value(fmt.encode_lower(&mut buffer));
+            builder
+                .append_value(fmt.encode_lower(&mut buffer))
+                .expect("Value always has 16 bytes");
         }
 
         Ok(ColumnarValue::Array(Arc::new(builder.finish())))