adriangb commented on code in PR #23217:
URL: https://github.com/apache/datafusion/pull/23217#discussion_r3732048293
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -762,62 +834,126 @@ where
.collect()
}
-/// Resolves struct field access to specific Parquet leaf column indices
+/// Returns the Parquet leaf column indices selected by the access tree.
+///
+/// # Matching
+///
+/// Iterates Parquet leaves in ascending order (`0..num_columns()`). For each
+/// leaf:
+///
+/// 1. **Root dispatch.** Look up the leaf's root index — the top-level Arrow
+/// column it belongs to — via `SchemaDescriptor::get_column_root_idx`. If
+/// that root is absent from the access tree (the filter never touched any
+/// field under it), skip the leaf without further work.
+///
+/// 2. **Path walk.** Otherwise, take the leaf's dotted column path
+/// (`col.path().parts()`), drop the first component (the root field name,
+/// already used in step 1), and walk the remaining components against the
+/// matching trie subtree via [`leaf_under_tree`].
///
-/// For every `StructFieldAccess`, finds the leaf columns in the Parquet schema
-/// whose path matches the struct root name + field path. This avoids reading
all
-/// leaves of a struct when only specific fields are needed
+/// 3. **Inclusion.** The leaf is added to the result iff the walk reaches a
+/// node with `selected_here = true` — either an ancestor along the
+/// descent (subsumption: a shallower access subsumes the leaf) or the
+/// terminal node reached at the end of the path (exact match).
+///
+/// # Returns
+///
+/// `Vec<usize>` of Parquet leaf column indices. The scan visits each leaf
+/// exactly once and pushes in iteration order, so the result is in ascending
+/// order and free of duplicates by construction — callers do not need to
+/// sort or dedup.
fn resolve_struct_field_leaves(
- accesses: &[StructFieldAccess],
- file_schema: &Schema,
+ access_tree: &StructAccessTree,
schema_descr: &SchemaDescriptor,
) -> Vec<usize> {
let mut leaf_indices = Vec::new();
- for access in accesses {
- let root_name = file_schema.field(access.root_index).name();
- let prefix = std::iter::once(root_name.as_str())
- .chain(access.field_path.iter().map(|p| p.as_str()))
- .collect::<Vec<_>>();
-
- for leaf_idx in 0..schema_descr.num_columns() {
- let col = schema_descr.column(leaf_idx);
- let col_path = col.path().parts();
-
- // A leaf matches if its path starts with our prefix.
- // e.g., prefix=["s", "value"] matches leaf path ["s", "value"]
- // prefix=["s", "outer"] matches ["s", "outer", "inner"]
- let leaf_matches_path = col_path.len() >= prefix.len()
- && col_path.iter().zip(prefix.iter()).all(|(a, b)| a == b);
-
- if leaf_matches_path {
- leaf_indices.push(leaf_idx);
- }
+ for leaf_idx in 0..schema_descr.num_columns() {
+ let root_idx = schema_descr.get_column_root_idx(leaf_idx);
+ let Some(root_node) = access_tree.roots.get(&root_idx) else {
+ continue;
+ };
+ // `parts()[0]` is the root field name; walk the rest against the tree.
+ let col = schema_descr.column(leaf_idx);
+ let path = col.path().parts();
+ if leaf_under_tree(root_node, &path[1..]) {
Review Comment:
could use `split_first()`?
##########
datafusion/datasource-parquet/src/projection_read_plan.rs:
##########
@@ -76,6 +76,83 @@ pub(crate) struct StructFieldAccess {
pub(crate) field_path: Vec<String>,
}
+/// Trie of nested struct accesses, keyed at the top by the root column index
in
+/// the file schema and then by field names down each access path.
+///
+/// # Example
+///
+/// For a filter expression
+///
+/// ```sql
+/// WHERE s['outer']['a'] > 10
+/// AND s['outer']['b'] < 20
+/// AND s['outer']['inner']['c'] IS NOT NULL
+/// ```
+///
+/// where `s` is column index `2` in the file schema, three accesses are
+/// recorded — all with `root_index = 2` and paths `["outer","a"]`,
+/// `["outer","b"]`, `["outer","inner","c"]`. They produce a trie in which
+/// the shared `"outer"` prefix is represented by a single intermediate node:
+///
+/// ```text
+/// roots:
+/// 2 ──► node { selected_here: false }
+/// children:
+/// "outer" ──► node { selected_here: false }
+/// children:
+/// "a" ──► { selected_here: true, children: {} }
+/// "b" ──► { selected_here: true, children: {} }
+/// "inner" ──► { selected_here: false,
+/// children: {
+/// "c" ──► { selected_here: true,
+/// children: {} }
+/// } }
+/// ```
+#[derive(Debug, Default)]
+struct StructAccessTree {
+ roots: BTreeMap<usize, StructAccessNode>,
+}
+
+/// One node in a [`StructAccessTree`].
+///
+/// `selected_here` is `true` when at least one access path terminates at this
+/// node. Duplicate paths are idempotent.
+#[derive(Debug, Default)]
+struct StructAccessNode {
+ children: BTreeMap<String, StructAccessNode>,
Review Comment:
Could these string keys be `'&str`?
--
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]
