Skip to content

Union the required leaves when a nested root is reached by multiple casts (or a cast plus get_field) instead of reading the whole column #24121

Description

@adriangb

Part of #24119.

Background

#24090 prunes Parquet leaves when a projected nested column is consumed through a narrowing cast. build_read_plan_with_cast_clipping partitions the referenced roots and, for a root consumed through a single cast, keeps only the leaves that cast's target names.

There are three cases where it currently gives up and reads the whole root instead, even though the set of leaves actually needed is computable:

1. Two different cast targets on the same root

SELECT CAST(s AS STRUCT(a)), CAST(s AS STRUCT(b)) FROM t

Both casts are recorded, their targets differ, and the code bails:

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

The correct read set is the union of the two clips, not the full column. (The full-read fallback here is a correctness fix from mbutrovich#1 — clipping to one target alone starved the other. Union is the optimization that fallback left on the table.)

2. A cast plus a separate get_field access on the same root

if struct_access_roots.contains(&root) {
    fallback_roots.insert(root);
    continue;
}

Same story: the leaves the cast names, unioned with the leaves the get_field paths reach. #24090's description notes this union was dropped deliberately in favour of the full-read fallback because it was untested — DefaultPhysicalExprAdapter never produces this shape (it routes get_field over a narrowed column through the same cast), though a custom PhysicalExprAdapter could.

3. Repeated casts with the same target

Already handled correctly (recognised as a repeat, nothing to do) — noted only for completeness.

What this issue asks for

Replace the full-read fallbacks in cases 1 and 2 with a union of the required leaf sets.

The blocker is not computing the union of leaf offsets — that is a set union over Vec<usize>. It is that ParquetReadPlan also carries the Arrow type the reader will emit for the kept leaves, and clip_for_cast returns exactly one (kept_offsets, pruned_type) pair per cast. A union needs the merged pruned type: the type the reader emits for the union of the leaf sets, which then has to satisfy both consumers (each cast's runtime cast_column, and each get_field).

Sketch:

  • Generalise clip_for_cast (or add a sibling) to accumulate into a shared "kept leaves" tree across multiple targets on one root, rather than returning a standalone result per cast.
  • Derive the emitted Arrow type from that merged tree in one pass, preserving the existing invariant that every clipped struct level keeps at least one leaf (the reader reconstructs ancestor validity from definition levels of surviving leaves).
  • Keep the total-fallback property: any shape the merge does not understand still degrades to a full read.

Test coverage to add

  • Two narrowing casts with disjoint field sets on one root → mask is the union, both projections produce correct values.
  • Two narrowing casts with overlapping field sets → union, no duplicate leaves.
  • Cast plus get_field on a disjoint subfield of the same root.
  • Nested case: casts that diverge below a shared List<Struct<...>> prefix.
  • A bytes_scanned assertion showing the union reads less than the full root (the existing assert_scan_prunes helper in datafusion/core/benches/parquet_nested_schema_pruning.rs is the model).

Notes

These shapes are rare in practice today, so this is an optimization on an already-correct fallback rather than a bug. Filing so it isn't lost.

Metadata

Metadata

Assignees

Labels

datasourceChanges to the datasource crateenhancementNew feature or requestperformanceMake DataFusion faster

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions