feat(parquet): observer-driven pruning in row-group and page filters

`RowGroupAccessPlanFilter::prune_by_statistics_with_observer` takes a
`PruningConjunction` plus `&mut dyn PruningObserver`. The untagged
`prune_by_statistics` wraps the predicate via `PruningConjunction::single`
and passes `NoopObserver` — one code path, two API shapes.

`PagePruningAccessPlanFilter`:
  * Predicates carry an optional caller `Tag` internally
    (`TaggedPagePredicate`); untagged callers see no behavior change.
  * `new_tagged(&[(Tag, expr)], &schema)` for tagged construction.
  * The page-index loop emits one observer event per leaf actually
    evaluated; leaves cut off by the existing `!selects_any`
    short-circuit are not observed (resolves reviewer Q2 on #22235:
    per-conjunct stats are not biased by predicates that never ran).
  * `prune_plan_with_page_index` and the new `prune_plan_with_observer`
    share a single body — no duplicated method as in the prior draft.

Fully-matched detection still consumes a flat `PruningPredicate` (it
inverts the combined expression); the observer path receives that
combined predicate explicitly and fires only for the per-leaf pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: adriangb

datafusion/datasource-parquet/src/opener/mod.rs

@@ -1103,13 +1103,15 @@ impl RowGroupsPrunedParquetOpen {
             && !access_plan.is_empty()
             && let Some(page_pruning_predicate) = page_pruning_predicate
         {
+            let mut obs = datafusion_pruning::NoopObserver;
             let page_pruning_result = page_pruning_predicate
                 .prune_plan_with_page_index_and_metrics(
                     access_plan,
                     &prepared.physical_file_schema,
                     reader_metadata.parquet_schema(),
                     file_metadata.as_ref(),
                     &prepared.file_metrics,
+                    &mut obs,
                 );
             access_plan = page_pruning_result.access_plan;
             ParquetFileMetrics::add_page_index_pages_skipped_by_fully_matched(

datafusion/datasource-parquet/src/page_filter.rs

@@ -31,7 +31,7 @@ use arrow::{
 use datafusion_common::ScalarValue;
 use datafusion_common::pruning::PruningStatistics;
 use datafusion_physical_expr::{PhysicalExpr, split_conjunction};
-use datafusion_pruning::PruningPredicate;
+use datafusion_pruning::{NoopObserver, PruningObserver, PruningPredicate, Tag};
 
 use log::{debug, trace};
 use parquet::arrow::arrow_reader::statistics::StatisticsConverter;
@@ -110,9 +110,20 @@ use parquet::{
 /// row selection that is added to the [`ParquetAccessPlan`].
 #[derive(Debug)]
 pub struct PagePruningAccessPlanFilter {
-    /// single column predicates (e.g. (`col = 5`) extracted from the overall
-    /// predicate. Must all be true for a row to be included in the result.
-    predicates: Vec<PruningPredicate>,
+    /// Single-column predicates (e.g. `col = 5`) extracted from the
+    /// overall predicate. Must all be true for a row to be included.
+    /// Each carries an optional caller [`Tag`]; when present, the
+    /// page-index evaluation loop fires `observer.on_leaf(tag, ..)`
+    /// after evaluating that predicate. Untagged predicates produce
+    /// `None` tags and no-op against `NoopObserver`.
+    predicates: Vec<TaggedPagePredicate>,
+}
+
+/// Single-column predicate paired with an optional caller tag.
+#[derive(Debug)]
+struct TaggedPagePredicate {
+    tag: Option<Tag>,
+    predicate: PruningPredicate,
 }
 
 /// Result of applying page-index pruning to a [`ParquetAccessPlan`].
@@ -137,40 +148,55 @@ impl PagePruningResult {
 
 impl PagePruningAccessPlanFilter {
     /// Create a new [`PagePruningAccessPlanFilter`] from a physical
-    /// expression.
+    /// expression. Predicates created this way have no caller tag.
     #[expect(clippy::needless_pass_by_value)]
     pub fn new(expr: &Arc<dyn PhysicalExpr>, schema: SchemaRef) -> Self {
-        // extract any single column predicates
         let predicates = split_conjunction(expr)
             .into_iter()
-            .filter_map(|predicate| {
-                let pp = match PruningPredicate::try_new(
-                    Arc::clone(predicate),
-                    Arc::clone(&schema),
-                ) {
-                    Ok(pp) => pp,
-                    Err(e) => {
-                        debug!("Ignoring error creating page pruning predicate: {e}");
-                        return None;
-                    }
-                };
-
-                if pp.always_true() {
-                    debug!("Ignoring always true page pruning predicate: {predicate}");
-                    return None;
-                }
-
-                if pp.required_columns().single_column().is_none() {
-                    debug!("Ignoring multi-column page pruning predicate: {predicate}");
-                    return None;
-                }
+            .filter_map(|predicate| Self::build_one(predicate, &schema, None))
+            .collect::<Vec<_>>();
+        Self { predicates }
+    }
 
-                Some(pp)
-            })
+    /// Like [`Self::new`], but each conjunct carries a caller-supplied
+    /// [`Tag`]. The page-index evaluation loop fires
+    /// `observer.on_leaf(tag, ..)` once per leaf actually evaluated;
+    /// leaves cut off by the existing AND short-circuit on row
+    /// selection (`!selects_any`) are not observed.
+    pub fn new_tagged(
+        tagged: &[(Tag, Arc<dyn PhysicalExpr>)],
+        schema: &SchemaRef,
+    ) -> Self {
+        let predicates = tagged
+            .iter()
+            .filter_map(|(tag, expr)| Self::build_one(expr, schema, Some(*tag)))
             .collect::<Vec<_>>();
         Self { predicates }
     }
 
+    fn build_one(
+        expr: &Arc<dyn PhysicalExpr>,
+        schema: &SchemaRef,
+        tag: Option<Tag>,
+    ) -> Option<TaggedPagePredicate> {
+        let pp = match PruningPredicate::try_new(Arc::clone(expr), Arc::clone(schema)) {
+            Ok(pp) => pp,
+            Err(e) => {
+                debug!("Ignoring error creating page pruning predicate: {e}");
+                return None;
+            }
+        };
+        if pp.always_true() {
+            debug!("Ignoring always true page pruning predicate: {expr}");
+            return None;
+        }
+        if pp.required_columns().single_column().is_none() {
+            debug!("Ignoring multi-column page pruning predicate: {expr}");
+            return None;
+        }
+        Some(TaggedPagePredicate { tag, predicate: pp })
+    }
+
     /// Returns an updated [`ParquetAccessPlan`] by applying predicates to the
     /// parquet page index, if any
     pub fn prune_plan_with_page_index(
@@ -181,26 +207,64 @@ impl PagePruningAccessPlanFilter {
         parquet_metadata: &ParquetMetaData,
         file_metrics: &ParquetFileMetrics,
     ) -> ParquetAccessPlan {
+        let mut obs = NoopObserver;
         self.prune_plan_with_page_index_and_metrics(
             access_plan,
             arrow_schema,
             parquet_schema,
             parquet_metadata,
             file_metrics,
+            &mut obs,
+        )
+        .access_plan
+    }
+
+    /// Tagged variant. Fires `observer.on_leaf(tag, mask)` once per
+    /// predicate that is actually evaluated against page-index
+    /// statistics. The `mask` is per-row-group at this layer:
+    /// `mask[i] = true` means row group `i` still has at least one
+    /// page that may match this leaf. Predicates skipped by the
+    /// existing per-row-group `!selects_any` short-circuit are not
+    /// observed — so per-conjunct stats are not biased by predicates
+    /// that never ran (resolves the reviewer's Q2 concern on
+    /// PR #22235).
+    pub fn prune_plan_with_observer<O>(
+        &self,
+        access_plan: ParquetAccessPlan,
+        arrow_schema: &Schema,
+        parquet_schema: &SchemaDescriptor,
+        parquet_metadata: &ParquetMetaData,
+        file_metrics: &ParquetFileMetrics,
+        observer: &mut O,
+    ) -> ParquetAccessPlan
+    where
+        O: PruningObserver + ?Sized,
+    {
+        self.prune_plan_with_page_index_and_metrics(
+            access_plan,
+            arrow_schema,
+            parquet_schema,
+            parquet_metadata,
+            file_metrics,
+            observer,
         )
         .access_plan
     }
 
     /// Returns an updated [`ParquetAccessPlan`] and metrics by applying predicates
     /// to the parquet page index, if any.
-    pub(crate) fn prune_plan_with_page_index_and_metrics(
+    pub(crate) fn prune_plan_with_page_index_and_metrics<O>(
         &self,
         mut access_plan: ParquetAccessPlan,
         arrow_schema: &Schema,
         parquet_schema: &SchemaDescriptor,
         parquet_metadata: &ParquetMetaData,
         file_metrics: &ParquetFileMetrics,
-    ) -> PagePruningResult {
+        observer: &mut O,
+    ) -> PagePruningResult
+    where
+        O: PruningObserver + ?Sized,
+    {
         // scoped timer updates on drop
         let _timer_guard = file_metrics.page_index_eval_time.timer();
         if self.predicates.is_empty() {
@@ -210,6 +274,16 @@ impl PagePruningAccessPlanFilter {
         let page_index_predicates = &self.predicates;
         let groups = parquet_metadata.row_groups();
 
+        // Per-leaf "did this row group still have any matching pages
+        // after the leaf alone?" mask. Built across all row groups,
+        // emitted via `observer.on_leaf` at the end so each leaf gets
+        // exactly one observation per call to this function. Leaves
+        // never evaluated (because an earlier conjunct emptied the
+        // running row selection — `!selects_any` break below) end up
+        // with `None` here and are intentionally not observed.
+        let mut per_leaf_mask: Vec<Option<Vec<bool>>> =
+            (0..page_index_predicates.len()).map(|_| None).collect();
+
         if groups.is_empty() {
             return PagePruningResult::new(access_plan, 0);
         }
@@ -262,7 +336,8 @@ impl PagePruningAccessPlanFilter {
             let mut matched_pages_in_group: HashSet<usize> =
                 HashSet::from_iter(0..total_pages_in_group);
 
-            for predicate in page_index_predicates {
+            for (leaf_idx, tagged) in page_index_predicates.iter().enumerate() {
+                let predicate = &tagged.predicate;
                 let Some(column) = predicate.required_columns().single_column() else {
                     debug!(
                         "Ignoring multi-column page pruning predicate: {:?}",
@@ -307,6 +382,15 @@ impl PagePruningAccessPlanFilter {
                     predicate.predicate_expr(),
                 );
 
+                // Per-leaf observation: this leaf ran for `row_group_index`
+                // and produced `selection`. Whether this row group is
+                // "kept" by the leaf is `selection.selects_any()`. The
+                // entry is created lazily so leaves never observed (due
+                // to short-circuit) stay `None`.
+                let mask = per_leaf_mask[leaf_idx]
+                    .get_or_insert_with(|| vec![false; groups.len()]);
+                mask[row_group_index] = selection.selects_any();
+
                 let matched_pages_indexes: HashSet<_> = pages
                     .into_iter()
                     .enumerate()
@@ -369,6 +453,18 @@ impl PagePruningAccessPlanFilter {
         file_metrics
             .page_index_pages_pruned
             .add_matched(total_pages_select);
+
+        // Emit one observer event per leaf that was actually
+        // evaluated against at least one row group. Leaves that the
+        // outer `!selects_any` short-circuit prevented from running
+        // stay `None` here and are correctly absent from the stats.
+        for (leaf_idx, mask_opt) in per_leaf_mask.into_iter().enumerate() {
+            if let Some(mask) = mask_opt {
+                let tag = page_index_predicates[leaf_idx].tag;
+                observer.on_leaf(tag, &mask);
+            }
+        }
+
         PagePruningResult::new(access_plan, total_pages_skipped_by_fully_matched)
     }