From df471a07828b12fc1fcbcb9fcc82a425257c7b1b Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 29 Sep 2026 21:25:14 +0700 Subject: [PATCH 1/2] fix: match container field ids that the INT96 coercion drops With spark.sql.parquet.fieldId.read.enabled, the native scan matches requested fields to file fields by the id on the Arrow schema DataFusion hands it. When a file has an INT96 column, DataFusion 55.1.0's INT96 coercion rebuilds struct, list and map fields without their metadata, so an id that sits only on a container is gone and the container is null filled, while Spark matches the id on the Parquet group and reads it. Spark writes timestamps as INT96 by default, so any id-bearing container next to a timestamp was affected. apache/datafusion#24790 keeps the metadata, but it is not in 55.1.0. Until Comet moves to a DataFusion release with it, the reader records the name and id of every field from the file's own Parquet schema in the footer it returns, the schema adapter puts the dropped container ids back on the physical schema, and the column expression relabels decoded arrays to match. The stamp is added only when field id matching is on, the file has an INT96 column and a container carries an id, and only to the metadata returned for that open, not the shared metadata cache. A key of the same name already in the file is removed. A file that would need the stamp while decryption properties are configured is refused, since its footer cannot be rewritten. A test drives DataFusion's INT96 coercion directly and fails once the coercion keeps container metadata, which is the signal to remove the workaround. --- .../user-guide/latest/compatibility/scans.md | 5 + native/core/src/parquet/cast_column.rs | 19 + .../eager_page_index_reader_factory.rs | 289 ++++++++++- native/core/src/parquet/field_id_stamp.rs | 484 ++++++++++++++++++ native/core/src/parquet/mod.rs | 1 + native/core/src/parquet/parquet_exec.rs | 189 ++++++- native/core/src/parquet/schema_adapter.rs | 52 ++ .../comet/parquet/ParquetReadSuite.scala | 111 ++++ 8 files changed, 1131 insertions(+), 19 deletions(-) create mode 100644 native/core/src/parquet/field_id_stamp.rs diff --git a/docs/source/user-guide/latest/compatibility/scans.md b/docs/source/user-guide/latest/compatibility/scans.md index 57363125331..1c31672f29d 100644 --- a/docs/source/user-guide/latest/compatibility/scans.md +++ b/docs/source/user-guide/latest/compatibility/scans.md @@ -78,6 +78,11 @@ The following limitations raise an error at scan time rather than falling back t duplicate names; inference across files can depend on merge order. Resolution is tracked in [#5884](https://github.com/apache/datafusion-comet/issues/5884), with mixed-type behavior in [#5964](https://github.com/apache/datafusion-comet/issues/5964). +- Reading an encrypted Parquet file with `spark.sql.parquet.fieldId.read.enabled=true` when a + struct, list or map carrying a field id sits in a file with an `INT96` timestamp. Comet + restores those field ids through the file footer, which it cannot rewrite while Parquet + decryption is configured ([#6131](https://github.com/apache/datafusion-comet/issues/6131)). + Set `spark.comet.scan.enabled=false` to read such a file with Spark. - Invalid UTF-8 bytes in `STRING` columns. Spark permits arbitrary byte sequences in a `STRING` column (for example from `CAST(X'C1' AS STRING)`), but Comet's native execution path is built on Arrow, whose string type is strictly UTF-8. Reading a Parquet file whose `STRING` column contains diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index ad7b909768b..2a23bf955ae 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -283,6 +283,25 @@ impl PhysicalExpr for CometCastColumnExpr { fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult { let value = self.expr.evaluate(batch)?; + // Under field-id matching the physical field may carry nested ids that + // `restore_field_ids` put back and the decoder, built from DataFusion's INT96-coerced + // schema, leaves off. Relabel the decoded array to the physical type first, so the + // conversion below resolves nested fields by those ids (#6131). + let physical_type = self.input_physical_field.data_type(); + let value = match value { + ColumnarValue::Array(array) + if self + .parquet_options + .as_ref() + .is_some_and(|options| options.use_field_id) + && array.data_type() != physical_type + && array.data_type().equals_datatype(physical_type) => + { + ColumnarValue::Array(relabel_array(array, physical_type)) + } + other => other, + }; + if self.target_field.has_valid_extension_type::() { return match value { ColumnarValue::Array(array) => Ok(ColumnarValue::Array(normalize_variant_array( 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 cadcf4c9f4e..446d29bdd0a 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -45,7 +45,7 @@ //! //! Filed upstream as apache/datafusion#23978. Once the opener merges its deferred page-index //! load back into `FileMetadataCache` instead of bypassing it, the eager policy can go, but the -//! factory cannot: `get_metadata` is the one per-file hook that sees the raw footer, and two +//! factory cannot: `get_metadata` is the one per-file hook that sees the raw footer, and three //! other things hang off it. //! //! The first is Spark's missing field id check. `ParquetReadSupport` refuses to open a file @@ -57,9 +57,15 @@ //! the INT96 coercion also rebuilt container fields without their metadata, so a struct id could //! vanish on the way to Arrow as well. //! -//! The second is the Variant footer rewrite, `with_spark_arrow_schema`, which replaces the -//! Arrow schema hint in the footer for scans that project Variant. +//! The second is the Variant footer rewrite, `spark_arrow_schema_key_values`, which replaces +//! the Arrow schema hint in the footer for scans that project Variant. +//! +//! The third is the field id stamp for scans that match by field id. Until the INT96 coercion +//! keeps container metadata, `field_id_stamp` records the raw ids in the returned footer so +//! the schema adapter can put them back (#6131). Both rewrites change only the metadata this +//! reader returns, never the copy in the shared cache. +use crate::parquet::field_id_stamp::{with_field_ids_stamp, StampChange, FIELD_IDS_KEY}; use arrow::datatypes::{DataType, FieldRef, Schema}; use async_trait::async_trait; use bytes::Bytes; @@ -181,6 +187,8 @@ pub struct EagerPageIndexReaderFactory { // Refuse a file whose Parquet schema carries no field id, as Spark's `ParquetReadSupport` // does when the requested schema carries one and `ignoreMissing` is not set. require_field_ids: bool, + // Stamp the raw field ids into the returned footer for a scan that matches by field id. + use_field_id: bool, } impl EagerPageIndexReaderFactory { @@ -210,6 +218,7 @@ impl EagerPageIndexReaderFactory { scan_io_metrics, spark_variant_schema: false, require_field_ids: false, + use_field_id: false, } } @@ -224,11 +233,19 @@ impl EagerPageIndexReaderFactory { self.require_field_ids = enabled; self } + + /// For a scan that matches by field id. When the INT96 coercion would drop the ids of a + /// file's struct, list and map fields, each reader stamps the raw ids into the footer it + /// returns. Off by default. + pub fn with_use_field_id(mut self, enabled: bool) -> Self { + self.use_field_id = enabled; + self + } } /// True when `node` or any node under it carries a field id, the way Spark's /// `containsFieldIds` answers it over the raw Parquet schema, message root included. -fn contains_field_ids(node: &ParquetType) -> bool { +pub(crate) fn contains_field_ids(node: &ParquetType) -> bool { node.get_basic_info().has_id() || (node.is_group() && node @@ -263,6 +280,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { metadata_size_hint, spark_variant_schema: self.spark_variant_schema, require_field_ids: self.require_field_ids, + use_field_id: self.use_field_id, })) } } @@ -279,6 +297,7 @@ struct EagerPageIndexReader { metadata_size_hint: Option, spark_variant_schema: bool, require_field_ids: bool, + use_field_id: bool, } // Arrow infers ENUM as Binary, losing the distinction from raw binary that Spark needs. @@ -366,10 +385,9 @@ fn spark_enum_schema(schema: &SchemaDescriptor) -> ParquetResult> /// 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. +/// Returns the new footer key-value list, or `None` when the footer needs no change. /// https://github.com/apache/datafusion-comet/issues/5477 -fn with_spark_arrow_schema(metadata: Arc) -> ParquetResult> { - let file = metadata.file_metadata(); +fn spark_arrow_schema_key_values(file: &FileMetaData) -> ParquetResult>> { let has_arrow_schema = file.key_value_metadata().is_some_and(|key_values| { key_values .iter() @@ -377,7 +395,7 @@ fn with_spark_arrow_schema(metadata: Arc) -> ParquetResult) -> ParquetResult) -> Arc { + let file = metadata.file_metadata(); let file = FileMetaData::new( file.version(), file.num_rows(), @@ -402,13 +426,13 @@ fn with_spark_arrow_schema(metadata: Arc) -> ParquetResult format!( + "its footer carries the key {FIELD_IDS_KEY}, which Comet reserves and \ + removes by rewriting the footer" + ), + StampChange::Added | StampChange::Unchanged => "a struct, list or map with \ + a field id sits in a file with an INT96 timestamp, and Comet restores \ + those ids by rewriting the footer" + .to_string(), + }; + return Err(ParquetError::General(format!( + "Cannot match Parquet field ids in {} while Parquet decryption properties \ + are configured: {reason}, which it cannot do for a decrypted read. Set \ + spark.comet.scan.enabled=false to read it with Spark", + object_meta.location + ))); } + Ok(with_key_values(&metadata, key_values)) } .boxed() } @@ -1183,6 +1239,200 @@ mod tests { assert!(metadata_for(false, without_ids).await.is_ok()); } + /// A one-row file whose struct `s` carries field id 1 and holds an INT96 leaf. + const INT96_STRUCT: &str = "message schema { optional group s = 1 { optional int96 ts; } }"; + + /// A one-row file of `message`, whose single leaf must be INT96, encrypted with `key` when + /// one is given and carrying `key_values` in its footer. + fn int96_file(message: &str, key: Option<&[u8]>, key_values: Option>) -> Vec { + use parquet::data_type::{Int96, Int96Type}; + use parquet::encryption::encrypt::FileEncryptionProperties; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + + let schema = Arc::new(parse_message_type(message).unwrap()); + let max_def_level = SchemaDescriptor::new(Arc::clone(&schema)) + .column(0) + .max_def_level(); + let mut properties = WriterProperties::builder().set_key_value_metadata(key_values); + if let Some(key) = key { + properties = properties.with_file_encryption_properties( + FileEncryptionProperties::builder(key.to_vec()) + .build() + .unwrap(), + ); + } + let properties = properties.build(); + let mut bytes = Vec::new(); + let mut writer = + SerializedFileWriter::new(&mut bytes, schema, Arc::new(properties)).unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut column = row_group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&[Int96::new()], Some(&[max_def_level]), None) + .unwrap(); + column.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); + bytes + } + + async fn put_file(store: &InMemory, location: &str, bytes: Vec) -> PartitionedFile { + let size = bytes.len() as u64; + store + .put(&Path::from(location), Bytes::from(bytes).into()) + .await + .unwrap(); + PartitionedFile::new(location.to_string(), size) + } + + fn field_ids_stamp_of(metadata: &ParquetMetaData) -> Option { + metadata + .file_metadata() + .key_value_metadata()? + .iter() + .find(|kv| kv.key == crate::parquet::field_id_stamp::FIELD_IDS_KEY) + .and_then(|kv| kv.value.clone()) + } + + /// A scan that matches by field id gets the stamp in the footer it is handed, other scans + /// do not, and the shared cache keeps the footer as written. + #[tokio::test] + async fn get_metadata_stamps_field_ids_only_when_matching_by_id() { + use datafusion::datasource::physical_plan::parquet::metadata::CachedParquetMetaData; + + let store = Arc::new(InMemory::new()); + let file = put_file( + &store, + "int96.parquet", + int96_file(INT96_STRUCT, None, None), + ) + .await; + let runtime = datafusion::execution::runtime_env::RuntimeEnv::default(); + let metadata_cache = runtime.cache_manager.get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let metadata_for = |use_field_id: bool| { + let factory = EagerPageIndexReaderFactory::new( + Arc::clone(&store) as Arc, + Arc::clone(&metadata_cache), + ScanIoSource::Local, + &metrics, + ) + .with_use_field_id(use_field_id); + let mut reader = factory + .create_reader(0, file.clone(), None, &metrics) + .unwrap(); + async move { reader.get_metadata(None).await.unwrap() } + }; + + assert_eq!( + field_ids_stamp_of(metadata_for(true).await.as_ref()).as_deref(), + Some(r#"[["s",1],["ts",null]]"#) + ); + assert_eq!(field_ids_stamp_of(metadata_for(false).await.as_ref()), None); + let cached = metadata_cache.get(&Path::from("int96.parquet")).unwrap(); + let cached = cached + .file_metadata + .as_any() + .downcast_ref::() + .unwrap() + .parquet_metadata(); + assert_eq!(field_ids_stamp_of(cached), None); + } + + /// With decryption properties configured, a rebuilt footer would lose the decryptor, so a + /// scan that would need the stamp refuses the file by name instead of reading the struct + /// as null. Other scans read it. + #[tokio::test] + async fn get_metadata_refuses_to_stamp_an_encrypted_file() { + let key = b"0123456789012345"; + let store = Arc::new(InMemory::new()); + let file = put_file( + &store, + "encrypted.parquet", + int96_file(INT96_STRUCT, Some(key), None), + ) + .await; + assert_refused_under_decryption( + &store, + file, + key, + "Cannot match Parquet field ids in encrypted.parquet while Parquet decryption \ + properties are configured: a struct, list or map with a field id", + ) + .await; + } + + /// A file that needs no stamp but carries the reserved key in its footer needs a rebuild + /// too, to remove the key, and is refused for that reason. + #[tokio::test] + async fn get_metadata_refuses_to_remove_a_carried_stamp_under_decryption() { + let key = b"0123456789012345"; + let store = Arc::new(InMemory::new()); + let forged = KeyValue::new(FIELD_IDS_KEY.to_string(), "[]".to_string()); + let file = put_file( + &store, + "carried.parquet", + int96_file( + "message schema { optional int96 ts = 1; }", + Some(key), + Some(vec![forged]), + ), + ) + .await; + assert_refused_under_decryption( + &store, + file, + key, + "Cannot match Parquet field ids in carried.parquet while Parquet decryption \ + properties are configured: its footer carries the key comet.parquet.field_ids", + ) + .await; + } + + /// With decryption properties set, a reader that matches by field id refuses `file` with an + /// error containing `expected`, and one that does not reads it. + async fn assert_refused_under_decryption( + store: &Arc, + file: PartitionedFile, + key: &[u8], + expected: &str, + ) { + use parquet::encryption::decrypt::FileDecryptionProperties; + + let runtime = datafusion::execution::runtime_env::RuntimeEnv::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let options = ArrowReaderOptions::new().with_file_decryption_properties( + FileDecryptionProperties::builder(key.to_vec()) + .build() + .unwrap(), + ); + let metadata_for = |use_field_id: bool| { + let factory = EagerPageIndexReaderFactory::new( + Arc::clone(store) as Arc, + runtime.cache_manager.get_file_metadata_cache(), + ScanIoSource::Local, + &metrics, + ) + .with_use_field_id(use_field_id); + let mut reader = factory + .create_reader(0, file.clone(), None, &metrics) + .unwrap(); + let options = options.clone(); + async move { reader.get_metadata(Some(&options)).await } + }; + + let err = metadata_for(true) + .await + .expect_err("the footer rebuild must be refused"); + assert!( + err.to_string().contains(expected), + "unexpected error: {err}" + ); + assert!(metadata_for(false).await.is_ok()); + } + #[test] fn variant_policy_preserves_footer_metadata_and_indexes() { let schema = Arc::new(Schema::new_with_metadata( @@ -1212,7 +1462,10 @@ mod tests { ) .unwrap(); let original = Arc::new(reader.metadata().clone()); - let rewritten = with_spark_arrow_schema(Arc::clone(&original)).unwrap(); + let key_values = spark_arrow_schema_key_values(original.file_metadata()) + .unwrap() + .unwrap(); + let rewritten = with_key_values(&original, key_values); assert!(original.column_index().is_some()); assert!(original.offset_index().is_some()); assert_eq!(rewritten.column_index(), original.column_index()); @@ -1232,9 +1485,9 @@ mod tests { 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() - )); + assert_eq!( + spark_arrow_schema_key_values(rewritten.file_metadata()).unwrap(), + None + ); } } diff --git a/native/core/src/parquet/field_id_stamp.rs b/native/core/src/parquet/field_id_stamp.rs new file mode 100644 index 00000000000..1b4a3f51bbc --- /dev/null +++ b/native/core/src/parquet/field_id_stamp.rs @@ -0,0 +1,484 @@ +// 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. + +//! Carries the Parquet field ids of struct, list and map fields past DataFusion's INT96 +//! coercion. +//! +//! When a file holds an INT96 leaf, DataFusion 55.1's `Int96Coercer` rebuilds every struct, +//! list and map field of the file's Arrow schema without its metadata, so `PARQUET:field_id` is +//! gone from those fields in the schema `SparkPhysicalExprAdapterFactory::create` receives, and +//! from the batches the decoder produces. Spark's `clipParquetGroupFields` matches the id on the +//! raw Parquet group, so a container whose id sits on the container alone read as null +//! (#6131). +//! +//! The coercion keeps the schema-level metadata, which parquet-rs fills from the footer's +//! key-value list. So `EagerPageIndexReader::get_metadata` records every field's name and id +//! under [`FIELD_IDS_KEY`] in the footer it returns, and the schema adapter puts the ids back +//! with [`restore_field_ids`]. Delete this module once DataFusion carries +//! apache/datafusion#24790; `int96_coercion_drops_container_ids` below fails at that point. + +use crate::parquet::eager_page_index_reader_factory::contains_field_ids; +use crate::parquet::parquet_support::field_id; +use arrow::datatypes::{DataType, FieldRef, Fields, Schema, SchemaRef}; +use parquet::arrow::{parquet_to_arrow_schema, PARQUET_FIELD_ID_META_KEY}; +use parquet::basic::Type as PhysicalType; +use parquet::errors::{ParquetError, Result as ParquetResult}; +use parquet::file::metadata::{FileMetaData, KeyValue}; +use parquet::schema::types::SchemaDescriptor; +use std::sync::Arc; + +/// Footer key, and so Arrow schema metadata key, of the stamp. Its value is a JSON array with +/// one `[name, id or null]` pair per Arrow field, in pre-order. +pub(crate) const FIELD_IDS_KEY: &str = "comet.parquet.field_ids"; + +/// The field types whose id the INT96 coercion drops. It rebuilds these three and clones every +/// other field, including the other list types, unchanged. +fn loses_id(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Struct(_) | DataType::List(_) | DataType::Map(_, _) + ) +} + +/// The child fields of `data_type`, in the order the walks below visit them. Every list type +/// and Map is descended, so a walk over a schema converted with an `ARROW:schema` hint visits +/// the same fields as one over the hint-free conversion. Dictionary is a leaf. +fn children(data_type: &DataType) -> &[FieldRef] { + match data_type { + DataType::Struct(fields) => fields, + DataType::List(child) + | DataType::LargeList(child) + | DataType::FixedSizeList(child, _) + | DataType::ListView(child) + | DataType::LargeListView(child) + | DataType::Map(child, _) => std::slice::from_ref(child), + _ => &[], + } +} + +fn collect_ids<'a>(field: &'a FieldRef, nodes: &mut Vec<(&'a str, Option)>) { + nodes.push((field.name(), field_id(field))); + for child in children(field.data_type()) { + collect_ids(child, nodes); + } +} + +/// The stamp for a file, or `None` when the coercion loses no id: the file has no INT96 leaf, +/// or none of its struct, list and map fields carries an id. The ids come from the hint-free +/// Arrow conversion, which takes them from the Parquet schema, as Spark does. +fn field_ids_stamp(schema: &SchemaDescriptor) -> ParquetResult> { + let has_int96 = schema + .columns() + .iter() + .any(|column| column.physical_type() == PhysicalType::INT96); + // Without an id below the message root no Arrow field carries one; skip the conversion. + // Leaves count too: a repeated primitive without a LIST annotation becomes a list that + // carries the primitive's id. + if !has_int96 + || !schema + .root_schema() + .get_fields() + .iter() + .any(|field| contains_field_ids(field)) + { + return Ok(None); + } + let arrow_schema = parquet_to_arrow_schema(schema, None)?; + if !arrow_schema.fields().iter().any(has_container_id) { + return Ok(None); + } + let mut nodes = Vec::new(); + for field in arrow_schema.fields() { + collect_ids(field, &mut nodes); + } + serde_json::to_string(&nodes) + .map(Some) + .map_err(|e| ParquetError::External(Box::new(e))) +} + +fn has_container_id(field: &FieldRef) -> bool { + (loses_id(field.data_type()) && field_id(field).is_some()) + || children(field.data_type()).iter().any(has_container_id) +} + +/// How [`with_field_ids_stamp`] changed the footer key-value list. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum StampChange { + /// The list is as it was. + Unchanged, + /// The file needs the stamp, which was added. + Added, + /// The file needs no stamp but its footer carried a [`FIELD_IDS_KEY`] entry, which was + /// removed. + Removed, +} + +/// The footer key-value list with this file's stamp in place, or `key_values` unchanged when +/// the file needs none, and what changed. `key_values` is the list an earlier footer rewrite +/// produced, if any; `None` in and out means the footer's own list stands. +/// +/// A [`FIELD_IDS_KEY`] entry in the footer's own list is removed, so that list, which parquet-rs +/// turns into the Arrow schema metadata, holds only a stamp computed here. parquet-rs also +/// merges the custom metadata of an `ARROW:schema` hint into the schema metadata, and a key +/// there survives when the file needs no stamp (a footer entry takes precedence over it). It +/// can do no more than the hint itself, which can already put `PARQUET:field_id` on any field: +/// [`restore_field_ids`] only adds ids to fields that lack one, after checking the field count, +/// every name and every id the schema carries. +pub(crate) fn with_field_ids_stamp( + file: &FileMetaData, + key_values: Option>, +) -> ParquetResult<(Option>, StampChange)> { + let stamp = field_ids_stamp(file.schema_descr())?; + let current = key_values + .as_deref() + .or(file.key_value_metadata().map(Vec::as_slice)) + .unwrap_or_default(); + let change = match (&stamp, current.iter().any(|kv| kv.key == FIELD_IDS_KEY)) { + (Some(_), _) => StampChange::Added, + (None, true) => StampChange::Removed, + (None, false) => return Ok((key_values, StampChange::Unchanged)), + }; + let mut key_values = + key_values.unwrap_or_else(|| file.key_value_metadata().cloned().unwrap_or_default()); + key_values.retain(|kv| kv.key != FIELD_IDS_KEY); + key_values.extend(stamp.map(|value| KeyValue::new(FIELD_IDS_KEY.to_string(), value))); + Ok((Some(key_values), change)) +} + +/// Put back the ids the INT96 coercion dropped from `schema`'s struct, list and map fields, +/// using the stamp in its metadata. Only a missing `PARQUET:field_id` is added, and only where +/// the coercion rebuilds: on a struct, list or map field whose ancestors are all structs, lists +/// or maps. The walk checks the stamp against the schema, field count, every name, and every id +/// the schema still carries, and on any mismatch, or without a stamp, returns `schema` itself. +/// +/// The stamp is trusted as computed from the file's Parquet schema because +/// `EagerPageIndexReaderFactory` removes any footer entry of that name before adding its own +/// (see [`with_field_ids_stamp`]). Only `init_datasource_exec` turns field-id matching on today, +/// and it always reads through that factory. +pub(crate) fn restore_field_ids(schema: &SchemaRef) -> SchemaRef { + let Some(stamp) = schema.metadata().get(FIELD_IDS_KEY) else { + return Arc::clone(schema); + }; + let nodes = match serde_json::from_str::)>>(stamp) { + Ok(nodes) => nodes, + Err(e) => { + log::debug!("Parquet field id stamp is not valid ({e}); ids not restored"); + return Arc::clone(schema); + } + }; + let mut next = 0; + let fields = schema + .fields() + .iter() + .map(|field| restore_field(field, &nodes, &mut next, true)) + .collect::>>(); + match fields { + Some(fields) if next == nodes.len() => { + Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone())) + } + _ => { + log::debug!( + "Parquet field id stamp with {} entries does not match the file schema of {} \ + root fields (last entry checked: {next}, {:?}); ids not restored", + nodes.len(), + schema.fields().len(), + next.checked_sub(1).and_then(|i| nodes.get(i)), + ); + Arc::clone(schema) + } + } +} + +/// Restore `field` and its subtree against the stamp entries from `next` on, or `None` on a +/// mismatch. `restorable` is false below a field the coercion leaves alone. +fn restore_field( + field: &FieldRef, + nodes: &[(String, Option)], + next: &mut usize, + restorable: bool, +) -> Option { + let (name, stamped_id) = nodes.get(*next)?; + *next += 1; + let current_id = field_id(field); + if name != field.name() || current_id.is_some_and(|id| Some(id) != *stamped_id) { + return None; + } + let restorable = restorable && loses_id(field.data_type()); + let mut restore = |child: &FieldRef| restore_field(child, nodes, next, restorable); + let data_type = match field.data_type() { + DataType::Struct(fields) => DataType::Struct( + fields + .iter() + .map(&mut restore) + .collect::>()?, + ), + DataType::List(child) => DataType::List(restore(child)?), + DataType::LargeList(child) => DataType::LargeList(restore(child)?), + DataType::FixedSizeList(child, size) => DataType::FixedSizeList(restore(child)?, *size), + DataType::ListView(child) => DataType::ListView(restore(child)?), + DataType::LargeListView(child) => DataType::LargeListView(restore(child)?), + DataType::Map(child, sorted) => DataType::Map(restore(child)?, *sorted), + other => other.clone(), + }; + let mut restored = field.as_ref().clone().with_data_type(data_type); + if let (true, None, Some(id)) = (restorable, current_id, stamped_id) { + let mut metadata = field.metadata().clone(); + metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string()); + restored = restored.with_metadata(metadata); + } + Some(Arc::new(restored)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{Field, TimeUnit}; + use datafusion::datasource::physical_plan::parquet::Int96Coercer; + use parquet::schema::parser::parse_message_type; + use std::collections::HashMap; + + /// A struct, a list and a map with ids on the containers and some leaves, next to INT96 + /// leaves. The map's `key_value` group never carries an Arrow id. + const NESTED: &str = "message schema { + optional group s = 1 { optional int32 a = 2; optional int96 ts; } + optional group l (LIST) = 3 { repeated group list { optional int96 element = 4; } } + optional group m (MAP) = 5 { + repeated group key_value { required binary key (STRING); optional int32 value = 6; } + } + }"; + + /// The ids of `NESTED`'s Arrow fields in pre-order: s, a, ts, l, element, m, key_value, + /// key, value. + const NESTED_IDS: [Option; 9] = [ + Some(1), + Some(2), + None, + Some(3), + Some(4), + Some(5), + None, + None, + Some(6), + ]; + + fn descriptor(message: &str) -> SchemaDescriptor { + SchemaDescriptor::new(Arc::new(parse_message_type(message).unwrap())) + } + + fn ids_in_pre_order(schema: &Schema) -> Vec> { + fn walk(field: &FieldRef, ids: &mut Vec>) { + ids.push(field_id(field)); + for child in children(field.data_type()) { + walk(child, ids); + } + } + let mut ids = Vec::new(); + for field in schema.fields() { + walk(field, &mut ids); + } + ids + } + + /// The Arrow schema DataFusion hands the schema adapter for `NESTED`, with `stamp` in its + /// metadata as parquet-rs copies it there from the footer. + fn coerced(stamp: Option) -> SchemaRef { + let descriptor = descriptor(NESTED); + let metadata: HashMap = stamp + .into_iter() + .map(|value| (FIELD_IDS_KEY.to_string(), value)) + .collect(); + let file_schema = parquet_to_arrow_schema(&descriptor, None) + .unwrap() + .with_metadata(metadata); + Arc::new( + Int96Coercer::new(&descriptor, &file_schema, &TimeUnit::Microsecond) + .with_timezone(Some(Arc::from("UTC"))) + .coerce() + .unwrap(), + ) + } + + fn stamp_of(nodes: &[(&str, Option)]) -> String { + serde_json::to_string(nodes).unwrap() + } + + #[test] + fn stamp_lists_every_field_in_pre_order() { + let stamp = field_ids_stamp(&descriptor(NESTED)).unwrap().unwrap(); + assert_eq!( + stamp, + r#"[["s",1],["a",2],["ts",null],["l",3],["element",4],["m",5],["key_value",null],["key",null],["value",6]]"# + ); + } + + /// A repeated primitive without a LIST annotation becomes a list that carries the + /// primitive's id, with an item named after it, although no group carries an id. + #[test] + fn stamp_covers_a_legacy_repeated_primitive() { + let message = "message schema { repeated int32 x = 5; optional int96 ts; }"; + assert_eq!( + field_ids_stamp(&descriptor(message)).unwrap().as_deref(), + Some(r#"[["x",5],["x",null],["ts",null]]"#) + ); + } + + #[test] + fn stamp_is_needed_only_for_container_ids_next_to_int96() { + let leaf_ids_only = "message schema { + optional group s { optional int32 a = 2; optional int96 ts = 3; } + }"; + let without_int96 = "message schema { + optional group s = 1 { optional int32 a; optional int64 ts; } + }"; + let list_group_id_only = "message schema { + optional group l (LIST) { repeated group list = 5 { optional int96 element; } } + }"; + for message in [leaf_ids_only, without_int96, list_group_id_only] { + assert_eq!(field_ids_stamp(&descriptor(message)).unwrap(), None); + } + } + + /// Pins DataFusion's behavior this module works around. When a DataFusion upgrade carries + /// apache/datafusion#24790, this fails: delete this module and its call sites in + /// `eager_page_index_reader_factory.rs`, `schema_adapter.rs` and `cast_column.rs`. + #[test] + fn int96_coercion_drops_container_ids() { + assert_eq!( + ids_in_pre_order(&coerced(None)), + vec![ + None, + Some(2), + None, + None, + Some(4), + None, + None, + None, + Some(6) + ] + ); + } + + #[test] + fn restore_puts_back_the_container_ids() { + let stamp = field_ids_stamp(&descriptor(NESTED)).unwrap(); + let restored = restore_field_ids(&coerced(stamp)); + assert_eq!(ids_in_pre_order(&restored), NESTED_IDS.to_vec()); + // Nothing but the ids changes. + let shape = |schema: &Schema| DataType::Struct(schema.fields().clone()); + assert!(shape(&restored).equals_datatype(&shape(&coerced(None)))); + } + + #[test] + fn restore_returns_the_schema_on_a_mismatch() { + let names = [ + "s", + "a", + "ts", + "l", + "element", + "m", + "key_value", + "key", + "value", + ]; + let nodes: Vec<(&str, Option)> = names.into_iter().zip(NESTED_IDS).collect(); + let mut renamed = nodes.clone(); + renamed[3].0 = "other"; + let mut leaf_disagrees = nodes.clone(); + leaf_disagrees[1].1 = Some(9); + let mut leaf_id_missing = nodes.clone(); + leaf_id_missing[1].1 = None; + let stamps = [ + stamp_of(&nodes[..8]), + stamp_of(&[nodes.as_slice(), &[("extra", None)]].concat()), + stamp_of(&renamed), + stamp_of(&leaf_disagrees), + stamp_of(&leaf_id_missing), + "not json".to_string(), + ]; + for stamp in stamps { + let schema = coerced(Some(stamp.clone())); + assert!( + Arc::ptr_eq(&restore_field_ids(&schema), &schema), + "stamp {stamp}" + ); + } + let schema = coerced(None); + assert!(Arc::ptr_eq(&restore_field_ids(&schema), &schema)); + } + + /// The coercion leaves a large list and everything below it alone, so a field there that + /// lacks an id lacks it for another reason and stays as it is. + #[test] + fn restore_skips_fields_the_coercion_leaves_alone() { + let item = Field::new("item", DataType::Struct(Fields::empty()), true); + let schema = Arc::new(Schema::new_with_metadata( + vec![Field::new("l", DataType::LargeList(Arc::new(item)), true)], + HashMap::from([( + FIELD_IDS_KEY.to_string(), + stamp_of(&[("l", Some(1)), ("item", Some(2))]), + )]), + )); + let restored = restore_field_ids(&schema); + assert_eq!(ids_in_pre_order(&restored), vec![None, None]); + } + + #[test] + fn stamp_replaces_one_the_file_carries() { + let forged = KeyValue::new(FIELD_IDS_KEY.to_string(), "[]".to_string()); + let other = KeyValue::new("other".to_string(), "kept".to_string()); + let file = |message: &str| { + FileMetaData::new( + 1, + 0, + None, + Some(vec![forged.clone(), other.clone()]), + Arc::new(descriptor(message)), + None, + ) + }; + + let (needs_stamp, change) = with_field_ids_stamp(&file(NESTED), None).unwrap(); + let needs_stamp = needs_stamp.unwrap(); + assert_eq!(change, StampChange::Added); + assert_eq!(needs_stamp.len(), 2); + assert_eq!(needs_stamp[0], other); + assert_eq!(needs_stamp[1].key, FIELD_IDS_KEY); + assert_eq!( + needs_stamp[1].value, + field_ids_stamp(&descriptor(NESTED)).unwrap() + ); + + let no_stamp = "message schema { optional int32 a = 1; }"; + assert_eq!( + with_field_ids_stamp(&file(no_stamp), None).unwrap(), + (Some(vec![other.clone()]), StampChange::Removed) + ); + let plain = FileMetaData::new(1, 0, None, None, Arc::new(descriptor(no_stamp)), None); + assert_eq!( + with_field_ids_stamp(&plain, None).unwrap(), + (None, StampChange::Unchanged) + ); + // A list an earlier rewrite produced passes through, with the stamp added when needed. + let earlier = vec![other.clone()]; + assert_eq!( + with_field_ids_stamp(&plain, Some(earlier.clone())).unwrap(), + (Some(earlier), StampChange::Unchanged) + ); + } +} diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index 7930320d148..9fd7f0fe6a8 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -24,5 +24,6 @@ pub mod schema_adapter; pub mod util; mod cast_column; +mod field_id_stamp; mod name_fold; pub(crate) mod objectstore; diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 5c5d26d3922..6c10fabb021 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -196,7 +196,8 @@ pub(crate) fn init_datasource_exec( parquet_source.metrics(), ) .with_spark_variant_schema(projects_variant) - .with_require_field_ids(require_field_ids), + .with_require_field_ids(require_field_ids) + .with_use_field_id(use_field_id), ); parquet_source = parquet_source.with_parquet_file_reader_factory(reader_factory); @@ -1338,6 +1339,192 @@ mod tests { assert!(scan_metric(&scan, "scan_io_metadata_bytes") >= bloom_bytes); } + /// Write `message` with two rows. Every leaf must be optional or repeated and holds one + /// value in each row, at its maximum definition level: `int32` leaves hold 1 and 2, `int96` + /// leaves the Unix epoch. + fn write_int96_file(message: &str) -> String { + use parquet::basic::Type as PhysicalType; + use parquet::data_type::{Int32Type, Int96, Int96Type}; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + use parquet::schema::types::SchemaDescriptor; + + let schema = Arc::new(parse_message_type(message).unwrap()); + let descriptor = SchemaDescriptor::new(Arc::clone(&schema)); + let filename = get_temp_filename().to_str().unwrap().to_string(); + let mut writer = + SerializedFileWriter::new(File::create(&filename).unwrap(), schema, Default::default()) + .unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + for column in descriptor.columns() { + let levels = [column.max_def_level(); 2]; + let repetition = [0; 2]; + let repetition = (column.max_rep_level() > 0).then_some(&repetition[..]); + let mut column_writer = row_group.next_column().unwrap().unwrap(); + match column.physical_type() { + PhysicalType::INT32 => column_writer.typed::().write_batch( + &[1, 2], + Some(&levels), + repetition, + ), + PhysicalType::INT96 => { + let mut epoch = Int96::new(); + epoch.set_data(0, 0, 2_440_588); + column_writer.typed::().write_batch( + &[epoch, epoch], + Some(&levels), + repetition, + ) + } + other => panic!("unexpected leaf type {other}"), + } + .unwrap(); + column_writer.close().unwrap(); + } + row_group.close().unwrap(); + writer.close().unwrap(); + filename + } + + /// Scan `filename` with Parquet field id matching on, as Spark's + /// `spark.sql.parquet.fieldId.read.enabled` asks, and return the single output batch. + async fn scan_with_field_ids(filename: String, required_schema: SchemaRef) -> RecordBatch { + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + required_schema, + None, + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![PartitionedFile::from_path(filename).unwrap()]], + None, + None, + None, + "UTC", + false, + false, + false, + false, + &session_ctx, + false, + true, + true, + ) + .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()); + batch + } + + fn with_field_id(field: Field, id: i32) -> Field { + field.with_metadata(HashMap::from([( + parquet::arrow::PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])) + } + + fn int96_timestamp_field() -> Field { + Field::new( + "ts", + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())), + true, + ) + } + + fn int_values(array: &dyn arrow::array::Array) -> Vec> { + use arrow::array::AsArray; + array + .as_primitive::() + .iter() + .collect() + } + + /// DataFusion's INT96 coercion rebuilds a struct holding an INT96 leaf without its field + /// id. The id sits only on the root struct here, so the root match by id has to see it. + #[tokio::test] + async fn field_id_on_a_root_struct_holding_int96_resolves() { + use arrow::array::AsArray; + + let filename = write_int96_file( + "message schema { optional group s = 1 { optional int32 a; optional int96 ts; } }", + ); + let children = vec![ + Field::new("a", DataType::Int32, true), + int96_timestamp_field(), + ]; + let required = Arc::new(Schema::new(vec![with_field_id( + Field::new("s", DataType::Struct(children.into()), true), + 1, + )])); + + let batch = scan_with_field_ids(filename, required).await; + let s = batch.column(0).as_struct(); + assert_eq!(int_values(s.column(0)), vec![Some(1), Some(2)]); + assert_eq!( + s.column(1) + .as_primitive::() + .iter() + .collect::>(), + vec![Some(0), Some(0)] + ); + } + + /// The same coercion drops the id of a struct nested in another one. The requested schema + /// renames the nested struct, so only its id can match it, and that match runs on each + /// decoded batch. + #[tokio::test] + async fn field_id_on_a_nested_struct_holding_int96_resolves() { + use arrow::array::AsArray; + + let filename = write_int96_file( + "message schema { optional group s = 1 { optional group inner = 2 { \ + optional int32 x = 3; optional int96 ts; } } }", + ); + let inner = vec![ + with_field_id(Field::new("x", DataType::Int32, true), 3), + int96_timestamp_field(), + ]; + let required = Arc::new(Schema::new(vec![with_field_id( + Field::new( + "s", + DataType::Struct( + vec![with_field_id( + Field::new("renamed_inner", DataType::Struct(inner.into()), true), + 2, + )] + .into(), + ), + true, + ), + 1, + )])); + + let batch = scan_with_field_ids(filename, required).await; + let inner = batch.column(0).as_struct().column(0).as_struct(); + assert_eq!(int_values(inner.column(0)), vec![Some(1), Some(2)]); + } + + /// A legacy repeated primitive without a LIST annotation reads as an Arrow list that + /// carries the primitive's id, and the coercion drops that id like any other list's. + #[tokio::test] + async fn field_id_on_a_legacy_repeated_primitive_next_to_int96_resolves() { + use arrow::array::AsArray; + + let filename = + write_int96_file("message schema { repeated int32 x = 5; optional int96 ts; }"); + let item = Arc::new(Field::new("element", DataType::Int32, true)); + let required = Arc::new(Schema::new(vec![with_field_id( + Field::new("x", DataType::List(item), true), + 5, + )])); + + let batch = scan_with_field_ids(filename, required).await; + let x = batch.column(0).as_list::(); + assert_eq!(int_values(x.values().as_ref()), vec![Some(1), Some(2)]); + assert_eq!(x.value_offsets(), &[0, 1, 2]); + } + #[tokio::test] async fn reports_projected_and_predicate_pruned_data_io() { let (filename, schema) = write_scan_io_fixture(); diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index c7ad0418578..9e03d9d8a89 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -16,6 +16,7 @@ // under the License. use crate::parquet::cast_column::CometCastColumnExpr; +use crate::parquet::field_id_stamp::restore_field_ids; use crate::parquet::name_fold::{fold_name, fold_names, fold_schema_names}; use crate::parquet::parquet_support::{ duplicate_parquet_field_error, field_id, field_names_with_id, match_struct_fields, @@ -862,6 +863,24 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { logical_file_schema: SchemaRef, physical_file_schema: SchemaRef, ) -> DataFusionResult> { + // DataFusion's INT96 coercion drops the field ids of struct, list and map fields. Put + // them back first, so every id match below sees the file's ids (#6131). A root column + // whose type changed decodes without the nested ids, so `rewrite` relabels it. + let (physical_file_schema, restored_columns) = if self.parquet_options.use_field_id { + let restored = restore_field_ids(&physical_file_schema); + let changed = restored + .fields() + .iter() + .zip(physical_file_schema.fields()) + .enumerate() + .filter(|(_, (after, before))| after.data_type() != before.data_type()) + .map(|(index, _)| index) + .collect(); + (restored, changed) + } else { + (physical_file_schema, HashSet::new()) + }; + // Remap physical schema field names to match logical names by Parquet field id // (when the logical schema carries IDs and `use_field_id` is set) and/or by // case-insensitive name match. The DefaultPhysicalExprAdapter uses exact name @@ -991,6 +1010,7 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { id_duplicate_roots, logical_folded, physical_folded, + restored_columns, })) } } @@ -1049,6 +1069,9 @@ struct SparkPhysicalExprAdapter { /// `physical_file_schema` field names pre-folded once, parallel to /// `physical_file_schema.fields()`. See `logical_folded`. physical_folded: Vec, + /// Indices of the physical columns whose type gained nested field ids from + /// `restore_field_ids`. The decoder still produces them without those ids. + restored_columns: HashSet, } impl PhysicalExprAdapter for SparkPhysicalExprAdapter { @@ -1117,6 +1140,7 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter { } }; let expr = self.wrap_direct_variant_column(expr)?; + let expr = self.wrap_restored_column(expr)?; // For case-insensitive mode: remap column names from logical back to // original physical names. The default adapter was given a remapped @@ -1177,6 +1201,34 @@ impl SparkPhysicalExprAdapter { checked_decoded_expr(physical_field.data_type(), cast) } + /// A bare column whose nested field ids were restored decodes without them, and the batch + /// the scan builds must match its schema exactly. `CometCastColumnExpr` relabels the + /// decoded array to the restored type. + fn wrap_restored_column( + &self, + expr: Arc, + ) -> DataFusionResult> { + let Some(column) = expr.downcast_ref::() else { + return Ok(expr); + }; + if !self.restored_columns.contains(&column.index()) { + return Ok(expr); + } + let Ok(logical_field) = self.logical_file_schema.field_with_name(column.name()) else { + return Ok(expr); + }; + let physical_field = Arc::clone(&self.physical_file_schema.fields()[column.index()]); + Ok(Arc::new( + CometCastColumnExpr::try_new( + expr, + 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 or Map). Uses `spark_parquet_convert` diff --git a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 9746d6553b7..cfc278ebac2 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala @@ -2425,6 +2425,117 @@ abstract class ParquetReadSuite extends CometTestBase { } } + // With the read flag on, the ids must also match. DataFusion's INT96 coercion rebuilds every + // struct, list and map of a file holding an INT96 timestamp without its field id, while Spark + // matches the id on the raw Parquet group, so a container whose id went missing would read as + // null. Each root column is read on its own, so a wrong answer names the column. The year 2500 + // timestamps sit outside the nanosecond range, which the INT96 decode has to keep. The last + // case writes TIMESTAMP_MICROS, which the coercion leaves alone. Before Spark 4.1 the + // vectorized reader rejects a nested struct renamed in the read schema, since its column + // vector compares the names, so that read is compared with Spark from 4.1 on. The pinned rows + // hold everywhere. + test( + "field ids on struct, list and map containers holding INT96 timestamps resolve with the " + + "read flag on") { + val ts = Timestamp.valueOf("2020-01-01 00:00:00") + val farTs = Timestamp.valueOf("2500-01-01 00:00:00") + def id(fieldId: Option[Int]): Metadata = fieldId.map(withId).getOrElse(Metadata.empty) + def leaves(aId: Option[Int] = None, tsId: Option[Int] = None): StructType = + new StructType() + .add("a", IntegerType, true, id(aId)) + .add("ts", TimestampType, true, id(tsId)) + def deep(inner: String, ids: Seq[Int]): StructType = new StructType().add( + inner, + new StructType().add("x", IntegerType, true, withId(ids(1))).add("ts", TimestampType), + true, + withId(ids(0))) + + val containerOnly = new StructType() + .add("s", leaves(), true, withId(1)) + .add("l", ArrayType(TimestampType), true, withId(2)) + .add("m", MapType(StringType, TimestampType), true, withId(3)) + val containerOnlyRows = Seq( + Row(Row(1, ts), Seq(ts, farTs), Map("k" -> ts)), + Row(Row(2, farTs), Seq(farTs), Map("f" -> farTs))) + + def withChildIds(base: Int): StructType = leaves(Some(base + 1), Some(base + 2)) + val containerAndChildren = new StructType() + .add("s", withChildIds(10), true, withId(1)) + .add("l", ArrayType(withChildIds(20)), true, withId(2)) + .add("m", MapType(StringType, withChildIds(30)), true, withId(3)) + val structRows = Seq( + Row(Row(1, ts), Seq(Row(2, ts), Row(3, farTs)), Map("k" -> Row(4, ts))), + Row(Row(5, farTs), Seq(Row(6, farTs)), Map("f" -> Row(7, farTs)))) + + def severalDepths(root: String, inner: String): StructType = new StructType() + .add(root, deep(inner, Seq(11, 12)), true, withId(1)) + .add("l", ArrayType(deep(inner, Seq(21, 22))), true, withId(2)) + .add("m", MapType(StringType, deep(inner, Seq(31, 32))), true, withId(3)) + val severalDepthsRows = Seq( + Row(Row(Row(1, ts)), Seq(Row(Row(2, farTs))), Map("k" -> Row(Row(3, ts)))), + Row(Row(Row(4, farTs)), Seq(Row(Row(5, ts))), Map("f" -> Row(Row(6, farTs))))) + + val farOnly = new StructType().add("s", leaves(), true, withId(1)) + val farOnlyRows = + Seq(Row(Row(1, farTs)), Row(Row(2, Timestamp.valueOf("9999-12-31 23:59:59")))) + + // (label, output timestamp type, write schema, read schema, rows) + val cases = Seq( + ("id only on the container", "INT96", containerOnly, containerOnly, containerOnlyRows), + ( + "ids on the container and its children", + "INT96", + containerAndChildren, + containerAndChildren, + structRows), + ( + "ids at several depths", + "INT96", + severalDepths("s", "inner"), + severalDepths("s", "inner"), + severalDepthsRows), + ( + "containers renamed in the read schema", + "INT96", + severalDepths("s", "inner"), + severalDepths("renamed", "renamed_inner"), + severalDepthsRows), + ("a timestamp outside the nanosecond range", "INT96", farOnly, farOnly, farOnlyRows), + ( + "TIMESTAMP_MICROS control", + "TIMESTAMP_MICROS", + containerOnly, + containerOnly, + containerOnlyRows)) + + for ((label, timestampType, writeSchema, readSchema, rows) <- cases) { + withSQLConf( + SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true", + SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> timestampType) { + withTempPath { dir => + spark + .createDataFrame(spark.sparkContext.parallelize(rows), writeSchema) + .write + .mode("overwrite") + .parquet(dir.getCanonicalPath) + readSchema.fieldNames.zipWithIndex.foreach { case (column, index) => + withClue(s"$label, column $column: ") { + def read(): DataFrame = + spark.read.schema(readSchema).parquet(dir.getCanonicalPath).select(column) + if (readSchema == writeSchema || isSpark41Plus) { + checkSparkAnswerAndOperator(read()) + } else { + checkCometOperators(stripAQEPlan(read().queryExecution.executedPlan)) + } + checkAnswer(read(), rows.map(row => Row(row.get(index)))) + } + } + } + } + } + } + // Spark checks each file on its own. A directory holding one file with ids and one without // raises on the second, and with `ignoreMissing` the file without ids reads as nulls because // no root field of it carries the requested id. Each side is written as one file. Spread over From 908eccbf2ab6b724f436a11dc0a2d59a8e0e935f Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 29 Sep 2026 22:21:45 +0700 Subject: [PATCH 2/2] fix: fall back to Spark for encrypted scans that match Parquet field ids Restoring the container field ids that the INT96 coercion drops means rewriting the footer the reader returns, which cannot be done when Parquet decryption properties are configured. Rather than failing such a scan in the native reader, CometScanRule now sends it to Spark at planning time when encryption is configured and the scan matches by field id, as it already does for native Variant scans. The planner cannot tell which files have INT96 or container ids, so every such scan falls back. The native refusal stays as a backstop. --- .../user-guide/latest/compatibility/scans.md | 10 ++--- .../apache/comet/rules/CometScanRule.scala | 13 +++++++ .../comet/rules/CometScanRuleSuite.scala | 39 +++++++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/scans.md b/docs/source/user-guide/latest/compatibility/scans.md index 1c31672f29d..b98604a8f8e 100644 --- a/docs/source/user-guide/latest/compatibility/scans.md +++ b/docs/source/user-guide/latest/compatibility/scans.md @@ -50,6 +50,11 @@ The following features are not supported and cause Comet to fall back to Spark: Comet Parquet scan regardless. - A read schema that repeats a Parquet field id, at the top level or within a struct, when `spark.sql.parquet.fieldId.read.enabled=true`. +- A read schema that carries Parquet field ids when `spark.sql.parquet.fieldId.read.enabled=true` + and Parquet encryption is configured. When a file has an `INT96` timestamp, Comet restores the + field ids of its structs, lists and maps through the file footer, which it cannot rewrite for a + decrypted read. Which files need this is not known when the query is planned, so every such + scan falls back ([#6131](https://github.com/apache/datafusion-comet/issues/6131)). The following limitation may produce incorrect results without falling back to Spark: @@ -78,11 +83,6 @@ The following limitations raise an error at scan time rather than falling back t duplicate names; inference across files can depend on merge order. Resolution is tracked in [#5884](https://github.com/apache/datafusion-comet/issues/5884), with mixed-type behavior in [#5964](https://github.com/apache/datafusion-comet/issues/5964). -- Reading an encrypted Parquet file with `spark.sql.parquet.fieldId.read.enabled=true` when a - struct, list or map carrying a field id sits in a file with an `INT96` timestamp. Comet - restores those field ids through the file footer, which it cannot rewrite while Parquet - decryption is configured ([#6131](https://github.com/apache/datafusion-comet/issues/6131)). - Set `spark.comet.scan.enabled=false` to read such a file with Spark. - Invalid UTF-8 bytes in `STRING` columns. Spark permits arbitrary byte sequences in a `STRING` column (for example from `CAST(X'C1' AS STRING)`), but Comet's native execution path is built on Arrow, whose string type is strictly UTF-8. Reading a Parquet file whose `STRING` column contains diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index 1b8c4cd15be..8e059156d59 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.getExistenceDefa import org.apache.spark.sql.comet.{CometBatchScanExec, CometScanExec} import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec} import org.apache.spark.sql.execution.datasources.HadoopFsRelation +import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.datasources.v2.csv.CSVScan import org.apache.spark.sql.internal.SQLConf @@ -347,6 +348,18 @@ case class CometScanRule(session: SparkSession) withFallbackReason(scanExec, "Native Parquet Variant scans do not support encryption") return None } + // The native reader restores container field ids that DataFusion's INT96 coercion drops by + // rewriting the footer, which it cannot do for a decrypted read. Whether a file needs that + // is unknown here, so fall back whenever the scan would match by field id (the same gate + // CometNativeScan uses to set useFieldId). + // https://github.com/apache/datafusion-comet/issues/6131 + if (encryptionEnabled(hadoopConf) && readFieldId(conf) && + ParquetUtils.hasFieldIds(scanExec.requiredSchema)) { + withFallbackReason( + scanExec, + "Native Parquet scans that match field ids do not support encryption") + return None + } // input_file_name, input_file_block_start, and input_file_block_length read from // InputFileBlockHolder, a thread-local set by Spark's FileScanRDD. The native DataFusion // scan does not use FileScanRDD, so these expressions would return empty/default values. diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala index 6bb4d6254a2..94c6dd24f14 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanRuleSuite.scala @@ -239,4 +239,43 @@ class CometScanRuleSuite extends CometTestBase { } } + test("encrypted scans that match Parquet field ids fall back to Spark") { + // The native reader restores container field ids that the INT96 coercion drops by rewriting + // the footer, which it cannot do for a decrypted read. The planner cannot see which files + // need that, so any encrypted read that matches by id goes to Spark. The file below holds a + // container id next to an INT96 timestamp, which the native reader would refuse. See + // https://github.com/apache/datafusion-comet/issues/6131. + val idOne = new MetadataBuilder().putLong("parquet.field.id", 1L).build() + val withIds = new StructType() + .add("s", new StructType().add("a", IntegerType).add("ts", TimestampType), true, idOne) + val withoutIds = new StructType() + .add("s", new StructType().add("a", IntegerType).add("ts", TimestampType)) + val rows = Seq(Row(Row(1, java.sql.Timestamp.valueOf("2020-01-01 00:00:00"))), Row(null)) + withSQLConf( + "parquet.crypto.factory.class" -> + "org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory", + "parquet.encryption.kms.client.class" -> + "org.apache.parquet.crypto.keytools.mocks.InMemoryKMS", + "parquet.encryption.key.list" -> "fieldIdKey: MDEyMzQ1Njc4OTAxMjM0NQ==", + "parquet.encryption.uniform.key" -> "fieldIdKey", + SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED.key -> "true", + SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "INT96") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark.createDataFrame(spark.sparkContext.parallelize(rows), withIds).write.parquet(path) + + withSQLConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") { + checkSparkAnswerAndFallbackReason( + spark.read.schema(withIds).parquet(path), + "Native Parquet scans that match field ids do not support encryption") + // Without ids in the read schema the scan matches by name, so it stays native. + checkSparkAnswerAndOperator(spark.read.schema(withoutIds).parquet(path)) + } + withSQLConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "false") { + checkSparkAnswerAndOperator(spark.read.schema(withIds).parquet(path)) + } + } + } + } + }