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 { + fn normalize_field(field: &FieldRef) -> Option { + 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::>(); + 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 { + 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::::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::::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::::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![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 { let mut keys = (0..30).map(|i| format!("k{i:02}")).collect::>(); 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 22a94d04a1..11dfd181a9 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 bytes::Bytes; use datafusion::common::Result as DFResult; use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; @@ -60,8 +61,12 @@ use futures::{FutureExt, TryFutureExt}; use object_store::{ObjectStore, ObjectStoreExt}; 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::{PageIndexPolicy, ParquetMetaData}; +use parquet::schema::types::{ColumnDescPtr, SchemaDescriptor}; use std::fmt::Debug; use std::ops::Range; use std::sync::Arc; @@ -70,6 +75,10 @@ use std::sync::Arc; pub struct EagerPageIndexReaderFactory { store: Arc, metadata_cache: Arc, + // 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 { @@ -77,8 +86,14 @@ impl EagerPageIndexReaderFactory { Self { store, metadata_cache, + 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 { @@ -101,6 +116,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { partitioned_file, metadata_cache: Arc::clone(&self.metadata_cache), metadata_size_hint, + spark_variant_schema: self.spark_variant_schema, })) } } @@ -114,6 +130,137 @@ struct EagerPageIndexReader { partitioned_file: PartitionedFile, metadata_cache: Arc, metadata_size_hint: Option, + 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 { + 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::>>()? + .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> { + 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::>>()?; + 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) -> ParquetResult> { + 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::>(); + 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 { @@ -154,17 +301,23 @@ impl AsyncFileReader for EagerPageIndexReader { let metadata_cache = Arc::clone(&self.metadata_cache); let store = Arc::clone(&self.store); let metadata_size_hint = self.metadata_size_hint; + 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); + 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 { options.map(|o| o.column_index_policy()) }; - DFParquetMetadata::new(store.as_ref(), &object_meta) + let metadata = DFParquetMetadata::new(store.as_ref(), &object_meta) .with_decryption_properties(file_decryption_properties) .with_file_metadata_cache(Some(metadata_cache)) .with_metadata_size_hint(metadata_size_hint) @@ -176,7 +329,12 @@ impl AsyncFileReader for EagerPageIndexReader { "Failed to fetch metadata for file {}: {e}", object_meta.location, )) - }) + })?; + if spark_variant_schema { + with_spark_arrow_schema(metadata) + } else { + Ok(metadata) + } } .boxed() } @@ -194,3 +352,72 @@ impl Drop for EagerPageIndexReader { .set_total(self.partitioned_file.object_meta.size as usize); } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{array::Int32Array, record_batch::RecordBatch}; + use parquet::{ + arrow::ArrowWriter, + file::{ + properties::{EnabledStatistics, WriterProperties}, + reader::FileReader, + serialized_reader::{ReadOptionsBuilder, SerializedFileReader}, + }, + }; + + #[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 8796cb2324..bccf3a2e5c 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -36,9 +36,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. @@ -145,6 +149,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::()); + 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 @@ -170,7 +183,8 @@ 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(); parquet_source = parquet_source.with_parquet_file_reader_factory(Arc::new( - EagerPageIndexReaderFactory::new(store, metadata_cache), + EagerPageIndexReaderFactory::new(store, metadata_cache) + .with_spark_variant_schema(projects_variant), )); // 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..3149460ecc --- /dev/null +++ b/native/core/src/parquet/parquet_exec/variant_tests.rs @@ -0,0 +1,348 @@ +// 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(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::>(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&metadata, None, None) + .unwrap(); + column.close().unwrap(); + + let mut column = row_group.next_column().unwrap().unwrap(); + column.typed::().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(), + 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(), + 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(), + 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::( + 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::>(); + 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::(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 = 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 = 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::()); } + #[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::().unwrap(); + let column = cast.children()[0].downcast_ref::().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