Skip to content

Commit 5cbf369

Browse files
committed
Remove necessity to implement arrow-rs' ExtensionType for DF extension types
1 parent 150b685 commit 5cbf369

8 files changed

Lines changed: 116 additions & 233 deletions

File tree

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

Lines changed: 8 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,12 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::error::_internal_err;
1920
use crate::types::extension::DFExtensionType;
2021
use arrow::array::{Array, Int8Array};
2122
use arrow::datatypes::DataType;
2223
use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult};
23-
use arrow_schema::ArrowError;
2424
use arrow_schema::extension::{Bool8, ExtensionType};
2525
use std::fmt::Write;
2626

@@ -30,32 +30,13 @@ use std::fmt::Write;
3030
#[derive(Debug, Clone)]
3131
pub struct DFBool8(Bool8);
3232

33-
impl ExtensionType for DFBool8 {
34-
const NAME: &'static str = Bool8::NAME;
35-
type Metadata = <Bool8 as ExtensionType>::Metadata;
36-
37-
fn metadata(&self) -> &Self::Metadata {
38-
self.0.metadata()
39-
}
40-
41-
fn serialize_metadata(&self) -> Option<String> {
42-
self.0.serialize_metadata()
43-
}
44-
45-
fn deserialize_metadata(
46-
metadata: Option<&str>,
47-
) -> Result<Self::Metadata, ArrowError> {
48-
Bool8::deserialize_metadata(metadata)
49-
}
50-
51-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
52-
self.0.supports_data_type(data_type)
53-
}
54-
55-
fn try_new(
33+
impl DFBool8 {
34+
/// Creates a new [`DFBool8`], validating that the storage type is compatible with the
35+
/// extension type.
36+
pub fn try_new(
5637
data_type: &DataType,
57-
metadata: Self::Metadata,
58-
) -> Result<Self, ArrowError> {
38+
metadata: <Bool8 as ExtensionType>::Metadata,
39+
) -> Result<Self> {
5940
Ok(Self(<Bool8 as ExtensionType>::try_new(
6041
data_type, metadata,
6142
)?))
@@ -75,7 +56,7 @@ impl DFExtensionType for DFBool8 {
7556
&self,
7657
array: &'fmt dyn Array,
7758
options: &FormatOptions<'fmt>,
78-
) -> crate::Result<Option<ArrayFormatter<'fmt>>> {
59+
) -> Result<Option<ArrayFormatter<'fmt>>> {
7960
if array.data_type() != &DataType::Int8 {
8061
return _internal_err!("Wrong array type for Bool8");
8162
}

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

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::types::extension::DFExtensionType;
1920
use arrow::datatypes::DataType;
20-
use arrow_schema::ArrowError;
2121
use arrow_schema::extension::{ExtensionType, FixedShapeTensor};
2222

2323
/// Defines the extension type logic for the canonical `arrow.fixed_shape_tensor` extension type.
@@ -34,32 +34,13 @@ pub struct DFFixedShapeTensor {
3434
storage_type: DataType,
3535
}
3636

37-
impl ExtensionType for DFFixedShapeTensor {
38-
const NAME: &'static str = FixedShapeTensor::NAME;
39-
type Metadata = <FixedShapeTensor as ExtensionType>::Metadata;
40-
41-
fn metadata(&self) -> &Self::Metadata {
42-
self.inner.metadata()
43-
}
44-
45-
fn serialize_metadata(&self) -> Option<String> {
46-
self.inner.serialize_metadata()
47-
}
48-
49-
fn deserialize_metadata(
50-
metadata: Option<&str>,
51-
) -> Result<Self::Metadata, ArrowError> {
52-
FixedShapeTensor::deserialize_metadata(metadata)
53-
}
54-
55-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
56-
self.inner.supports_data_type(data_type)
57-
}
58-
59-
fn try_new(
37+
impl DFFixedShapeTensor {
38+
/// Creates a new [`DFFixedShapeTensor`], validating that the storage type is compatible with
39+
/// the extension type.
40+
pub fn try_new(
6041
data_type: &DataType,
61-
metadata: Self::Metadata,
62-
) -> Result<Self, ArrowError> {
42+
metadata: <FixedShapeTensor as ExtensionType>::Metadata,
43+
) -> Result<Self> {
6344
Ok(Self {
6445
inner: <FixedShapeTensor as ExtensionType>::try_new(data_type, metadata)?,
6546
storage_type: data_type.clone(),

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

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::types::extension::DFExtensionType;
1920
use arrow::datatypes::DataType;
20-
use arrow_schema::ArrowError;
2121
use arrow_schema::extension::{ExtensionType, Json};
2222

2323
/// Defines the extension type logic for the canonical `arrow.json` extension type.
@@ -29,32 +29,13 @@ pub struct DFJson {
2929
storage_type: DataType,
3030
}
3131

32-
impl ExtensionType for DFJson {
33-
const NAME: &'static str = Json::NAME;
34-
type Metadata = <Json as ExtensionType>::Metadata;
35-
36-
fn metadata(&self) -> &Self::Metadata {
37-
self.inner.metadata()
38-
}
39-
40-
fn serialize_metadata(&self) -> Option<String> {
41-
self.inner.serialize_metadata()
42-
}
43-
44-
fn deserialize_metadata(
45-
metadata: Option<&str>,
46-
) -> Result<Self::Metadata, ArrowError> {
47-
Json::deserialize_metadata(metadata)
48-
}
49-
50-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
51-
self.inner.supports_data_type(data_type)
52-
}
53-
54-
fn try_new(
32+
impl DFJson {
33+
/// Creates a new [`DFJson`], validating that the storage type is compatible with the
34+
/// extension type.
35+
pub fn try_new(
5536
data_type: &DataType,
56-
metadata: Self::Metadata,
57-
) -> Result<Self, ArrowError> {
37+
metadata: <Json as ExtensionType>::Metadata,
38+
) -> Result<Self> {
5839
Ok(Self {
5940
inner: <Json as ExtensionType>::try_new(data_type, metadata)?,
6041
storage_type: data_type.clone(),

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

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::types::extension::DFExtensionType;
1920
use arrow::datatypes::DataType;
20-
use arrow_schema::ArrowError;
2121
use arrow_schema::extension::{ExtensionType, Opaque};
2222

2323
/// Defines the extension type logic for the canonical `arrow.opaque` extension type.
@@ -29,32 +29,13 @@ pub struct DFOpaque {
2929
storage_type: DataType,
3030
}
3131

32-
impl ExtensionType for DFOpaque {
33-
const NAME: &'static str = Opaque::NAME;
34-
type Metadata = <Opaque as ExtensionType>::Metadata;
35-
36-
fn metadata(&self) -> &Self::Metadata {
37-
self.inner.metadata()
38-
}
39-
40-
fn serialize_metadata(&self) -> Option<String> {
41-
self.inner.serialize_metadata()
42-
}
43-
44-
fn deserialize_metadata(
45-
metadata: Option<&str>,
46-
) -> Result<Self::Metadata, ArrowError> {
47-
Opaque::deserialize_metadata(metadata)
48-
}
49-
50-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
51-
self.inner.supports_data_type(data_type)
52-
}
53-
54-
fn try_new(
32+
impl DFOpaque {
33+
/// Creates a new [`DFOpaque`], validating that the storage type is compatible with the
34+
/// extension type.
35+
pub fn try_new(
5536
data_type: &DataType,
56-
metadata: Self::Metadata,
57-
) -> Result<Self, ArrowError> {
37+
metadata: <Opaque as ExtensionType>::Metadata,
38+
) -> Result<Self> {
5839
Ok(Self {
5940
inner: <Opaque as ExtensionType>::try_new(data_type, metadata)?,
6041
storage_type: data_type.clone(),

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

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::ScalarValue;
1920
use crate::error::_internal_err;
2021
use crate::types::extension::DFExtensionType;
@@ -39,32 +40,13 @@ pub struct DFTimestampWithOffset {
3940
storage_type: DataType,
4041
}
4142

42-
impl ExtensionType for DFTimestampWithOffset {
43-
const NAME: &'static str = TimestampWithOffset::NAME;
44-
type Metadata = <TimestampWithOffset as ExtensionType>::Metadata;
45-
46-
fn metadata(&self) -> &Self::Metadata {
47-
self.inner.metadata()
48-
}
49-
50-
fn serialize_metadata(&self) -> Option<String> {
51-
self.inner.serialize_metadata()
52-
}
53-
54-
fn deserialize_metadata(
55-
metadata: Option<&str>,
56-
) -> Result<Self::Metadata, ArrowError> {
57-
TimestampWithOffset::deserialize_metadata(metadata)
58-
}
59-
60-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
61-
self.inner.supports_data_type(data_type)
62-
}
63-
64-
fn try_new(
43+
impl DFTimestampWithOffset {
44+
/// Creates a new [`DFTimestampWithOffset`], validating that the storage type is compatible with
45+
/// the extension type.
46+
pub fn try_new(
6547
data_type: &DataType,
66-
metadata: Self::Metadata,
67-
) -> Result<Self, ArrowError> {
48+
metadata: <TimestampWithOffset as ExtensionType>::Metadata,
49+
) -> Result<Self> {
6850
Ok(Self {
6951
inner: <TimestampWithOffset as ExtensionType>::try_new(data_type, metadata)?,
7052
storage_type: data_type.clone(),
@@ -85,7 +67,7 @@ impl DFExtensionType for DFTimestampWithOffset {
8567
&self,
8668
array: &'fmt dyn Array,
8769
options: &FormatOptions<'fmt>,
88-
) -> crate::Result<Option<ArrayFormatter<'fmt>>> {
70+
) -> Result<Option<ArrayFormatter<'fmt>>> {
8971
if array.data_type() != &self.storage_type {
9072
return _internal_err!(
9173
"Unexpected data type for TimestampWithOffset: {}",

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

Lines changed: 9 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,12 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::error::_internal_err;
1920
use crate::types::extension::DFExtensionType;
2021
use arrow::array::{Array, FixedSizeBinaryArray};
2122
use arrow::datatypes::DataType;
2223
use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult};
23-
use arrow_schema::ArrowError;
2424
use arrow_schema::extension::{ExtensionType, Uuid};
2525
use std::fmt::Write;
2626
use uuid::Bytes;
@@ -31,32 +31,13 @@ use uuid::Bytes;
3131
#[derive(Debug, Clone)]
3232
pub struct DFUuid(Uuid);
3333

34-
impl ExtensionType for DFUuid {
35-
const NAME: &'static str = Uuid::NAME;
36-
type Metadata = <Uuid as ExtensionType>::Metadata;
37-
38-
fn metadata(&self) -> &Self::Metadata {
39-
self.0.metadata()
40-
}
41-
42-
fn serialize_metadata(&self) -> Option<String> {
43-
self.0.serialize_metadata()
44-
}
45-
46-
fn deserialize_metadata(
47-
metadata: Option<&str>,
48-
) -> Result<Self::Metadata, ArrowError> {
49-
Uuid::deserialize_metadata(metadata)
50-
}
51-
52-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
53-
self.0.supports_data_type(data_type)
54-
}
55-
56-
fn try_new(
34+
impl DFUuid {
35+
/// Creates a new [`DFUuid`], validating that the storage type is compatible with the
36+
/// extension type.
37+
pub fn try_new(
5738
data_type: &DataType,
58-
metadata: Self::Metadata,
59-
) -> Result<Self, ArrowError> {
39+
metadata: <Uuid as ExtensionType>::Metadata,
40+
) -> Result<Self> {
6041
Ok(Self(<Uuid as ExtensionType>::try_new(data_type, metadata)?))
6142
}
6243
}
@@ -74,7 +55,7 @@ impl DFExtensionType for DFUuid {
7455
&self,
7556
array: &'fmt dyn Array,
7657
options: &FormatOptions<'fmt>,
77-
) -> crate::Result<Option<ArrayFormatter<'fmt>>> {
58+
) -> Result<Option<ArrayFormatter<'fmt>>> {
7859
if array.data_type() != &DataType::FixedSizeBinary(16) {
7960
return _internal_err!("Wrong array type for Uuid");
8061
}
@@ -116,6 +97,7 @@ impl DisplayIndex for UuidValueDisplayIndex<'_> {
11697
mod tests {
11798
use super::*;
11899
use crate::ScalarValue;
100+
use arrow_schema::ArrowError;
119101

120102
#[test]
121103
pub fn test_pretty_print_uuid() -> Result<(), ArrowError> {

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

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::Result;
1819
use crate::types::extension::DFExtensionType;
1920
use arrow::datatypes::DataType;
20-
use arrow_schema::ArrowError;
2121
use arrow_schema::extension::{ExtensionType, VariableShapeTensor};
2222

2323
/// Defines the extension type logic for the canonical `arrow.variable_shape_tensor` extension type.
@@ -32,32 +32,13 @@ pub struct DFVariableShapeTensor {
3232
storage_type: DataType,
3333
}
3434

35-
impl ExtensionType for DFVariableShapeTensor {
36-
const NAME: &'static str = VariableShapeTensor::NAME;
37-
type Metadata = <VariableShapeTensor as ExtensionType>::Metadata;
38-
39-
fn metadata(&self) -> &Self::Metadata {
40-
self.inner.metadata()
41-
}
42-
43-
fn serialize_metadata(&self) -> Option<String> {
44-
self.inner.serialize_metadata()
45-
}
46-
47-
fn deserialize_metadata(
48-
metadata: Option<&str>,
49-
) -> Result<Self::Metadata, ArrowError> {
50-
VariableShapeTensor::deserialize_metadata(metadata)
51-
}
52-
53-
fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
54-
self.inner.supports_data_type(data_type)
55-
}
56-
57-
fn try_new(
35+
impl DFVariableShapeTensor {
36+
/// Creates a new [`DFVariableShapeTensor`], validating that the storage type is compatible with
37+
/// the extension type.
38+
pub fn try_new(
5839
data_type: &DataType,
59-
metadata: Self::Metadata,
60-
) -> Result<Self, ArrowError> {
40+
metadata: <VariableShapeTensor as ExtensionType>::Metadata,
41+
) -> Result<Self> {
6142
Ok(Self {
6243
inner: <VariableShapeTensor as ExtensionType>::try_new(data_type, metadata)?,
6344
storage_type: data_type.clone(),

0 commit comments

Comments
 (0)