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/11 16:48:48 UTC
(datafusion-comet) branch main updated: feat: adapt Parquet storage for Variant projection (#5794)
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 8b818b53bd feat: adapt Parquet storage for Variant projection (#5794)
8b818b53bd is described below
commit 8b818b53bd6a1b6512cbea890132537c5bcfda98
Author: Peter Lee <pe...@gmail.com>
AuthorDate: Sat Sep 12 00:48:42 2026 +0800
feat: adapt Parquet storage for Variant projection (#5794)
---
native/core/Cargo.toml | 2 +-
native/core/src/parquet/cast_column/variant.rs | 101 +++++-
.../core/src/parquet/cast_column/variant/tests.rs | 168 ++++++++++
.../src/parquet/eager_page_index_reader_factory.rs | 226 ++++++++++++-
native/core/src/parquet/parquet_exec.rs | 28 +-
.../core/src/parquet/parquet_exec/variant_tests.rs | 351 +++++++++++++++++++++
native/core/src/parquet/schema_adapter.rs | 87 +++--
7 files changed, 922 insertions(+), 41 deletions(-)
diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml
index 9ebb69b2f6..20ea283653 100644
--- a/native/core/Cargo.toml
+++ b/native/core/Cargo.toml
@@ -38,7 +38,7 @@ publish = false
arrow = { workspace = true }
base64 = "0.23.0"
bytes = { workspace = true }
-parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
+parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "arrow_canonical_extension_types", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
futures = { workspace = true }
mimalloc = { version = "*", default-features = false, optional = true }
tikv-jemallocator = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls"] }
diff --git a/native/core/src/parquet/cast_column/variant.rs b/native/core/src/parquet/cast_column/variant.rs
index ce2c75a878..4e727ecab3 100644
--- a/native/core/src/parquet/cast_column/variant.rs
+++ b/native/core/src/parquet/cast_column/variant.rs
@@ -21,8 +21,8 @@ use arrow::{
StructArray,
},
buffer::NullBuffer,
- compute::cast,
- datatypes::{DataType, FieldRef},
+ compute::{cast, cast_with_options},
+ datatypes::{DataType, FieldRef, TimeUnit, DECIMAL128_MAX_PRECISION},
error::ArrowError,
};
use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -59,6 +59,7 @@ pub(super) fn normalize_variant_array(
// 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 array = normalize_variant_storage(array)?;
let variant = VariantArray::try_new(array.as_ref())?;
let prepared = prepare_variant_for_unshredding(&variant)?;
let unshredded = unshred_variant(&prepared)?;
@@ -74,6 +75,102 @@ pub(super) fn normalize_variant_array(
)?))
}
+/// Arrow Variant compute rejects some storage types that Spark's Parquet reader accepts.
+/// Choose supported types recursively for encoded, unsigned, decimal, timestamp, and fixed
+/// binary/list children before reconstructing the whole value.
+/// https://github.com/apache/datafusion-comet/issues/5477
+fn normalize_variant_type(data_type: &DataType) -> Option<DataType> {
+ fn normalize_field(field: &FieldRef) -> Option<FieldRef> {
+ normalize_variant_type(field.data_type())
+ .map(|data_type| Arc::new(field.as_ref().clone().with_data_type(data_type)))
+ }
+
+ match data_type {
+ DataType::Dictionary(_, value_type) => {
+ Some(normalize_variant_type(value_type).unwrap_or_else(|| value_type.as_ref().clone()))
+ }
+ DataType::UInt8 => Some(DataType::Int16),
+ DataType::UInt16 => Some(DataType::Int32),
+ DataType::UInt32 => Some(DataType::Int64),
+ // Spark reads Parquet UINT_64 as Decimal(20, 0). This is lossless for the full range and
+ // preserves values larger than i64::MAX for Variant decimal encoding.
+ DataType::UInt64 => Some(DataType::Decimal128(20, 0)),
+ // Arrow chooses Decimal256 from the physical byte width, but Spark's DecimalType is
+ // precision-based and stores every supported precision (<= 38) in 128 bits.
+ DataType::Decimal256(precision, scale) if *precision <= DECIMAL128_MAX_PRECISION => {
+ Some(DataType::Decimal128(*precision, *scale))
+ }
+ DataType::Timestamp(TimeUnit::Millisecond, timezone) => {
+ Some(DataType::Timestamp(TimeUnit::Microsecond, timezone.clone()))
+ }
+ DataType::FixedSizeBinary(_) => Some(DataType::Binary),
+ DataType::FixedSizeList(field, _) => Some(DataType::List(
+ normalize_field(field).unwrap_or_else(|| Arc::clone(field)),
+ )),
+ DataType::List(field) => normalize_field(field).map(DataType::List),
+ DataType::LargeList(field) => normalize_field(field).map(DataType::LargeList),
+ DataType::ListView(field) => normalize_field(field).map(DataType::ListView),
+ DataType::LargeListView(field) => normalize_field(field).map(DataType::LargeListView),
+ DataType::Struct(fields) => {
+ let mut changed = false;
+ let fields = fields
+ .iter()
+ .map(|field| match normalize_field(field) {
+ Some(field) => {
+ changed = true;
+ field
+ }
+ None => Arc::clone(field),
+ })
+ .collect::<Vec<_>>();
+ changed.then(|| DataType::Struct(fields.into()))
+ }
+ _ => None,
+ }
+}
+
+fn contains_uuid_extension(data_type: &DataType) -> bool {
+ fn field_contains_uuid(field: &FieldRef) -> bool {
+ (field.data_type() == &DataType::FixedSizeBinary(16)
+ && field.extension_type_name() == Some("arrow.uuid"))
+ || contains_uuid_extension(field.data_type())
+ }
+
+ match data_type {
+ DataType::Struct(fields) => fields.iter().any(field_contains_uuid),
+ DataType::List(field)
+ | DataType::LargeList(field)
+ | DataType::ListView(field)
+ | DataType::LargeListView(field)
+ | DataType::FixedSizeList(field, _)
+ | DataType::Map(field, _) => field_contains_uuid(field),
+ DataType::Dictionary(_, value_type) => contains_uuid_extension(value_type),
+ _ => false,
+ }
+}
+
+/// Arrow Variant compute cannot consume every storage type Spark reads. Decode and cast those
+/// children before validation, preserving UUID rejection and reporting conversion overflow.
+/// https://github.com/apache/datafusion-comet/issues/5477
+fn normalize_variant_storage(array: &ArrayRef) -> DataFusionResult<ArrayRef> {
+ if contains_uuid_extension(array.data_type()) {
+ return Err(DataFusionError::Execution(
+ "Parquet UUID is not supported as a shredded Variant child".to_string(),
+ ));
+ }
+ let Some(data_type) = normalize_variant_type(array.data_type()) else {
+ return Ok(Arc::clone(array));
+ };
+ Ok(cast_with_options(
+ array.as_ref(),
+ &data_type,
+ &arrow::compute::CastOptions {
+ safe: false,
+ ..Default::default()
+ },
+ )?)
+}
+
/// 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
diff --git a/native/core/src/parquet/cast_column/variant/tests.rs b/native/core/src/parquet/cast_column/variant/tests.rs
index 698bf675c7..46acc2b56a 100644
--- a/native/core/src/parquet/cast_column/variant/tests.rs
+++ b/native/core/src/parquet/cast_column/variant/tests.rs
@@ -39,6 +39,174 @@ fn target_field(nullable: bool) -> FieldRef {
)
}
+#[test]
+fn normalize_encoded_storage_and_unsigned_extremes() {
+ use arrow::array::{
+ DictionaryArray, Int8Array, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+ };
+ use arrow::datatypes::Int8Type;
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::Null);
+ let (metadata, _) = builder.finish();
+ let metadata: ArrayRef = Arc::new(BinaryArray::from(vec![metadata.as_slice()]));
+ let encoded_metadata: ArrayRef =
+ Arc::new(DictionaryArray::<Int8Type>::try_new(Int8Array::from(vec![0]), metadata).unwrap());
+ let cases: Vec<(ArrayRef, i128)> = vec![
+ (
+ Arc::new(UInt8Array::from(vec![u8::MAX])),
+ i128::from(u8::MAX),
+ ),
+ (
+ Arc::new(UInt16Array::from(vec![u16::MAX])),
+ i128::from(u16::MAX),
+ ),
+ (
+ Arc::new(UInt32Array::from(vec![u32::MAX])),
+ i128::from(u32::MAX),
+ ),
+ (
+ Arc::new(UInt64Array::from(vec![u64::MAX])),
+ i128::from(u64::MAX),
+ ),
+ ];
+ for (typed, expected) in cases {
+ let typed: ArrayRef = Arc::new(
+ DictionaryArray::<Int8Type>::try_new(Int8Array::from(vec![0]), typed).unwrap(),
+ );
+ let input: ArrayRef = Arc::new(StructArray::new(
+ vec![
+ Field::new("typed_value", typed.data_type().clone(), false),
+ Field::new("metadata", encoded_metadata.data_type().clone(), false),
+ ]
+ .into(),
+ vec![typed, Arc::clone(&encoded_metadata)],
+ None,
+ ));
+ let output = normalize_variant_array(&input, &target_field(false)).unwrap();
+ let output = VariantArray::try_new(output.as_ref()).unwrap();
+ let value = output.value(0);
+ if let Variant::Decimal16(decimal) = value {
+ assert_eq!(decimal.integer(), expected);
+ assert_eq!(decimal.scale(), 0);
+ } else {
+ assert_eq!(i128::from(value.as_int64().unwrap()), expected);
+ }
+ }
+
+ let mut builder = VariantArrayBuilder::new(1);
+ builder.append_variant(Variant::from(42_i64));
+ let base = builder.build();
+ let encoded_value: ArrayRef = Arc::new(
+ DictionaryArray::<Int8Type>::try_new(
+ Int8Array::from(vec![0]),
+ Arc::clone(base.value_column()),
+ )
+ .unwrap(),
+ );
+ let input: ArrayRef = Arc::new(StructArray::new(
+ vec![
+ Field::new("value", encoded_value.data_type().clone(), false),
+ Field::new("metadata", encoded_metadata.data_type().clone(), false),
+ ]
+ .into(),
+ vec![encoded_value, encoded_metadata],
+ None,
+ ));
+ let output = normalize_variant_array(&input, &target_field(false)).unwrap();
+ assert_eq!(
+ VariantArray::try_new(output.as_ref())
+ .unwrap()
+ .value(0)
+ .as_int64(),
+ Some(42)
+ );
+}
+
+#[test]
+fn normalize_fixed_storage_and_checked_timestamps() {
+ use arrow::array::{FixedSizeBinaryArray, TimestampMillisecondArray};
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::Null);
+ let (metadata, _) = builder.finish();
+ let normalize = |typed: ArrayRef| {
+ let input: ArrayRef = Arc::new(StructArray::new(
+ vec![
+ Field::new("metadata", DataType::Binary, false),
+ Field::new("typed_value", typed.data_type().clone(), false),
+ ]
+ .into(),
+ vec![
+ Arc::new(BinaryArray::from(vec![metadata.as_slice()])),
+ typed,
+ ],
+ None,
+ ));
+ normalize_variant_array(&input, &target_field(false))
+ };
+ for width in [3, 16] {
+ let bytes = vec![5; width];
+ let typed = FixedSizeBinaryArray::try_from_iter([bytes.as_slice()].into_iter()).unwrap();
+ let output = normalize(Arc::new(typed)).unwrap();
+ assert_eq!(
+ VariantArray::try_new(output.as_ref()).unwrap().value(0),
+ Variant::Binary(&bytes)
+ );
+ }
+ assert!(normalize(Arc::new(TimestampMillisecondArray::from(vec![i64::MAX]))).is_err());
+ assert!(normalize(Arc::new(TimestampMillisecondArray::from(vec![123]))).is_ok());
+}
+
+#[test]
+fn normalize_fixed_size_list_and_reject_uuid() {
+ use arrow::array::{FixedSizeBinaryArray, FixedSizeListArray, UInt16Array};
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::Null);
+ let (metadata, _) = builder.finish();
+ let wrap = |field: Field, typed: ArrayRef| -> ArrayRef {
+ Arc::new(StructArray::new(
+ vec![Field::new("metadata", DataType::Binary, false), field].into(),
+ vec![
+ Arc::new(BinaryArray::from(vec![metadata.as_slice()])),
+ typed,
+ ],
+ None,
+ ))
+ };
+ let elements: ArrayRef = Arc::new(StructArray::new(
+ vec![Field::new("typed_value", DataType::UInt16, false)].into(),
+ vec![Arc::new(UInt16Array::from(vec![1, u16::MAX]))],
+ None,
+ ));
+ let list: ArrayRef = Arc::new(FixedSizeListArray::new(
+ Arc::new(Field::new("item", elements.data_type().clone(), false)),
+ 2,
+ elements,
+ None,
+ ));
+ let input = wrap(
+ Field::new("typed_value", list.data_type().clone(), false),
+ list,
+ );
+ let output = normalize_variant_array(&input, &target_field(false)).unwrap();
+ let output = VariantArray::try_new(output.as_ref()).unwrap();
+ let Variant::List(list) = output.value(0) else {
+ panic!("expected list")
+ };
+ assert_eq!(
+ list.iter()
+ .map(|value| value.as_int64().unwrap())
+ .collect::<Vec<_>>(),
+ vec![1, 65535]
+ );
+
+ let uuid: ArrayRef =
+ Arc::new(FixedSizeBinaryArray::try_from_iter([[0_u8; 16]].into_iter()).unwrap());
+ let field = Field::new("typed_value", uuid.data_type().clone(), false)
+ .with_metadata([("ARROW:extension:name".to_string(), "arrow.uuid".to_string())].into());
+ let error = normalize_variant_array(&wrap(field, uuid), &target_field(false)).unwrap_err();
+ assert!(error.to_string().contains("Parquet UUID"));
+}
+
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());
diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs
index 59c68a3f86..d89a177283 100644
--- a/native/core/src/parquet/eager_page_index_reader_factory.rs
+++ b/native/core/src/parquet/eager_page_index_reader_factory.rs
@@ -46,6 +46,7 @@
//! Filed upstream as apache/datafusion#23978. Revert this once the opener merges its deferred
//! page-index load back into `FileMetadataCache` instead of bypassing it.
+use arrow::datatypes::{DataType, FieldRef, Schema};
use async_trait::async_trait;
use bytes::Bytes;
use datafusion::common::Result as DFResult;
@@ -69,10 +70,14 @@ use object_store::{
};
use parquet::arrow::arrow_reader::ArrowReaderOptions;
use parquet::arrow::async_reader::AsyncFileReader;
-use parquet::errors::ParquetError;
+use parquet::arrow::{encode_arrow_schema, parquet_to_arrow_schema, ARROW_SCHEMA_META_KEY};
+use parquet::basic::{ConvertedType, LogicalType};
+use parquet::errors::{ParquetError, Result as ParquetResult};
+use parquet::file::metadata::{FileMetaData, KeyValue, ParquetMetaDataBuilder};
use parquet::file::metadata::{
FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader,
};
+use parquet::schema::types::{ColumnDescPtr, SchemaDescriptor};
use std::fmt::{Debug, Display, Formatter};
use std::ops::Range;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
@@ -154,6 +159,10 @@ pub struct EagerPageIndexReaderFactory {
store: Arc<dyn ObjectStore>,
metadata_cache: Arc<FileMetadataCache>,
scan_io_metrics: Arc<ScanIoMetrics>,
+ // Arrow schema hints and ENUM inference can change Spark's Variant interpretation.
+ // Enable the footer workaround only for scans that project Variant.
+ // https://github.com/apache/datafusion-comet/issues/5477
+ spark_variant_schema: bool,
}
impl EagerPageIndexReaderFactory {
@@ -181,8 +190,14 @@ impl EagerPageIndexReaderFactory {
store,
metadata_cache,
scan_io_metrics,
+ spark_variant_schema: false,
}
}
+
+ pub fn with_spark_variant_schema(mut self, enabled: bool) -> Self {
+ self.spark_variant_schema = enabled;
+ self
+ }
}
impl ParquetFileReaderFactory for EagerPageIndexReaderFactory {
@@ -209,6 +224,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory {
partitioned_file,
metadata_cache: Arc::clone(&self.metadata_cache),
metadata_size_hint,
+ spark_variant_schema: self.spark_variant_schema,
}))
}
}
@@ -223,6 +239,137 @@ struct EagerPageIndexReader {
partitioned_file: PartitionedFile,
metadata_cache: Arc<FileMetadataCache>,
metadata_size_hint: Option<usize>,
+ spark_variant_schema: bool,
+}
+
+// Arrow infers ENUM as Binary, losing the distinction from raw binary that Spark needs.
+// Inspect the Parquet annotation so Variant reconstruction can retain ENUM as a string.
+// https://github.com/apache/datafusion-comet/issues/5477
+fn is_enum_column(column: &ColumnDescPtr) -> bool {
+ matches!(column.logical_type_ref(), Some(LogicalType::Enum))
+ || column.converted_type() == ConvertedType::ENUM
+}
+
+fn spark_enum_field(
+ field: &FieldRef,
+ columns: &[ColumnDescPtr],
+ column_index: &mut usize,
+) -> ParquetResult<FieldRef> {
+ let rewrite =
+ |field: &FieldRef, data_type| Arc::new(field.as_ref().clone().with_data_type(data_type));
+ let data_type = match field.data_type() {
+ DataType::Struct(fields) => DataType::Struct(
+ fields
+ .iter()
+ .map(|field| spark_enum_field(field, columns, column_index))
+ .collect::<ParquetResult<Vec<_>>>()?
+ .into(),
+ ),
+ DataType::List(child) => DataType::List(spark_enum_field(child, columns, column_index)?),
+ DataType::LargeList(child) => {
+ DataType::LargeList(spark_enum_field(child, columns, column_index)?)
+ }
+ DataType::FixedSizeList(child, size) => {
+ DataType::FixedSizeList(spark_enum_field(child, columns, column_index)?, *size)
+ }
+ DataType::ListView(child) => {
+ DataType::ListView(spark_enum_field(child, columns, column_index)?)
+ }
+ DataType::LargeListView(child) => {
+ DataType::LargeListView(spark_enum_field(child, columns, column_index)?)
+ }
+ DataType::Map(child, sorted) => {
+ DataType::Map(spark_enum_field(child, columns, column_index)?, *sorted)
+ }
+ _ => {
+ let column = columns.get(*column_index).ok_or_else(|| {
+ ParquetError::General(
+ "Arrow schema contains more leaves than the Parquet schema".to_string(),
+ )
+ })?;
+ *column_index += 1;
+ if is_enum_column(column) {
+ DataType::Utf8
+ } else {
+ return Ok(Arc::clone(field));
+ }
+ }
+ };
+ Ok(rewrite(field, data_type))
+}
+
+/// Arrow maps Parquet ENUM to Binary, while Spark reads it as String. Supply a schema hint
+/// that changes only ENUM leaves so Variant reconstruction preserves Spark's interpretation.
+/// https://github.com/apache/datafusion-comet/issues/5477
+fn spark_enum_schema(schema: &SchemaDescriptor) -> ParquetResult<Option<Schema>> {
+ let columns = schema.columns();
+ if !columns.iter().any(is_enum_column) {
+ return Ok(None);
+ }
+
+ let arrow_schema = parquet_to_arrow_schema(schema, None)?;
+ let mut column_index = 0;
+ let fields = arrow_schema
+ .fields()
+ .iter()
+ .map(|field| spark_enum_field(field, columns, &mut column_index))
+ .collect::<ParquetResult<Vec<_>>>()?;
+ if column_index != columns.len() {
+ return Err(ParquetError::General(
+ "Parquet schema contains more leaves than the Arrow schema".to_string(),
+ ));
+ }
+ Ok(Some(Schema::new_with_metadata(
+ fields,
+ arrow_schema.metadata().clone(),
+ )))
+}
+
+/// Arrow restores advisory `ARROW:schema` types that can differ from Spark's physical Parquet
+/// interpretation. Replace that hint with physical inference and the ENUM string mapping.
+/// Rebuild only the returned metadata; the shared cache retains the original footer.
+/// https://github.com/apache/datafusion-comet/issues/5477
+fn with_spark_arrow_schema(metadata: Arc<ParquetMetaData>) -> ParquetResult<Arc<ParquetMetaData>> {
+ let file = metadata.file_metadata();
+ let has_arrow_schema = file.key_value_metadata().is_some_and(|key_values| {
+ key_values
+ .iter()
+ .any(|key_value| key_value.key == ARROW_SCHEMA_META_KEY)
+ });
+ let enum_schema = spark_enum_schema(file.schema_descr())?;
+ if !has_arrow_schema && enum_schema.is_none() {
+ return Ok(metadata);
+ }
+
+ let mut key_values = file
+ .key_value_metadata()
+ .into_iter()
+ .flatten()
+ .filter(|key_value| key_value.key != ARROW_SCHEMA_META_KEY)
+ .cloned()
+ .collect::<Vec<_>>();
+ if let Some(schema) = enum_schema {
+ key_values.push(KeyValue {
+ key: ARROW_SCHEMA_META_KEY.to_string(),
+ value: Some(encode_arrow_schema(&schema)),
+ });
+ }
+
+ let file = FileMetaData::new(
+ file.version(),
+ file.num_rows(),
+ file.created_by().map(str::to_owned),
+ Some(key_values),
+ file.schema_descr_ptr(),
+ file.column_orders().cloned(),
+ );
+ Ok(Arc::new(
+ ParquetMetaDataBuilder::new(file)
+ .set_row_groups(metadata.row_groups().to_vec())
+ .set_column_index(metadata.column_index().cloned())
+ .set_offset_index(metadata.offset_index().cloned())
+ .build(),
+ ))
}
impl AsyncFileReader for EagerPageIndexReader {
@@ -291,11 +438,17 @@ impl AsyncFileReader for EagerPageIndexReader {
let store = Arc::clone(&self.store);
let metadata_size_hint = self.metadata_size_hint;
let scan_io_metrics = Arc::clone(&self.scan_io_metrics);
+ let spark_variant_schema = self.spark_variant_schema;
async move {
let file_decryption_properties = options
.and_then(|o| o.file_decryption_properties())
.map(Arc::clone);
let cache_enabled = file_decryption_properties.is_none();
+ if spark_variant_schema && file_decryption_properties.is_some() {
+ return Err(ParquetError::General(
+ "Projected Variant with Parquet encryption requires Spark fallback".to_string(),
+ ));
+ }
let page_index_policy = if file_decryption_properties.is_none() {
Some(PageIndexPolicy::Optional)
} else {
@@ -344,7 +497,12 @@ impl AsyncFileReader for EagerPageIndexReader {
metadata_store.record_valid_footer();
}
- metadata
+ let metadata = metadata?;
+ if spark_variant_schema {
+ with_spark_arrow_schema(metadata)
+ } else {
+ Ok(metadata)
+ }
}
.boxed()
}
@@ -672,7 +830,16 @@ impl Drop for EagerPageIndexReader {
#[cfg(test)]
mod tests {
use super::*;
+ use arrow::{array::Int32Array, record_batch::RecordBatch};
use object_store::memory::InMemory;
+ use parquet::{
+ arrow::ArrowWriter,
+ file::{
+ properties::{EnabledStatistics, WriterProperties},
+ reader::FileReader,
+ serialized_reader::{ReadOptionsBuilder, SerializedFileReader},
+ },
+ };
#[derive(Debug)]
struct RecordingRangeStore {
@@ -824,4 +991,59 @@ mod tests {
if remote { 6 } else { 0 }
);
}
+
+ #[test]
+ fn variant_policy_preserves_footer_metadata_and_indexes() {
+ let schema = Arc::new(Schema::new_with_metadata(
+ vec![arrow::datatypes::Field::new("id", DataType::Int32, false)],
+ [("application".to_string(), "keep".to_string())].into(),
+ ));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+ )
+ .unwrap();
+ let file = tempfile::NamedTempFile::new().unwrap();
+ let props = WriterProperties::builder()
+ .set_key_value_metadata(Some(vec![KeyValue::new(
+ "application".to_string(),
+ "keep".to_string(),
+ )]))
+ .set_statistics_enabled(EnabledStatistics::Page)
+ .set_data_page_row_count_limit(1)
+ .build();
+ let mut writer = ArrowWriter::try_new(file.reopen().unwrap(), schema, Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ let reader = SerializedFileReader::new_with_options(
+ file.reopen().unwrap(),
+ ReadOptionsBuilder::new().with_page_index().build(),
+ )
+ .unwrap();
+ let original = Arc::new(reader.metadata().clone());
+ let rewritten = with_spark_arrow_schema(Arc::clone(&original)).unwrap();
+ assert!(original.column_index().is_some());
+ assert!(original.offset_index().is_some());
+ assert_eq!(rewritten.column_index(), original.column_index());
+ assert_eq!(rewritten.offset_index(), original.offset_index());
+ assert_eq!(rewritten.row_groups(), original.row_groups());
+ assert_eq!(
+ rewritten.file_metadata().column_orders(),
+ original.file_metadata().column_orders()
+ );
+ assert!(original
+ .file_metadata()
+ .key_value_metadata()
+ .unwrap()
+ .iter()
+ .any(|entry| entry.key == ARROW_SCHEMA_META_KEY));
+ assert_eq!(
+ rewritten.file_metadata().key_value_metadata().unwrap(),
+ &vec![KeyValue::new("application".to_string(), "keep".to_string())]
+ );
+ assert!(Arc::ptr_eq(
+ &rewritten,
+ &with_spark_arrow_schema(Arc::clone(&rewritten)).unwrap()
+ ));
+ }
}
diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs
index ff611dd3d1..60a1027f97 100644
--- a/native/core/src/parquet/parquet_exec.rs
+++ b/native/core/src/parquet/parquet_exec.rs
@@ -37,9 +37,13 @@ use datafusion::prelude::SessionContext;
use datafusion::scalar::ScalarValue;
use datafusion_comet_spark_expr::EvalMode;
use datafusion_datasource::TableSchema;
+use parquet::variant::VariantType;
use std::collections::HashMap;
use std::sync::Arc;
+#[cfg(test)]
+mod variant_tests;
+
/// Initializes a DataSourceExec plan with a ParquetSource for Comet's native Parquet scan.
///
/// `required_schema`: Schema to be projected by the scan.
@@ -147,6 +151,15 @@ pub(crate) fn init_datasource_exec(
.with_table_parquet_options(table_parquet_options)
.with_metadata_size_hint(512 * 1024); // Same as DataFusion's default
+ let projects_variant = required_schema
+ .fields()
+ .iter()
+ .any(|field| field.has_valid_extension_type::<VariantType>());
+ if projects_variant && encryption_enabled {
+ return Err(ExecutionError::GeneralError(
+ "Projected Variant with Parquet encryption requires Spark fallback".to_string(),
+ ));
+ }
if encryption_enabled {
parquet_source = parquet_source.with_encryption_factory(
session_ctx
@@ -173,12 +186,15 @@ pub(crate) fn init_datasource_exec(
let store = runtime_env.object_store(&object_store_url)?;
let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache();
let scan_io_source = scan_io_source(object_store_backend);
- let reader_factory = Arc::new(EagerPageIndexReaderFactory::new(
- store,
- metadata_cache,
- scan_io_source,
- parquet_source.metrics(),
- ));
+ let reader_factory = Arc::new(
+ EagerPageIndexReaderFactory::new(
+ store,
+ metadata_cache,
+ scan_io_source,
+ parquet_source.metrics(),
+ )
+ .with_spark_variant_schema(projects_variant),
+ );
parquet_source = parquet_source.with_parquet_file_reader_factory(reader_factory);
// Route data filters through `try_pushdown_filters` rather than calling
diff --git a/native/core/src/parquet/parquet_exec/variant_tests.rs b/native/core/src/parquet/parquet_exec/variant_tests.rs
new file mode 100644
index 0000000000..641529c40d
--- /dev/null
+++ b/native/core/src/parquet/parquet_exec/variant_tests.rs
@@ -0,0 +1,351 @@
+// 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::datatypes::i256;
+use arrow::{
+ array::{Array, ArrayRef, BinaryArray, Date64Array, Decimal256Array, StructArray},
+ datatypes::{DataType, Fields, Schema},
+ record_batch::RecordBatch,
+};
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename;
+use futures::StreamExt;
+use parquet::{
+ arrow::{arrow_writer::ArrowWriterOptions, ArrowWriter},
+ basic::{LogicalType, Repetition, Type as PhysicalType},
+ data_type::{
+ ByteArray, ByteArrayType, DataType as ParquetDataType, FixedLenByteArray,
+ FixedLenByteArrayType,
+ },
+ file::{properties::WriterProperties, writer::SerializedFileWriter},
+ schema::types::{Type as ParquetType, TypePtr},
+ variant::{Variant, VariantArray, VariantBuilder, VariantDecimal16},
+};
+use std::{fs::File, path::PathBuf};
+fn required_variant_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new(
+ "v",
+ DataType::Struct(Fields::from(vec![
+ Field::new("value", DataType::Binary, false),
+ Field::new("metadata", DataType::Binary, false),
+ ])),
+ false,
+ )
+ .with_extension_type(VariantType)]))
+}
+
+async fn write_and_scan_shredded_variant(
+ typed_value: ArrayRef,
+ coerce_types: bool,
+) -> VariantArray {
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::Null);
+ let (metadata, _) = builder.finish();
+ let metadata: ArrayRef = Arc::new(BinaryArray::from(vec![Some(metadata.as_slice())]));
+ 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![metadata, typed_value],
+ None,
+ )
+ .unwrap(),
+ );
+ let file_schema = Arc::new(Schema::new(vec![Field::new(
+ "v",
+ physical.data_type().clone(),
+ false,
+ )
+ .with_extension_type(VariantType)]));
+ let batch = RecordBatch::try_new(Arc::clone(&file_schema), vec![physical]).unwrap();
+
+ let filename = get_temp_filename();
+ let file = File::create(&filename).unwrap();
+ let properties = WriterProperties::builder()
+ .set_coerce_types(coerce_types)
+ .build();
+ let mut writer = ArrowWriter::try_new_with_options(
+ file,
+ file_schema,
+ ArrowWriterOptions::new().with_properties(properties),
+ )
+ .unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+
+ scan_variant_file(filename).await
+}
+
+fn write_variant_typed_value<T: ParquetDataType>(typed_value: TypePtr, values: &[T::T]) -> PathBuf {
+ let filename = get_temp_filename();
+ let file = File::create(&filename).unwrap();
+ let metadata = Arc::new(
+ ParquetType::primitive_type_builder("metadata", PhysicalType::BYTE_ARRAY)
+ .with_repetition(Repetition::REQUIRED)
+ .build()
+ .unwrap(),
+ );
+ let variant = Arc::new(
+ ParquetType::group_type_builder("v")
+ .with_repetition(Repetition::REQUIRED)
+ .with_logical_type(Some(LogicalType::variant(None)))
+ .with_fields(vec![metadata, typed_value])
+ .build()
+ .unwrap(),
+ );
+ let schema = Arc::new(
+ ParquetType::group_type_builder("schema")
+ .with_fields(vec![variant])
+ .build()
+ .unwrap(),
+ );
+ let mut writer = SerializedFileWriter::new(file, schema, Default::default()).unwrap();
+ let mut row_group = writer.next_row_group().unwrap();
+
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::Null);
+ let (metadata, _) = builder.finish();
+ let metadata = (0..values.len())
+ .map(|_| ByteArray::from(metadata.clone()))
+ .collect::<Vec<_>>();
+ let mut column = row_group.next_column().unwrap().unwrap();
+ column
+ .typed::<ByteArrayType>()
+ .write_batch(&metadata, None, None)
+ .unwrap();
+ column.close().unwrap();
+
+ let mut column = row_group.next_column().unwrap().unwrap();
+ column.typed::<T>().write_batch(values, None, None).unwrap();
+ column.close().unwrap();
+ row_group.close().unwrap();
+ writer.close().unwrap();
+ filename
+}
+
+async fn scan_variant_file(filename: PathBuf) -> VariantArray {
+ let partitioned_file =
+ PartitionedFile::from_path(filename.to_string_lossy().into_owned()).unwrap();
+ let session_ctx = Arc::new(SessionContext::new());
+ let scan = init_datasource_exec(
+ required_variant_schema(),
+ None,
+ None,
+ ObjectStoreUrl::local_filesystem(),
+ ObjectStoreBackend::Local,
+ vec![vec![partitioned_file]],
+ None,
+ None,
+ None,
+ "UTC",
+ true,
+ false,
+ false,
+ false,
+ &session_ctx,
+ false,
+ false,
+ false,
+ )
+ .unwrap();
+ let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap();
+ let batch = stream.next().await.unwrap().unwrap();
+ assert!(stream.next().await.is_none());
+ VariantArray::try_new(batch.column(0).as_ref()).unwrap()
+}
+
+#[tokio::test]
+async fn unread_variant_does_not_override_arrow_schema_hint() {
+ let variant = required_variant_schema().field(0).clone();
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("d", DataType::Date64, false),
+ variant,
+ ]));
+ let mut builder = VariantBuilder::new();
+ builder.append_value(Variant::Null);
+ let (metadata, value) = builder.finish();
+ let physical = StructArray::new(
+ match schema.field(1).data_type() {
+ DataType::Struct(fields) => fields.clone(),
+ _ => unreachable!(),
+ },
+ vec![
+ Arc::new(BinaryArray::from(vec![value.as_slice()])),
+ Arc::new(BinaryArray::from(vec![metadata.as_slice()])),
+ ],
+ None,
+ );
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Date64Array::from(vec![86_400_000])),
+ Arc::new(physical),
+ ],
+ )
+ .unwrap();
+ let file = tempfile::NamedTempFile::new().unwrap();
+ let mut writer = ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ let required = Arc::new(Schema::new(vec![Field::new("d", DataType::Date64, false)]));
+ let session = Arc::new(SessionContext::new());
+ let scan = init_datasource_exec(
+ required,
+ None,
+ None,
+ ObjectStoreUrl::local_filesystem(),
+ ObjectStoreBackend::Local,
+ vec![vec![PartitionedFile::from_path(
+ file.path().to_string_lossy().into_owned(),
+ )
+ .unwrap()]],
+ None,
+ None,
+ None,
+ "UTC",
+ true,
+ false,
+ false,
+ false,
+ &session,
+ false,
+ false,
+ false,
+ )
+ .unwrap();
+ let mut stream = scan.execute(0, session.task_ctx()).unwrap();
+ let output = stream.next().await.unwrap().unwrap();
+ assert_eq!(output.column(0).data_type(), &DataType::Date64);
+ assert_eq!(output.num_columns(), 1);
+}
+
+#[test]
+fn encrypted_projected_variant_is_rejected_before_reader_creation() {
+ let session = Arc::new(SessionContext::new());
+ let result = init_datasource_exec(
+ required_variant_schema(),
+ None,
+ None,
+ ObjectStoreUrl::local_filesystem(),
+ ObjectStoreBackend::Local,
+ vec![],
+ None,
+ None,
+ None,
+ "UTC",
+ true,
+ false,
+ false,
+ false,
+ &session,
+ true,
+ false,
+ false,
+ );
+ assert!(result
+ .unwrap_err()
+ .to_string()
+ .contains("requires Spark fallback"));
+}
+
+#[tokio::test]
+async fn variant_scan_uses_parquet_physical_types_instead_of_arrow_schema_hints() {
+ let decimal: ArrayRef = Arc::new(
+ Decimal256Array::from(vec![i256::from_i128(123)])
+ .with_precision_and_scale(38, 2)
+ .unwrap(),
+ );
+ let output = write_and_scan_shredded_variant(decimal, false).await;
+ assert_eq!(
+ output.value(0),
+ Variant::Decimal16(VariantDecimal16::try_new(123, 2).unwrap())
+ );
+
+ let date64: ArrayRef = Arc::new(Date64Array::from(vec![86_400_000]));
+ let output = write_and_scan_shredded_variant(Arc::clone(&date64), false).await;
+ assert_eq!(output.value(0).as_int64(), Some(86_400_000));
+
+ let output = write_and_scan_shredded_variant(date64, true).await;
+ let Variant::Date(date) = output.value(0) else {
+ panic!("expected DATE-annotated physical value")
+ };
+ assert_eq!(date.to_string(), "1970-01-02");
+}
+
+#[tokio::test]
+async fn variant_scan_preserves_parquet_enum_string_and_binary_semantics() {
+ for (logical_type, expected_string) in [
+ (Some(LogicalType::Enum), true),
+ (Some(LogicalType::String), true),
+ (None, false),
+ ] {
+ let typed_value = Arc::new(
+ ParquetType::primitive_type_builder("typed_value", PhysicalType::BYTE_ARRAY)
+ .with_repetition(Repetition::REQUIRED)
+ .with_logical_type(logical_type)
+ .build()
+ .unwrap(),
+ );
+ let filename = write_variant_typed_value::<ByteArrayType>(
+ typed_value,
+ &[ByteArray::from(b"red".to_vec())],
+ );
+
+ let output = scan_variant_file(filename).await;
+ if expected_string {
+ assert_eq!(output.value(0).as_string(), Some("red"));
+ } else {
+ assert_eq!(output.value(0), Variant::Binary(b"red"));
+ }
+ }
+}
+
+#[tokio::test]
+async fn variant_scan_reads_wide_physical_decimal_as_decimal128() {
+ for width in [17, 32] {
+ let values = [123_i128, -123_i128]
+ .into_iter()
+ .map(|value| {
+ let mut bytes = vec![if value.is_negative() { 0xff } else { 0 }; width];
+ bytes[width - 16..].copy_from_slice(&value.to_be_bytes());
+ FixedLenByteArray::from(bytes)
+ })
+ .collect::<Vec<_>>();
+ let typed_value = Arc::new(
+ ParquetType::primitive_type_builder("typed_value", PhysicalType::FIXED_LEN_BYTE_ARRAY)
+ .with_repetition(Repetition::REQUIRED)
+ .with_logical_type(Some(LogicalType::decimal(2, 38)))
+ .with_length(width as i32)
+ .with_precision(38)
+ .with_scale(2)
+ .build()
+ .unwrap(),
+ );
+ let filename = write_variant_typed_value::<FixedLenByteArrayType>(typed_value, &values);
+
+ let output = scan_variant_file(filename).await;
+ for (index, value) in [123, -123].into_iter().enumerate() {
+ assert_eq!(
+ output.value(index),
+ Variant::Decimal16(VariantDecimal16::try_new(value, 2).unwrap())
+ );
+ }
+ }
+}
diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs
index 004cf93ebf..818916f126 100644
--- a/native/core/src/parquet/schema_adapter.rs
+++ b/native/core/src/parquet/schema_adapter.rs
@@ -254,29 +254,16 @@ fn remap_physical_schema(
let logical_folded = fold_schema_names(logical_schema, case_sensitive);
let physical_folded = fold_schema_names(physical_schema, case_sensitive);
- // Folded names of ID-bearing logical fields whose ID is not present in the file. Any physical
- // field that shares one of these names must be renamed to something the
- // `DefaultPhysicalExprAdapter` cannot name-match, otherwise the read would silently fall
- // through to a name match. Spark's `matchIdField` solves the same problem with
- // `generateFakeColumnName` (see `ParquetReadSupport.scala`).
- let unmatched_id_logical_folded: HashSet<String> = if should_match_by_id {
- logical_schema
- .fields()
- .iter()
- .enumerate()
- .filter_map(|(j, lf)| {
- parse_field_id(lf).and_then(|id| {
- if id_to_phys_names.contains_key(&id) {
- None
- } else {
- Some(logical_folded[j].clone())
- }
- })
- })
- .collect()
- } else {
- HashSet::new()
- };
+ // All ID-bearing targets resolve by ID, even when a different file column has the
+ // requested name. Hide that shadowing column after giving its own ID match precedence.
+ let id_logical_folded: HashSet<&String> = logical_schema
+ .fields()
+ .iter()
+ .zip(&logical_folded)
+ .filter(|(field, _)| should_match_by_id && parse_field_id(field).is_some())
+ .map(|(_, name)| name)
+ .collect();
+ let mut occupied_names = HashSet::new();
let mut fake_counter: usize = 0;
let mut name_map: HashMap<String, String> = HashMap::new();
@@ -305,13 +292,19 @@ fn remap_physical_schema(
}
}
- // Block accidental name match for ID-bearing logical fields whose ID is missing
- // from the file. Mirrors Spark's `generateFakeColumnName` in `matchIdField`.
- if should_match_by_id
- && unmatched_id_logical_folded.contains(&physical_folded[phys_idx])
- {
- fake_counter += 1;
- let fake_name = format!("__comet_unmatched_field_id_{}", fake_counter);
+ // Block accidental name matches for ID-bearing targets, whether their ID was
+ // missing or resolved to a different physical column.
+ if should_match_by_id && id_logical_folded.contains(&physical_folded[phys_idx]) {
+ if fake_counter == 0 {
+ occupied_names.extend(logical_folded.iter().chain(&physical_folded).cloned());
+ }
+ let fake_name = loop {
+ fake_counter += 1;
+ let name = format!("__comet_unmatched_field_id_{}", fake_counter);
+ if occupied_names.insert(name.clone()) {
+ break name;
+ }
+ };
return Arc::new(
Field::new(fake_name, field.data_type().clone(), field.is_nullable())
.with_metadata(field.metadata().clone()),
@@ -2167,6 +2160,40 @@ mod test {
.has_valid_extension_type::<VariantType>());
}
+ #[test]
+ fn variant_field_id_wins_over_a_shadowing_name() {
+ let storage = DataType::Struct(Fields::from(vec![
+ Field::new("value", DataType::Binary, false),
+ Field::new("metadata", DataType::Binary, false),
+ ]));
+ for case_sensitive in [true, false] {
+ let logical = Arc::new(Schema::new(vec![Field::new("v", storage.clone(), true)
+ .with_metadata(id_meta("1"))
+ .with_extension_type(VariantType)]));
+ let physical = Arc::new(Schema::new(vec![
+ Field::new(
+ if case_sensitive { "v" } else { "V" },
+ DataType::Binary,
+ true,
+ )
+ .with_metadata(id_meta("2")),
+ Field::new("other", storage.clone(), true).with_metadata(id_meta("1")),
+ Field::new("__comet_unmatched_field_id_1", DataType::Binary, true),
+ ]));
+ let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false);
+ options.case_sensitive = case_sensitive;
+ options.use_field_id = true;
+ let adapter = SparkPhysicalExprAdapterFactory::new(options, None)
+ .create(logical, Arc::clone(&physical))
+ .unwrap();
+ let rewritten = adapter.rewrite(Arc::new(Column::new("v", 0))).unwrap();
+ let cast = rewritten.downcast_ref::<CometCastColumnExpr>().unwrap();
+ let column = cast.children()[0].downcast_ref::<Column>().unwrap();
+ assert_eq!(column.name(), "other");
+ assert_eq!(column.index(), 1);
+ }
+ }
+
/// #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