unify parquet read plan

Author: friendlymatthew

datafusion/datasource-parquet/src/row_filter.rs

@@ -137,7 +137,7 @@ impl DatafusionArrowPredicate {
         time: metrics::Time,
     ) -> Result<Self> {
         let physical_expr =
-            reassign_expr_columns(candidate.expr, &candidate.filter_schema)?;
+            reassign_expr_columns(candidate.expr, &candidate.read_plan.projected_schema)?;
 
         Ok(Self {
             physical_expr,
@@ -146,7 +146,7 @@ impl DatafusionArrowPredicate {
             // can properly project and filter nested structures.
             projection_mask: ProjectionMask::leaves(
                 metadata.file_metadata().schema_descr(),
-                candidate.projection.leaf_indices.iter().copied(),
+                candidate.read_plan.leaf_indices.iter().copied(),
             ),
             rows_pruned,
             rows_matched,
@@ -199,22 +199,22 @@ pub(crate) struct FilterCandidate {
     required_bytes: usize,
     /// Can this filter use an index (e.g. a page index) to prune rows?
     can_use_index: bool,
-    /// Column indices into the parquet file schema required to evaluate this filter.
-    projection: LeafProjection,
-    /// The Arrow schema containing only the columns required by this filter,
-    /// projected from the file's Arrow schema.
-    filter_schema: SchemaRef,
+    /// The resolved Parquet read plan (leaf indices + projected schema).
+    read_plan: ParquetReadPlan,
 }
 
-/// Projection specification for nested columns using Parquet leaf column indices.
+/// The result of resolving which Parquet leaf columns and Arrow schema fields
+/// are needed to evaluate an expression against a Parquet file
 ///
-/// For nested types like List and Struct, Parquet stores data in leaf columns
-/// (the primitive fields). This struct tracks which leaf columns are needed
-/// to evaluate a filter expression.
+/// This is the shared output of the column resolution pipeline used by both
+/// the row fiter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s
 #[derive(Debug, Clone)]
-struct LeafProjection {
-    /// Leaf column indices in the Parquet schema descriptor.
-    leaf_indices: Vec<usize>,
+pub(crate) struct ParquetReadPlan {
+    /// Leaf column indices in the Parquet schema descriptor to decode
+    pub leaf_indices: Vec<usize>,
+    /// 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`.
@@ -246,41 +246,18 @@ impl FilterCandidateBuilder {
     /// * `Ok(None)` if the expression cannot be used as an ArrowFilter
     /// * `Err(e)` if an error occurs while building the candidate
     pub fn build(self, metadata: &ParquetMetaData) -> Result<Option<FilterCandidate>> {
-        let Some(required_columns) = pushdown_columns(&self.expr, &self.file_schema)?
-        else {
-            return Ok(None);
-        };
-
         let schema_descr = metadata.file_metadata().schema_descr();
-        let root_indices: Vec<_> =
-            required_columns.required_columns.into_iter().collect();
-
-        let mut leaf_indices = leaf_indices_for_roots(&root_indices, schema_descr);
-
-        let struct_leaf_indices = resolve_struct_field_leaves(
-            &required_columns.struct_field_accesses,
-            &self.file_schema,
-            schema_descr,
-        );
-        leaf_indices.extend_from_slice(&struct_leaf_indices);
-        leaf_indices.sort_unstable();
-        leaf_indices.dedup();
-
-        let projected_schema = build_filter_schema(
-            &self.file_schema,
-            &root_indices,
-            &required_columns.struct_field_accesses,
-        );
-
-        let required_bytes = size_of_columns(&leaf_indices, metadata)?;
-        let can_use_index = columns_sorted(&leaf_indices, metadata)?;
+        let read_plan =
+            match build_parquet_read_plan(&self.expr, &self.file_schema, schema_descr)? {
+                Some(plan) => plan,
+                None => return Ok(None),
+            };
 
         Ok(Some(FilterCandidate {
             expr: self.expr,
-            required_bytes,
-            can_use_index,
-            projection: LeafProjection { leaf_indices },
-            filter_schema: projected_schema,
+            required_bytes: size_of_columns(&read_plan.leaf_indices, metadata)?,
+            can_use_index: columns_sorted(&read_plan.leaf_indices, metadata)?,
+            read_plan,
         }))
     }
 }
@@ -562,17 +539,64 @@ fn pushdown_columns(
     Ok((!checker.prevents_pushdown()).then(|| checker.into_sorted_columns()))
 }
 
-fn leaf_indices_for_roots(
-    root_indices: &[usize],
+/// Resolves which Parquet leaf columns and Arrow schema fields are needed
+/// to evaluate `expr` against a Parquet file
+///
+/// Returns `Ok(Some(plan))` when the expression can be evaluated using only
+/// pushdown-compatible columns. `Ok(None)` when it can not (it references
+/// whole struct columns or columns missing from disk)
+///
+/// Note: this is a shared entry point used by both row filter construction and
+/// the opener's projection logic
+pub(crate) fn build_parquet_read_plan(
+    expr: &Arc<dyn PhysicalExpr>,
+    file_schema: &Schema,
     schema_descr: &SchemaDescriptor,
-) -> Vec<usize> {
+) -> Result<Option<ParquetReadPlan>> {
+    let Some(required_columns) = pushdown_columns(expr, file_schema)? else {
+        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(
+        &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 projected_schema = build_filter_schema(
+        file_schema,
+        root_indices,
+        &required_columns.struct_field_accesses,
+    );
+
+    Ok(Some(ParquetReadPlan {
+        leaf_indices,
+        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.iter().copied().collect();
+    let root_set: BTreeSet<_> = root_indices.into_iter().collect();
 
     (0..schema_descr.num_columns())
         .filter(|leaf_idx| {
@@ -983,7 +1007,7 @@ mod test {
             .expect("building candidate")
             .expect("list pushdown should be supported");
 
-        assert_eq!(candidate.projection.leaf_indices, vec![list_index]);
+        assert_eq!(candidate.read_plan.leaf_indices, vec![list_index]);
     }
 
     #[test]
@@ -1441,7 +1465,7 @@ mod test {
 
         // col_b is Parquet leaf 3 (shifted by struct_col's two children).
         assert_eq!(
-            candidate.projection.leaf_indices,
+            candidate.read_plan.leaf_indices,
             vec![3],
             "leaf_indices should be [3] for col_b"
         );
@@ -1591,7 +1615,7 @@ mod test {
         // The filter accesses only s.value, so only Parquet leaf 1 is needed.
         // Leaf 2 (s.label) is not read, reducing unnecessary I/O.
         assert_eq!(
-            candidate.projection.leaf_indices,
+            candidate.read_plan.leaf_indices,
             vec![1],
             "leaf_indices should contain only the accessed struct field leaf"
         );
@@ -1714,7 +1738,7 @@ mod test {
 
         // Only s.outer.inner (leaf 2) should be projected,
         assert_eq!(
-            candidate.projection.leaf_indices,
+            candidate.read_plan.leaf_indices,
             vec![2],
             "leaf_indices should be [2] for s.outer.inner, skipping sibling and cousin leaves"
         );