From 775dfe40fbcdedefb4d47ec9ef12899558334d41 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Thu, 13 Aug 2026 01:23:18 +0300 Subject: [PATCH] Clip nested wrappers in Parquet schema pruning. --- .../src/nested_schema_pruning.rs | 431 ++++++++++++++---- 1 file changed, 337 insertions(+), 94 deletions(-) diff --git a/datafusion/datasource-parquet/src/nested_schema_pruning.rs b/datafusion/datasource-parquet/src/nested_schema_pruning.rs index 9768282c3bab0..85efc5ca2e121 100644 --- a/datafusion/datasource-parquet/src/nested_schema_pruning.rs +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -36,19 +36,25 @@ //! //! # Safety of clipping //! -//! The runtime cast for nested types -//! ([`datafusion_common::nested_struct::cast_column`]) consumes source struct -//! children exclusively by looking up the *target* field names, recursively -//! through list wrappers. Physical subtrees not named by the target are -//! provably dead: removing them from the read cannot change the cast's -//! output. That holds for *any* +//! The runtime cast path for nested types consumes source struct children +//! exclusively by looking up the *target* field names, recursively through +//! matching list, list-view, fixed-size-list, and dictionary wrappers. +//! Physical subtrees not named by the target are provably dead: removing +//! them from the read cannot change the cast's output. That holds for *any* //! [`CastExpr`](datafusion_physical_expr::expressions::CastExpr) over a //! nested type, not just the ones the schema adapter inserts: -//! `ColumnarValue::cast_to` routes every -//! cast for which +//! `ColumnarValue::cast_to` routes every cast for which //! [`requires_nested_struct_cast`](datafusion_common::nested_struct::requires_nested_struct_cast) -//! holds, the same predicate the projection analysis gates on, through -//! `cast_column`. +//! holds, the same predicate the projection analysis gates on, through the +//! name-based [`datafusion_common::nested_struct::cast_column`]. The wrapper +//! pairs clipped here must therefore stay exactly the pairs that predicate +//! recurses through. A wrapper outside the set — run-end-encoded, for +//! example — is cast by Arrow's kernel instead, and Arrow's struct cast +//! falls back to *positional* field mapping when a target name is missing +//! from the source, so a physical child the target does not name can still +//! feed a (renamed) target field. Clipping it could change the cast's +//! output; extend `requires_nested_struct_cast` and `cast_column` first if +//! a new wrapper is to become clippable. //! //! Struct-level nullability is preserved because the Parquet reader //! reconstructs ancestor validity from the definition levels of any surviving @@ -64,25 +70,25 @@ //! [`clip_for_cast`] detects the empty level and declines to clip. //! //! The clip is *total*: any type shape it does not understand (maps, -//! dictionaries, wrapper-kind mismatches, ...) keeps all of its leaves, so -//! the worst case is today's behavior of reading the full column. Map values -//! are deliberately not clipped: the runtime cast routes maps through Arrow's -//! positional struct cast, which requires all children to be present. Nor are -//! `ListView`/`LargeListView`/`Dictionary` wrappers clipped here, even though -//! `cast_column` does recurse through them by name. That is a conservative -//! choice (safe, since the worst case is still just a full read) left as a -//! candidate follow-up rather than something this module currently handles. +//! run-end-encoded wrappers, wrapper-kind or fixed-size-list-width +//! mismatches, ...) keeps all of its leaves, so the worst case is today's +//! behavior of reading the full column. Map values are deliberately not +//! clipped: the runtime cast routes maps through Arrow's positional struct +//! cast, which requires all children to be present. Run-end-encoded values +//! are deliberately not clipped either, per the routing argument above. +//! Matching single-child wrappers are rebuilt with their physical outer +//! parameters around the clipped child because that is the type the reader +//! emits before the target cast runs. use std::collections::HashMap; use std::sync::Arc; use arrow::datatypes::{DataType, Field, FieldRef, Fields}; -/// The single child type one level of container nesting wraps, or `None` for -/// a type this module does not descend through (leaves, `Struct`, `Map`, and -/// wrapper kinds this module intentionally does not clip, see the module -/// doc). Shared by [`count_leaves`] and [`contains_struct`], which otherwise -/// need to agree on the exact same set of container variants. +/// The single logical child type one level of container nesting wraps. `Map` +/// and `RunEndEncoded` are included because [`count_leaves`] and +/// [`contains_struct`] must descend through them for leaf accounting, even +/// though [`clip_type`] deliberately keeps both opaque. fn nested_child(dt: &DataType) -> Option<&DataType> { match dt { DataType::List(f) @@ -243,8 +249,44 @@ fn clip_type( ); DataType::LargeList(field_with_type(p_item, pruned)) } - // Anything else, leaf pairs, wrapper-kind mismatches, maps, - // dictionaries, fixed-size lists, views, is kept wholesale. + (DataType::ListView(p_item), DataType::ListView(t_item)) => { + let pruned = clip_type( + p_item.data_type(), + t_item.data_type(), + next_leaf, + kept, + unclippable, + ); + DataType::ListView(field_with_type(p_item, pruned)) + } + (DataType::LargeListView(p_item), DataType::LargeListView(t_item)) => { + let pruned = clip_type( + p_item.data_type(), + t_item.data_type(), + next_leaf, + kept, + unclippable, + ); + DataType::LargeListView(field_with_type(p_item, pruned)) + } + ( + DataType::FixedSizeList(p_item, p_size), + DataType::FixedSizeList(t_item, t_size), + ) if p_size == t_size => { + let pruned = clip_type( + p_item.data_type(), + t_item.data_type(), + next_leaf, + kept, + unclippable, + ); + DataType::FixedSizeList(field_with_type(p_item, pruned), *p_size) + } + (DataType::Dictionary(p_key, p_value), DataType::Dictionary(_, t_value)) => { + let pruned = clip_type(p_value, t_value, next_leaf, kept, unclippable); + DataType::Dictionary(p_key.clone(), Box::new(pruned)) + } + // All other shapes are opaque and kept wholesale; see the module docs. _ => keep_all_leaves(physical, next_leaf, kept), } } @@ -287,6 +329,12 @@ pub(crate) fn field_with_type(field: &Field, data_type: DataType) -> FieldRef { mod tests { use super::*; + use arrow::array::{Array, ArrayRef, Int64Array, StringArray, StructArray}; + use arrow::buffer::NullBuffer; + use arrow::record_batch::RecordBatch; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::{ArrowWriter, ProjectionMask}; + fn utf8(name: &str) -> Field { Field::new(name, DataType::Utf8, true) } @@ -326,10 +374,9 @@ mod tests { let dict = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); assert_eq!(count_leaves(&dict), 1); - // Wrapper kinds must be descended through, not counted as one leaf. - // A dictionary or run-end-encoded *value* that is itself a struct has - // as many leaves as the struct: counting it as 1 would misalign every - // later leaf index in the mask. + // A dictionary or run-end-encoded *value* that is itself a struct + // has as many leaves as the struct; counting the wrapper as one leaf + // would misalign every later leaf index in the mask. assert_eq!( count_leaves(&DataType::Dictionary( Box::new(DataType::Int32), @@ -496,6 +543,98 @@ mod tests { assert_eq!(emitted, DataType::LargeList(item(vec![int64("x")]))); } + #[test] + fn clip_list_views_of_struct() { + let item = |fields| Arc::new(Field::new("item", struct_of(fields), true)); + + let physical = DataType::ListView(item(vec![int64("x"), utf8("pad")])); + let target = DataType::ListView(item(vec![int64("x")])); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, DataType::ListView(item(vec![int64("x")]))); + + let physical = DataType::LargeListView(item(vec![int64("x"), utf8("pad")])); + let target = DataType::LargeListView(item(vec![int64("x")])); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, DataType::LargeListView(item(vec![int64("x")]))); + } + + #[test] + fn clip_fixed_size_list_of_struct() { + let item = |fields| Arc::new(Field::new("item", struct_of(fields), true)); + let physical = DataType::FixedSizeList(item(vec![int64("x"), utf8("pad")]), 4); + let target = DataType::FixedSizeList(item(vec![int64("x")]), 4); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, DataType::FixedSizeList(item(vec![int64("x")]), 4)); + } + + #[test] + fn no_clip_on_fixed_size_list_width_mismatch() { + let item = |fields| Arc::new(Field::new("item", struct_of(fields), true)); + let physical = DataType::FixedSizeList(item(vec![int64("x"), utf8("pad")]), 4); + let target = DataType::FixedSizeList(item(vec![int64("x")]), 3); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + #[test] + fn clip_dictionary_struct_value() { + let physical = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(struct_of(vec![int64("x"), utf8("pad")])), + ); + let target = DataType::Dictionary( + Box::new(DataType::Int64), + Box::new(struct_of(vec![int64("x")])), + ); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted, + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(struct_of(vec![int64("x")])) + ) + ); + } + + /// Run-end-encoded wrappers remain opaque; see the module documentation. + #[test] + fn no_clip_on_run_end_encoded() { + let run_ends = Arc::new(Field::new("run_ends", DataType::Int32, false)); + let values = |fields| Arc::new(Field::new("values", struct_of(fields), true)); + let physical = DataType::RunEndEncoded( + Arc::clone(&run_ends), + values(vec![int64("x"), utf8("pad")]), + ); + let target = DataType::RunEndEncoded(run_ends, values(vec![int64("x")])); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Production-reachable nested case for the opaque run-end-encoded fallback. + #[test] + fn no_clip_inside_run_end_encoded() { + let ree = |value_fields| { + DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", struct_of(value_fields), true)), + ) + }; + let physical = struct_of(vec![ + ree(vec![int64("x"), utf8("pad")]).into_field("r"), + utf8("dropped"), + ]); + let target = struct_of(vec![ree(vec![int64("x")]).into_field("r")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + // Both REE leaves survive (x=0, pad=1); only the sibling is dropped. + assert_eq!(kept, vec![0, 1]); + assert_eq!( + emitted, + struct_of(vec![ree(vec![int64("x"), utf8("pad")]).into_field("r")]) + ); + } + /// Wrapper-kind mismatch cannot be clipped. #[test] fn no_clip_on_wrapper_mismatch() { @@ -621,11 +760,6 @@ mod tests { /// type the reader emits, rather than surviving as an empty struct. #[test] fn reader_drops_struct_child_with_no_selected_leaves() { - use arrow::array::{ArrayRef, Int64Array, StructArray}; - use arrow::record_batch::RecordBatch; - use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; - use parquet::arrow::{ArrowWriter, ProjectionMask}; - let inner_fields = Fields::from(vec![int64("a"), int64("b")]); let outer_fields = Fields::from(vec![ Field::new("inner", DataType::Struct(inner_fields.clone()), true), @@ -672,84 +806,57 @@ mod tests { ); } - /// Pins the arrow-rs behavior this module relies on: selecting a subset - /// of leaves under a `List` column with `ProjectionMask::leaves` - /// makes the reader emit exactly the type predicted by [`clip_for_cast`], - /// and null list rows / null struct elements survive (their validity is - /// reconstructed from the surviving leaves' definition levels). - #[test] - fn arrow_reader_emits_clipped_type_for_masked_list_struct() { - use arrow::array::{ - Array, ArrayRef, Int64Array, ListArray, StringArray, StructArray, - }; - use arrow::buffer::{NullBuffer, OffsetBuffer}; - use arrow::record_batch::RecordBatch; - use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; - use parquet::arrow::{ArrowWriter, ProjectionMask}; - + /// Item struct `{x, y, pad}` with the three elements + /// `[{1, "a", "p0"}, NULL, {3, "c", "p2"}]` shared by the reader pin + /// tests below, whose clip target `{x, y}` drops `pad` (leaf 2). + fn pin_item_values() -> (FieldRef, StructArray) { let item_fields = Fields::from(vec![int64("x"), utf8("y"), utf8("pad")]); - let item_field = Arc::new(Field::new( - "item", - DataType::Struct(item_fields.clone()), - true, - )); - let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( - "events", - DataType::List(Arc::clone(&item_field)), - true, - )])); - - // 3 elements; element 1 is a NULL struct. Rows: [e0, e1], NULL, [e2]. let columns: Vec = vec![ Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])), Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])), Arc::new(StringArray::from(vec![Some("p0"), None, Some("p2")])), ]; - let struct_validity = NullBuffer::from(vec![true, false, true]); - let values = StructArray::new(item_fields, columns, Some(struct_validity)); - let list_validity = NullBuffer::from(vec![true, false, true]); - let events = ListArray::new( - item_field, - OffsetBuffer::from_lengths([2, 0, 1]), - Arc::new(values), - Some(list_validity), - ); - let batch = - RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(events)]).unwrap(); - + let validity = NullBuffer::from(vec![true, false, true]); + let values = StructArray::new(item_fields.clone(), columns, Some(validity)); + ( + Arc::new(Field::new("item", DataType::Struct(item_fields), true)), + values, + ) + } + + /// Roundtrip `batch`'s single column through Parquet reading only the + /// leaves [`clip_for_cast`] keeps for `target`, asserting the reader + /// emits exactly the predicted type. Returns the batch read back. + fn roundtrip_with_clipped_mask( + batch: &RecordBatch, + target: &DataType, + expected_kept: &[usize], + ) -> RecordBatch { let file = tempfile::NamedTempFile::new().unwrap(); let mut writer = - ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap(); - writer.write(&batch).unwrap(); + ArrowWriter::try_new(file.reopen().unwrap(), batch.schema(), None).unwrap(); + writer.write(batch).unwrap(); writer.close().unwrap(); - // Clip to the narrow target {x, y}. let physical = batch.schema().field(0).data_type().clone(); - let target = list_of(struct_of(vec![int64("x"), utf8("y")])); - let (kept, predicted_type) = clip_for_cast(&physical, &target).unwrap(); - assert_eq!(kept, vec![0, 1]); + let (kept, predicted_type) = clip_for_cast(&physical, target).unwrap(); + assert_eq!(kept, expected_kept); let builder = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap(); let mask = ProjectionMask::leaves(builder.parquet_schema(), kept.iter().copied()); let reader = builder.with_projection(mask).build().unwrap(); - let out: Vec = reader.map(|b| b.unwrap()).collect(); - assert_eq!(out.len(), 1); - let out = &out[0]; - - // Emitted type matches the prediction. + let mut batches: Vec = reader.map(|b| b.unwrap()).collect(); + assert_eq!(batches.len(), 1); + let out = batches.remove(0); assert_eq!(out.schema().field(0).data_type(), &predicted_type); + out + } - // Null semantics survive the clip. - let events = out.column(0).as_any().downcast_ref::().unwrap(); - assert!(events.is_valid(0)); - assert!(events.is_null(1)); - assert!(events.is_valid(2)); - let structs = events - .values() - .as_any() - .downcast_ref::() - .unwrap(); + /// Assert the [`pin_item_values`] elements survived the masked read: + /// struct validity `[valid, NULL, valid]` reconstructed from the + /// surviving leaves' definition levels, and `x` values intact. + fn assert_pin_item_values(structs: &StructArray) { assert_eq!(structs.len(), 3); assert!(structs.is_valid(0)); assert!(structs.is_null(1)); @@ -763,6 +870,142 @@ mod tests { assert_eq!(x.value(2), 3); } + /// Pins the arrow-rs behavior this module relies on: selecting a subset + /// of leaves under a `List` column with `ProjectionMask::leaves` + /// makes the reader emit exactly the type predicted by [`clip_for_cast`], + /// and null list rows / null struct elements survive (their validity is + /// reconstructed from the surviving leaves' definition levels). + #[test] + fn arrow_reader_emits_clipped_type_for_masked_list_struct() { + use arrow::array::ListArray; + use arrow::buffer::OffsetBuffer; + + let (item_field, values) = pin_item_values(); + // Rows: [e0, e1], NULL, [e2]. + let events = ListArray::new( + Arc::clone(&item_field), + OffsetBuffer::from_lengths([2, 0, 1]), + Arc::new(values), + Some(NullBuffer::from(vec![true, false, true])), + ); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "events", + DataType::List(item_field), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(events)]).unwrap(); + + let target = list_of(struct_of(vec![int64("x"), utf8("y")])); + let out = roundtrip_with_clipped_mask(&batch, &target, &[0, 1]); + + let events = out.column(0).as_any().downcast_ref::().unwrap(); + assert!(events.is_valid(0)); + assert!(events.is_null(1)); + assert!(events.is_valid(2)); + let structs = events + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_pin_item_values(structs); + } + + /// Pins the reader contract for `ListView`. + #[test] + fn arrow_reader_emits_clipped_type_for_masked_list_view_struct() { + use arrow::array::ListViewArray; + + let (item_field, values) = pin_item_values(); + // Rows: [e0, e1], NULL, [e2]. + let events = ListViewArray::new( + Arc::clone(&item_field), + vec![0_i32, 0, 2].into(), + vec![2_i32, 0, 1].into(), + Arc::new(values), + Some(NullBuffer::from(vec![true, false, true])), + ); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "events", + DataType::ListView(item_field), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(events)]).unwrap(); + + let target = DataType::ListView(Arc::new(Field::new( + "item", + struct_of(vec![int64("x"), utf8("y")]), + true, + ))); + roundtrip_with_clipped_mask(&batch, &target, &[0, 1]); + } + + /// Pins the reader contract for `FixedSizeList`, whose reader + /// path is distinct from the variable-size lists': the type is rebuilt + /// from the embedded Arrow schema hint, and null rows carry hidden child + /// slots. + #[test] + fn arrow_reader_emits_clipped_type_for_masked_fixed_size_list_struct() { + use arrow::array::FixedSizeListArray; + + // 3 rows of width 2; row 1 is NULL (its child slots are hidden) and + // element 1 of row 0 is a NULL struct. + let item_fields = Fields::from(vec![int64("x"), utf8("y"), utf8("pad")]); + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![ + Some(1), + None, + None, + None, + Some(3), + Some(4), + ])), + Arc::new(StringArray::from(vec![ + Some("a"), + None, + None, + None, + Some("c"), + Some("d"), + ])), + Arc::new(StringArray::from(vec![ + Some("p0"), + None, + None, + None, + Some("p2"), + Some("p3"), + ])), + ]; + let struct_validity = + NullBuffer::from(vec![true, false, false, false, true, true]); + let values = + StructArray::new(item_fields.clone(), columns, Some(struct_validity)); + let item_field = + Arc::new(Field::new("item", DataType::Struct(item_fields), true)); + let events = FixedSizeListArray::new( + Arc::clone(&item_field), + 2, + Arc::new(values), + Some(NullBuffer::from(vec![true, false, true])), + ); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "events", + DataType::FixedSizeList(item_field, 2), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(events)]).unwrap(); + + let target = DataType::FixedSizeList( + Arc::new(Field::new( + "item", + struct_of(vec![int64("x"), utf8("y")]), + true, + )), + 2, + ); + roundtrip_with_clipped_mask(&batch, &target, &[0, 1]); + } + trait IntoField { fn into_field(self, name: &str) -> Field; }