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-22346-2d1d53e2afe17f2aa35be649e85ebe611ee23497 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b4739e5031f16e1f0c151e52c976973f99ab8642 Author: Adrian Garcia Badaracco <[email protected]> AuthorDate: Mon May 18 21:10:02 2026 -0700 refactor(parquet-datasource): split opener.rs into an opener/ module (#22346) ## Which issue does this PR close? Relates to the discussion in #22024 about the Parquet datasource crate becoming hard to navigate. Split out of #22156, which bundled several code-motion moves into one PR — this is one of three smaller, independently-reviewable PRs that replace it. ## Rationale for this change `opener.rs` had grown to ~2,700 LOC, bundling several distinct responsibilities into one file. That makes it hard to read and hard to review changes in isolation. This PR is **pure code motion**: no behavior change and no public API change. ## What changes are included in this PR? Splits `opener.rs` into an `opener/` directory module: - `opener/early_stop.rs` — `EarlyStoppingStream`, the dynamic-filter early-termination wrapper applied at the end of `build_stream`. - `opener/encryption.rs` — `EncryptionContext` and the `ParquetMorselizer::get_encryption_context` helpers, isolating the `#[cfg(feature = "parquet_encryption")]` gating that previously bled through the main file. `opener.rs` becomes `opener/mod.rs`. Note: #22156 originally also extracted an `opener/push_decoder_stream.rs`. That move is now obsolete — #22289 has since extracted `PushDecoderStreamState` into `push_decoder.rs` — so it is dropped here. ## Are these changes tested? Yes, covered by existing tests. `cargo test -p datafusion-datasource-parquet --all-features` (122 passing) and `cargo clippy -p datafusion-datasource-parquet --all-targets --all-features -- -D warnings` both pass. ## Are there any user-facing changes? No. `opener` was already a private module; this only reorganizes files inside the crate. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> --- .../datasource-parquet/src/opener/early_stop.rs | 107 ++++++++++++++ .../datasource-parquet/src/opener/encryption.rs | 104 +++++++++++++ .../src/{opener.rs => opener/mod.rs} | 161 +-------------------- 3 files changed, 219 insertions(+), 153 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/early_stop.rs b/datafusion/datasource-parquet/src/opener/early_stop.rs new file mode 100644 index 0000000000..75749d2840 --- /dev/null +++ b/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()) + } + } + } +} diff --git a/datafusion/datasource-parquet/src/opener/encryption.rs b/datafusion/datasource-parquet/src/opener/encryption.rs new file mode 100644 index 0000000000..b725198237 --- /dev/null +++ b/datafusion/datasource-parquet/src/opener/encryption.rs @@ -0,0 +1,104 @@ +// 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. + +//! Encryption context used during Parquet file open. +//! +//! Isolated here so the `#[cfg(feature = "parquet_encryption")]` gating does +//! not pollute the rest of the opener module. + +#[cfg(feature = "parquet_encryption")] +use std::sync::Arc; + +use datafusion_common::Result; +#[cfg(feature = "parquet_encryption")] +use datafusion_common::config::EncryptionFactoryOptions; +#[cfg(feature = "parquet_encryption")] +use datafusion_common::encryption::FileDecryptionProperties; +#[cfg(feature = "parquet_encryption")] +use datafusion_execution::parquet_encryption::EncryptionFactory; + +use super::ParquetMorselizer; + +#[derive(Default)] +pub(super) struct EncryptionContext { + #[cfg(feature = "parquet_encryption")] + file_decryption_properties: Option<Arc<FileDecryptionProperties>>, + #[cfg(feature = "parquet_encryption")] + encryption_factory: Option<(Arc<dyn EncryptionFactory>, EncryptionFactoryOptions)>, +} + +#[cfg(feature = "parquet_encryption")] +impl EncryptionContext { + fn new( + file_decryption_properties: Option<Arc<FileDecryptionProperties>>, + encryption_factory: Option<( + Arc<dyn EncryptionFactory>, + EncryptionFactoryOptions, + )>, + ) -> Self { + Self { + file_decryption_properties, + encryption_factory, + } + } + + pub(super) async fn get_file_decryption_properties( + &self, + file_location: &object_store::path::Path, + ) -> Result<Option<Arc<FileDecryptionProperties>>> { + match &self.file_decryption_properties { + Some(file_decryption_properties) => { + Ok(Some(Arc::clone(file_decryption_properties))) + } + None => match &self.encryption_factory { + Some((encryption_factory, encryption_config)) => Ok(encryption_factory + .get_file_decryption_properties(encryption_config, file_location) + .await?), + None => Ok(None), + }, + } + } +} + +#[cfg(not(feature = "parquet_encryption"))] +#[expect(dead_code)] +impl EncryptionContext { + pub(super) async fn get_file_decryption_properties( + &self, + _file_location: &object_store::path::Path, + ) -> Result< + Option<std::sync::Arc<datafusion_common::encryption::FileDecryptionProperties>>, + > { + Ok(None) + } +} + +impl ParquetMorselizer { + #[cfg(feature = "parquet_encryption")] + pub(super) fn get_encryption_context(&self) -> EncryptionContext { + EncryptionContext::new( + self.file_decryption_properties.clone(), + self.encryption_factory.clone(), + ) + } + + #[cfg(not(feature = "parquet_encryption"))] + #[expect(dead_code)] + pub(super) fn get_encryption_context(&self) -> EncryptionContext { + EncryptionContext::default() + } +} diff --git a/datafusion/datasource-parquet/src/opener.rs b/datafusion/datasource-parquet/src/opener/mod.rs similarity index 96% rename from datafusion/datasource-parquet/src/opener.rs rename to datafusion/datasource-parquet/src/opener/mod.rs index 11cf786a3d..e5929fd43f 100644 --- a/datafusion/datasource-parquet/src/opener.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -17,6 +17,12 @@ //! [`ParquetMorselizer`] state machines for opening Parquet files +mod early_stop; +mod encryption; + +use self::early_stop::EarlyStoppingStream; +#[cfg(feature = "parquet_encryption")] +use self::encryption::EncryptionContext; use crate::access_plan::PreparedAccessPlan; use crate::page_filter::PagePruningAccessPlanFilter; use crate::push_decoder::{DecoderBuilderConfig, PushDecoderStreamState}; @@ -36,11 +42,10 @@ use std::collections::{HashMap, VecDeque}; use std::fmt; use std::future::Future; use std::mem; -use std::pin::Pin; use std::sync::Arc; -use std::task::{Context, Poll}; use arrow::datatypes::{SchemaRef, TimeUnit}; +#[cfg(feature = "parquet_encryption")] use datafusion_common::encryption::FileDecryptionProperties; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, Result, ScalarValue, Statistics, exec_err}; @@ -53,7 +58,6 @@ use datafusion_physical_expr_common::physical_expr::{ use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, - PruningMetrics, }; use datafusion_pruning::{FilePruner, PruningPredicate, build_pruning_predicate}; @@ -61,9 +65,7 @@ use datafusion_pruning::{FilePruner, PruningPredicate, build_pruning_predicate}; use datafusion_common::config::EncryptionFactoryOptions; #[cfg(feature = "parquet_encryption")] use datafusion_execution::parquet_encryption::EncryptionFactory; -use futures::{ - FutureExt, Stream, StreamExt, future::BoxFuture, ready, stream::BoxStream, -}; +use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream}; use log::debug; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; @@ -1322,153 +1324,6 @@ fn constant_value_from_stats( None } -/// 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). -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 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()) - } - } - } -} - -#[derive(Default)] -struct EncryptionContext { - #[cfg(feature = "parquet_encryption")] - file_decryption_properties: Option<Arc<FileDecryptionProperties>>, - #[cfg(feature = "parquet_encryption")] - encryption_factory: Option<(Arc<dyn EncryptionFactory>, EncryptionFactoryOptions)>, -} - -#[cfg(feature = "parquet_encryption")] -impl EncryptionContext { - fn new( - file_decryption_properties: Option<Arc<FileDecryptionProperties>>, - encryption_factory: Option<( - Arc<dyn EncryptionFactory>, - EncryptionFactoryOptions, - )>, - ) -> Self { - Self { - file_decryption_properties, - encryption_factory, - } - } - - async fn get_file_decryption_properties( - &self, - file_location: &object_store::path::Path, - ) -> Result<Option<Arc<FileDecryptionProperties>>> { - match &self.file_decryption_properties { - Some(file_decryption_properties) => { - Ok(Some(Arc::clone(file_decryption_properties))) - } - None => match &self.encryption_factory { - Some((encryption_factory, encryption_config)) => Ok(encryption_factory - .get_file_decryption_properties(encryption_config, file_location) - .await?), - None => Ok(None), - }, - } - } -} - -#[cfg(not(feature = "parquet_encryption"))] -#[expect(dead_code)] -impl EncryptionContext { - async fn get_file_decryption_properties( - &self, - _file_location: &object_store::path::Path, - ) -> Result<Option<Arc<FileDecryptionProperties>>> { - Ok(None) - } -} - -impl ParquetMorselizer { - #[cfg(feature = "parquet_encryption")] - fn get_encryption_context(&self) -> EncryptionContext { - EncryptionContext::new( - self.file_decryption_properties.clone(), - self.encryption_factory.clone(), - ) - } - - #[cfg(not(feature = "parquet_encryption"))] - #[expect(dead_code)] - fn get_encryption_context(&self) -> EncryptionContext { - EncryptionContext::default() - } -} - /// Return the initial [`ParquetAccessPlan`] /// /// If the user has supplied one as an extension, use that --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
