Split tensor extension types into multiple files

Author: tobixdev

…src/types/canonical_extensions/tensor.rs …nonical_extensions/fixed_shape_tensor.rsdatafusion/common/src/types/canonical_extensions/tensor.rs renamed to datafusion/common/src/types/canonical_extensions/fixed_shape_tensor.rs datafusion/common/src/types/canonical_extensions/tensor.rs renamed to datafusion/common/src/types/canonical_extensions/fixed_shape_tensor.rs

@@ -18,7 +18,7 @@
 use crate::types::extension::DFExtensionType;
 use arrow::datatypes::DataType;
 use arrow_schema::ArrowError;
-use arrow_schema::extension::{ExtensionType, FixedShapeTensor, VariableShapeTensor};
+use arrow_schema::extension::{ExtensionType, FixedShapeTensor};
 
 /// Defines the extension type logic for the canonical `arrow.fixed_shape_tensor` extension type.
 ///
@@ -76,58 +76,3 @@ impl DFExtensionType for DFFixedShapeTensor {
         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()
-    }
-}

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

@@ -21,15 +21,17 @@
 // under the License.
 
 mod bool8;
+mod fixed_shape_tensor;
 mod json;
 mod opaque;
-mod tensor;
 mod timestamp_with_offset;
 mod uuid;
+mod variable_shape_tensor;
 
 pub use bool8::DFBool8;
+pub use fixed_shape_tensor::DFFixedShapeTensor;
 pub use json::DFJson;
 pub use opaque::DFOpaque;
-pub use tensor::{DFFixedShapeTensor, DFVariableShapeTensor};
 pub use timestamp_with_offset::DFTimestampWithOffset;
 pub use uuid::DFUuid;
+pub use variable_shape_tensor::DFVariableShapeTensor;

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

@@ -186,8 +186,10 @@ fn format_offset(minutes: i16) -> String {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use arrow::array::{StructArray, TimestampSecondArray};
     use arrow::datatypes::{Field, Fields};
     use chrono::{TimeZone, Utc};
+    use std::sync::Arc;
 
     #[test]
     fn test_pretty_print_timestamp_with_offset() -> Result<(), ArrowError> {

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

@@ -0,0 +1,76 @@
+// 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, VariableShapeTensor};
+
+/// 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()
+    }
+}