From a01dfe0645978a358ce54c6ebc3199fe7233ebe2 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:34:40 -0500 Subject: [PATCH 1/2] refactor: collapse repeated root-level plan construction into a helper `build_projection_read_plan` had three copies of the same `ProjectionMask::roots` + `Schema::project` + `ParquetReadPlan` sequence, one per early-return path. Extract `root_level_plan`. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01BsRGsN18YBCji5iFTQmWpr --- .../src/projection_read_plan.rs | 65 ++++++++----------- 1 file changed, 28 insertions(+), 37 deletions(-) diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs index 46caed661614a..154aeeb3e130e 100644 --- a/datafusion/datasource-parquet/src/projection_read_plan.rs +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -363,18 +363,7 @@ pub(crate) fn build_projection_read_plan( 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, - }; + return root_level_plan(&root_indices, file_schema, schema_descr); } // secondary fast path: if the schema has no struct columns, we can skip @@ -393,19 +382,7 @@ pub(crate) fn build_projection_read_plan( 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, - }; + return root_level_plan(&root_indices, file_schema, schema_descr); } let mut all_root_indices = Vec::new(); @@ -426,18 +403,7 @@ pub(crate) fn build_projection_read_plan( // 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, - }; + return root_level_plan(&all_root_indices, file_schema, schema_descr); } let leaf_indices = { @@ -465,6 +431,31 @@ pub(crate) fn build_projection_read_plan( } } +/// Builds a [`ParquetReadPlan`] that decodes whole root columns. +/// +/// `root_indices` must be sorted, deduplicated indices into `file_schema`. Every +/// leaf below each root is decoded, and the projected schema keeps each root +/// field's full type. Callers that need to decode only some leaves of a struct +/// root must build the plan from leaf indices instead. +fn root_level_plan( + root_indices: &[usize], + file_schema: &Schema, + schema_descr: &SchemaDescriptor, +) -> ParquetReadPlan { + 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"), + ); + + ParquetReadPlan { + projection_mask, + projected_schema, + } +} + pub(crate) fn leaf_indices_for_roots( root_indices: I, schema_descr: &SchemaDescriptor, From 022f6658115a2469f0d26c91439949588d5d74ca Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:39:42 -0500 Subject: [PATCH 2/2] refactor: share leaf-level read plan assembly with the row filter `build_projection_read_plan` and `row_filter::build_parquet_read_plan` ran the same chain to build a leaf-level plan: leaf_indices_for_roots, resolve_struct_field_leaves, sort/dedup, ProjectionMask::leaves, build_filter_schema. They differed only in the row filter also sizing the leaves and returning an Option. Extract `assemble_read_plan`, which returns the plan plus the resolved leaf indices so the row filter can compute `required_bytes` from them. The three helpers it subsumes are now private to the module. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01BsRGsN18YBCji5iFTQmWpr --- .../src/projection_read_plan.rs | 62 +++++++++++++------ .../datasource-parquet/src/row_filter.rs | 31 ++-------- 2 files changed, 46 insertions(+), 47 deletions(-) diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs index 154aeeb3e130e..c9d8beab1466d 100644 --- a/datafusion/datasource-parquet/src/projection_read_plan.rs +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -406,29 +406,51 @@ pub(crate) fn build_projection_read_plan( return root_level_plan(&all_root_indices, file_schema, schema_descr); } - 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); + let (read_plan, _leaf_indices) = assemble_read_plan( + &all_root_indices, + &all_struct_accesses, + file_schema, + schema_descr, + ); - out.extend_from_slice(&struct_leaf_indices); - out.sort_unstable(); - out.dedup(); + read_plan +} - out - }; +/// Builds a leaf-level [`ParquetReadPlan`] covering `root_indices` in full plus +/// the individual leaves reached by `struct_field_accesses`. +/// +/// `root_indices` must be sorted, deduplicated indices into `file_schema`. +/// +/// Also returns the resolved Parquet leaf indices, sorted and deduplicated, so +/// callers can size the columns the decoder will read. +pub(crate) fn assemble_read_plan( + root_indices: &[usize], + struct_field_accesses: &[StructFieldAccess], + file_schema: &Schema, + schema_descr: &SchemaDescriptor, +) -> (ParquetReadPlan, Vec) { + let mut leaf_indices = + leaf_indices_for_roots(root_indices.iter().copied(), schema_descr); + leaf_indices.extend_from_slice(&resolve_struct_field_leaves( + struct_field_accesses, + file_schema, + schema_descr, + )); + leaf_indices.sort_unstable(); + leaf_indices.dedup(); 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, - } + build_filter_schema(file_schema, root_indices, struct_field_accesses); + + ( + ParquetReadPlan { + projection_mask, + projected_schema, + }, + leaf_indices, + ) } /// Builds a [`ParquetReadPlan`] that decodes whole root columns. @@ -456,7 +478,7 @@ fn root_level_plan( } } -pub(crate) fn leaf_indices_for_roots( +fn leaf_indices_for_roots( root_indices: I, schema_descr: &SchemaDescriptor, ) -> Vec @@ -482,7 +504,7 @@ where /// 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( +fn resolve_struct_field_leaves( accesses: &[StructFieldAccess], file_schema: &Schema, schema_descr: &SchemaDescriptor, @@ -521,7 +543,7 @@ pub(crate) fn resolve_struct_field_leaves( /// 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( +fn build_filter_schema( file_schema: &Schema, regular_indices: &[usize], struct_field_accesses: &[StructFieldAccess], diff --git a/datafusion/datasource-parquet/src/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs index 4505f7fd62c91..a375e6611e004 100644 --- a/datafusion/datasource-parquet/src/row_filter.rs +++ b/datafusion/datasource-parquet/src/row_filter.rs @@ -86,8 +86,7 @@ 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, + ParquetReadPlan, PushdownChecker, PushdownColumns, assemble_read_plan, }; /// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform @@ -269,38 +268,16 @@ pub(crate) fn build_parquet_read_plan( return Ok(None); }; - let root_indices = &required_columns.required_columns; - - let mut leaf_indices = - leaf_indices_for_roots(root_indices.iter().copied(), schema_descr); - - let struct_leaf_indices = resolve_struct_field_leaves( + let (read_plan, leaf_indices) = assemble_read_plan( + &required_columns.required_columns, &required_columns.struct_field_accesses, file_schema, schema_descr, ); - leaf_indices.extend_from_slice(&struct_leaf_indices); - leaf_indices.sort_unstable(); - leaf_indices.dedup(); let required_bytes = size_of_columns(&leaf_indices, metadata)?; - let projection_mask = - ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied()); - - let projected_schema = build_filter_schema( - file_schema, - root_indices, - &required_columns.struct_field_accesses, - ); - - Ok(Some(( - ParquetReadPlan { - projection_mask, - projected_schema, - }, - required_bytes, - ))) + Ok(Some((read_plan, required_bytes))) } /// Checks if a predicate expression can be pushed down to the parquet decoder.