This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24130-e7e037dbeb6e578f13a6ba16d528366e888c253a in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit bf1c17ae6b41a0aa8fe95a7b20273aefb47d3b10 Author: Goutam Adwant <[email protected]> AuthorDate: Sun Aug 16 01:13:03 2026 +0000 feat: union nested Parquet leaves across projection accesses (#24130) ## Which issue does this PR close? - Closes #24121. ## Rationale for this change When one projected nested column is consumed through multiple narrowing casts, or through both a narrowing cast and `get_field`, the Parquet read plan currently falls back to reading the entire root column. The leaves required by those consumers can be combined, so reading the full root performs unnecessary I/O. ## What changes are included in this PR? - Accumulate and deduplicate the leaf offsets required by every narrowing cast on a root column. - Add leaves required by `get_field` accesses to the same union. - Derive the Arrow type emitted by the reader from the merged leaf set while retaining a full-read fallback for unsupported partial wrapper shapes. - Add coverage for disjoint, overlapping, and repeated casts, cast plus `get_field`, and casts that diverge below `List<Struct>`. ## Are these changes tested? Yes. The following checks pass: - `cargo fmt --all -- --check` - `cargo test -p datafusion-datasource-parquet --lib` - `cargo test --profile=ci --test sqllogictests -- parquet_nested_schema_pruning.slt` - `cargo clippy -p datafusion-datasource-parquet --all-targets --all-features -- -D warnings` The SQL logic test also verifies that the merged projections return the expected values and scan fewer bytes than a full-root read. ## Are there any user-facing changes? No API or SQL behavior changes. Queries with multiple nested accesses to the same Parquet root can read fewer leaf columns. --- .../src/nested_schema_pruning.rs | 109 +++++++++++ .../datasource-parquet/src/projection_read_plan.rs | 213 ++++++++++++++------- .../test_files/parquet_nested_schema_pruning.slt | 44 ++++- 3 files changed, 287 insertions(+), 79 deletions(-) diff --git a/datafusion/datasource-parquet/src/nested_schema_pruning.rs b/datafusion/datasource-parquet/src/nested_schema_pruning.rs index 9768282c3b..9d1534eff6 100644 --- a/datafusion/datasource-parquet/src/nested_schema_pruning.rs +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -128,6 +128,35 @@ pub(crate) fn clip_for_cast( Some((kept, pruned_type)) } +/// Rebuilds `physical` for a sorted, deduplicated subset of its leaf offsets. +/// +/// This is the second half of merging several nested accesses to one root: +/// callers union the offsets required by each access, then use this function +/// to derive the Arrow type the Parquet reader emits for that union. Partial +/// projection is supported through the same struct, list, and large-list +/// shapes as [`clip_for_cast`]. Any partial selection below another wrapper +/// returns `None`, preserving the total-fallback property of cast clipping. +pub(crate) fn type_for_leaf_subset( + physical: &DataType, + kept: &[usize], +) -> Option<DataType> { + let total = count_leaves(physical); + if kept.is_empty() + || kept.len() >= total + || kept.windows(2).any(|pair| pair[0] >= pair[1]) + || kept.last().is_some_and(|last| *last >= total) + { + return None; + } + + let mut next_leaf = 0; + let projected = project_type_for_leaves(physical, kept, &mut next_leaf); + if projected.is_ok() { + debug_assert_eq!(next_leaf, total, "leaf accounting must cover the type"); + } + projected.ok().flatten() +} + /// Number of Parquet leaf columns a (Parquet-derived) Arrow type occupies. pub(crate) fn count_leaves(dt: &DataType) -> usize { match dt { @@ -249,6 +278,54 @@ fn clip_type( } } +/// Rebuilds one physical type subtree for `kept` leaf offsets. +/// +/// `Ok(None)` means no selected leaf lies below this subtree. `Err(())` means +/// only part of an unsupported wrapper was selected, so the caller must fall +/// back to reading the full root. +fn project_type_for_leaves( + physical: &DataType, + kept: &[usize], + next_leaf: &mut usize, +) -> Result<Option<DataType>, ()> { + match physical { + DataType::Struct(fields) => { + let mut projected = Vec::with_capacity(fields.len()); + for field in fields { + if let Some(dt) = + project_type_for_leaves(field.data_type(), kept, next_leaf)? + { + projected.push(field_with_type(field, dt)); + } + } + Ok((!projected.is_empty()).then(|| DataType::Struct(projected.into()))) + } + DataType::List(item) => { + Ok(project_type_for_leaves(item.data_type(), kept, next_leaf)? + .map(|projected| DataType::List(field_with_type(item, projected)))) + } + DataType::LargeList(item) => { + Ok(project_type_for_leaves(item.data_type(), kept, next_leaf)? + .map(|projected| DataType::LargeList(field_with_type(item, projected)))) + } + _ => { + let start = *next_leaf; + let end = start + count_leaves(physical); + let selected_start = kept.partition_point(|leaf| *leaf < start); + let selected_end = kept.partition_point(|leaf| *leaf < end); + let selected = selected_end - selected_start; + *next_leaf = end; + if selected == 0 { + Ok(None) + } else if selected == end - start { + Ok(Some(physical.clone())) + } else { + Err(()) + } + } + } +} + /// Keep every leaf of `dt` (no pruning below this point); returns `dt` /// unchanged since nothing was clipped. fn keep_all_leaves( @@ -455,6 +532,38 @@ mod tests { assert_eq!(emitted, list_of(struct_of(vec![int64("x"), utf8("y")]))); } + /// A union assembled from two casts below a shared List<Struct> prefix + /// reconstructs one emitted type in physical field order. + #[test] + fn type_for_leaf_subset_unions_nested_cast_leaves() { + let physical = list_of(struct_of(vec![int64("x"), utf8("y"), utf8("pad")])); + let emitted = type_for_leaf_subset(&physical, &[0, 1]).unwrap(); + assert_eq!(emitted, list_of(struct_of(vec![int64("x"), utf8("y")]))); + } + + /// A complete leaf set has nothing to prune, so callers must retain the + /// original root type rather than treating it as a clipped projection. + #[test] + fn type_for_leaf_subset_rejects_complete_selection() { + let physical = struct_of(vec![int64("x"), utf8("y"), utf8("pad")]); + assert!(type_for_leaf_subset(&physical, &[0, 1, 2]).is_none()); + } + + /// Partial selection below a wrapper cast clipping does not understand + /// must fall back instead of predicting a type the reader may not emit. + #[test] + fn type_for_leaf_subset_rejects_unsupported_partial_wrapper() { + let physical = DataType::FixedSizeList( + Arc::new(Field::new( + "item", + struct_of(vec![int64("x"), utf8("pad")]), + true, + )), + 2, + ); + assert!(type_for_leaf_subset(&physical, &[0]).is_none()); + } + /// Two levels of `list<struct>` nesting, the inner one also narrowed, /// the `events: array<struct<..., items: array<struct<...>>>>` shape /// reported in `datafusion-comet#4859`, where a sibling struct field at diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs index 0381808511..c00478f562 100644 --- a/datafusion/datasource-parquet/src/projection_read_plan.rs +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -45,6 +45,7 @@ use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; use crate::nested_schema_pruning::{ CastColumnAccess, clip_for_cast, contains_struct, count_leaves, field_with_type, + type_for_leaf_subset, }; /// The result of resolving which Parquet leaf columns and Arrow schema fields @@ -577,22 +578,16 @@ pub(crate) fn build_projection_read_plan( /// - roots referenced as whole columns keep every leaf and their full /// physical field (whole-column reads take precedence; cast accesses on /// such roots were already dropped by the caller); -/// - roots consumed through a cast, and not also through a `get_field` -/// access on the same root, keep only the leaves the cast target names -/// (see `crate::nested_schema_pruning`); +/// - roots consumed through one or more casts keep the union of the leaves +/// their targets name (see `crate::nested_schema_pruning`); +/// - a root consumed through both casts and `get_field` accesses keeps the +/// union of both access kinds; /// - roots consumed only through `get_field` accesses keep the union of the /// leaves those accesses reach, as before; /// - any other referenced root, a cast that can't be safely clipped (see -/// `nested_schema_pruning::clip_for_cast`), a root reached by two casts -/// with *different* targets (a projection can consume the same column -/// through more than one narrowing cast, e.g. -/// `SELECT CAST(s AS STRUCT(a)), CAST(s AS STRUCT(b)) FROM t`; clipping to -/// either target alone would starve the other), or a root reached by both a -/// cast and a `get_field` access (not produced by -/// `DefaultPhysicalExprAdapter`, which always routes a `get_field` over a -/// narrowed column through the same cast rather than a separate access, -/// but a custom `PhysicalExprAdapter` could in principle inject both), -/// falls back to a full read of that root. +/// `nested_schema_pruning::clip_for_cast`), or a merged leaf set whose +/// emitted Arrow type can't be derived safely, falls back to a full read of +/// that root. fn build_read_plan_with_cast_clipping( file_schema: &Schema, schema_descr: &SchemaDescriptor, @@ -601,45 +596,22 @@ fn build_read_plan_with_cast_clipping( cast_accesses: &[CastColumnAccess], ) -> ParquetReadPlan { let whole_roots: BTreeSet<usize> = whole_root_indices.iter().copied().collect(); - let struct_access_roots: BTreeSet<usize> = - struct_accesses.iter().map(|a| a.root_index).collect(); // Every referenced root's Parquet leaves, grouped in one pass over the // schema descriptor rather than one `leaf_indices_for_roots` scan per // root (this function may look up several roots). let leaves_by_root = leaves_grouped_by_root(schema_descr); - // Root -> (absolute kept leaf indices, cast-clipped Arrow type) for - // roots successfully clipped via a cast. - let mut clipped_by_root: BTreeMap<usize, (Vec<usize>, DataType)> = BTreeMap::new(); + // Root -> relative leaf offsets required by every narrowing cast on that + // root. A set makes repeated and overlapping targets a natural union. + let mut kept_offsets_by_root: BTreeMap<usize, BTreeSet<usize>> = BTreeMap::new(); // Roots with a cast access that must fall back to a full read. let mut fallback_roots: BTreeSet<usize> = BTreeSet::new(); - // The cast target already clipped for a root, so a second cast on the - // same root can be recognised as either a repeat (same target: nothing to - // do) or a conflict (different target: neither clip is valid on its own). - let mut clipped_target_by_root: BTreeMap<usize, &DataType> = BTreeMap::new(); for access in cast_accesses { let root = access.root_index; if whole_roots.contains(&root) || fallback_roots.contains(&root) { continue; } - if let Some(previous) = clipped_target_by_root.get(&root) { - if **previous != access.target_type { - // The projection consumes this root through two different - // narrowing casts. Each cast only needs its own leaves, but - // the mask is per column: clipping to the first target would - // silently null-fill whatever the second one needs. Read the - // whole root instead. - clipped_by_root.remove(&root); - clipped_target_by_root.remove(&root); - fallback_roots.insert(root); - } - continue; - } - if struct_access_roots.contains(&root) { - fallback_roots.insert(root); - continue; - } let physical_type = file_schema.field(root).data_type(); let root_leaves = leaves_by_root.get(&root).map_or(&[][..], Vec::as_slice); @@ -654,33 +626,73 @@ fn build_read_plan_with_cast_clipping( } match clip_for_cast(physical_type, &access.target_type) { - Some((kept_offsets, pruned_type)) => { - let start = root_leaves[0]; - let absolute = kept_offsets.into_iter().map(|o| start + o).collect(); - clipped_by_root.insert(root, (absolute, pruned_type)); - clipped_target_by_root.insert(root, &access.target_type); + Some((kept_offsets, _pruned_type)) => { + kept_offsets_by_root + .entry(root) + .or_default() + .extend(kept_offsets); } // Nothing prunable for this cast: every leaf is consumed. None => { + kept_offsets_by_root.remove(&root); fallback_roots.insert(root); } } } + // Add leaves reached through `get_field` to cast roots. The resolver + // returns absolute Parquet leaf indices; convert them back to offsets in + // their root so they can share the same union as cast clipping. + let struct_access_tree = StructAccessTree::from_accesses(struct_accesses); + for leaf in resolve_struct_field_leaves(&struct_access_tree, schema_descr) { + let root = schema_descr.get_column_root_idx(leaf); + if !kept_offsets_by_root.contains_key(&root) { + continue; + } + let Some(offset) = leaves_by_root + .get(&root) + .and_then(|root_leaves| root_leaves.binary_search(&leaf).ok()) + else { + kept_offsets_by_root.remove(&root); + fallback_roots.insert(root); + continue; + }; + kept_offsets_by_root + .get_mut(&root) + .expect("root presence checked above") + .insert(offset); + } + + // Derive the reader's one emitted Arrow type from each merged leaf set. + // Any unsupported partial wrapper retains the total fallback guarantee. + let mut clipped_by_root: BTreeMap<usize, (Vec<usize>, DataType)> = BTreeMap::new(); + for (root, kept_offsets) in kept_offsets_by_root { + if fallback_roots.contains(&root) { + continue; + } + let physical_type = file_schema.field(root).data_type(); + let root_leaves = leaves_by_root.get(&root).map_or(&[][..], Vec::as_slice); + let kept_offsets = kept_offsets.into_iter().collect::<Vec<_>>(); + let Some(pruned_type) = type_for_leaf_subset(physical_type, &kept_offsets) else { + fallback_roots.insert(root); + continue; + }; + let absolute = kept_offsets + .into_iter() + .map(|offset| root_leaves[offset]) + .collect(); + clipped_by_root.insert(root, (absolute, pruned_type)); + } + // `get_field` accesses on roots not already read in full (as a whole // column, or as a cast that fell back) keep the existing (non-cast) leaf // resolution. let get_field_accesses: Vec<StructFieldAccess> = struct_accesses .iter() .filter(|a| { - // A root carrying a `get_field` access is put into - // `fallback_roots` before any clip is attempted (see the loop - // above), so it can never also be clipped. Assert that rather - // than re-testing it here, so a future reordering trips the - // assert instead of silently changing which leaves are read. - debug_assert!(!clipped_by_root.contains_key(&a.root_index)); !whole_roots.contains(&a.root_index) && !fallback_roots.contains(&a.root_index) + && !clipped_by_root.contains_key(&a.root_index) }) .cloned() .collect(); @@ -1358,10 +1370,9 @@ mod test { ); } - /// Once conflicting cast targets have demoted a root to a full read, a - /// *third* cast on it must not resurrect the clip. + /// A repeated third cast does not duplicate leaves or widen the union. #[test] - fn build_projection_read_plan_keeps_full_read_after_a_third_cast() { + fn build_projection_read_plan_keeps_union_after_a_third_cast() { let (file_schema, metadata) = write_id_struct_file(); let schema_descr = metadata.file_metadata().schema_descr(); @@ -1374,10 +1385,19 @@ mod test { assert_eq!( read_plan.projection_mask, - ProjectionMask::leaves(schema_descr, [1, 2, 3]) + ProjectionMask::leaves(schema_descr, [1, 2]) ); let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); - assert_eq!(s_field.data_type(), file_schema.field(1).data_type()); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into() + ) + ); } /// A whole-column reference wins over a `get_field` access on the same @@ -1490,12 +1510,10 @@ mod test { ); } - /// Two casts on the same root with *different* targets cannot both be - /// served by one mask: clipping to either target alone would null-fill - /// whatever the other one needs (or fail its runtime struct-compatibility - /// check outright). Read the whole root instead. + /// Two casts on the same root with disjoint targets share the union of + /// their leaves, while an unreferenced sibling remains pruned. #[test] - fn build_projection_read_plan_falls_back_on_conflicting_cast_targets() { + fn build_projection_read_plan_unions_disjoint_cast_targets() { let (file_schema, metadata) = write_id_struct_file(); let schema_descr = metadata.file_metadata().schema_descr(); @@ -1515,11 +1533,68 @@ mod test { assert_eq!( read_plan.projection_mask, - ProjectionMask::leaves(schema_descr, [1, 2, 3]), - "every leaf of `s` must be read so both casts see their fields" + ProjectionMask::leaves(schema_descr, [1, 2]), + "the union must serve both casts without reading `s.pad`" ); let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); - assert_eq!(s_field.data_type(), file_schema.field(1).data_type()); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into() + ) + ); + } + + /// A union that covers every leaf falls back to the full root type. + #[test] + fn build_projection_read_plan_falls_back_for_complete_cast_union() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![ + cast_to_struct("s", 1, vec![("value", DataType::Int32)]), + cast_to_struct( + "s", + 1, + vec![("label", DataType::Utf8), ("pad", DataType::Utf8)], + ), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1, 2, 3]) + ); + assert_eq!( + read_plan.projected_schema.field_with_name("s").unwrap(), + file_schema.field(1) + ); + } + + /// Overlapping cast targets deduplicate their shared leaves. + #[test] + fn build_projection_read_plan_unions_overlapping_cast_targets() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![ + cast_to_struct("s", 1, vec![("value", DataType::Int32)]), + cast_to_struct( + "s", + 1, + vec![("value", DataType::Int32), ("label", DataType::Utf8)], + ), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1, 2]) + ); } /// The struct fast-path gate looks at the *projected* columns, not at @@ -1546,12 +1621,10 @@ mod test { assert_eq!(read_plan.projected_schema.fields().len(), 1); } - /// A root reached by both a narrowing cast and a `get_field` access (not - /// producible by `DefaultPhysicalExprAdapter`, but a custom - /// `PhysicalExprAdapter` could inject both) falls back to a full read of - /// that root rather than attempting to union the two leaf sets. + /// A root reached by both a narrowing cast and a disjoint `get_field` + /// access shares the union of their leaves. #[test] - fn build_projection_read_plan_falls_back_when_cast_and_get_field_share_a_root() { + fn build_projection_read_plan_unions_cast_and_get_field_on_one_root() { let (file_schema, metadata) = write_id_struct_file(); let schema_descr = metadata.file_metadata().schema_descr(); @@ -1575,8 +1648,7 @@ mod test { let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); - // Every leaf of `s` is read (full fallback), not just value/label. - let expected_mask = ProjectionMask::leaves(schema_descr, [1, 2, 3]); + let expected_mask = ProjectionMask::leaves(schema_descr, [1, 2]); assert_eq!(read_plan.projection_mask, expected_mask); let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); @@ -1586,7 +1658,6 @@ mod test { vec![ Arc::new(Field::new("value", DataType::Int32, false)), Arc::new(Field::new("label", DataType::Utf8, false)), - Arc::new(Field::new("pad", DataType::Utf8, false)), ] .into() ), diff --git a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt index d936a89beb..0aa171824f 100644 --- a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -473,10 +473,8 @@ set datafusion.execution.parquet.pushdown_filters = false; # Query-level casts. `ProjectionExec` is merged into the scan, so a `CAST` # written in the query reaches the same read-plan analysis as an # adapter-inserted one — including one column consumed through two *different* -# cast targets, which no single clipped read can serve. Clipping to one -# target's leaves would leave the other cast reading a struct that is missing -# the fields it names, which `cast_column` either null-fills (wrong results) -# or, for disjoint targets, rejects outright. +# cast targets. The read plan unions the targets' leaves so both casts receive +# the fields they name without reading unreferenced siblings. # # `exact` infers its schema from the file, so no adapter cast is interposed # and the casts below are the only ones the scan sees. @@ -515,6 +513,25 @@ FROM exact ORDER BY id; {x: 200} {x: 200, y: s2} NULL NULL +# A cast plus a separate get_field access on the same physical root. +query ?T +SELECT CAST(s AS STRUCT<x BIGINT>) AS q0, s['y'] AS q1 +FROM exact ORDER BY id; +---- +{x: 100} s1 +{x: 200} s2 +NULL NULL + +# The targets diverge below a shared ARRAY<STRUCT<...>> prefix. +query ?? +SELECT CAST(events AS ARRAY<STRUCT<x BIGINT>>) AS q0, + CAST(events AS ARRAY<STRUCT<y VARCHAR>>) AS q1 +FROM exact ORDER BY id; +---- +[{x: 10}] [{y: a1}] +[{x: 20}, {x: 21}] [{y: b1}, {y: b2}] +NULL NULL + # Repeated identical targets still clip. query ?? SELECT CAST(s AS STRUCT<x BIGINT>) AS q0, @@ -534,14 +551,25 @@ FROM exact ORDER BY id; {x: 200} {x: 200, y: s2, pad: sp2} NULL NULL -# The conflicting-target fallback reads the whole column -- exactly what a -# scan with no clipping at all reads, and never more. These two must match: -# the first falls back, the second never clips in the first place. +# The disjoint targets read their two-leaf union and skip `pad`. query TT explain analyze select CAST(s AS STRUCT<x BIGINT>) AS q0, CAST(s AS STRUCT<pad VARCHAR>) AS q1 from exact; ---- -Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=219<slt:ignore>] +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=148<slt:ignore>] + +# The cast/get_field union likewise skips `pad`. +query TT +explain analyze select CAST(s AS STRUCT<x BIGINT>) AS q0, s['y'] AS q1 from exact; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=146<slt:ignore>] + +# The nested union reads x/y but skips both pad leaves. +query TT +explain analyze select CAST(events AS ARRAY<STRUCT<x BIGINT>>) AS q0, CAST(events AS ARRAY<STRUCT<y VARCHAR>>) AS q1 from exact; +---- +Plan with Metrics DataSourceExec: <slt:ignore>metrics=[output_rows=3, <slt:ignore>bytes_scanned=172<slt:ignore>] +# A full-root baseline remains larger than every union above. query TT explain analyze select s from exact; ---- --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
