adriangb commented on code in PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#discussion_r3823593778
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -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.
+ 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()) =>
+ {
+ 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));
Review 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.
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -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.
Review 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.
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -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.
+ 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()) =>
Review 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.
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -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.
+ 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()) =>
+ {
+ 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));
+ 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]));
Review 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?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
