From c2f6aa81eeab0e98b238b722b661647969c048cf Mon Sep 17 00:00:00 2001 From: sdf-jkl Date: Tue, 27 Oct 2026 15:13:43 -0400 Subject: [PATCH] feat: Support the canonical Parquet Variant extension type --- Cargo.lock | 47 +++++++++++++++ Cargo.toml | 1 + datafusion/common/Cargo.toml | 1 + .../src/types/canonical_extensions/mod.rs | 2 + .../canonical_extensions/parquet_variant.rs | 57 +++++++++++++++++++ datafusion/expr/Cargo.toml | 1 + datafusion/expr/src/registry.rs | 12 +++- 7 files changed, 120 insertions(+), 1 deletion(-) create mode 100644 datafusion/common/src/types/canonical_extensions/parquet_variant.rs diff --git a/Cargo.lock b/Cargo.lock index cbf3c06465bef..371b211a6d77e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1927,6 +1927,7 @@ dependencies = [ "num-traits", "object_store", "parquet", + "parquet-variant-compute", "rand 0.9.4", "recursive", "sqlparser", @@ -2201,6 +2202,7 @@ dependencies = [ "indexmap 2.14.2", "insta", "itertools 0.15.0", + "parquet-variant-compute", "recursive", "serde", "serde_json", @@ -4698,6 +4700,51 @@ dependencies = [ "zstd", ] +[[package]] +name = "parquet-variant" +version = "60.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae8443f2224cbf220f0457a70ad8f41e62cbde25c2c0d0d7c6dfebdc389894b8" +dependencies = [ + "arrow-schema", + "chrono", + "half", + "indexmap 2.14.2", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-compute" +version = "60.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cd02f2a545490a8d7c26f9f740627291e3c1b68272aa1f9bbb76e28d3ad5d61" +dependencies = [ + "arrow", + "arrow-schema", + "chrono", + "half", + "indexmap 2.14.2", + "num-traits", + "parquet-variant", + "parquet-variant-json", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "60.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d83b1566b90235a000b3795081af4aa62843669e9a9e1a48cb9b63b23d48bfb3" +dependencies = [ + "arrow-schema", + "base64 0.23.1", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + [[package]] name = "parse-display" version = "0.9.1" diff --git a/Cargo.toml b/Cargo.toml index 3070ed33f3231..ed96d7f6d61f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -186,6 +186,7 @@ parquet = { version = "60.0.0", default-features = false, features = [ "async", "object_store", ] } +parquet-variant-compute = { version = "60.0.0" } pbjson = { version = "0.9.0" } pbjson-types = "0.9" percent-encoding = "2.3" diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 9ee199fe82f28..ba6e0fbb8289d 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -84,6 +84,7 @@ log = { workspace = true } num-traits = { workspace = true } object_store = { workspace = true, optional = true } parquet = { workspace = true, optional = true, default-features = true } +parquet-variant-compute = { workspace = true } recursive = { workspace = true, optional = true } sqlparser = { workspace = true, optional = true } tokio = { workspace = true } diff --git a/datafusion/common/src/types/canonical_extensions/mod.rs b/datafusion/common/src/types/canonical_extensions/mod.rs index 2d74d0669d213..a3f6da5314248 100644 --- a/datafusion/common/src/types/canonical_extensions/mod.rs +++ b/datafusion/common/src/types/canonical_extensions/mod.rs @@ -19,6 +19,7 @@ mod bool8; mod fixed_shape_tensor; mod json; mod opaque; +mod parquet_variant; mod timestamp_with_offset; mod uuid; mod variable_shape_tensor; @@ -27,6 +28,7 @@ pub use bool8::DFBool8; pub use fixed_shape_tensor::DFFixedShapeTensor; pub use json::DFJson; pub use opaque::DFOpaque; +pub use parquet_variant::DFVariant; pub use timestamp_with_offset::DFTimestampWithOffset; pub use uuid::DFUuid; pub use variable_shape_tensor::DFVariableShapeTensor; diff --git a/datafusion/common/src/types/canonical_extensions/parquet_variant.rs b/datafusion/common/src/types/canonical_extensions/parquet_variant.rs new file mode 100644 index 0000000000000..5bb93d6e3b377 --- /dev/null +++ b/datafusion/common/src/types/canonical_extensions/parquet_variant.rs @@ -0,0 +1,57 @@ +// 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::Result; +use crate::types::extension::DFExtensionType; +use arrow::datatypes::DataType; +use arrow_schema::extension::ExtensionType; +use parquet_variant_compute::VariantType; + +/// Defines the extension type logic for the canonical `arrow.parquet.variant` extension type. +/// +/// See [`DFExtensionType`] for information on DataFusion's extension type mechanism. See also +/// [`VariantType`] for the implementation of arrow-rs, which this type uses internally. +/// +/// +#[derive(Debug, Clone)] +pub struct DFVariant { + storage_type: DataType, +} + +impl DFVariant { + /// Creates a new [`DFVariant`], validating that the storage type is compatible with the + /// extension type. + pub fn try_new( + data_type: &DataType, + metadata: ::Metadata, + ) -> Result { + VariantType::try_new(data_type, metadata)?; + Ok(Self { + storage_type: data_type.clone(), + }) + } +} + +impl DFExtensionType for DFVariant { + fn storage_type(&self) -> DataType { + self.storage_type.clone() + } + + fn serialize_metadata(&self) -> Option { + VariantType.serialize_metadata() + } +} diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 45b66ee89885e..eec77892ecd94 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -64,6 +64,7 @@ datafusion-proto-common = { workspace = true, optional = true } datafusion-proto-models = { workspace = true, optional = true } indexmap = { workspace = true } itertools = { workspace = true } +parquet-variant-compute = { workspace = true } recursive = { workspace = true, optional = true } serde = { workspace = true } serde_json = { workspace = true } diff --git a/datafusion/expr/src/registry.rs b/datafusion/expr/src/registry.rs index 4b9744d9573b6..d0f9d0bc749e4 100644 --- a/datafusion/expr/src/registry.rs +++ b/datafusion/expr/src/registry.rs @@ -29,9 +29,10 @@ use arrow_schema::extension::{ }; use datafusion_common::types::{ DFBool8, DFExtensionTypeRef, DFFixedShapeTensor, DFJson, DFOpaque, - DFTimestampWithOffset, DFUuid, DFVariableShapeTensor, + DFTimestampWithOffset, DFUuid, DFVariableShapeTensor, DFVariant, }; use datafusion_common::{HashMap, Result, not_impl_err, plan_datafusion_err}; +use parquet_variant_compute::VariantType; use std::collections::HashSet; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, RwLock}; @@ -493,6 +494,15 @@ impl MemoryExtensionTypeRegistry { Bool8::deserialize_metadata(metadata)?, )?)) }), + ExtensionTypeRegistration::new_arc( + VariantType::NAME, + |storage_type, metadata| { + Ok(Arc::new(DFVariant::try_new( + storage_type, + VariantType::deserialize_metadata(metadata)?, + )?)) + }, + ), ExtensionTypeRegistration::new_arc( TimestampWithOffset::NAME, |storage_type, metadata| {