diff --git a/docs/source/user-guide/latest/compatibility/scans.md b/docs/source/user-guide/latest/compatibility/scans.md index 57363125331..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: 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/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index f015ed70f83..faa0d8f19ea 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/parquet/ParquetReadSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/ParquetReadSuite.scala index 9b6bceeef25..e0074508ab7 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 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)) + } + } + } + } + }