Skip to content

refactor: Simplify and optimize Parquet projection planning for nested casts - #24522

Open
jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:refactor/parquet-projection-read-plan
Open

refactor: Simplify and optimize Parquet projection planning for nested casts#24522
jayzhan211 wants to merge 1 commit into
apache:mainfrom
jayzhan211:refactor/parquet-projection-read-plan

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

Rationale for this change

build_read_plan_with_cast_clipping previously tracked projection decisions
across 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_field projections.

This PR simplifies that planning path while preserving its conservative
fallback behavior.

What changes are included in this PR?

  • Represent each projected root with one explicit state: a full read or a set
    of selected leaf offsets.
  • Merge cast and get_field leaf requirements into the same per-root state.
  • Preserve full-root reads whenever partial projection cannot be proven safe.
  • Build the projected schema and leaf mask in one ordered finalization loop.
  • Avoid an additional struct-leaf resolution pass and several temporary
    collections.
  • Remove the final sort and deduplication because root ordering, descriptor
    ordering, and BTreeSet offsets already produce sorted unique leaf indices.
  • Add comments explaining the ordering and fallback invariants.

Are these changes tested?

Yes. Existing tests cover full and partial projections, repeated and overlapping
casts, cast and get_field combinations, fallback behavior, and stale column
indices.

The following checks passed:

  • cargo fmt --all --check
  • cargo clippy -p datafusion-datasource-parquet --lib -- -D warnings
  • cargo test -p datafusion-datasource-parquet --lib
  • cargo test -p datafusion-sqllogictest --test sqllogictests -- parquet_nested_schema_pruning

Are there any user-facing changes?

No. This is an internal performance and readability improvement with no public
API or intended query behavior changes.

@github-actions github-actions Bot added the datasource Changes to the datasource crate label Aug 20, 2026
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.30769% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.30%. Comparing base (40c208e) to head (7819b01).
⚠️ Report is 6 commits behind head on main.

Files with missing lines Patch % Lines
...ion/datasource-parquet/src/projection_read_plan.rs 92.30% 2 Missing and 1 partial ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@adriangb adriangb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some small changes requested. Once fixed (or rejected) feel free to merge.

Comment on lines +651 to +652
// 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()) =>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

datasource Changes to the datasource crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants