You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "jayzhan211 (via GitHub)" <gi...@apache.org> on 2026/08/20 14:31:42 UTC
[PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
jayzhan211 opened a new pull request, #24522:
URL: https://github.com/apache/datafusion/pull/24522
## 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.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#discussion_r3830090084
##########
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();
Review Comment:
It would be useful to add a decoder-level regression test for a root accessed through both a cast and `get_field`. The test could install the generated projection mask in `ParquetRecordBatchReader`, verify that the emitted batch schema matches the read plan's projected schema, and then evaluate both expressions. The existing coverage checks the planned mask and schema, while this would also confirm that the decoder output is compatible with the resulting read plan.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#discussion_r3835076392
##########
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:
`assemble_read_plan` updated
##########
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:
remove the previous check
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#issuecomment-5358145649
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/24522?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `92.30769%` with `3 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 81.30%. Comparing base ([`40c208e`](https://app.codecov.io/gh/apache/datafusion/commit/40c208e5d65b931ef3cc3e121e85e228045a1a17?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`7819b01`](https://app.codecov.io/gh/apache/datafusion/commit/7819b01140e09dffe7f9b124a30fc9e79310afb9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 6 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/24522?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...ion/datasource-parquet/src/projection\_read\_plan.rs](https://app.codecov.io/gh/apache/datafusion/pull/24522?src=pr&el=tree&filepath=datafusion%2Fdatasource-parquet%2Fsrc%2Fprojection_read_plan.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9kYXRhc291cmNlLXBhcnF1ZXQvc3JjL3Byb2plY3Rpb25fcmVhZF9wbGFuLnJz) | 92.30% | [2 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/24522?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ 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
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/24522?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#discussion_r3835077323
##########
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:
remove this assertion and add another one
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#discussion_r3823809763
##########
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:
If it fixes a bug can we get a regression test?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#issuecomment-5379596387
Thanks @kosiew @adriangb @mbutrovich for the review
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #24522:
URL: https://github.com/apache/datafusion/pull/24522#discussion_r3835075296
##########
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:
Test added
##########
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();
Review Comment:
Test added
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: Simplify and optimize Parquet projection planning for nested casts [datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 merged PR #24522:
URL: https://github.com/apache/datafusion/pull/24522
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org