Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jayzhan211 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
160 changes: 64 additions & 96 deletions datafusion/datasource-parquet/src/projection_read_plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -571,6 +571,14 @@ pub(crate) fn build_projection_read_plan(
read_plan
}

/// Leaf selection accumulated for one projected root column.
enum RootRead {
/// Decode every leaf and preserve the physical Arrow field.
Full,
/// Decode these leaf offsets relative to the start of this root.
Partial(BTreeSet<usize>),
}

/// Builds a [`ParquetReadPlan`] when at least one projected root column is
/// consumed through a cast to a narrower nested type.
///
Expand All @@ -595,21 +603,20 @@ fn build_read_plan_with_cast_clipping(
struct_accesses: &[StructFieldAccess],
cast_accesses: &[CastColumnAccess],
) -> ParquetReadPlan {
let whole_roots: BTreeSet<usize> = whole_root_indices.iter().copied().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 -> 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();
// Keep one decision per root. The ordered map also determines the output
// schema order, which must match the Parquet reader's root order.
let mut root_reads: BTreeMap<usize, RootRead> = whole_root_indices
.iter()
.map(|root| (*root, RootRead::Full))
.collect();

for access in cast_accesses {
let root = access.root_index;
if whole_roots.contains(&root) || fallback_roots.contains(&root) {
if matches!(root_reads.get(&root), Some(RootRead::Full)) {
continue;
}

Expand All @@ -621,130 +628,91 @@ fn build_read_plan_with_cast_clipping(
// arrow schema). If not, never risk a wrong mask: read the whole
// root.
if root_leaves.len() != count_leaves(physical_type) {
fallback_roots.insert(root);
root_reads.insert(root, RootRead::Full);
continue;
}

match clip_for_cast(physical_type, &access.target_type) {
Some((kept_offsets, _pruned_type)) => {
kept_offsets_by_root
if let RootRead::Partial(offsets) = root_reads
.entry(root)
.or_default()
.extend(kept_offsets);
.or_insert_with(|| RootRead::Partial(BTreeSet::new()))
{
offsets.extend(kept_offsets);
}
}
// Nothing prunable for this cast: every leaf is consumed.
None => {
kept_offsets_by_root.remove(&root);
fallback_roots.insert(root);
root_reads.insert(root, RootRead::Full);
}
}
}

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

@adriangb adriangb Aug 20, 2026

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.

@mbutrovich mbutrovich Aug 20, 2026

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?

for access in struct_accesses {
root_reads
.entry(access.root_index)
.or_insert_with(|| RootRead::Partial(BTreeSet::new()));
}

// The resolver returns absolute Parquet leaf indices. Convert each selected
// leaf to a root-relative offset so casts and field accesses share one union.
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) {
let Some(RootRead::Partial(offsets)) = root_reads.get_mut(&root) else {
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);
root_reads.insert(root, RootRead::Full);
continue;
};
let absolute = kept_offsets
.into_iter()
.map(|offset| root_leaves[offset])
.collect();
clipped_by_root.insert(root, (absolute, pruned_type));
offsets.insert(offset);
}

// `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| {
!whole_roots.contains(&a.root_index)
&& !fallback_roots.contains(&a.root_index)
&& !clipped_by_root.contains_key(&a.root_index)
})
.cloned()
.collect();

let mut leaf_indices: Vec<usize> = Vec::new();
let mut fields: BTreeMap<usize, Arc<Field>> = BTreeMap::new();

for root in whole_roots.iter().chain(fallback_roots.iter()) {
// A root with no parquet leaves contributes nothing to the mask;
// `ProjectionMask::roots` handles that case the same way, so match it
// rather than indexing and panicking.
if let Some(leaves) = leaves_by_root.get(root) {
leaf_indices.extend(leaves.iter().copied());
let mut fields = Vec::with_capacity(root_reads.len());
for (root, read) in root_reads {
let field = file_schema.field(root);
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()) =>

@adriangb adriangb Aug 20, 2026

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.

{
let offsets = offsets.into_iter().collect::<Vec<_>>();
if let Some(projected_type) =
type_for_leaf_subset(field.data_type(), &offsets)
{
leaf_indices
.extend(offsets.into_iter().map(|offset| root_leaves[offset]));
fields.push(field_with_type(field, projected_type));

@adriangb adriangb Aug 20, 2026

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.

continue;
}
}
RootRead::Full | RootRead::Partial(_) => {}
}
fields.insert(*root, Arc::new(file_schema.field(*root).clone()));
}

for (&root, (kept, pruned_type)) in &clipped_by_root {
leaf_indices.extend(kept.iter().copied());
fields.insert(
root,
field_with_type(file_schema.field(root), pruned_type.clone()),
);
// Full reads and unsupported/empty partial reads preserve the physical
// field. A root with no Parquet leaves contributes only its Arrow field.
leaf_indices.extend(root_leaves.iter().copied());
fields.push(Arc::new(field.clone()));
}

if !get_field_accesses.is_empty() {
let get_field_tree = StructAccessTree::from_accesses(&get_field_accesses);
leaf_indices.extend(resolve_struct_field_leaves(&get_field_tree, schema_descr));
let get_field_schema = build_filter_schema(file_schema, &[], &get_field_tree);
let get_field_roots: BTreeSet<usize> =
get_field_accesses.iter().map(|a| a.root_index).collect();
// `build_filter_schema` emits one field per accessed root in
// ascending root order, which is the order `get_field_roots` iterates
// in, so the two line up positionally. Pairing them beats looking each
// one up by name: no repeated linear scans, and no ambiguity if two
// roots happen to share a name.
debug_assert_eq!(get_field_roots.len(), get_field_schema.fields().len());
for (root, field) in get_field_roots.iter().zip(get_field_schema.fields()) {
fields.insert(*root, Arc::clone(field));
}
}

leaf_indices.sort_unstable();
leaf_indices.dedup();
// `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]));

@adriangb adriangb Aug 20, 2026

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?


ParquetReadPlan {
projection_mask: ProjectionMask::leaves(
schema_descr,
leaf_indices.iter().copied(),
),
projected_schema: Arc::new(Schema::new_with_metadata(
fields.into_values().collect::<Vec<_>>(),
fields,
file_schema.metadata().clone(),
)),
}
Expand Down