You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/09 16:55:32 UTC
(datafusion) 01/01: refactor: extract parquet projection read plan into its own module (#23396)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23396-4f63ada24056d9ec9b8decc04f128512bcdd735e
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit a3091beeb1a930b4b0fd91ef57987fec1548417c
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Thu Jul 9 16:55:09 2026 +0000
refactor: extract parquet projection read plan into its own module (#23396)
## Which issue does this PR close?
- N/A — pure code movement, split out of a larger change (see
rationale). First of a 3-PR stack: this refactor, #23397 (benchmark),
and the feature PR (nested schema pruning for the parquet reader,
upstream work for apache/datafusion-comet#4859).
## Rationale for this change
`row_filter.rs` is 2,100+ lines and contains two distinct concerns:
row-filter/`ArrowPredicate` construction, and the shared "expressions →
parquet leaf indices → `ProjectionMask` + projected schema" resolution
used by both the row filter and the opener's projection handling. This
PR moves the second concern into its own module,
`projection_read_plan.rs`, so upcoming changes to projection-mask
derivation can be reviewed without wading through the row-filter
machinery.
This is mergeable as standalone cleanup regardless of whether the
follow-up feature is accepted.
## What changes are included in this PR?
A move-only refactor (best reviewed with `git diff --color-moved`):
- `ParquetReadPlan`, `StructFieldAccess`, `build_projection_read_plan`,
`leaf_indices_for_roots`, `resolve_struct_field_leaves`,
`build_filter_schema`, `prune_struct_type` and the path-grouping helpers
move from `row_filter.rs` to a new `projection_read_plan.rs`, together
with their test.
- `PushdownChecker` / `PushdownColumns` become `pub(crate)` so the new
module can keep using them.
- `decoder_projection.rs` imports `build_projection_read_plan` from its
new home.
No behavior change; net +57 lines (module docs, imports, visibility).
## Are these changes tested?
Covered by existing tests (the moved unit test plus the
`datafusion-datasource-parquet` and core parquet integration suites,
which all pass unchanged).
## Are there any user-facing changes?
No.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_01KuMaRtFSPDQesuzjN5Koyd
---------
Co-authored-by: Claude Fable 5 <no...@anthropic.com>
---
.../datasource-parquet/src/decoder_projection.rs | 2 +-
datafusion/datasource-parquet/src/mod.rs | 1 +
.../datasource-parquet/src/projection_read_plan.rs | 745 +++++++++++++++++++++
datafusion/datasource-parquet/src/row_filter.rs | 717 +-------------------
4 files changed, 756 insertions(+), 709 deletions(-)
diff --git a/datafusion/datasource-parquet/src/decoder_projection.rs b/datafusion/datasource-parquet/src/decoder_projection.rs
index 192600ce7a..89fdc01af4 100644
--- a/datafusion/datasource-parquet/src/decoder_projection.rs
+++ b/datafusion/datasource-parquet/src/decoder_projection.rs
@@ -45,7 +45,7 @@ use parquet::arrow::ProjectionMask;
use parquet::schema::types::SchemaDescriptor;
use crate::opener::{VirtualColumnsState, append_fields};
-use crate::row_filter::build_projection_read_plan;
+use crate::projection_read_plan::build_projection_read_plan;
/// Per-file decoder projection: the [`ProjectionMask`] installed on the
/// parquet decoder, plus the per-batch transform that maps the decoder's
diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs
index e6e372cd78..25b79a6188 100644
--- a/datafusion/datasource-parquet/src/mod.rs
+++ b/datafusion/datasource-parquet/src/mod.rs
@@ -32,6 +32,7 @@ pub mod metadata;
mod metrics;
mod opener;
mod page_filter;
+mod projection_read_plan;
mod push_decoder;
mod reader;
mod row_filter;
diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs
new file mode 100644
index 0000000000..46caed6616
--- /dev/null
+++ b/datafusion/datasource-parquet/src/projection_read_plan.rs
@@ -0,0 +1,745 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Resolution of expressions against a Parquet file's schema into a
+//! [`ParquetReadPlan`]: the leaf-level [`ProjectionMask`] to install on the
+//! decoder plus the Arrow schema the decoder will emit under that mask.
+//!
+//! This is shared by the opener's projection handling (via
+//! [`build_projection_read_plan`]) and row-filter construction (via
+//! [`crate::row_filter`]), which both need to translate column and struct
+//! field references into Parquet leaf indices. [`PushdownChecker`], the
+//! expression traversal that discovers those references, lives here as well
+//! so that [`crate::row_filter`] depends on this module and not vice versa.
+
+use std::collections::{BTreeMap, BTreeSet};
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use parquet::arrow::ProjectionMask;
+use parquet::schema::types::SchemaDescriptor;
+
+use datafusion_common::Result;
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
+use datafusion_functions::core::file_row_index::FileRowIndexFunc;
+use datafusion_functions::core::getfield::GetFieldFunc;
+use datafusion_physical_expr::expressions::{Column, Literal};
+use datafusion_physical_expr::utils::collect_columns;
+use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr};
+
+/// The result of resolving which Parquet leaf columns and Arrow schema fields
+/// are needed to evaluate an expression against a Parquet file
+///
+/// This is the shared output of the column resolution pipeline used by both
+/// the row filter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s
+#[derive(Debug, Clone)]
+pub(crate) struct ParquetReadPlan {
+ /// Projection mask built from leaf column indices in the Parquet schema.
+ /// Using a `ProjectionMask` directly (rather than raw indices) prevents
+ /// bugs from accidentally mixing up root vs leaf indices.
+ pub projection_mask: ProjectionMask,
+ /// The projected Arrow schema containing only the columns/fields required
+ /// Struct types are pruned to include only the accessed sub-fields
+ pub projected_schema: SchemaRef,
+}
+
+/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`.
+///
+/// This allows the row filter to project only the specific Parquet leaf columns
+/// needed by the filter, rather than all leaves of the struct.
+#[derive(Debug, Clone)]
+pub(crate) struct StructFieldAccess {
+ /// Arrow root column index of the struct in the file schema.
+ pub(crate) root_index: usize,
+ /// Field names forming the path into the struct.
+ /// e.g., `["value"]` for `s['value']`, `["outer", "inner"]` for `s['outer']['inner']`.
+ pub(crate) field_path: Vec<String>,
+}
+
+/// Traverses a `PhysicalExpr` tree to determine if any column references would
+/// prevent the expression from being pushed down to the parquet decoder.
+///
+/// An expression cannot be pushed down if it references:
+/// - Unsupported nested columns (whole struct references or list fields that are
+/// not covered by the supported predicate set)
+/// - Columns that don't exist in the file schema
+///
+/// Struct field access via `get_field` is supported when the resolved leaf type
+/// is primitive (e.g. `get_field(struct_col, 'field') > 5`).
+pub(crate) struct PushdownChecker<'schema> {
+ /// Does the expression require any non-primitive columns (like structs)?
+ non_primitive_columns: bool,
+ /// Does the expression reference any columns not present in the file schema?
+ projected_columns: bool,
+ /// Does the expression references a ScalarUDF that requires some rewrite
+ /// and therefore can't be pushed down into the row-filter.
+ has_unpushable_udfs: bool,
+ /// Indices into the file schema of columns required to evaluate the expression.
+ /// Does not include struct columns accessed via `get_field`.
+ required_columns: Vec<usize>,
+ /// Struct field accesses via `get_field`.
+ struct_field_accesses: Vec<StructFieldAccess>,
+ /// Whether nested list columns are supported by the predicate semantics.
+ allow_list_columns: bool,
+ /// The Arrow schema of the parquet file.
+ file_schema: &'schema Schema,
+}
+
+impl<'schema> PushdownChecker<'schema> {
+ pub(crate) fn new(file_schema: &'schema Schema, allow_list_columns: bool) -> Self {
+ Self {
+ non_primitive_columns: false,
+ projected_columns: false,
+ has_unpushable_udfs: false,
+ required_columns: Vec::new(),
+ struct_field_accesses: Vec::new(),
+ allow_list_columns,
+ file_schema,
+ }
+ }
+
+ /// Checks whether a struct's root column exists in the file schema and, if so,
+ /// records its index so the entire struct is decoded for filter evaluation.
+ ///
+ /// This is called when we see a `get_field` expression that resolves to a
+ /// primitive leaf type. We only need the *root* column index because the
+ /// Parquet reader decodes all leaves of a struct together.
+ ///
+ /// # Example
+ ///
+ /// Given file schema `{a: Int32, s: Struct(foo: Utf8, bar: Int64)}` and the
+ /// expression `get_field(s, 'foo') = 'hello'`:
+ ///
+ /// - `column_name` = `"s"` (the root struct column)
+ /// - `file_schema.index_of("s")` returns `1`
+ /// - We push `1` into `required_columns`
+ /// - Return `None` (no issue — traversal continues in the caller)
+ ///
+ /// If `"s"` is not in the file schema (e.g. a projected-away column), we set
+ /// `projected_columns = true` and return `Jump` to skip the subtree.
+ fn check_struct_field_column(
+ &mut self,
+ column_name: &str,
+ field_path: Vec<String>,
+ ) -> Option<TreeNodeRecursion> {
+ let Ok(idx) = self.file_schema.index_of(column_name) else {
+ self.projected_columns = true;
+ return Some(TreeNodeRecursion::Jump);
+ };
+
+ self.struct_field_accesses.push(StructFieldAccess {
+ root_index: idx,
+ field_path,
+ });
+
+ None
+ }
+
+ fn check_single_column(&mut self, column_name: &str) -> Option<TreeNodeRecursion> {
+ let idx = match self.file_schema.index_of(column_name) {
+ Ok(idx) => idx,
+ Err(_) => {
+ // Column does not exist in the file schema, so we can't push this down.
+ self.projected_columns = true;
+ return Some(TreeNodeRecursion::Jump);
+ }
+ };
+
+ // Duplicates are handled by dedup() in into_sorted_columns()
+ self.required_columns.push(idx);
+ let data_type = self.file_schema.field(idx).data_type();
+
+ if DataType::is_nested(data_type) {
+ self.handle_nested_type(data_type)
+ } else {
+ None
+ }
+ }
+
+ /// Determines whether a nested data type can be pushed down to Parquet decoding.
+ ///
+ /// Returns `Some(TreeNodeRecursion::Jump)` if the nested type prevents pushdown,
+ /// `None` if the type is supported and pushdown can continue.
+ fn handle_nested_type(&mut self, data_type: &DataType) -> Option<TreeNodeRecursion> {
+ if self.is_nested_type_supported(data_type) {
+ None
+ } else {
+ // Block pushdown for unsupported nested types:
+ // - Structs (regardless of predicate support)
+ // - Lists without supported predicates
+ self.non_primitive_columns = true;
+ Some(TreeNodeRecursion::Jump)
+ }
+ }
+
+ /// Checks if a nested data type is supported for list column pushdown.
+ ///
+ /// List columns are only supported if:
+ /// 1. The data type is a list variant (List, LargeList, or FixedSizeList)
+ /// 2. The expression contains supported list predicates (e.g., array_has_all)
+ fn is_nested_type_supported(&self, data_type: &DataType) -> bool {
+ let is_list = matches!(
+ data_type,
+ DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _)
+ );
+ self.allow_list_columns && is_list
+ }
+
+ #[inline]
+ pub(crate) fn prevents_pushdown(&self) -> bool {
+ self.non_primitive_columns || self.projected_columns || self.has_unpushable_udfs
+ }
+
+ /// Consumes the checker and returns sorted, deduplicated column indices
+ /// wrapped in a `PushdownColumns` struct.
+ ///
+ /// This method sorts the column indices and removes duplicates. The sort
+ /// is required because downstream code relies on column indices being in
+ /// ascending order for correct schema projection.
+ pub(crate) fn into_sorted_columns(mut self) -> PushdownColumns {
+ self.required_columns.sort_unstable();
+ self.required_columns.dedup();
+ PushdownColumns {
+ required_columns: self.required_columns,
+ struct_field_accesses: self.struct_field_accesses,
+ }
+ }
+}
+
+impl TreeNodeVisitor<'_> for PushdownChecker<'_> {
+ type Node = Arc<dyn PhysicalExpr>;
+
+ fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
+ // Handle struct field access like `s['foo']['bar'] > 10`.
+ //
+ // DataFusion represents nested field access as `get_field(Column("s"), "foo")`
+ // (or chained: `get_field(get_field(Column("s"), "foo"), "bar")`).
+ //
+ // We intercept the outermost `get_field` on the way *down* the tree so
+ // the visitor never reaches the raw `Column("s")` node. Without this,
+ // `check_single_column` would see that `s` is a Struct and reject it.
+ //
+ // The strategy:
+ // 1. Match `get_field` whose first arg is a `Column` (the struct root).
+ // 2. Check that the *resolved* return type is primitive — meaning we've
+ // drilled all the way to a leaf (e.g. `s['foo']` → Utf8).
+ // 3. Record the root column index via `check_struct_field_column` and
+ // return `Jump` to skip visiting the children (the Column and the
+ // literal field-name args), since we've already handled them.
+ //
+ // If the return type is still nested (e.g. `s['nested_struct']` → Struct),
+ // we fall through and let normal traversal continue, which will
+ // eventually reject the expression when it hits the struct Column.
+ if let Some(func) =
+ ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(node.as_ref())
+ {
+ let args = func.args();
+
+ if let Some(column) = args.first().and_then(|a| a.downcast_ref::<Column>()) {
+ // for Map columns, get_field performs a runtime key lookup rather than a
+ // schema-level field access so the entire Map column must be read,
+ // we skip the struct field optimization and defer to normal Column traversal
+ let is_map_column = self
+ .file_schema
+ .index_of(column.name())
+ .ok()
+ .map(|idx| {
+ matches!(
+ self.file_schema.field(idx).data_type(),
+ DataType::Map(_, _)
+ )
+ })
+ .unwrap_or(false);
+
+ let return_type = func.return_type();
+
+ if !is_map_column
+ && (!DataType::is_nested(return_type)
+ || self.is_nested_type_supported(return_type))
+ {
+ // if any field name argument is not a string literal we cannot
+ // determine the exact leaf path, so we fall back to reading the
+ // entire struct root column
+ let field_path = args[1..]
+ .iter()
+ .map(|arg| {
+ arg.downcast_ref::<Literal>().and_then(|lit| {
+ lit.value().try_as_str().flatten().map(|s| s.to_string())
+ })
+ })
+ .collect();
+
+ match field_path {
+ Some(path) => {
+ if let Some(recursion) =
+ self.check_struct_field_column(column.name(), path)
+ {
+ return Ok(recursion);
+ }
+ }
+ None => {
+ // Could not resolve field path — fall back to
+ // reading the entire struct root column.
+ if let Some(recursion) =
+ self.check_single_column(column.name())
+ {
+ return Ok(recursion);
+ }
+ }
+ }
+
+ return Ok(TreeNodeRecursion::Jump);
+ }
+ }
+ }
+
+ if let Some(column) = node.downcast_ref::<Column>()
+ && let Some(recursion) = self.check_single_column(column.name())
+ {
+ return Ok(recursion);
+ }
+
+ if ScalarFunctionExpr::try_downcast_func::<FileRowIndexFunc>(node.as_ref())
+ .is_some()
+ {
+ self.has_unpushable_udfs = true;
+ return Ok(TreeNodeRecursion::Jump);
+ }
+
+ Ok(TreeNodeRecursion::Continue)
+ }
+}
+
+/// Result of checking which columns are required for filter pushdown.
+#[derive(Debug)]
+pub(crate) struct PushdownColumns {
+ /// Sorted, unique column indices into the file schema required to evaluate
+ /// the filter expression. Must be in ascending order for correct schema
+ /// projection matching. Does not include struct columns accessed via `get_field`.
+ pub(crate) required_columns: Vec<usize>,
+ /// Struct field accesses via `get_field`. Each entry records the root struct
+ /// column index and the field path being accessed.
+ pub(crate) struct_field_accesses: Vec<StructFieldAccess>,
+}
+
+/// Builds a unified [`ParquetReadPlan`] for a set of projection expressions
+///
+/// Unlike [`crate::row_filter::build_parquet_read_plan`] (which is used for
+/// filter pushdown and returns `None` when an expression references
+/// unsupported nested types or missing columns), this function always
+/// succeeds. It collects every column that *can* be resolved in the file and
+/// produces a leaf-level projection mask. Columns missing from the file are
+/// silently skipped since the projection layer handles those by inserting
+/// nulls.
+pub(crate) fn build_projection_read_plan(
+ exprs: impl IntoIterator<Item = Arc<dyn PhysicalExpr>>,
+ file_schema: &Schema,
+ schema_descr: &SchemaDescriptor,
+) -> ParquetReadPlan {
+ // fast path: if every expression is a plain Column reference, skip all
+ // struct analysis and use root-level projection directly
+ let exprs = exprs.into_iter().collect::<Vec<_>>();
+ let all_plain_columns = exprs.iter().all(|e| e.downcast_ref::<Column>().is_some());
+
+ if all_plain_columns {
+ let mut root_indices: Vec<usize> = exprs
+ .iter()
+ .map(|e| e.downcast_ref::<Column>().unwrap().index())
+ .collect();
+ root_indices.sort_unstable();
+ root_indices.dedup();
+
+ let projection_mask =
+ ProjectionMask::roots(schema_descr, root_indices.iter().copied());
+ let projected_schema = Arc::new(
+ file_schema
+ .project(&root_indices)
+ .expect("valid column indices"),
+ );
+
+ return ParquetReadPlan {
+ projection_mask,
+ projected_schema,
+ };
+ }
+
+ // secondary fast path: if the schema has no struct columns, we can skip
+ // PushdownChecker traversal and use root-level projection
+ let has_struct_columns = file_schema
+ .fields()
+ .iter()
+ .any(|f| matches!(f.data_type(), DataType::Struct(_)));
+
+ if !has_struct_columns {
+ let mut root_indices = exprs
+ .into_iter()
+ .flat_map(|e| collect_columns(&e).into_iter().map(|col| col.index()))
+ .collect::<Vec<_>>();
+
+ root_indices.sort_unstable();
+ root_indices.dedup();
+
+ let projection_mask =
+ ProjectionMask::roots(schema_descr, root_indices.iter().copied());
+
+ let projected_schema = Arc::new(
+ file_schema
+ .project(&root_indices)
+ .expect("valid column indices"),
+ );
+
+ return ParquetReadPlan {
+ projection_mask,
+ projected_schema,
+ };
+ }
+
+ let mut all_root_indices = Vec::new();
+ let mut all_struct_accesses = Vec::new();
+
+ for expr in exprs {
+ let mut checker = PushdownChecker::new(file_schema, true);
+ let _ = expr.visit(&mut checker);
+ let columns = checker.into_sorted_columns();
+
+ all_root_indices.extend_from_slice(&columns.required_columns);
+ all_struct_accesses.extend(columns.struct_field_accesses);
+ }
+
+ all_root_indices.sort_unstable();
+ all_root_indices.dedup();
+
+ // when no struct field accesses were found, fall back to root-level projection
+ // to match the performance of the simple path
+ if all_struct_accesses.is_empty() {
+ let projection_mask =
+ ProjectionMask::roots(schema_descr, all_root_indices.iter().copied());
+ let projected_schema = Arc::new(
+ file_schema
+ .project(&all_root_indices)
+ .expect("valid column indices"),
+ );
+
+ return ParquetReadPlan {
+ projection_mask,
+ projected_schema,
+ };
+ }
+
+ let leaf_indices = {
+ let mut out =
+ leaf_indices_for_roots(all_root_indices.iter().copied(), schema_descr);
+ let struct_leaf_indices =
+ resolve_struct_field_leaves(&all_struct_accesses, file_schema, schema_descr);
+
+ out.extend_from_slice(&struct_leaf_indices);
+ out.sort_unstable();
+ out.dedup();
+
+ out
+ };
+
+ let projection_mask =
+ ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied());
+
+ let projected_schema =
+ build_filter_schema(file_schema, &all_root_indices, &all_struct_accesses);
+
+ ParquetReadPlan {
+ projection_mask,
+ projected_schema,
+ }
+}
+
+pub(crate) fn leaf_indices_for_roots<I>(
+ root_indices: I,
+ schema_descr: &SchemaDescriptor,
+) -> Vec<usize>
+where
+ I: IntoIterator<Item = usize>,
+{
+ // Always map root (Arrow) indices to Parquet leaf indices via the schema
+ // descriptor. Arrow root indices only equal Parquet leaf indices when the
+ // schema has no group columns (Struct, Map, etc.); when group columns
+ // exist, their children become separate leaves and shift all subsequent
+ // leaf indices.
+ let root_set: BTreeSet<_> = root_indices.into_iter().collect();
+
+ (0..schema_descr.num_columns())
+ .filter(|leaf_idx| {
+ root_set.contains(&schema_descr.get_column_root_idx(*leaf_idx))
+ })
+ .collect()
+}
+
+/// Resolves struct field access to specific Parquet leaf column indices
+///
+/// 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
+pub(crate) fn resolve_struct_field_leaves(
+ accesses: &[StructFieldAccess],
+ file_schema: &Schema,
+ 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);
+ }
+ }
+ }
+
+ leaf_indices
+}
+
+/// Builds a filter schema that includes only the fields actually accessed by the
+/// filter expression.
+///
+/// For regular (non-struct) columns, the full field type is used.
+/// For struct columns accessed via `get_field`, a pruned struct type is created
+/// containing only the fields along the access path. Note: it must match the schema
+/// that the Parquet reader produces when projecting specific struct leaves
+pub(crate) fn build_filter_schema(
+ file_schema: &Schema,
+ regular_indices: &[usize],
+ struct_field_accesses: &[StructFieldAccess],
+) -> SchemaRef {
+ let regular_set: BTreeSet<usize> = regular_indices.iter().copied().collect();
+ let paths_by_root = group_access_paths_by_root(struct_field_accesses);
+
+ let all_indices = regular_indices
+ .iter()
+ .copied()
+ .chain(paths_by_root.keys().copied())
+ .collect::<BTreeSet<_>>();
+
+ let fields = all_indices
+ .iter()
+ .map(|&idx| {
+ let field = file_schema.field(idx);
+
+ // if this column appears as a regular (whole-column) reference,
+ // keep the full type
+ //
+ // Pruning is only valid when the column is accessed exclusively
+ // through struct field accesses
+ if regular_set.contains(&idx) {
+ return Arc::new(field.clone());
+ }
+
+ let Some(field_paths) = paths_by_root.get(&idx) else {
+ return Arc::new(field.clone());
+ };
+
+ let pruned_data_type = prune_struct_type(field.data_type(), field_paths);
+ Arc::new(Field::new(
+ field.name(),
+ pruned_data_type,
+ field.is_nullable(),
+ ))
+ })
+ .collect::<Vec<_>>();
+
+ Arc::new(Schema::new_with_metadata(
+ fields,
+ file_schema.metadata().clone(),
+ ))
+}
+
+/// Groups struct field access paths once for the root schema level.
+///
+/// Each map entry contains the complete field paths accessed below a root
+/// column. Recursive pruning groups these paths by their next component at each
+/// nested struct level.
+fn group_access_paths_by_root(
+ struct_field_accesses: &[StructFieldAccess],
+) -> BTreeMap<usize, Vec<&[String]>> {
+ let mut paths_by_root: BTreeMap<usize, Vec<&[String]>> = BTreeMap::new();
+ for StructFieldAccess {
+ root_index,
+ field_path,
+ } in struct_field_accesses
+ {
+ paths_by_root
+ .entry(*root_index)
+ .or_default()
+ .push(field_path.as_slice());
+ }
+
+ paths_by_root
+}
+
+/// Groups access paths once for the current struct level.
+///
+/// The map key is the field name at this level. The map value is the list of
+/// remaining path suffixes below that field. An empty suffix means the access
+/// path terminates at that field, so the full field must be preserved.
+fn group_paths_by_next_field<'a>(
+ paths: &'a [&'a [String]],
+) -> BTreeMap<&'a str, Vec<&'a [String]>> {
+ let mut paths_by_field: BTreeMap<&str, Vec<&[String]>> = BTreeMap::new();
+ for path in paths {
+ if let Some((field, sub_path)) = path.split_first() {
+ paths_by_field
+ .entry(field.as_str())
+ .or_default()
+ .push(sub_path);
+ }
+ }
+
+ paths_by_field
+}
+
+fn prune_struct_type(dt: &DataType, paths: &[&[String]]) -> DataType {
+ let DataType::Struct(fields) = dt else {
+ return dt.clone();
+ };
+
+ let paths_by_field = group_paths_by_next_field(paths);
+
+ let pruned_fields = fields
+ .iter()
+ .filter_map(|f| {
+ let sub_paths = paths_by_field.get(f.name().as_str())?;
+
+ let out = if sub_paths.iter().any(|sub| sub.is_empty()) {
+ // Leaf of access path — keep the field as-is.
+ Arc::clone(f)
+ } else {
+ // Recurse into nested struct.
+ let pruned = prune_struct_type(f.data_type(), sub_paths);
+ Arc::new(Field::new(f.name(), pruned, f.is_nullable()))
+ };
+
+ Some(out)
+ })
+ .collect::<Vec<_>>();
+
+ DataType::Struct(pruned_fields.into())
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+ use Column as PhysicalColumn;
+ use arrow::array::{Int32Array, RecordBatch, StringArray, StructArray};
+ use arrow::datatypes::Fields;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{Expr, col};
+ use datafusion_functions::core::get_field;
+ use datafusion_physical_expr::planner::logical2physical;
+ use parquet::arrow::ArrowWriter;
+ use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+ use tempfile::NamedTempFile;
+
+ #[test]
+ fn projection_read_plan_preserves_full_struct() {
+ // Schema: id (Int32), s (Struct{value: Int32, label: Utf8})
+ // Parquet leaves: id=0, s.value=1, s.label=2
+ let struct_fields: Fields = vec![
+ Arc::new(Field::new("value", DataType::Int32, false)),
+ Arc::new(Field::new("label", DataType::Utf8, false)),
+ ]
+ .into();
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("s", DataType::Struct(struct_fields.clone()), false),
+ ]));
+
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2, 3])),
+ Arc::new(StructArray::new(
+ struct_fields,
+ vec![
+ Arc::new(Int32Array::from(vec![10, 20, 30])) as _,
+ Arc::new(StringArray::from(vec!["a", "b", "c"])) as _,
+ ],
+ None,
+ )),
+ ],
+ )
+ .unwrap();
+
+ let file = NamedTempFile::new().expect("temp file");
+ let mut writer =
+ ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None)
+ .expect("writer");
+ writer.write(&batch).expect("write batch");
+ writer.close().expect("close writer");
+
+ let reader_file = file.reopen().expect("reopen file");
+ let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file)
+ .expect("reader builder");
+ let metadata = builder.metadata().clone();
+ let file_schema = builder.schema().clone();
+ let schema_descr = metadata.file_metadata().schema_descr();
+
+ // Simulate SELECT * output projection: Column("id") and Column("s")
+ // Plus a get_field(s, 'value') expression from the pushed-down filter
+ let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
+ Arc::new(PhysicalColumn::new("id", 0)),
+ Arc::new(PhysicalColumn::new("s", 1)),
+ logical2physical(
+ &get_field().call(vec![
+ col("s"),
+ Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None),
+ ]),
+ &file_schema,
+ ),
+ ];
+
+ let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr);
+
+ // The projected schema must have the FULL struct type because Column("s")
+ // is in the projection. It should NOT be narrowed to Struct{value: Int32}.
+ let s_field = read_plan.projected_schema.field_with_name("s").unwrap();
+ assert_eq!(
+ s_field.data_type(),
+ &DataType::Struct(
+ vec![
+ Arc::new(Field::new("value", DataType::Int32, false)),
+ Arc::new(Field::new("label", DataType::Utf8, false)),
+ ]
+ .into()
+ ),
+ );
+
+ // all 3 Parquet leaves should be in the projection mask
+ let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]);
+ assert_eq!(read_plan.projection_mask, expected_mask,);
+ }
+}
diff --git a/datafusion/datasource-parquet/src/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs
index ef0478f315..4505f7fd62 100644
--- a/datafusion/datasource-parquet/src/row_filter.rs
+++ b/datafusion/datasource-parquet/src/row_filter.rs
@@ -65,32 +65,30 @@
//! - `WHERE s['value'] > 5` — pushed down (accesses a primitive leaf)
//! - `WHERE s IS NOT NULL` — not pushed down (references the whole struct)
-use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use arrow::array::BooleanArray;
-use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::RecordBatch;
-use datafusion_functions::core::file_row_index::FileRowIndexFunc;
-use datafusion_functions::core::getfield::GetFieldFunc;
use parquet::arrow::ProjectionMask;
use parquet::arrow::arrow_reader::{ArrowPredicate, RowFilter};
use parquet::file::metadata::ParquetMetaData;
-use parquet::schema::types::SchemaDescriptor;
use datafusion_common::Result;
use datafusion_common::cast::as_boolean_array;
-use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
-use datafusion_physical_expr::ScalarFunctionExpr;
-use datafusion_physical_expr::expressions::{Column, Literal};
-use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns};
+use datafusion_common::tree_node::TreeNode;
+use datafusion_physical_expr::utils::reassign_expr_columns;
use datafusion_physical_expr::{PhysicalExpr, split_conjunction};
use datafusion_physical_plan::metrics;
use super::ParquetFileMetrics;
use super::supported_predicates::supports_list_predicates;
+use crate::projection_read_plan::{
+ ParquetReadPlan, PushdownChecker, PushdownColumns, build_filter_schema,
+ leaf_indices_for_roots, resolve_struct_field_leaves,
+};
/// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform
/// row-level filtering during parquet decoding.
@@ -189,22 +187,6 @@ pub(crate) struct FilterCandidate {
read_plan: ParquetReadPlan,
}
-/// The result of resolving which Parquet leaf columns and Arrow schema fields
-/// are needed to evaluate an expression against a Parquet file
-///
-/// This is the shared output of the column resolution pipeline used by both
-/// the row filter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s
-#[derive(Debug, Clone)]
-pub(crate) struct ParquetReadPlan {
- /// Projection mask built from leaf column indices in the Parquet schema.
- /// Using a `ProjectionMask` directly (rather than raw indices) prevents
- /// bugs from accidentally mixing up root vs leaf indices.
- pub projection_mask: ProjectionMask,
- /// The projected Arrow schema containing only the columns/fields required
- /// Struct types are pruned to include only the accessed sub-fields
- pub projected_schema: SchemaRef,
-}
-
/// Helper to build a `FilterCandidate`.
///
/// This will do several things:
@@ -246,289 +228,6 @@ impl FilterCandidateBuilder {
}
}
-/// Traverses a `PhysicalExpr` tree to determine if any column references would
-/// prevent the expression from being pushed down to the parquet decoder.
-///
-/// An expression cannot be pushed down if it references:
-/// - Unsupported nested columns (whole struct references or list fields that are
-/// not covered by the supported predicate set)
-/// - Columns that don't exist in the file schema
-///
-/// Struct field access via `get_field` is supported when the resolved leaf type
-/// is primitive (e.g. `get_field(struct_col, 'field') > 5`).
-struct PushdownChecker<'schema> {
- /// Does the expression require any non-primitive columns (like structs)?
- non_primitive_columns: bool,
- /// Does the expression reference any columns not present in the file schema?
- projected_columns: bool,
- /// Does the expression references a ScalarUDF that requires some rewrite
- /// and therefore can't be pushed down into the row-filter.
- has_unpushable_udfs: bool,
- /// Indices into the file schema of columns required to evaluate the expression.
- /// Does not include struct columns accessed via `get_field`.
- required_columns: Vec<usize>,
- /// Struct field accesses via `get_field`.
- struct_field_accesses: Vec<StructFieldAccess>,
- /// Whether nested list columns are supported by the predicate semantics.
- allow_list_columns: bool,
- /// The Arrow schema of the parquet file.
- file_schema: &'schema Schema,
-}
-
-impl<'schema> PushdownChecker<'schema> {
- fn new(file_schema: &'schema Schema, allow_list_columns: bool) -> Self {
- Self {
- non_primitive_columns: false,
- projected_columns: false,
- has_unpushable_udfs: false,
- required_columns: Vec::new(),
- struct_field_accesses: Vec::new(),
- allow_list_columns,
- file_schema,
- }
- }
-
- /// Checks whether a struct's root column exists in the file schema and, if so,
- /// records its index so the entire struct is decoded for filter evaluation.
- ///
- /// This is called when we see a `get_field` expression that resolves to a
- /// primitive leaf type. We only need the *root* column index because the
- /// Parquet reader decodes all leaves of a struct together.
- ///
- /// # Example
- ///
- /// Given file schema `{a: Int32, s: Struct(foo: Utf8, bar: Int64)}` and the
- /// expression `get_field(s, 'foo') = 'hello'`:
- ///
- /// - `column_name` = `"s"` (the root struct column)
- /// - `file_schema.index_of("s")` returns `1`
- /// - We push `1` into `required_columns`
- /// - Return `None` (no issue — traversal continues in the caller)
- ///
- /// If `"s"` is not in the file schema (e.g. a projected-away column), we set
- /// `projected_columns = true` and return `Jump` to skip the subtree.
- fn check_struct_field_column(
- &mut self,
- column_name: &str,
- field_path: Vec<String>,
- ) -> Option<TreeNodeRecursion> {
- let Ok(idx) = self.file_schema.index_of(column_name) else {
- self.projected_columns = true;
- return Some(TreeNodeRecursion::Jump);
- };
-
- self.struct_field_accesses.push(StructFieldAccess {
- root_index: idx,
- field_path,
- });
-
- None
- }
-
- fn check_single_column(&mut self, column_name: &str) -> Option<TreeNodeRecursion> {
- let idx = match self.file_schema.index_of(column_name) {
- Ok(idx) => idx,
- Err(_) => {
- // Column does not exist in the file schema, so we can't push this down.
- self.projected_columns = true;
- return Some(TreeNodeRecursion::Jump);
- }
- };
-
- // Duplicates are handled by dedup() in into_sorted_columns()
- self.required_columns.push(idx);
- let data_type = self.file_schema.field(idx).data_type();
-
- if DataType::is_nested(data_type) {
- self.handle_nested_type(data_type)
- } else {
- None
- }
- }
-
- /// Determines whether a nested data type can be pushed down to Parquet decoding.
- ///
- /// Returns `Some(TreeNodeRecursion::Jump)` if the nested type prevents pushdown,
- /// `None` if the type is supported and pushdown can continue.
- fn handle_nested_type(&mut self, data_type: &DataType) -> Option<TreeNodeRecursion> {
- if self.is_nested_type_supported(data_type) {
- None
- } else {
- // Block pushdown for unsupported nested types:
- // - Structs (regardless of predicate support)
- // - Lists without supported predicates
- self.non_primitive_columns = true;
- Some(TreeNodeRecursion::Jump)
- }
- }
-
- /// Checks if a nested data type is supported for list column pushdown.
- ///
- /// List columns are only supported if:
- /// 1. The data type is a list variant (List, LargeList, or FixedSizeList)
- /// 2. The expression contains supported list predicates (e.g., array_has_all)
- fn is_nested_type_supported(&self, data_type: &DataType) -> bool {
- let is_list = matches!(
- data_type,
- DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _)
- );
- self.allow_list_columns && is_list
- }
-
- #[inline]
- fn prevents_pushdown(&self) -> bool {
- self.non_primitive_columns || self.projected_columns || self.has_unpushable_udfs
- }
-
- /// Consumes the checker and returns sorted, deduplicated column indices
- /// wrapped in a `PushdownColumns` struct.
- ///
- /// This method sorts the column indices and removes duplicates. The sort
- /// is required because downstream code relies on column indices being in
- /// ascending order for correct schema projection.
- fn into_sorted_columns(mut self) -> PushdownColumns {
- self.required_columns.sort_unstable();
- self.required_columns.dedup();
- PushdownColumns {
- required_columns: self.required_columns,
- struct_field_accesses: self.struct_field_accesses,
- }
- }
-}
-
-impl TreeNodeVisitor<'_> for PushdownChecker<'_> {
- type Node = Arc<dyn PhysicalExpr>;
-
- fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
- // Handle struct field access like `s['foo']['bar'] > 10`.
- //
- // DataFusion represents nested field access as `get_field(Column("s"), "foo")`
- // (or chained: `get_field(get_field(Column("s"), "foo"), "bar")`).
- //
- // We intercept the outermost `get_field` on the way *down* the tree so
- // the visitor never reaches the raw `Column("s")` node. Without this,
- // `check_single_column` would see that `s` is a Struct and reject it.
- //
- // The strategy:
- // 1. Match `get_field` whose first arg is a `Column` (the struct root).
- // 2. Check that the *resolved* return type is primitive — meaning we've
- // drilled all the way to a leaf (e.g. `s['foo']` → Utf8).
- // 3. Record the root column index via `check_struct_field_column` and
- // return `Jump` to skip visiting the children (the Column and the
- // literal field-name args), since we've already handled them.
- //
- // If the return type is still nested (e.g. `s['nested_struct']` → Struct),
- // we fall through and let normal traversal continue, which will
- // eventually reject the expression when it hits the struct Column.
- if let Some(func) =
- ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(node.as_ref())
- {
- let args = func.args();
-
- if let Some(column) = args.first().and_then(|a| a.downcast_ref::<Column>()) {
- // for Map columns, get_field performs a runtime key lookup rather than a
- // schema-level field access so the entire Map column must be read,
- // we skip the struct field optimization and defer to normal Column traversal
- let is_map_column = self
- .file_schema
- .index_of(column.name())
- .ok()
- .map(|idx| {
- matches!(
- self.file_schema.field(idx).data_type(),
- DataType::Map(_, _)
- )
- })
- .unwrap_or(false);
-
- let return_type = func.return_type();
-
- if !is_map_column
- && (!DataType::is_nested(return_type)
- || self.is_nested_type_supported(return_type))
- {
- // try to resolve all field name arguments to strinrg literals
- // if any argument is not a string literal, we can not determine the exact
- // leaf path so we fall back to reading the entire struct root column
- let field_path = args[1..]
- .iter()
- .map(|arg| {
- arg.downcast_ref::<Literal>().and_then(|lit| {
- lit.value().try_as_str().flatten().map(|s| s.to_string())
- })
- })
- .collect();
-
- match field_path {
- Some(path) => {
- if let Some(recursion) =
- self.check_struct_field_column(column.name(), path)
- {
- return Ok(recursion);
- }
- }
- None => {
- // Could not resolve field path — fall back to
- // reading the entire struct root column.
- if let Some(recursion) =
- self.check_single_column(column.name())
- {
- return Ok(recursion);
- }
- }
- }
-
- return Ok(TreeNodeRecursion::Jump);
- }
- }
- }
-
- if let Some(column) = node.downcast_ref::<Column>()
- && let Some(recursion) = self.check_single_column(column.name())
- {
- return Ok(recursion);
- }
-
- if ScalarFunctionExpr::try_downcast_func::<FileRowIndexFunc>(node.as_ref())
- .is_some()
- {
- self.has_unpushable_udfs = true;
- return Ok(TreeNodeRecursion::Jump);
- }
-
- Ok(TreeNodeRecursion::Continue)
- }
-}
-
-/// Describes the nested column behavior for filter pushdown.
-///
-/// This enum makes explicit the different states a predicate can be in
-/// with respect to nested column handling during Parquet decoding.
-/// Result of checking which columns are required for filter pushdown.
-#[derive(Debug)]
-struct PushdownColumns {
- /// Sorted, unique column indices into the file schema required to evaluate
- /// the filter expression. Must be in ascending order for correct schema
- /// projection matching. Does not include struct columns accessed via `get_field`.
- required_columns: Vec<usize>,
- /// Struct field accesses via `get_field`. Each entry records the root struct
- /// column index and the field path being accessed.
- struct_field_accesses: Vec<StructFieldAccess>,
-}
-
-/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`.
-///
-/// This allows the row filter to project only the specific Parquet leaf columns
-/// needed by the filter, rather than all leaves of the struct.
-#[derive(Debug, Clone)]
-struct StructFieldAccess {
- /// Arrow root column index of the struct in the file schema.
- root_index: usize,
- /// Field names forming the path into the struct.
- /// e.g., `["value"]` for `s['value']`, `["outer", "inner"]` for `s['outer']['inner']`.
- field_path: Vec<String>,
-}
-
/// Checks if a given expression can be pushed down to the parquet decoder.
///
/// Returns `Some(PushdownColumns)` if the expression can be pushed down,
@@ -604,323 +303,6 @@ pub(crate) fn build_parquet_read_plan(
)))
}
-/// Builds a unified [`ParquetReadPlan`] for a set of projection expressions
-///
-/// Unlike [`build_parquet_read_plan`] (which is used for filter pushdown and
-/// returns `None` when an expression references unsupported nested types or
-/// missing columns), this function always succeeds. It collects every column
-/// that *can* be resolved in the file and produces a leaf-level projection
-/// mask. Columns missing from the file are silently skipped since the projection
-/// layer handles those by inserting nulls.
-pub(crate) fn build_projection_read_plan(
- exprs: impl IntoIterator<Item = Arc<dyn PhysicalExpr>>,
- file_schema: &Schema,
- schema_descr: &SchemaDescriptor,
-) -> ParquetReadPlan {
- // fast path: if every expression is a plain Column reference, skip all
- // struct analysis and use root-level projection directly
- let exprs = exprs.into_iter().collect::<Vec<_>>();
- let all_plain_columns = exprs.iter().all(|e| e.downcast_ref::<Column>().is_some());
-
- if all_plain_columns {
- let mut root_indices: Vec<usize> = exprs
- .iter()
- .map(|e| e.downcast_ref::<Column>().unwrap().index())
- .collect();
- root_indices.sort_unstable();
- root_indices.dedup();
-
- let projection_mask =
- ProjectionMask::roots(schema_descr, root_indices.iter().copied());
- let projected_schema = Arc::new(
- file_schema
- .project(&root_indices)
- .expect("valid column indices"),
- );
-
- return ParquetReadPlan {
- projection_mask,
- projected_schema,
- };
- }
-
- // secondary fast path: if the schema has no struct columns, we can skip
- // PushdownChecker traversal and use root-level projection
- let has_struct_columns = file_schema
- .fields()
- .iter()
- .any(|f| matches!(f.data_type(), DataType::Struct(_)));
-
- if !has_struct_columns {
- let mut root_indices = exprs
- .into_iter()
- .flat_map(|e| collect_columns(&e).into_iter().map(|col| col.index()))
- .collect::<Vec<_>>();
-
- root_indices.sort_unstable();
- root_indices.dedup();
-
- let projection_mask =
- ProjectionMask::roots(schema_descr, root_indices.iter().copied());
-
- let projected_schema = Arc::new(
- file_schema
- .project(&root_indices)
- .expect("valid column indices"),
- );
-
- return ParquetReadPlan {
- projection_mask,
- projected_schema,
- };
- }
-
- let mut all_root_indices = Vec::new();
- let mut all_struct_accesses = Vec::new();
-
- for expr in exprs {
- let mut checker = PushdownChecker::new(file_schema, true);
- let _ = expr.visit(&mut checker);
- let columns = checker.into_sorted_columns();
-
- all_root_indices.extend_from_slice(&columns.required_columns);
- all_struct_accesses.extend(columns.struct_field_accesses);
- }
-
- all_root_indices.sort_unstable();
- all_root_indices.dedup();
-
- // when no struct field accesses were found, fall back to root-level projection
- // to match the performance of the simple path
- if all_struct_accesses.is_empty() {
- let projection_mask =
- ProjectionMask::roots(schema_descr, all_root_indices.iter().copied());
- let projected_schema = Arc::new(
- file_schema
- .project(&all_root_indices)
- .expect("valid column indices"),
- );
-
- return ParquetReadPlan {
- projection_mask,
- projected_schema,
- };
- }
-
- let leaf_indices = {
- let mut out =
- leaf_indices_for_roots(all_root_indices.iter().copied(), schema_descr);
- let struct_leaf_indices =
- resolve_struct_field_leaves(&all_struct_accesses, file_schema, schema_descr);
-
- out.extend_from_slice(&struct_leaf_indices);
- out.sort_unstable();
- out.dedup();
-
- out
- };
-
- let projection_mask =
- ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied());
-
- let projected_schema =
- build_filter_schema(file_schema, &all_root_indices, &all_struct_accesses);
-
- ParquetReadPlan {
- projection_mask,
- projected_schema,
- }
-}
-
-fn leaf_indices_for_roots<I>(
- root_indices: I,
- schema_descr: &SchemaDescriptor,
-) -> Vec<usize>
-where
- I: IntoIterator<Item = usize>,
-{
- // Always map root (Arrow) indices to Parquet leaf indices via the schema
- // descriptor. Arrow root indices only equal Parquet leaf indices when the
- // schema has no group columns (Struct, Map, etc.); when group columns
- // exist, their children become separate leaves and shift all subsequent
- // leaf indices.
- // Struct columns are unsupported.
- let root_set: BTreeSet<_> = root_indices.into_iter().collect();
-
- (0..schema_descr.num_columns())
- .filter(|leaf_idx| {
- root_set.contains(&schema_descr.get_column_root_idx(*leaf_idx))
- })
- .collect()
-}
-
-/// Resolves struct field access to specific Parquet leaf column indices
-///
-/// 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
-fn resolve_struct_field_leaves(
- accesses: &[StructFieldAccess],
- file_schema: &Schema,
- 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"]
-
- // a leaf matches if its path starts with our prefix
- // for example: 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);
- }
- }
- }
-
- leaf_indices
-}
-
-/// Builds a filter schema that includes only the fields actually accessed by the
-/// filter expression.
-///
-/// For regular (non-struct) columns, the full field type is used.
-/// For struct columns accessed via `get_field`, a pruned struct type is created
-/// containing only the fields along the access path. Note: it must match the schema
-/// that the Parquet reader produces when projecting specific struct leaves
-fn build_filter_schema(
- file_schema: &Schema,
- regular_indices: &[usize],
- struct_field_accesses: &[StructFieldAccess],
-) -> SchemaRef {
- let regular_set: BTreeSet<usize> = regular_indices.iter().copied().collect();
- let paths_by_root = group_access_paths_by_root(struct_field_accesses);
-
- let all_indices = regular_indices
- .iter()
- .copied()
- .chain(paths_by_root.keys().copied())
- .collect::<BTreeSet<_>>();
-
- let fields = all_indices
- .iter()
- .map(|&idx| {
- let field = file_schema.field(idx);
-
- // if this column appears as a regular (whole-column) reference,
- // keep the full type
- //
- // Pruning is only valid when the column is accessed exclusively
- // through struct field accesses
- if regular_set.contains(&idx) {
- return Arc::new(field.clone());
- }
-
- let Some(field_paths) = paths_by_root.get(&idx) else {
- return Arc::new(field.clone());
- };
-
- let pruned_data_type = prune_struct_type(field.data_type(), field_paths);
- Arc::new(Field::new(
- field.name(),
- pruned_data_type,
- field.is_nullable(),
- ))
- })
- .collect::<Vec<_>>();
-
- Arc::new(Schema::new_with_metadata(
- fields,
- file_schema.metadata().clone(),
- ))
-}
-
-/// Groups struct field access paths once for the root schema level.
-///
-/// Each map entry contains the complete field paths accessed below a root
-/// column. Recursive pruning groups these paths by their next component at each
-/// nested struct level.
-fn group_access_paths_by_root(
- struct_field_accesses: &[StructFieldAccess],
-) -> BTreeMap<usize, Vec<&[String]>> {
- let mut paths_by_root: BTreeMap<usize, Vec<&[String]>> = BTreeMap::new();
- for StructFieldAccess {
- root_index,
- field_path,
- } in struct_field_accesses
- {
- paths_by_root
- .entry(*root_index)
- .or_default()
- .push(field_path.as_slice());
- }
-
- paths_by_root
-}
-
-/// Groups access paths once for the current struct level.
-///
-/// The map key is the field name at this level. The map value is the list of
-/// remaining path suffixes below that field. An empty suffix means the access
-/// path terminates at that field, so the full field must be preserved.
-fn group_paths_by_next_field<'a>(
- paths: &'a [&'a [String]],
-) -> BTreeMap<&'a str, Vec<&'a [String]>> {
- let mut paths_by_field: BTreeMap<&str, Vec<&[String]>> = BTreeMap::new();
- for path in paths {
- if let Some((field, sub_path)) = path.split_first() {
- paths_by_field
- .entry(field.as_str())
- .or_default()
- .push(sub_path);
- }
- }
-
- paths_by_field
-}
-
-fn prune_struct_type(dt: &DataType, paths: &[&[String]]) -> DataType {
- let DataType::Struct(fields) = dt else {
- return dt.clone();
- };
-
- let paths_by_field = group_paths_by_next_field(paths);
-
- let pruned_fields = fields
- .iter()
- .filter_map(|f| {
- let sub_paths = paths_by_field.get(f.name().as_str())?;
-
- let out = if sub_paths.iter().any(|sub| sub.is_empty()) {
- // Leaf of access path — keep the field as-is.
- Arc::clone(f)
- } else {
- // Recurse into nested struct.
- let pruned = prune_struct_type(f.data_type(), sub_paths);
- Arc::new(Field::new(f.name(), pruned, f.is_nullable()))
- };
-
- Some(out)
- })
- .collect::<Vec<_>>();
-
- DataType::Struct(pruned_fields.into())
-}
-
/// Checks if a predicate expression can be pushed down to the parquet decoder.
///
/// Returns `true` if all columns referenced by the expression:
@@ -1167,7 +549,7 @@ impl<'a> RowFilterGenerator<'a> {
#[cfg(test)]
mod test {
use super::*;
- use arrow::datatypes::Fields;
+ use arrow::datatypes::{DataType, Fields};
use datafusion_common::ScalarValue;
use arrow::array::{
@@ -1182,6 +564,7 @@ mod test {
use datafusion_functions_nested::expr_fn::{
array_has, array_has_all, array_has_any, make_array,
};
+ use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::planner::logical2physical;
use datafusion_physical_expr_adapter::{
DefaultPhysicalExprAdapterFactory, PhysicalExprAdapterFactory,
@@ -1194,8 +577,6 @@ mod test {
use parquet::file::reader::{FileReader, SerializedFileReader};
use tempfile::NamedTempFile;
- use datafusion_physical_expr::expressions::Column as PhysicalColumn;
-
// List predicates used by the decoder should be accepted for pushdown
#[test]
fn test_filter_candidate_builder_supports_list_types() {
@@ -2052,86 +1433,6 @@ mod test {
assert_eq!(file_metrics.pushdown_rows_matched.value(), 2);
}
- #[test]
- fn projection_read_plan_preserves_full_struct() {
- // Schema: id (Int32), s (Struct{value: Int32, label: Utf8})
- // Parquet leaves: id=0, s.value=1, s.label=2
- let struct_fields: Fields = vec![
- Arc::new(Field::new("value", DataType::Int32, false)),
- Arc::new(Field::new("label", DataType::Utf8, false)),
- ]
- .into();
-
- let schema = Arc::new(Schema::new(vec![
- Field::new("id", DataType::Int32, false),
- Field::new("s", DataType::Struct(struct_fields.clone()), false),
- ]));
-
- let batch = RecordBatch::try_new(
- Arc::clone(&schema),
- vec![
- Arc::new(Int32Array::from(vec![1, 2, 3])),
- Arc::new(StructArray::new(
- struct_fields,
- vec![
- Arc::new(Int32Array::from(vec![10, 20, 30])) as _,
- Arc::new(StringArray::from(vec!["a", "b", "c"])) as _,
- ],
- None,
- )),
- ],
- )
- .unwrap();
-
- let file = NamedTempFile::new().expect("temp file");
- let mut writer =
- ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None)
- .expect("writer");
- writer.write(&batch).expect("write batch");
- writer.close().expect("close writer");
-
- let reader_file = file.reopen().expect("reopen file");
- let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file)
- .expect("reader builder");
- let metadata = builder.metadata().clone();
- let file_schema = builder.schema().clone();
- let schema_descr = metadata.file_metadata().schema_descr();
-
- // Simulate SELECT * output projection: Column("id") and Column("s")
- // Plus a get_field(s, 'value') expression from the pushed-down filter
- let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
- Arc::new(PhysicalColumn::new("id", 0)),
- Arc::new(PhysicalColumn::new("s", 1)),
- logical2physical(
- &get_field().call(vec![
- col("s"),
- Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None),
- ]),
- &file_schema,
- ),
- ];
-
- let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr);
-
- // The projected schema must have the FULL struct type because Column("s")
- // is in the projection. It should NOT be narrowed to Struct{value: Int32}.
- let s_field = read_plan.projected_schema.field_with_name("s").unwrap();
- assert_eq!(
- s_field.data_type(),
- &DataType::Struct(
- vec![
- Arc::new(Field::new("value", DataType::Int32, false)),
- Arc::new(Field::new("label", DataType::Utf8, false)),
- ]
- .into()
- ),
- );
-
- // all3 Parquet leaves should be in the projection mask
- let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]);
- assert_eq!(read_plan.projection_mask, expected_mask,);
- }
-
/// Sanity check that the given expression could be evaluated against the given schema without any errors.
/// This will fail if the expression references columns that are not in the schema or if the types of the columns are incompatible, etc.
fn check_expression_can_evaluate_against_schema(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org