Add draft for more extension types

Author: tobixdev

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

@@ -20,18 +20,55 @@ use crate::types::extension::DFExtensionType;
 use arrow::array::{Array, Int8Array};
 use arrow::datatypes::DataType;
 use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult};
+use arrow_schema::ArrowError;
+use arrow_schema::extension::{Bool8, ExtensionType};
 use std::fmt::Write;
 
-/// Defines the extension type logic for the canonical `arrow.uuid` extension type.
+/// Defines the extension type logic for the canonical `arrow.bool8` extension type.
 ///
 /// See [`DFExtensionType`] for information on DataFusion's extension type mechanism.
-impl DFExtensionType for arrow_schema::extension::Bool8 {
+#[derive(Debug, Clone)]
+pub struct DFBool8(Bool8);
+
+impl ExtensionType for DFBool8 {
+    const NAME: &'static str = Bool8::NAME;
+    type Metadata = <Bool8 as ExtensionType>::Metadata;
+
+    fn metadata(&self) -> &Self::Metadata {
+        self.0.metadata()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.0.serialize_metadata()
+    }
+
+    fn deserialize_metadata(
+        metadata: Option<&str>,
+    ) -> Result<Self::Metadata, ArrowError> {
+        Bool8::deserialize_metadata(metadata)
+    }
+
+    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
+        self.0.supports_data_type(data_type)
+    }
+
+    fn try_new(
+        data_type: &DataType,
+        metadata: Self::Metadata,
+    ) -> Result<Self, ArrowError> {
+        Ok(Self(<Bool8 as ExtensionType>::try_new(
+            data_type, metadata,
+        )?))
+    }
+}
+
+impl DFExtensionType for DFBool8 {
     fn storage_type(&self) -> DataType {
         DataType::Int8
     }
 
     fn serialize_metadata(&self) -> Option<String> {
-        None
+        self.0.serialize_metadata()
     }
 
     fn create_array_formatter<'fmt>(
@@ -82,7 +119,7 @@ mod tests {
     pub fn test_pretty_bool8() {
         let values = Int8Array::from_iter([Some(0), Some(1), Some(-20), None]);
 
-        let extension_type = arrow_schema::extension::Bool8 {};
+        let extension_type = DFBool8(Bool8 {});
         let formatter = extension_type
             .create_array_formatter(&values, &FormatOptions::default().with_null("NULL"))
             .unwrap()

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

@@ -0,0 +1,78 @@
+// 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.  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::extension::DFExtensionType;
+use arrow::datatypes::DataType;
+use arrow_schema::ArrowError;
+use arrow_schema::extension::{ExtensionType, Json};
+
+/// Defines the extension type logic for the canonical `arrow.json` extension type.
+///
+/// See [`DFExtensionType`] for information on DataFusion's extension type mechanism.
+#[derive(Debug, Clone)]
+pub struct DFJson {
+    inner: Json,
+    storage_type: DataType,
+}
+
+impl ExtensionType for DFJson {
+    const NAME: &'static str = Json::NAME;
+    type Metadata = <Json as ExtensionType>::Metadata;
+
+    fn metadata(&self) -> &Self::Metadata {
+        self.inner.metadata()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+
+    fn deserialize_metadata(
+        metadata: Option<&str>,
+    ) -> Result<Self::Metadata, ArrowError> {
+        Json::deserialize_metadata(metadata)
+    }
+
+    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
+        self.inner.supports_data_type(data_type)
+    }
+
+    fn try_new(
+        data_type: &DataType,
+        metadata: Self::Metadata,
+    ) -> Result<Self, ArrowError> {
+        Ok(Self {
+            inner: <Json as ExtensionType>::try_new(data_type, metadata)?,
+            storage_type: data_type.clone(),
+        })
+    }
+}
+
+impl DFExtensionType for DFJson {
+    fn storage_type(&self) -> DataType {
+        self.storage_type.clone()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+}

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

@@ -4,7 +4,12 @@
 // 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
+// with the License.  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
 //
@@ -16,4 +21,15 @@
 // under the License.
 
 mod bool8;
+mod json;
+mod opaque;
+mod tensor;
+mod timestamp_with_offset;
 mod uuid;
+
+pub use bool8::DFBool8;
+pub use json::DFJson;
+pub use opaque::DFOpaque;
+pub use tensor::{DFFixedShapeTensor, DFVariableShapeTensor};
+pub use timestamp_with_offset::DFTimestampWithOffset;
+pub use uuid::DFUuid;

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

@@ -0,0 +1,78 @@
+// 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.  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::extension::DFExtensionType;
+use arrow::datatypes::DataType;
+use arrow_schema::ArrowError;
+use arrow_schema::extension::{ExtensionType, Opaque};
+
+/// Defines the extension type logic for the canonical `arrow.opaque` extension type.
+///
+/// See [`DFExtensionType`] for information on DataFusion's extension type mechanism.
+#[derive(Debug, Clone)]
+pub struct DFOpaque {
+    inner: Opaque,
+    storage_type: DataType,
+}
+
+impl ExtensionType for DFOpaque {
+    const NAME: &'static str = Opaque::NAME;
+    type Metadata = <Opaque as ExtensionType>::Metadata;
+
+    fn metadata(&self) -> &Self::Metadata {
+        self.inner.metadata()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+
+    fn deserialize_metadata(
+        metadata: Option<&str>,
+    ) -> Result<Self::Metadata, ArrowError> {
+        Opaque::deserialize_metadata(metadata)
+    }
+
+    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
+        self.inner.supports_data_type(data_type)
+    }
+
+    fn try_new(
+        data_type: &DataType,
+        metadata: Self::Metadata,
+    ) -> Result<Self, ArrowError> {
+        Ok(Self {
+            inner: <Opaque as ExtensionType>::try_new(data_type, metadata)?,
+            storage_type: data_type.clone(),
+        })
+    }
+}
+
+impl DFExtensionType for DFOpaque {
+    fn storage_type(&self) -> DataType {
+        self.storage_type.clone()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+}

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

@@ -0,0 +1,133 @@
+// 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::extension::DFExtensionType;
+use arrow::datatypes::DataType;
+use arrow_schema::ArrowError;
+use arrow_schema::extension::{ExtensionType, FixedShapeTensor, VariableShapeTensor};
+
+/// Defines the extension type logic for the canonical `arrow.fixed_shape_tensor` extension type.
+///
+/// See [`DFExtensionType`] for information on DataFusion's extension type mechanism.
+#[derive(Debug, Clone)]
+pub struct DFFixedShapeTensor {
+    inner: FixedShapeTensor,
+    /// The storage type of the tensor.
+    ///
+    /// While we could reconstruct the storage type from the inner [`FixedShapeTensor`], we may
+    /// choose a different name for the field within the [`DataType::FixedSizeList`] which can
+    /// cause problems down the line (e.g., checking for equality).
+    storage_type: DataType,
+}
+
+impl ExtensionType for DFFixedShapeTensor {
+    const NAME: &'static str = FixedShapeTensor::NAME;
+    type Metadata = <FixedShapeTensor as ExtensionType>::Metadata;
+
+    fn metadata(&self) -> &Self::Metadata {
+        self.inner.metadata()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+
+    fn deserialize_metadata(
+        metadata: Option<&str>,
+    ) -> Result<Self::Metadata, ArrowError> {
+        FixedShapeTensor::deserialize_metadata(metadata)
+    }
+
+    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
+        self.inner.supports_data_type(data_type)
+    }
+
+    fn try_new(
+        data_type: &DataType,
+        metadata: Self::Metadata,
+    ) -> Result<Self, ArrowError> {
+        Ok(Self {
+            inner: <FixedShapeTensor as ExtensionType>::try_new(data_type, metadata)?,
+            storage_type: data_type.clone(),
+        })
+    }
+}
+
+impl DFExtensionType for DFFixedShapeTensor {
+    fn storage_type(&self) -> DataType {
+        self.storage_type.clone()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+}
+
+/// Defines the extension type logic for the canonical `arrow.variable_shape_tensor` extension type.
+///
+/// See [`DFExtensionType`] for information on DataFusion's extension type mechanism.
+#[derive(Debug, Clone)]
+pub struct DFVariableShapeTensor {
+    inner: VariableShapeTensor,
+    /// While we could reconstruct the storage type from the inner [`VariableShapeTensor`], we may
+    /// choose a different name for the field within the [`DataType::List`] which can cause problems
+    /// down the line (e.g., checking for equality).
+    storage_type: DataType,
+}
+
+impl ExtensionType for DFVariableShapeTensor {
+    const NAME: &'static str = VariableShapeTensor::NAME;
+    type Metadata = <VariableShapeTensor as ExtensionType>::Metadata;
+
+    fn metadata(&self) -> &Self::Metadata {
+        self.inner.metadata()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+
+    fn deserialize_metadata(
+        metadata: Option<&str>,
+    ) -> Result<Self::Metadata, ArrowError> {
+        VariableShapeTensor::deserialize_metadata(metadata)
+    }
+
+    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
+        self.inner.supports_data_type(data_type)
+    }
+
+    fn try_new(
+        data_type: &DataType,
+        metadata: Self::Metadata,
+    ) -> Result<Self, ArrowError> {
+        Ok(Self {
+            inner: <VariableShapeTensor as ExtensionType>::try_new(data_type, metadata)?,
+            storage_type: data_type.clone(),
+        })
+    }
+}
+
+impl DFExtensionType for DFVariableShapeTensor {
+    fn storage_type(&self) -> DataType {
+        self.storage_type.clone()
+    }
+
+    fn serialize_metadata(&self) -> Option<String> {
+        self.inner.serialize_metadata()
+    }
+}