refactor: Simplify and optimize Parquet projection planning for nested casts - #24522
refactor: Simplify and optimize Parquet projection planning for nested casts#24522jayzhan211 wants to merge 1 commit into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24522 +/- ##
==========================================
+ Coverage 81.29% 81.30% +0.01%
==========================================
Files 1116 1117 +1
Lines 395460 395879 +419
Branches 395460 395879 +419
==========================================
+ Hits 321486 321877 +391
- Misses 55167 55180 +13
- Partials 18807 18822 +15 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
adriangb
left a comment
There was a problem hiding this comment.
Some small changes requested. Once fixed (or rejected) feel free to merge.
| // Add every `get_field` root before resolving leaves. If an access matches | ||
| // no leaf, finalization safely falls back to a full read for that root. |
There was a problem hiding this comment.
This isn't a pure refactor — it fixes a bug. Routing get_field-only roots through type_for_leaf_subset instead of build_filter_schema changes what happens when an access path matches no Parquet leaf. With a cast on root a and a b['nonexistent'] access on root b:
| projected schema | mask | |
|---|---|---|
| before | a: Struct<p>, b: Struct() |
[0] |
| after | a: Struct<p>, b: Struct<m,n> |
[0, 2, 3] |
The old path emitted an empty struct with zero leaves selected, a schema the reader can't produce. The fallback here is correct, but nothing tests it; the private access() helper makes a regression test ~20 lines. Worth calling out in the description too, so this isn't reviewed as a no-op.
While here: the doc comment at L594 still says get_field-only roots behave "as before", which is no longer true.
There was a problem hiding this comment.
If it fixes a bug can we get a regression test?
| { | ||
| leaf_indices | ||
| .extend(offsets.into_iter().map(|offset| root_leaves[offset])); | ||
| fields.push(field_with_type(field, projected_type)); |
There was a problem hiding this comment.
field_with_type preserves root field metadata; assemble_read_plan still goes through build_filter_schema, which rebuilds roots with Field::new(...) and drops it. Same access, same leaves, same type:
cast path : b = Struct("m": Int32) meta={"k": "v"}
plain path: b = Struct("m": Int32) meta={}
So a column's projected field now depends on whether an unrelated column carries a narrowing cast. This side is the right behavior: worth switching build_filter_schema to field_with_type too, here or as a follow-up.
| let root_leaves = leaves_by_root.get(&root).map_or(&[][..], Vec::as_slice); | ||
| match read { | ||
| RootRead::Partial(offsets) | ||
| if root_leaves.len() == count_leaves(field.data_type()) => |
There was a problem hiding this comment.
This guard is redundant for cast roots (already checked in the loop above) but is newly applied to get_field-only roots, which never had it. This is a safe direction since it just falls back to a full read. But it's a silent narrowing; worth a comment saying so.
| // `root_reads` visits roots in schema order, every root's leaves were | ||
| // collected in descriptor order, and partial offsets are a `BTreeSet`. | ||
| // Therefore the final mask is already sorted and deduplicated. | ||
| debug_assert!(leaf_indices.windows(2).all(|pair| pair[0] < pair[1])); |
There was a problem hiding this comment.
Does this even matter?ProjectionMask::leaves just flips booleans in a vec![false; num_columns], so order and duplicates are irrelevant. Dropping the sort_unstable/dedup is fine regardless. The invariant that does carry weight is that fields is pushed in ascending root order to match the reader's output. Consider asserting/documenting that instead?
Rationale for this change
build_read_plan_with_cast_clippingpreviously tracked projection decisionsacross several overlapping collections. This made it difficult to understand
whether each root column required a full or partial read and required additional
work to combine cast and
get_fieldprojections.This PR simplifies that planning path while preserving its conservative
fallback behavior.
What changes are included in this PR?
of selected leaf offsets.
get_fieldleaf requirements into the same per-root state.collections.
ordering, and
BTreeSetoffsets already produce sorted unique leaf indices.Are these changes tested?
Yes. Existing tests cover full and partial projections, repeated and overlapping
casts, cast and
get_fieldcombinations, fallback behavior, and stale columnindices.
The following checks passed:
cargo fmt --all --checkcargo clippy -p datafusion-datasource-parquet --lib -- -D warningscargo test -p datafusion-datasource-parquet --libcargo test -p datafusion-sqllogictest --test sqllogictests -- parquet_nested_schema_pruningAre there any user-facing changes?
No. This is an internal performance and readability improvement with no public
API or intended query behavior changes.