refactor(parquet): split opener.rs into a module and add ParquetAccessPlanOptimizer trait

Two complementary changes that keep the Parquet opener manageable as more
pruning strategies stack up:

1. Convert `opener.rs` into an `opener/` module. `EarlyStoppingStream` and
   `EncryptionContext` move to sibling files (`early_stop.rs`, `encryption.rs`).
   Pure code motion: no public API change, no behavior change.

2. Add `ParquetAccessPlanOptimizer`, an extension trait for refining a
   `ParquetAccessPlan` during file open. The opener invokes registered
   user-supplied optimizers after each built-in pruning stage
   (`AfterMetadata`, `AfterBloomFilters`, `BeforeBuildStream`).
   `ParquetSource::with_access_plan_optimizer` is the registration point.

The built-in pruning passes (range, statistics, bloom filter, limit, page
index) are deliberately *not* migrated to the trait in this PR; they share
a `is_fully_matched` side-state with `prune_by_limit` that isn't yet
modeled on the trait surface, so migrating them would require either
plumbing that state through the trait or losing the
limit-via-fully-matched optimization for built-in passes. Tracking that
as follow-up.

Author: adriangb

datafusion/datasource-parquet/src/access_plan_optimizer.rs

@@ -0,0 +1,170 @@
+// 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.
+
+//! Extension point for refining a [`ParquetAccessPlan`] during file open.
+//!
+//! The Parquet opener narrows down which row groups (and which rows within
+//! them) it will read through a fixed sequence of built-in passes:
+//!
+//! - file-range pruning,
+//! - row-group statistics pruning,
+//! - bloom-filter pruning,
+//! - limit-based pruning,
+//! - page-index pruning.
+//!
+//! Each pass operates on a [`ParquetAccessPlan`]. The
+//! [`ParquetAccessPlanOptimizer`] trait exposes that pipeline as an extension
+//! point so external crates can append additional passes — sampling, custom
+//! statistics, user-defined Parquet indexes, etc. — without having to fork
+//! the opener.
+//!
+//! User-supplied optimizers are invoked **after** the built-in passes for the
+//! corresponding [`OptimizerStage`]. They can read the access plan and the
+//! pruning context, then return a (possibly narrowed) plan; they cannot
+//! widen access beyond what the built-ins produced.
+
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_datasource::{FileRange, PartitionedFile};
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use parquet::file::metadata::ParquetMetaData;
+
+use crate::ParquetAccessPlan;
+use crate::ParquetFileMetrics;
+use crate::page_filter::PagePruningAccessPlanFilter;
+use arrow::datatypes::SchemaRef;
+use datafusion_pruning::PruningPredicate;
+
+/// Stage at which a [`ParquetAccessPlanOptimizer`] runs during file open.
+///
+/// Each stage has access to a different subset of file metadata, reflecting
+/// the order in which the opener loads it. See [`AccessPlanContext`] for
+/// the fields available at each stage.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum OptimizerStage {
+    /// Runs after the Parquet footer and page index (when enabled) are
+    /// loaded, and after the built-in file-range and row-group-statistics
+    /// passes have refined the plan. Bloom filters have **not** been
+    /// loaded at this point.
+    AfterMetadata,
+    /// Runs after the bloom filters for surviving row groups have been
+    /// loaded, and after the built-in bloom-filter pruning pass.
+    AfterBloomFilters,
+    /// Runs at the end of the pruning pipeline, just before the stream
+    /// is built, after the built-in limit and page-index passes.
+    BeforeBuildStream,
+}
+
+/// Read-only context passed to a [`ParquetAccessPlanOptimizer`].
+///
+/// Some fields are only populated at certain [`OptimizerStage`]s; see each
+/// field's docs.
+#[derive(Debug)]
+pub struct AccessPlanContext<'a> {
+    /// Execution partition index for the scan.
+    pub partition_index: usize,
+    /// The file being opened.
+    pub partitioned_file: &'a PartitionedFile,
+    /// Optional byte range restricting which part of the file to read.
+    pub file_range: Option<&'a FileRange>,
+    /// Schema of the file after type coercions (the schema the parquet
+    /// reader will produce, before projection).
+    pub physical_file_schema: &'a SchemaRef,
+    /// Loaded Parquet metadata, including page index when enabled.
+    pub file_metadata: &'a ParquetMetaData,
+    /// Raw predicate applied to this scan, if any.
+    pub predicate: Option<&'a Arc<dyn PhysicalExpr>>,
+    /// Row-group-level pruning predicate derived from `predicate`.
+    pub pruning_predicate: Option<&'a Arc<PruningPredicate>>,
+    /// Page-index pruning predicate derived from `predicate`.
+    pub page_pruning_predicate: Option<&'a Arc<PagePruningAccessPlanFilter>>,
+    /// Outer query limit, if any.
+    pub limit: Option<usize>,
+    /// Whether the query requires the original row order to be preserved.
+    pub preserve_order: bool,
+    /// Per-file metrics. Optimizers may emit to these counters.
+    pub file_metrics: &'a ParquetFileMetrics,
+    /// Current optimizer stage.
+    pub stage: OptimizerStage,
+}
+
+/// Trait for narrowing a [`ParquetAccessPlan`] during Parquet file open.
+///
+/// The opener invokes registered optimizers after each built-in pruning
+/// stage (see [`OptimizerStage`]). An optimizer that does not apply at the
+/// current stage should return the plan unchanged.
+///
+/// # Example
+///
+/// A sampling optimizer that keeps a fraction of the surviving row groups:
+///
+/// ```ignore
+/// use std::sync::Arc;
+/// use datafusion_common::Result;
+/// use datafusion_datasource_parquet::{
+///     access_plan_optimizer::{
+///         AccessPlanContext, OptimizerStage, ParquetAccessPlanOptimizer,
+///     },
+///     ParquetAccessPlan, RowGroupAccess,
+/// };
+///
+/// #[derive(Debug)]
+/// struct SampleHalf;
+///
+/// impl ParquetAccessPlanOptimizer for SampleHalf {
+///     fn stage(&self) -> OptimizerStage { OptimizerStage::BeforeBuildStream }
+///
+///     fn optimize(
+///         &self,
+///         _ctx: &AccessPlanContext<'_>,
+///         mut plan: ParquetAccessPlan,
+///     ) -> Result<ParquetAccessPlan> {
+///         for (idx, count) in (0..plan.len()).step_by(2).zip(std::iter::repeat(())) {
+///             let _ = count;
+///             plan.skip(idx);
+///         }
+///         Ok(plan)
+///     }
+/// }
+/// ```
+pub trait ParquetAccessPlanOptimizer: Debug + Send + Sync {
+    /// At which stage of the opener this optimizer runs.
+    fn stage(&self) -> OptimizerStage;
+
+    /// Refine `plan` for `ctx`. Returning `plan` unchanged is the no-op.
+    fn optimize(
+        &self,
+        ctx: &AccessPlanContext<'_>,
+        plan: ParquetAccessPlan,
+    ) -> Result<ParquetAccessPlan>;
+}
+
+/// Run all `optimizers` whose [`OptimizerStage`] matches `ctx.stage`.
+pub(crate) fn run_stage(
+    optimizers: &[Arc<dyn ParquetAccessPlanOptimizer>],
+    ctx: &AccessPlanContext<'_>,
+    mut plan: ParquetAccessPlan,
+) -> Result<ParquetAccessPlan> {
+    for opt in optimizers {
+        if opt.stage() == ctx.stage {
+            plan = opt.optimize(ctx, plan)?;
+        }
+    }
+    Ok(plan)
+}

datafusion/datasource-parquet/src/mod.rs

@@ -25,6 +25,7 @@
 #![cfg_attr(test, allow(clippy::needless_pass_by_value))]
 
 pub mod access_plan;
+pub mod access_plan_optimizer;
 pub mod file_format;
 pub mod metadata;
 mod metrics;
@@ -39,6 +40,9 @@ mod supported_predicates;
 mod writer;
 
 pub use access_plan::{ParquetAccessPlan, RowGroupAccess};
+pub use access_plan_optimizer::{
+    AccessPlanContext, OptimizerStage, ParquetAccessPlanOptimizer,
+};
 pub use file_format::*;
 pub use metrics::ParquetFileMetrics;
 pub use page_filter::PagePruningAccessPlanFilter;

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

@@ -0,0 +1,107 @@
+// 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.
+
+//! [`EarlyStoppingStream`] terminates a Parquet file scan when a dynamic
+//! filter narrows after the scan has already started.
+
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use arrow::array::RecordBatch;
+use datafusion_common::Result;
+use datafusion_physical_plan::metrics::PruningMetrics;
+use datafusion_pruning::FilePruner;
+use futures::{Stream, StreamExt, ready};
+
+/// Wraps an inner RecordBatchStream and a [`FilePruner`]
+///
+/// This can terminate the scan early when some dynamic filters is updated after
+/// the scan starts, so we discover after the scan starts that the file can be
+/// pruned (can't have matching rows).
+pub(super) struct EarlyStoppingStream<S> {
+    /// Has the stream finished processing? All subsequent polls will return
+    /// None
+    done: bool,
+    file_pruner: FilePruner,
+    files_ranges_pruned_statistics: PruningMetrics,
+    /// The inner stream
+    inner: S,
+}
+
+impl<S> EarlyStoppingStream<S> {
+    pub(super) fn new(
+        stream: S,
+        file_pruner: FilePruner,
+        files_ranges_pruned_statistics: PruningMetrics,
+    ) -> Self {
+        Self {
+            done: false,
+            inner: stream,
+            file_pruner,
+            files_ranges_pruned_statistics,
+        }
+    }
+}
+
+impl<S> EarlyStoppingStream<S>
+where
+    S: Stream<Item = Result<RecordBatch>> + Unpin,
+{
+    fn check_prune(&mut self, input: Result<RecordBatch>) -> Result<Option<RecordBatch>> {
+        let batch = input?;
+
+        // Since dynamic filters may have been updated, see if we can stop
+        // reading this stream entirely.
+        if self.file_pruner.should_prune()? {
+            self.files_ranges_pruned_statistics.add_pruned(1);
+            // Previously this file range has been counted as matched
+            self.files_ranges_pruned_statistics.subtract_matched(1);
+            self.done = true;
+            Ok(None)
+        } else {
+            // Return the adapted batch
+            Ok(Some(batch))
+        }
+    }
+}
+
+impl<S> Stream for EarlyStoppingStream<S>
+where
+    S: Stream<Item = Result<RecordBatch>> + Unpin,
+{
+    type Item = Result<RecordBatch>;
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Self::Item>> {
+        if self.done {
+            return Poll::Ready(None);
+        }
+        match ready!(self.inner.poll_next_unpin(cx)) {
+            None => {
+                // input done
+                self.done = true;
+                Poll::Ready(None)
+            }
+            Some(input_batch) => {
+                let output = self.check_prune(input_batch);
+                Poll::Ready(output.transpose())
+            }
+        }
+    }
+}