You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/09/08 16:05:58 UTC
(datafusion-comet) branch main updated: feat: normalize marked Variant arrays at the native Parquet boundary (#5715)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new e696f53bec feat: normalize marked Variant arrays at the native Parquet boundary (#5715)
e696f53bec is described below
commit e696f53beca01b3f242165290c6e1d28f78c8414
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Wed Sep 9 00:05:54 2026 +0800
feat: normalize marked Variant arrays at the native Parquet boundary (#5715)
---
native/core/src/parquet/cast_column.rs | 16 +
native/core/src/parquet/cast_column/variant.rs | 596 +++++++++++++++++++++
.../core/src/parquet/cast_column/variant/tests.rs | 550 +++++++++++++++++++
native/core/src/parquet/schema_adapter.rs | 98 +++-
4 files changed, 1257 insertions(+), 3 deletions(-)
diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs
index 231f88c507..8c72455f4d 100644
--- a/native/core/src/parquet/cast_column.rs
+++ b/native/core/src/parquet/cast_column.rs
@@ -14,6 +14,9 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
+mod variant;
+
+use self::variant::normalize_variant_array;
use arrow::{
array::{make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray},
compute::CastOptions,
@@ -26,6 +29,7 @@ use datafusion::common::format::DEFAULT_CAST_OPTIONS;
use datafusion::common::{DataFusionError, Result as DataFusionResult};
use datafusion::logical_expr::ColumnarValue;
use datafusion::physical_expr::PhysicalExpr;
+use parquet::variant::VariantType;
use std::{
fmt::{self, Display},
hash::Hash,
@@ -243,6 +247,18 @@ impl PhysicalExpr for CometCastColumnExpr {
fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> {
let value = self.expr.evaluate(batch)?;
+ if self.target_field.has_valid_extension_type::<VariantType>() {
+ return match value {
+ ColumnarValue::Array(array) => Ok(ColumnarValue::Array(normalize_variant_array(
+ &array,
+ &self.target_field,
+ )?)),
+ ColumnarValue::Scalar(_) => Err(DataFusionError::Execution(
+ "Variant Parquet projection requires an array".to_string(),
+ )),
+ };
+ }
+
// Use == (PartialEq) instead of equals_datatype because equals_datatype
// ignores field names in nested types (Struct, List, Map). We need to detect
// when field names differ (e.g., Struct("a","b") vs Struct("c","d")) so that
diff --git a/native/core/src/parquet/cast_column/variant.rs b/native/core/src/parquet/cast_column/variant.rs
new file mode 100644
index 0000000000..ce2c75a878
--- /dev/null
+++ b/native/core/src/parquet/cast_column/variant.rs
@@ -0,0 +1,596 @@
+// 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 arrow::{
+ array::{
+ make_array, Array, ArrayRef, AsArray, BinaryArray, BinaryBuilder, ListLikeArray,
+ StructArray,
+ },
+ buffer::NullBuffer,
+ compute::cast,
+ datatypes::{DataType, FieldRef},
+ error::ArrowError,
+};
+use datafusion::common::{DataFusionError, Result as DataFusionResult};
+use parquet::variant::{
+ unshred_variant, ListBuilder, MetadataBuilder, ObjectBuilder, ParentState,
+ ReadOnlyMetadataBuilder, ValueBuilder, Variant, VariantArray, VariantMetadata,
+};
+use std::{
+ panic::{catch_unwind, AssertUnwindSafe},
+ sync::Arc,
+};
+
+pub(super) fn normalize_variant_array(
+ array: &ArrayRef,
+ target_field: &FieldRef,
+) -> DataFusionResult<ArrayRef> {
+ let DataType::Struct(fields) = target_field.data_type() else {
+ return Err(DataFusionError::Execution(
+ "Variant extension field must use Struct storage".to_string(),
+ ));
+ };
+ if fields.len() != 2
+ || fields[0].name() != "value"
+ || fields[1].name() != "metadata"
+ || fields
+ .iter()
+ .any(|field| field.data_type() != &DataType::Binary)
+ {
+ return Err(DataFusionError::Execution(
+ "Variant output must contain Binary children [value, metadata]".to_string(),
+ ));
+ }
+
+ // VariantArray resolves metadata/value/typed_value by name, so the reader's child order is
+ // irrelevant. Legacy Spark residuals must be put in Arrow order before the single upstream
+ // unshred call; the whole output is then put back in the order expected by released Spark 4.
+ let variant = VariantArray::try_new(array.as_ref())?;
+ let prepared = prepare_variant_for_unshredding(&variant)?;
+ let unshredded = unshred_variant(&prepared)?;
+ let value = unshredded.value_column();
+ let value = cast(value.as_ref(), &DataType::Binary)?;
+ let metadata = cast(unshredded.metadata_column().as_ref(), &DataType::Binary)?;
+ let value = reorder_variant_values(&value, &metadata, unshredded.inner().nulls())?;
+
+ Ok(Arc::new(StructArray::try_new(
+ fields.clone(),
+ vec![value, metadata],
+ unshredded.inner().nulls().cloned(),
+ )?))
+}
+
+/// Arrow validates every residual `value` while unshredding. Spark versions before SPARK-58949
+/// wrote object keys in Java UTF-16 order, so rewrite every reachable legacy residual to Arrow's
+/// UTF-8 order before unshredding. `metadata_rows` carries each root metadata row through nested
+/// lists.
+fn rewrite_shredding_state(
+ state: &StructArray,
+ metadata: &BinaryArray,
+ metadata_rows: &[Option<usize>],
+) -> DataFusionResult<(ArrayRef, bool)> {
+ if state.len() != metadata_rows.len() {
+ return Err(DataFusionError::Execution(
+ "Variant shredding state and metadata row mapping have different lengths".to_string(),
+ ));
+ }
+
+ let active_rows = metadata_rows
+ .iter()
+ .enumerate()
+ .map(|(index, row)| state.is_valid(index).then_some(*row).flatten())
+ .collect::<Vec<_>>();
+ let mut fields = state.fields().iter().cloned().collect::<Vec<_>>();
+ let mut columns = state.columns().to_vec();
+ let mut changed = false;
+
+ if let Some(index) = fields.iter().position(|field| field.name() == "value") {
+ let (value, value_changed) =
+ rewrite_residual_values(&columns[index], metadata, &active_rows)?;
+ if value_changed {
+ fields[index] = Arc::new(
+ fields[index]
+ .as_ref()
+ .clone()
+ .with_data_type(value.data_type().clone()),
+ );
+ columns[index] = value;
+ changed = true;
+ }
+ }
+
+ if let Some(index) = fields
+ .iter()
+ .position(|field| field.name() == "typed_value")
+ {
+ let typed_rows = active_rows
+ .iter()
+ .enumerate()
+ .map(|(row, metadata)| columns[index].is_valid(row).then_some(*metadata).flatten())
+ .collect::<Vec<_>>();
+ let (typed_value, typed_changed) =
+ rewrite_typed_value(&columns[index], metadata, &typed_rows)?;
+ if typed_changed {
+ fields[index] = Arc::new(
+ fields[index]
+ .as_ref()
+ .clone()
+ .with_data_type(typed_value.data_type().clone()),
+ );
+ columns[index] = typed_value;
+ changed = true;
+ }
+ }
+
+ if !changed {
+ return Ok((Arc::new(state.clone()), false));
+ }
+ Ok((
+ Arc::new(StructArray::try_new(
+ fields.into(),
+ columns,
+ state.nulls().cloned(),
+ )?),
+ true,
+ ))
+}
+
+fn rewrite_residual_values(
+ value: &ArrayRef,
+ metadata: &BinaryArray,
+ metadata_rows: &[Option<usize>],
+) -> DataFusionResult<(ArrayRef, bool)> {
+ let binary = cast(value.as_ref(), &DataType::Binary)?;
+ let binary = binary.as_binary::<i32>();
+ let mut output: Option<BinaryBuilder> = None;
+
+ for (index, metadata_row) in metadata_rows.iter().enumerate() {
+ if binary.is_null(index) {
+ if let Some(output) = &mut output {
+ output.append_null();
+ }
+ continue;
+ }
+ let Some(metadata_row) = metadata_row else {
+ if let Some(output) = &mut output {
+ output.append_value(binary.value(index));
+ }
+ continue;
+ };
+ if metadata.is_null(*metadata_row) {
+ return Err(DataFusionError::Execution(format!(
+ "Variant metadata is null at row {metadata_row}"
+ )));
+ }
+
+ let rebuilt = catch_unwind(AssertUnwindSafe(
+ || -> Result<Option<Vec<u8>>, ArrowError> {
+ let metadata = VariantMetadata::try_new(metadata.value(*metadata_row))?;
+ let variant = Variant::new_with_metadata(metadata.clone(), binary.value(index));
+ if is_compatible_variant(&variant, VariantObjectKeyOrder::ArrowUtf8) {
+ return Ok(None);
+ }
+ if !is_compatible_variant(&variant, VariantObjectKeyOrder::SparkUtf16) {
+ return Err(ArrowError::InvalidArgumentError(
+ "Variant residual is neither UTF-8 nor Spark UTF-16 ordered".to_string(),
+ ));
+ }
+ Ok(Some(variant_bytes(
+ &metadata,
+ variant,
+ VariantObjectKeyOrder::ArrowUtf8,
+ )?))
+ },
+ ))
+ .map_err(|_| {
+ DataFusionError::Execution(format!("Invalid Variant residual at row {metadata_row}"))
+ })??;
+ if rebuilt.is_some() && output.is_none() {
+ output = Some(binary_prefix_builder(binary, index));
+ }
+ if let Some(output) = &mut output {
+ output.append_value(rebuilt.as_deref().unwrap_or_else(|| binary.value(index)));
+ }
+ }
+
+ if let Some(mut output) = output {
+ Ok((Arc::new(output.finish()), true))
+ } else {
+ Ok((Arc::clone(value), false))
+ }
+}
+
+fn list_metadata_rows<L: ListLikeArray>(
+ list: &L,
+ parent_rows: &[Option<usize>],
+) -> DataFusionResult<Vec<Option<usize>>> {
+ let mut child_rows = vec![None; list.values().len()];
+ for (index, metadata_row) in parent_rows.iter().enumerate() {
+ let Some(metadata_row) = metadata_row else {
+ continue;
+ };
+ for child_index in list.element_range(index) {
+ match child_rows[child_index] {
+ Some(existing) if existing != *metadata_row => {
+ return Err(DataFusionError::Execution(
+ "A shared Variant list child refers to different metadata rows".to_string(),
+ ));
+ }
+ _ => child_rows[child_index] = Some(*metadata_row),
+ }
+ }
+ }
+ Ok(child_rows)
+}
+
+fn rewrite_list_typed_value<L: ListLikeArray>(
+ array: &ArrayRef,
+ list: &L,
+ metadata: &BinaryArray,
+ metadata_rows: &[Option<usize>],
+) -> DataFusionResult<(ArrayRef, bool)> {
+ let child_rows = list_metadata_rows(list, metadata_rows)?;
+ let values = list.values().as_struct_opt().ok_or_else(|| {
+ DataFusionError::Execution(format!(
+ "Invalid shredded Variant list values: expected Struct, got {}",
+ list.values().data_type()
+ ))
+ })?;
+ let (values, changed) = rewrite_shredding_state(values, metadata, &child_rows)?;
+ if !changed {
+ return Ok((Arc::clone(array), false));
+ }
+
+ let data_type = match array.data_type() {
+ DataType::List(field) => DataType::List(Arc::new(
+ field
+ .as_ref()
+ .clone()
+ .with_data_type(values.data_type().clone()),
+ )),
+ DataType::LargeList(field) => DataType::LargeList(Arc::new(
+ field
+ .as_ref()
+ .clone()
+ .with_data_type(values.data_type().clone()),
+ )),
+ DataType::ListView(field) => DataType::ListView(Arc::new(
+ field
+ .as_ref()
+ .clone()
+ .with_data_type(values.data_type().clone()),
+ )),
+ DataType::LargeListView(field) => DataType::LargeListView(Arc::new(
+ field
+ .as_ref()
+ .clone()
+ .with_data_type(values.data_type().clone()),
+ )),
+ data_type => {
+ return Err(DataFusionError::Execution(format!(
+ "Expected a Variant list, got {data_type}"
+ )));
+ }
+ };
+ let data = array
+ .to_data()
+ .into_builder()
+ .data_type(data_type)
+ .child_data(vec![values.to_data()])
+ .build()?;
+ Ok((make_array(data), true))
+}
+
+fn rewrite_typed_value(
+ typed_value: &ArrayRef,
+ metadata: &BinaryArray,
+ metadata_rows: &[Option<usize>],
+) -> DataFusionResult<(ArrayRef, bool)> {
+ match typed_value.data_type() {
+ DataType::Struct(_) => {
+ let object = typed_value.as_struct();
+ let mut fields = object.fields().iter().cloned().collect::<Vec<_>>();
+ let mut columns = object.columns().to_vec();
+ let mut changed = false;
+ for (index, column) in object.columns().iter().enumerate() {
+ let child = column.as_struct_opt().ok_or_else(|| {
+ DataFusionError::Execution(format!(
+ "Invalid shredded Variant object field '{}': expected Struct, got {}",
+ fields[index].name(),
+ column.data_type()
+ ))
+ })?;
+ let (child, child_changed) =
+ rewrite_shredding_state(child, metadata, metadata_rows)?;
+ if child_changed {
+ fields[index] = Arc::new(
+ fields[index]
+ .as_ref()
+ .clone()
+ .with_data_type(child.data_type().clone()),
+ );
+ columns[index] = child;
+ changed = true;
+ }
+ }
+ if !changed {
+ return Ok((Arc::clone(typed_value), false));
+ }
+ Ok((
+ Arc::new(StructArray::try_new(
+ fields.into(),
+ columns,
+ object.nulls().cloned(),
+ )?),
+ true,
+ ))
+ }
+ DataType::List(_) => rewrite_list_typed_value(
+ typed_value,
+ typed_value.as_list::<i32>(),
+ metadata,
+ metadata_rows,
+ ),
+ DataType::LargeList(_) => rewrite_list_typed_value(
+ typed_value,
+ typed_value.as_list::<i64>(),
+ metadata,
+ metadata_rows,
+ ),
+ DataType::ListView(_) => rewrite_list_typed_value(
+ typed_value,
+ typed_value.as_list_view::<i32>(),
+ metadata,
+ metadata_rows,
+ ),
+ DataType::LargeListView(_) => rewrite_list_typed_value(
+ typed_value,
+ typed_value.as_list_view::<i64>(),
+ metadata,
+ metadata_rows,
+ ),
+ _ => Ok((Arc::clone(typed_value), false)),
+ }
+}
+
+fn prepare_variant_for_unshredding(variant: &VariantArray) -> DataFusionResult<VariantArray> {
+ if variant.typed_value_column().is_none() {
+ return Ok(variant.clone());
+ }
+
+ let metadata = cast(variant.metadata_column().as_ref(), &DataType::Binary)?;
+ let metadata = metadata.as_binary::<i32>();
+ let metadata_rows = (0..variant.len())
+ .map(|index| variant.inner().is_valid(index).then_some(index))
+ .collect::<Vec<_>>();
+ let (array, changed) = rewrite_shredding_state(variant.inner(), metadata, &metadata_rows)?;
+ if changed {
+ Ok(VariantArray::try_new(array.as_ref())?)
+ } else {
+ Ok(variant.clone())
+ }
+}
+
+/// Supplies sort-only field names whose Rust ordering matches Java `String.compareTo` ordering.
+/// Field IDs still come from the original metadata dictionary.
+#[derive(Debug)]
+struct SparkMetadataBuilder<'a, 'm> {
+ metadata: &'a VariantMetadata<'m>,
+ sort_keys: Vec<String>,
+}
+
+impl<'a, 'm> SparkMetadataBuilder<'a, 'm> {
+ fn new(metadata: &'a VariantMetadata<'m>) -> Self {
+ let sort_keys = metadata
+ .iter()
+ .map(|field_name| {
+ field_name
+ .encode_utf16()
+ .map(|unit| char::from_u32(0x10000 + u32::from(unit)).unwrap())
+ .collect()
+ })
+ .collect();
+ Self {
+ metadata,
+ sort_keys,
+ }
+ }
+}
+
+impl MetadataBuilder for SparkMetadataBuilder<'_, '_> {
+ fn try_upsert_field_name(&mut self, field_name: &str) -> Result<u32, ArrowError> {
+ self.metadata
+ .get_entry(field_name)
+ .map(|(field_id, _)| field_id)
+ .ok_or_else(|| {
+ ArrowError::InvalidArgumentError(format!(
+ "Field name '{field_name}' not found in metadata dictionary"
+ ))
+ })
+ }
+
+ fn field_name(&self, field_id: usize) -> &str {
+ &self.sort_keys[field_id]
+ }
+
+ fn num_field_names(&self) -> usize {
+ self.metadata.len()
+ }
+
+ fn truncate_field_names(&mut self, new_size: usize) {
+ debug_assert_eq!(self.metadata.len(), new_size);
+ }
+
+ fn finish(&mut self) -> usize {
+ self.metadata.size()
+ }
+}
+
+#[derive(Clone, Copy)]
+enum VariantObjectKeyOrder {
+ ArrowUtf8,
+ SparkUtf16,
+}
+
+fn is_compatible_variant(variant: &Variant<'_, '_>, order: VariantObjectKeyOrder) -> bool {
+ match variant {
+ Variant::Object(object) => {
+ let mut previous = None;
+ object.iter().all(|(name, value)| {
+ let ordered = previous
+ .map(|previous: &str| match order {
+ VariantObjectKeyOrder::ArrowUtf8 => previous <= name,
+ VariantObjectKeyOrder::SparkUtf16 => {
+ previous.encode_utf16().cmp(name.encode_utf16())
+ != std::cmp::Ordering::Greater
+ }
+ })
+ .unwrap_or(true);
+ previous = Some(name);
+ ordered && is_compatible_variant(&value, order)
+ })
+ }
+ Variant::List(list) => list
+ .iter()
+ .all(|value| is_compatible_variant(&value, order)),
+ _ => true,
+ }
+}
+
+fn variant_bytes(
+ metadata: &VariantMetadata<'_>,
+ variant: Variant<'_, '_>,
+ order: VariantObjectKeyOrder,
+) -> Result<Vec<u8>, ArrowError> {
+ let mut value_builder = ValueBuilder::new();
+ match variant {
+ Variant::Object(object) => {
+ let mut metadata_builder: Box<dyn MetadataBuilder> = match order {
+ VariantObjectKeyOrder::ArrowUtf8 => {
+ Box::new(ReadOnlyMetadataBuilder::new(metadata))
+ }
+ VariantObjectKeyOrder::SparkUtf16 => Box::new(SparkMetadataBuilder::new(metadata)),
+ };
+ let mut builder = ObjectBuilder::new(
+ ParentState::variant(&mut value_builder, metadata_builder.as_mut()),
+ matches!(order, VariantObjectKeyOrder::ArrowUtf8),
+ );
+ for (name, value) in object.iter() {
+ let value = variant_bytes(metadata, value, order)?;
+ builder
+ .try_insert_bytes(name, Variant::new_with_metadata(metadata.clone(), &value))?;
+ }
+ builder.finish();
+ }
+ Variant::List(list) => {
+ let mut metadata_builder = ReadOnlyMetadataBuilder::new(metadata);
+ let mut builder = ListBuilder::new(
+ ParentState::variant(&mut value_builder, &mut metadata_builder),
+ matches!(order, VariantObjectKeyOrder::ArrowUtf8),
+ );
+ for value in list.iter() {
+ let value = variant_bytes(metadata, value, order)?;
+ builder.append_value_bytes(Variant::new_with_metadata(metadata.clone(), &value));
+ }
+ builder.finish();
+ }
+ variant => {
+ let mut metadata_builder = ReadOnlyMetadataBuilder::new(metadata);
+ ValueBuilder::try_append_variant(
+ ParentState::variant(&mut value_builder, &mut metadata_builder),
+ variant,
+ )?;
+ }
+ }
+ Ok(value_builder.into_inner())
+}
+
+/// Released Spark 4 profiles search object fields in Java UTF-16 order. Convert whole-value output
+/// to that order until #5474 can remove this rewrite after every supported profile includes
+/// SPARK-58949. Values already in the requested order remain byte-for-byte unchanged.
+/// https://github.com/apache/datafusion-comet/issues/5474
+fn reorder_variant_values(
+ value: &ArrayRef,
+ metadata: &ArrayRef,
+ parent_nulls: Option<&NullBuffer>,
+) -> DataFusionResult<ArrayRef> {
+ let original = value;
+ let value = value.as_binary::<i32>();
+ let metadata = metadata.as_binary::<i32>();
+ let mut output: Option<BinaryBuilder> = None;
+
+ for index in 0..value.len() {
+ if parent_nulls.is_some_and(|nulls| nulls.is_null(index)) {
+ if value.is_valid(index) && output.is_none() {
+ output = Some(binary_prefix_builder(value, index));
+ }
+ if let Some(output) = &mut output {
+ output.append_null();
+ }
+ continue;
+ }
+ if value.is_null(index) {
+ return Err(DataFusionError::Execution(format!(
+ "Variant value is null at row {index}"
+ )));
+ }
+ if metadata.is_null(index) {
+ return Err(DataFusionError::Execution(format!(
+ "Variant metadata is null at row {index}"
+ )));
+ }
+
+ let rebuilt = catch_unwind(AssertUnwindSafe(
+ || -> Result<Option<Vec<u8>>, ArrowError> {
+ let metadata = VariantMetadata::try_new(metadata.value(index))?;
+ let variant = Variant::new_with_metadata(metadata.clone(), value.value(index));
+ if is_compatible_variant(&variant, VariantObjectKeyOrder::SparkUtf16) {
+ return Ok(None);
+ }
+ Ok(Some(variant_bytes(
+ &metadata,
+ variant,
+ VariantObjectKeyOrder::SparkUtf16,
+ )?))
+ },
+ ))
+ .map_err(|_| {
+ DataFusionError::Execution(format!("Invalid Variant value at row {index}"))
+ })??;
+ if rebuilt.is_some() && output.is_none() {
+ output = Some(binary_prefix_builder(value, index));
+ }
+ if let Some(output) = &mut output {
+ output.append_value(rebuilt.as_deref().unwrap_or_else(|| value.value(index)));
+ }
+ }
+
+ match output {
+ Some(mut output) => Ok(Arc::new(output.finish())),
+ None => Ok(Arc::clone(original)),
+ }
+}
+
+// Called only at the first changed row; the unchanged prefix is copied once.
+fn binary_prefix_builder(value: &BinaryArray, end: usize) -> BinaryBuilder {
+ let mut builder = BinaryBuilder::new();
+ builder.extend(value.iter().take(end));
+ builder
+}
+
+#[cfg(test)]
+mod tests;
diff --git a/native/core/src/parquet/cast_column/variant/tests.rs b/native/core/src/parquet/cast_column/variant/tests.rs
new file mode 100644
index 0000000000..698bf675c7
--- /dev/null
+++ b/native/core/src/parquet/cast_column/variant/tests.rs
@@ -0,0 +1,550 @@
+// 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 super::*;
+use arrow::{
+ array::{Int64Array, ListArray},
+ buffer::OffsetBuffer,
+ datatypes::{Field, Fields},
+};
+use parquet::variant::{
+ shred_variant, VariantArrayBuilder, VariantBuilder, VariantBuilderExt, VariantType,
+};
+
+fn target_field(nullable: bool) -> FieldRef {
+ Arc::new(
+ Field::new(
+ "v",
+ DataType::Struct(Fields::from(vec![
+ Field::new("value", DataType::Binary, false),
+ Field::new("metadata", DataType::Binary, false),
+ ])),
+ nullable,
+ )
+ .with_extension_type(VariantType),
+ )
+}
+
+fn unicode_object_keys() -> Vec<String> {
+ let mut keys = (0..30).map(|i| format!("k{i:02}")).collect::<Vec<_>>();
+ keys.push("\u{e000}".to_string());
+ keys.push("๐".to_string());
+ keys
+}
+
+fn assert_spark_unicode_variant(variant: Variant<'_, '_>) {
+ let Variant::Object(object) = variant else {
+ panic!("expected object")
+ };
+ let fields = object.iter().collect::<Vec<_>>();
+ assert_eq!(fields.len(), 32);
+ assert_eq!(fields[30].0, "๐");
+ assert_eq!(fields[31].0, "\u{e000}");
+ assert_eq!(object.get("๐").unwrap().as_int64(), Some(531));
+}
+
+fn assert_spark_unicode_output(output: &StructArray) {
+ let value = output.column(0).as_binary::<i32>();
+ let metadata = output.column(1).as_binary::<i32>();
+ assert_spark_unicode_variant(Variant::new(metadata.value(0), value.value(0)));
+}
+
+#[test]
+fn normalize_full_shredding_reorders_children_and_preserves_parent_nulls() {
+ let mut builder = VariantArrayBuilder::new(3);
+ builder.append_variant(Variant::from(1_i64));
+ builder.append_null();
+ builder.append_variant(Variant::from(3_i64));
+ let base = builder.build();
+ let metadata = Arc::clone(base.metadata_column());
+ let typed_value: ArrayRef = Arc::new(Int64Array::from(vec![Some(10), None, Some(30)]));
+ let physical: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![
+ Field::new("typed_value", DataType::Int64, true),
+ Field::new("metadata", metadata.data_type().clone(), false),
+ ]),
+ vec![typed_value, metadata],
+ base.inner().nulls().cloned(),
+ )
+ .unwrap(),
+ );
+
+ let output = normalize_variant_array(&physical, &target_field(true)).unwrap();
+ let output = output.as_struct();
+ assert_eq!(
+ output
+ .fields()
+ .iter()
+ .map(|field| field.name().as_str())
+ .collect::<Vec<_>>(),
+ ["value", "metadata"]
+ );
+ assert!(output
+ .columns()
+ .iter()
+ .all(|column| column.data_type() == &DataType::Binary));
+ assert!(output.is_null(1));
+
+ let variant = VariantArray::try_new(output).unwrap();
+ assert_eq!(variant.value(0), Variant::from(10_i64));
+ assert_eq!(variant.value(2), Variant::from(30_i64));
+}
+
+#[test]
+fn normalize_fully_shredded_object_orders_for_spark() {
+ let keys = unicode_object_keys();
+ let mut builder = VariantBuilder::new().with_field_names(keys.iter().map(String::as_str));
+ builder.append_value(Variant::Null);
+ let (metadata, _) = builder.finish();
+ let mut fields = Vec::with_capacity(keys.len());
+ let mut columns = Vec::with_capacity(keys.len());
+ for (index, key) in keys.iter().enumerate() {
+ let state = StructArray::try_new(
+ Fields::from(vec![Field::new("typed_value", DataType::Int64, false)]),
+ vec![Arc::new(Int64Array::from(vec![if key == "๐" {
+ 531
+ } else {
+ index as i64
+ }]))],
+ None,
+ )
+ .unwrap();
+ fields.push(Field::new(key, state.data_type().clone(), false));
+ columns.push(Arc::new(state) as ArrayRef);
+ }
+ let typed_value: ArrayRef =
+ Arc::new(StructArray::try_new(fields.into(), columns, None).unwrap());
+ let physical: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![
+ Field::new("metadata", DataType::Binary, false),
+ Field::new("typed_value", typed_value.data_type().clone(), false),
+ ]),
+ vec![
+ Arc::new(BinaryArray::from(vec![Some(metadata.as_slice())])),
+ typed_value,
+ ],
+ None,
+ )
+ .unwrap(),
+ );
+
+ let output = normalize_variant_array(&physical, &target_field(false)).unwrap();
+ assert_spark_unicode_output(output.as_struct());
+}
+
+#[test]
+fn canonical_and_shredded_values_normalize_equally() {
+ let mut builder = VariantArrayBuilder::new(6);
+ builder.new_object().with_field("known", 1_i64).finish();
+ builder
+ .new_object()
+ .with_field("known", 2_i64)
+ .with_field("extra", 3_i64)
+ .finish();
+ builder
+ .new_list()
+ .with_value(4_i64)
+ .with_value(Variant::Null)
+ .finish();
+ builder.append_variant(Variant::from(5_i64));
+ builder.append_variant(Variant::Null);
+ builder.append_null();
+ let canonical = builder.build();
+ let shredded = shred_variant(
+ &canonical,
+ &DataType::Struct(Fields::from(vec![Field::new(
+ "known",
+ DataType::Int64,
+ true,
+ )])),
+ )
+ .unwrap();
+
+ let prepared = prepare_variant_for_unshredding(&shredded).unwrap();
+ assert!(Arc::ptr_eq(
+ shredded.value_column(),
+ prepared.value_column()
+ ));
+ assert!(Arc::ptr_eq(
+ shredded.typed_value_column().unwrap(),
+ prepared.typed_value_column().unwrap(),
+ ));
+
+ let normalize = |array: &VariantArray| {
+ let array: ArrayRef = Arc::new(array.inner().clone());
+ let output = normalize_variant_array(&array, &target_field(true)).unwrap();
+ VariantArray::try_new(output.as_ref()).unwrap()
+ };
+ let canonical = normalize(&canonical);
+ let shredded = normalize(&shredded);
+
+ for index in 0..canonical.len() {
+ assert_eq!(canonical.is_null(index), shredded.is_null(index));
+ if canonical.is_valid(index) {
+ assert_eq!(canonical.value(index), shredded.value(index));
+ }
+ }
+}
+
+#[test]
+fn normalize_unshredded_variant_orders_for_spark_and_is_idempotent() {
+ let keys = unicode_object_keys();
+ let mut builder = VariantBuilder::new().with_field_names(keys.iter().map(String::as_str));
+ let mut object = builder.new_object();
+ for (index, key) in keys.iter().enumerate() {
+ object.insert(key, if key == "๐" { 531 } else { index as i64 });
+ }
+ object.finish();
+ let (metadata, value) = builder.finish();
+
+ let physical: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![
+ Field::new("metadata", DataType::Binary, false),
+ Field::new("value", DataType::Binary, false),
+ ]),
+ vec![
+ Arc::new(BinaryArray::from(vec![Some(metadata.as_slice())])),
+ Arc::new(BinaryArray::from(vec![Some(value.as_slice())])),
+ ],
+ None,
+ )
+ .unwrap(),
+ );
+
+ let first = normalize_variant_array(&physical, &target_field(false)).unwrap();
+ assert_spark_unicode_output(first.as_struct());
+ let first_value = first.as_struct().column(0).as_binary::<i32>().value(0);
+
+ let second = normalize_variant_array(&first, &target_field(false)).unwrap();
+ assert_spark_unicode_output(second.as_struct());
+ assert_eq!(
+ second.as_struct().column(0).as_binary::<i32>().value(0),
+ first_value
+ );
+ assert_eq!(
+ second
+ .as_struct()
+ .column(0)
+ .as_binary::<i32>()
+ .values()
+ .as_ptr(),
+ first
+ .as_struct()
+ .column(0)
+ .as_binary::<i32>()
+ .values()
+ .as_ptr()
+ );
+}
+
+#[test]
+fn unchanged_values_reuse_buffers_and_still_validate() {
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::from("unchanged payload"));
+ let (metadata, bytes) = builder.finish();
+ let values: ArrayRef = Arc::new(BinaryArray::from(vec![
+ None,
+ Some(bytes.as_slice()),
+ Some(bytes.as_slice()),
+ ]));
+ // Exercise a sliced array: prefix offsets need not start at zero.
+ let values = values.slice(1, 2);
+ let metadata: ArrayRef = Arc::new(BinaryArray::from(vec![metadata.as_slice(); 2]));
+ let output = reorder_variant_values(&values, &metadata, None).unwrap();
+ assert!(Arc::ptr_eq(&values, &output));
+ let target = target_field(false);
+ let DataType::Struct(fields) = target.data_type() else {
+ unreachable!()
+ };
+ let input: ArrayRef = Arc::new(StructArray::new(
+ fields.clone(),
+ vec![Arc::clone(&values), Arc::clone(&metadata)],
+ None,
+ ));
+ let normalized = normalize_variant_array(&input, &target).unwrap();
+ for (input, output) in input
+ .as_struct()
+ .columns()
+ .iter()
+ .zip(normalized.as_struct().columns())
+ {
+ assert_eq!(
+ input.as_binary::<i32>().values().as_ptr(),
+ output.as_binary::<i32>().values().as_ptr()
+ );
+ }
+ let (output, changed) =
+ rewrite_residual_values(&values, metadata.as_binary::<i32>(), &[Some(0), None]).unwrap();
+ assert!(!changed);
+ assert!(Arc::ptr_eq(&values, &output));
+
+ let nulls = NullBuffer::from(vec![true, false]);
+ let output = reorder_variant_values(&values, &metadata, Some(&nulls)).unwrap();
+ assert_eq!(output.as_binary::<i32>().value(0), bytes);
+ assert!(output.is_null(1));
+ let again = reorder_variant_values(&output, &metadata, Some(&nulls)).unwrap();
+ assert!(Arc::ptr_eq(&output, &again));
+
+ let missing_metadata: ArrayRef = Arc::new(BinaryArray::from(vec![
+ Some(metadata.as_binary::<i32>().value(0)),
+ None,
+ ]));
+ assert!(reorder_variant_values(&values, &missing_metadata, None).is_err());
+ assert!(rewrite_residual_values(
+ &values,
+ missing_metadata.as_binary::<i32>(),
+ &[Some(0), Some(1)],
+ )
+ .is_err());
+}
+
+#[test]
+fn lazy_rewrites_preserve_prefix_nulls_and_suffix() {
+ let keys = unicode_object_keys();
+ let mut builder = VariantBuilder::new().with_field_names(keys.iter().map(String::as_str));
+ let mut object = builder.new_object();
+ for (index, key) in keys.iter().enumerate() {
+ object.insert(key, index as i64);
+ }
+ object.finish();
+ let (metadata, canonical) = builder.finish();
+ let values: ArrayRef = Arc::new(BinaryArray::from(vec![
+ Some(canonical.as_slice()),
+ None,
+ Some(canonical.as_slice()),
+ None,
+ ]));
+ let metadata: ArrayRef = Arc::new(BinaryArray::from(vec![metadata.as_slice(); 4]));
+ let parents = NullBuffer::from(vec![true, false, true, false]);
+ let legacy = reorder_variant_values(&values, &metadata, Some(&parents)).unwrap();
+ let mixed: ArrayRef = Arc::new(BinaryArray::from(vec![
+ Some(canonical.as_slice()),
+ None,
+ Some(legacy.as_binary::<i32>().value(2)),
+ None,
+ ]));
+ let (output, changed) = rewrite_residual_values(
+ &mixed,
+ metadata.as_binary::<i32>(),
+ &[Some(0), None, Some(2), None],
+ )
+ .unwrap();
+ assert!(changed);
+ let output = output.as_binary::<i32>();
+ assert_eq!(output.value(0), canonical);
+ assert!(output.is_null(1));
+ assert!(output.is_null(3));
+ assert_eq!(
+ Variant::new(metadata.as_binary::<i32>().value(2), output.value(2)),
+ Variant::new(metadata.as_binary::<i32>().value(2), &canonical),
+ );
+}
+
+// Run explicitly with --ignored --nocapture; fixture construction is outside the timed loop.
+#[test]
+#[ignore]
+fn benchmark_variant_buffer_reuse() {
+ use std::{hint::black_box, time::Instant};
+ let rows = 4096;
+ let payload = "x".repeat(4096);
+ let mut strings = VariantArrayBuilder::new(rows);
+ let mut objects = VariantArrayBuilder::new(rows);
+ for _ in 0..rows {
+ strings.append_variant(Variant::from(payload.as_str()));
+ objects
+ .new_object()
+ .with_field("known", 1_i64)
+ .with_field("payload", payload.as_str())
+ .finish();
+ }
+ let strings = strings.build();
+ let objects = objects.build();
+ let shredded = shred_variant(
+ &objects,
+ &DataType::Struct(Fields::from(vec![Field::new(
+ "known",
+ DataType::Int64,
+ true,
+ )])),
+ )
+ .unwrap();
+ let target = target_field(false);
+ let DataType::Struct(fields) = target.data_type() else {
+ unreachable!()
+ };
+ let strings: ArrayRef = Arc::new(StructArray::new(
+ fields.clone(),
+ vec![
+ cast(strings.value_column().as_ref(), &DataType::Binary).unwrap(),
+ cast(strings.metadata_column().as_ref(), &DataType::Binary).unwrap(),
+ ],
+ None,
+ ));
+ let shredded: ArrayRef = Arc::new(shredded.into_inner());
+ for (name, input) in [("canonical", strings), ("partially_shredded", shredded)] {
+ for _ in 0..3 {
+ black_box(normalize_variant_array(&input, &target).unwrap());
+ }
+ let start = Instant::now();
+ for _ in 0..30 {
+ black_box(normalize_variant_array(black_box(&input), &target).unwrap());
+ }
+ eprintln!(
+ "{name}: {:.3} ms/batch, {rows} rows, 4096-byte payload",
+ start.elapsed().as_secs_f64() * 1000.0 / 30.0
+ );
+ }
+}
+
+#[test]
+fn normalize_partially_shredded_legacy_residual() {
+ let keys = unicode_object_keys();
+ let mut builder = VariantBuilder::new().with_field_names(keys.iter().map(String::as_str));
+ let mut object = builder.new_object();
+ for (index, key) in keys.iter().enumerate().skip(1) {
+ object.insert(key, if key == "๐" { 531 } else { index as i64 });
+ }
+ object.finish();
+ let (metadata, value) = builder.finish();
+ let metadata_array: ArrayRef = Arc::new(BinaryArray::from(vec![Some(metadata.as_slice())]));
+ let value_array: ArrayRef = Arc::new(BinaryArray::from(vec![Some(value.as_slice())]));
+ let legacy_value = reorder_variant_values(&value_array, &metadata_array, None).unwrap();
+
+ let shredded: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![Field::new("typed_value", DataType::Int64, false)]),
+ vec![Arc::new(Int64Array::from(vec![0]))],
+ None,
+ )
+ .unwrap(),
+ );
+ let typed_value: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![Field::new("k00", shredded.data_type().clone(), false)]),
+ vec![shredded],
+ None,
+ )
+ .unwrap(),
+ );
+ let physical: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![
+ Field::new("metadata", DataType::Binary, false),
+ Field::new("value", DataType::Binary, true),
+ Field::new("typed_value", typed_value.data_type().clone(), true),
+ ]),
+ vec![metadata_array, legacy_value, typed_value],
+ None,
+ )
+ .unwrap(),
+ );
+
+ let output = normalize_variant_array(&physical, &target_field(false)).unwrap();
+ assert_spark_unicode_output(output.as_struct());
+}
+
+#[test]
+fn normalize_nested_list_residuals_use_their_root_metadata() {
+ fn legacy_row(keys: &[&str]) -> (Vec<u8>, Vec<u8>) {
+ let mut builder = VariantBuilder::new().with_field_names(keys.iter().copied());
+ let mut object = builder.new_object();
+ for (index, key) in keys.iter().enumerate() {
+ object.insert(key, index as i64);
+ }
+ object.finish();
+ let (metadata, value) = builder.finish();
+ let metadata_array: ArrayRef = Arc::new(BinaryArray::from(vec![Some(metadata.as_slice())]));
+ let value: ArrayRef = Arc::new(BinaryArray::from(vec![Some(value.as_slice())]));
+ let value = reorder_variant_values(&value, &metadata_array, None).unwrap();
+ (metadata, value.as_binary::<i32>().value(0).to_vec())
+ }
+
+ let (metadata0, value0) = legacy_row(&["a", "\u{e000}", "๐"]);
+ let (metadata1, value1) = legacy_row(&["b", "zz", "\u{ffff}", "๐"]);
+ let states: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![Field::new("value", DataType::Binary, true)]),
+ vec![Arc::new(BinaryArray::from(vec![
+ Some(value0.as_slice()),
+ Some(value1.as_slice()),
+ ]))],
+ None,
+ )
+ .unwrap(),
+ );
+ let list: ArrayRef = Arc::new(
+ ListArray::try_new(
+ Arc::new(Field::new("element", states.data_type().clone(), false)),
+ OffsetBuffer::new(vec![0, 1, 2].into()),
+ states,
+ None,
+ )
+ .unwrap(),
+ );
+ let metadata: ArrayRef = Arc::new(BinaryArray::from(vec![
+ Some(metadata0.as_slice()),
+ Some(metadata1.as_slice()),
+ ]));
+ let physical: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![
+ Field::new("metadata", DataType::Binary, false),
+ Field::new("typed_value", list.data_type().clone(), false),
+ ]),
+ vec![metadata, list],
+ None,
+ )
+ .unwrap(),
+ );
+
+ let output = normalize_variant_array(&physical, &target_field(false)).unwrap();
+ let output = VariantArray::try_new(output.as_ref()).unwrap();
+ for (index, key) in ["๐", "๐"].into_iter().enumerate() {
+ let Variant::List(list) = output.value(index) else {
+ panic!("expected list")
+ };
+ let Variant::Object(object) = list.get(0).unwrap() else {
+ panic!("expected object")
+ };
+ assert_eq!(object.get(key).unwrap().as_int64(), Some(index as i64 + 2));
+ }
+}
+
+#[test]
+fn normalize_null_parent_ignores_empty_children() {
+ let physical: ArrayRef = Arc::new(
+ StructArray::try_new(
+ Fields::from(vec![
+ Field::new("value", DataType::Binary, false),
+ Field::new("metadata", DataType::Binary, false),
+ ]),
+ vec![
+ Arc::new(BinaryArray::from(vec![Some(&b""[..])])),
+ Arc::new(BinaryArray::from(vec![Some(&b""[..])])),
+ ],
+ Some(NullBuffer::from(vec![false])),
+ )
+ .unwrap(),
+ );
+
+ let output = normalize_variant_array(&physical, &target_field(true)).unwrap();
+ assert!(output.is_null(0));
+ assert!(output.as_struct().column(0).is_null(0));
+}
diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs
index 4ac716a435..9926bde135 100644
--- a/native/core/src/parquet/schema_adapter.rs
+++ b/native/core/src/parquet/schema_adapter.rs
@@ -34,7 +34,7 @@ use datafusion_physical_expr_adapter::{
replace_columns_with_literals, DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter,
PhysicalExprAdapterFactory,
};
-use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+use parquet::{arrow::PARQUET_FIELD_ID_META_KEY, variant::VariantType};
use std::collections::{HashMap, HashSet};
use std::fmt::{self, Display};
use std::hash::{Hash, Hasher};
@@ -680,6 +680,7 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter {
self.wrap_all_type_mismatches(expr)?
}
};
+ let expr = self.wrap_direct_variant_column(expr)?;
// For case-insensitive mode: remap column names from logical back to
// original physical names. The default adapter was given a remapped
@@ -708,6 +709,37 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter {
}
impl SparkPhysicalExprAdapter {
+ /// The default adapter leaves an identical physical/logical Field as a bare Column. Variant
+ /// still needs normalization because a canonical unshredded value may already have that exact
+ /// marked layout.
+ fn wrap_direct_variant_column(
+ &self,
+ expr: Arc<dyn PhysicalExpr>,
+ ) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
+ let Some(column) = expr.downcast_ref::<Column>() else {
+ return Ok(expr);
+ };
+ let Ok(logical_field) = self.logical_file_schema.field_with_name(column.name()) else {
+ return Ok(expr);
+ };
+ if !logical_field.has_valid_extension_type::<VariantType>() {
+ return Ok(expr);
+ }
+ let Some(physical_field) = self.physical_file_schema.fields().get(column.index()) else {
+ return Ok(expr);
+ };
+
+ Ok(Arc::new(
+ CometCastColumnExpr::try_new(
+ expr,
+ Arc::clone(physical_field),
+ Arc::new(logical_field.clone()),
+ None,
+ )?
+ .with_parquet_options(self.parquet_options.clone()),
+ ))
+ }
+
/// Wrap ALL Column expressions that have type mismatches with CometCastColumnExpr.
/// This is the fallback path when the default adapter fails (e.g., for complex
/// nested type casts like List<Struct> or Map). Uses `spark_parquet_convert`
@@ -753,7 +785,9 @@ impl SparkPhysicalExprAdapter {
Arc::clone(&e)
};
- if logical_field.data_type() != physical_field.data_type() {
+ if logical_field.has_valid_extension_type::<VariantType>()
+ || logical_field.data_type() != physical_field.data_type()
+ {
// Mirror the same string/binary -> non-string/binary rejection in
// `replace_with_spark_cast`; this branch is reached when the default
// adapter rejected the cast and we'd otherwise build a CometCastColumnExpr
@@ -812,6 +846,22 @@ impl SparkPhysicalExprAdapter {
};
let physical_type = input_field.data_type();
+ if cast
+ .target_field()
+ .has_valid_extension_type::<VariantType>()
+ {
+ let comet_cast: Arc<dyn PhysicalExpr> = Arc::new(
+ CometCastColumnExpr::try_new(
+ child,
+ input_field,
+ Arc::clone(cast.target_field()),
+ None,
+ )?
+ .with_parquet_options(self.parquet_options.clone()),
+ );
+ return Ok(Transformed::yes(comet_cast));
+ }
+
// Identity cast: DataFusion's default adapter inserts a CastExpr
// whenever the logical and physical Arrow Fields differ in any
// attribute (data type, nullability, or metadata), so with identical
@@ -1297,6 +1347,7 @@ impl PhysicalExpr for RejectOnNonEmpty {
#[cfg(test)]
mod test {
+ use crate::parquet::cast_column::CometCastColumnExpr;
use crate::parquet::parquet_support::SparkParquetOptions;
use crate::parquet::schema_adapter::{
is_pure_structural_narrowing, SparkPhysicalExprAdapterFactory,
@@ -1307,7 +1358,7 @@ mod test {
Int64Array, StringArray, TimestampMicrosecondArray,
};
use arrow::datatypes::SchemaRef;
- use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::datatypes::{DataType, Field, Fields, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::common::DataFusionError;
use datafusion::datasource::listing::PartitionedFile;
@@ -1315,6 +1366,7 @@ mod test {
use datafusion::datasource::source::DataSourceExec;
use datafusion::execution::object_store::ObjectStoreUrl;
use datafusion::execution::TaskContext;
+ use datafusion::physical_expr::expressions::Column;
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_plan::ExecutionPlan;
use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename;
@@ -1323,6 +1375,7 @@ mod test {
use futures::StreamExt;
use parquet::arrow::ArrowWriter;
use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+ use parquet::variant::VariantType;
use std::collections::HashMap;
use std::fs::File;
use std::sync::Arc;
@@ -2075,6 +2128,45 @@ mod test {
);
}
+ #[test]
+ fn marked_variant_column_is_wrapped_after_unicode_name_remap() {
+ let storage = DataType::Struct(Fields::from(vec![
+ Field::new("value", DataType::Binary, false),
+ Field::new("metadata", DataType::Binary, false),
+ ]));
+ let logical = Arc::new(Schema::new(vec![Field::new(
+ "mรผnchen",
+ storage.clone(),
+ true,
+ )
+ .with_extension_type(VariantType)]));
+ let physical = Arc::new(Schema::new(vec![
+ Field::new("MรNCHEN", storage, true).with_extension_type(VariantType)
+ ]));
+ let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false);
+ options.case_sensitive = false;
+ let adapter = SparkPhysicalExprAdapterFactory::new(options, None)
+ .create(Arc::clone(&logical), Arc::clone(&physical))
+ .unwrap();
+
+ let expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("mรผnchen", 0));
+ let rewritten = adapter.rewrite(expr).unwrap();
+ let cast = rewritten
+ .downcast_ref::<CometCastColumnExpr>()
+ .expect("marked Variant must retain its normalization wrapper");
+ assert_eq!(
+ cast.children()[0]
+ .downcast_ref::<Column>()
+ .expect("normalization input must remain a column")
+ .name(),
+ "MรNCHEN"
+ );
+ assert!(rewritten
+ .return_field(&physical)
+ .unwrap()
+ .has_valid_extension_type::<VariantType>());
+ }
+
/// #4859 investigation: for a pure structural narrowing of a nested column (dropping
/// unrequested struct fields, no leaf-type promotion, no tz relabeling, no field-id
/// matching, case-sensitive), does Comet's `CometCastColumnExpr` (via
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org