From fe49ae7e64dc423c7e6c9fc7098ef075dcc7c044 Mon Sep 17 00:00:00 2001 From: mitchsw Date: Mon, 6 Jul 2026 13:36:21 -0400 Subject: [PATCH] Decouple FileSource trait from concrete FileScanConfig DataSource impl, via FileSourceArgs FileSource::create_file_opener and create_morselizer took a &FileScanConfig, tying the file-format engines to that one concrete DataSource. A custom DataSource could not reuse an existing FileSource (e.g. ParquetSource) without first constructing a FileScanConfig. Introduce FileSourceArgs, a small freely-constructible struct carrying the open-time inputs a FileSource needs (object_store, batch_size, limit, preserve_order, file_compression_type, expr_adapter_factory), and have both methods take &FileSourceArgs instead of &FileScanConfig. FileScanConfig builds it inline in open_with_args; a custom DataSource can build it the same way. Also remove FileSource::with_batch_size: batch size now flows through FileSourceArgs, so sources no longer store it or guard it with an expect(). This opens the door for custom DataSource implementations to reuse FileSource without depending on FileScanConfig. Updates the Parquet/Csv/Json/Arrow/Avro sources, the test sources, the csv/json opener example, and adds a 55.0.0 upgrade note. --- .../custom_data_source/csv_json_opener.rs | 13 ++- .../core/src/datasource/physical_plan/mod.rs | 2 +- .../physical_optimizer/pushdown_utils.rs | 24 ++---- datafusion/datasource-arrow/src/source.rs | 37 +++------ datafusion/datasource-avro/src/source.rs | 53 ++++++------ datafusion/datasource-csv/src/source.rs | 40 ++++------ datafusion/datasource-json/src/source.rs | 22 ++--- datafusion/datasource-parquet/src/source.rs | 33 +++----- datafusion/datasource/src/file.rs | 78 ++++++++++++++++-- .../datasource/src/file_scan_config/mod.rs | 38 ++++----- datafusion/datasource/src/test_util.rs | 10 +-- .../library-user-guide/upgrading/55.0.0.md | 80 +++++++++++++++++++ 12 files changed, 251 insertions(+), 179 deletions(-) diff --git a/datafusion-examples/examples/custom_data_source/csv_json_opener.rs b/datafusion-examples/examples/custom_data_source/csv_json_opener.rs index 51c0e2167053e..a672fbb542d6c 100644 --- a/datafusion-examples/examples/custom_data_source/csv_json_opener.rs +++ b/datafusion-examples/examples/custom_data_source/csv_json_opener.rs @@ -28,7 +28,7 @@ use datafusion::{ listing::PartitionedFile, object_store::ObjectStoreUrl, physical_plan::{ - CsvSource, FileSource, FileStreamBuilder, JsonOpener, JsonSource, + CsvSource, FileSourceArgs, FileStreamBuilder, JsonOpener, JsonSource, }, }, error::Result, @@ -66,20 +66,17 @@ async fn csv_opener() -> Result<()> { let source = CsvSource::new(Arc::clone(&schema)) .with_csv_options(options) - .with_comment(Some(b'#')) - .with_batch_size(8192); + .with_comment(Some(b'#')); let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source.into()) .with_projection_indices(Some(vec![0, 1]))? .with_limit(Some(5)) .with_file(PartitionedFile::new(csv_path.display().to_string(), 10)) .build(); - let opener = - scan_config - .file_source() - .create_file_opener(object_store, &scan_config, 0)?; + let args = FileSourceArgs::new(object_store).with_batch_size(8192); + let opener = scan_config.file_source().create_file_opener(&args, 0)?; let mut result = vec![]; let metrics = ExecutionPlanMetricsSet::new(); diff --git a/datafusion/core/src/datasource/physical_plan/mod.rs b/datafusion/core/src/datasource/physical_plan/mod.rs index 8e4855afa66bb..64c6fdd32349e 100644 --- a/datafusion/core/src/datasource/physical_plan/mod.rs +++ b/datafusion/core/src/datasource/physical_plan/mod.rs @@ -39,7 +39,7 @@ pub use json::{JsonOpener, JsonSource}; pub use arrow::{ArrowOpener, ArrowSource}; pub use csv::{CsvOpener, CsvSource}; -pub use datafusion_datasource::file::FileSource; +pub use datafusion_datasource::file::{FileSource, FileSourceArgs}; pub use datafusion_datasource::file_groups::FileGroup; pub use datafusion_datasource::file_groups::FileGroupPartitioner; pub use datafusion_datasource::file_scan_config::{ diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 2ffd1899b3c1d..7ecef5fd91f37 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -20,9 +20,12 @@ use arrow::{array::RecordBatch, compute::concat_batches}; use datafusion::{datasource::object_store::ObjectStoreUrl, physical_plan::PhysicalExpr}; use datafusion_common::{Result, config::ConfigOptions, internal_err}; use datafusion_datasource::{ - PartitionedFile, file::FileSource, file_scan_config::FileScanConfig, - file_scan_config::FileScanConfigBuilder, file_stream::FileOpenFuture, - file_stream::FileOpener, source::DataSourceExec, + PartitionedFile, + file::{FileSource, FileSourceArgs}, + file_scan_config::FileScanConfigBuilder, + file_stream::FileOpenFuture, + file_stream::FileOpener, + source::DataSourceExec, }; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr_common::physical_expr::fmt_sql; @@ -40,7 +43,6 @@ use datafusion_physical_plan::{ }; use futures::StreamExt; use futures::{FutureExt, Stream}; -use object_store::ObjectStore; use std::{ fmt::{Display, Formatter}, pin::Pin, @@ -102,7 +104,6 @@ impl FileOpener for TestOpener { pub struct TestSource { support: bool, predicate: Option>, - batch_size: Option, batches: Vec, metrics: ExecutionPlanMetricsSet, projection: Option, @@ -117,7 +118,6 @@ impl TestSource { metrics: ExecutionPlanMetricsSet::new(), batches, predicate: None, - batch_size: None, projection: None, table_schema, } @@ -127,13 +127,12 @@ impl TestSource { impl FileSource for TestSource { fn create_file_opener( &self, - _object_store: Arc, - _base_config: &FileScanConfig, + args: &FileSourceArgs, _partition: usize, ) -> Result> { Ok(Arc::new(TestOpener { batches: self.batches.clone(), - batch_size: self.batch_size, + batch_size: Some(args.batch_size), projection: self.projection.clone(), predicate: self.predicate.clone(), })) @@ -143,13 +142,6 @@ impl FileSource for TestSource { self.predicate.clone() } - fn with_batch_size(&self, batch_size: usize) -> Arc { - Arc::new(TestSource { - batch_size: Some(batch_size), - ..self.clone() - }) - } - fn metrics(&self) -> &ExecutionPlanMetricsSet { &self.metrics } diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 59c020c779ca2..672a79c4e7f22 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -41,7 +41,7 @@ use arrow::ipc::reader::{FileDecoder, FileReader, StreamReader}; use datafusion_common::error::Result; use datafusion_common::exec_datafusion_err; use datafusion_datasource::PartitionedFile; -use datafusion_datasource::file::FileSource; +use datafusion_datasource::file::{FileSource, FileSourceArgs}; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -290,10 +290,10 @@ impl ArrowSource { impl FileSource for ArrowSource { fn create_file_opener( &self, - object_store: Arc, - _base_config: &FileScanConfig, + args: &FileSourceArgs, _partition: usize, ) -> Result> { + let object_store = Arc::clone(&args.object_store); let split_projection = self.projection.clone(); let opener: Arc = match self.format { @@ -313,10 +313,6 @@ impl FileSource for ArrowSource { ) } - fn with_batch_size(&self, _batch_size: usize) -> Arc { - Arc::new(Self { ..self.clone() }) - } - fn metrics(&self) -> &ExecutionPlanMetricsSet { &self.metrics } @@ -484,13 +480,8 @@ mod tests { Arc::new(ArrowSource::new_file_source(schema)) }; - let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - source.clone(), - ) - .build(); - - let file_opener = source.create_file_opener(object_store, &scan_config, 0)?; + let file_opener = + source.create_file_opener(&FileSourceArgs::new(object_store), 0)?; let mut stream = file_opener.open(partitioned_file)?.await?; assert!(stream.next().await.is_some()); @@ -526,13 +517,8 @@ mod tests { let source = Arc::new(ArrowSource::new_file_source(schema)); - let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - source.clone(), - ) - .build(); - - let file_opener = source.create_file_opener(object_store, &scan_config, 0)?; + let file_opener = + source.create_file_opener(&FileSourceArgs::new(object_store), 0)?; let mut stream = file_opener.open(partitioned_file)?.await?; assert!(stream.next().await.is_some()); @@ -567,13 +553,8 @@ mod tests { let source = Arc::new(ArrowSource::new_stream_file_source(schema)); - let scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::local_filesystem(), - source.clone(), - ) - .build(); - - let file_opener = source.create_file_opener(object_store, &scan_config, 0)?; + let file_opener = + source.create_file_opener(&FileSourceArgs::new(object_store), 0)?; let result = file_opener.open(partitioned_file); assert!(result.is_err()); diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index e3be9d8a401d0..fa820dd2c8e86 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -23,21 +23,17 @@ use arrow::datatypes::{Schema, SchemaRef}; use arrow_avro::reader::{Reader, ReaderBuilder}; use datafusion_common::error::Result; use datafusion_datasource::TableSchema; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file::{FileSource, FileSourceArgs}; use datafusion_datasource::file_stream::FileOpener; use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use datafusion_physical_expr_adapter::BatchAdapterFactory; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::projection::ProjectionExprs; -use object_store::ObjectStore; - /// AvroSource holds the extra configuration that is necessary for opening avro files #[derive(Clone)] pub struct AvroSource { table_schema: TableSchema, - batch_size: Option, projection: SplitProjection, metrics: ExecutionPlanMetricsSet, } @@ -49,7 +45,6 @@ impl AvroSource { Self { projection: SplitProjection::unprojected(&table_schema), table_schema, - batch_size: None, metrics: ExecutionPlanMetricsSet::new(), } } @@ -58,9 +53,9 @@ impl AvroSource { &self, reader: R, projection: Option>, + batch_size: usize, ) -> Result> { - let mut builder = ReaderBuilder::new() - .with_batch_size(self.batch_size.expect("Batch size must set before open")); + let mut builder = ReaderBuilder::new().with_batch_size(batch_size); if let Some(projection) = projection { builder = builder.with_projection(projection); } @@ -114,13 +109,13 @@ impl AvroSource { impl FileSource for AvroSource { fn create_file_opener( &self, - object_store: Arc, - _base_config: &FileScanConfig, + args: &FileSourceArgs, _partition: usize, ) -> Result> { let mut opener = Arc::new(private::AvroOpener { config: Arc::new(self.clone()), - object_store, + object_store: Arc::clone(&args.object_store), + batch_size: args.batch_size, }) as Arc; opener = ProjectionOpener::try_new( self.projection.clone(), @@ -134,12 +129,6 @@ impl FileSource for AvroSource { &self.table_schema } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - fn try_pushdown_projection( &self, projection: &ProjectionExprs, @@ -183,12 +172,14 @@ mod private { pub struct AvroOpener { pub config: Arc, pub object_store: Arc, + pub batch_size: usize, } impl FileOpener for AvroOpener { fn open(&self, partitioned_file: PartitionedFile) -> Result { let object_store = Arc::clone(&self.object_store); let config = Arc::clone(&self.config); + let batch_size = self.batch_size; let projected_file_schema = config.projected_file_schema(); Ok(Box::pin(async move { @@ -199,15 +190,21 @@ mod private { GetResultPayload::File(mut file, _) => { // Probe the writer schema first so logical projected columns can be // translated to the writer-schema ordinals expected by `arrow-avro`. - let probe_reader = - config.open(BufReader::new(file.try_clone()?), None)?; + let probe_reader = config.open( + BufReader::new(file.try_clone()?), + None, + batch_size, + )?; let writer_projection = config.writer_projection_for_schema( probe_reader.schema().as_ref(), projected_file_schema.as_ref(), ); file.rewind()?; - let reader = - config.open(BufReader::new(file), writer_projection)?; + let reader = config.open( + BufReader::new(file), + writer_projection, + batch_size, + )?; let batch_adapter = BatchAdapterFactory::new(Arc::clone(&projected_file_schema)) .make_adapter(&reader.schema())?; @@ -222,14 +219,20 @@ mod private { let bytes = r.bytes().await?; // As above, inspect the writer schema before constructing the real // reader so `with_projection` can use writer-schema ordinals. - let probe_reader = - config.open(BufReader::new(bytes.clone().reader()), None)?; + let probe_reader = config.open( + BufReader::new(bytes.clone().reader()), + None, + batch_size, + )?; let writer_projection = config.writer_projection_for_schema( probe_reader.schema().as_ref(), projected_file_schema.as_ref(), ); - let reader = config - .open(BufReader::new(bytes.reader()), writer_projection)?; + let reader = config.open( + BufReader::new(bytes.reader()), + writer_projection, + batch_size, + )?; let batch_adapter = BatchAdapterFactory::new(Arc::clone(&projected_file_schema)) .make_adapter(&reader.schema())?; diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 25ec311880405..c6ea81d8afa5e 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -35,8 +35,7 @@ use arrow::csv; use datafusion_common::config::CsvOptions; use datafusion_common::{DataFusionError, Result, exec_datafusion_err}; use datafusion_common_runtime::JoinSet; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file::{FileSource, FileSourceArgs}; use datafusion_execution::TaskContext; use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use datafusion_physical_plan::{ @@ -85,7 +84,6 @@ use tokio::io::AsyncWriteExt; #[derive(Debug, Clone)] pub struct CsvSource { options: CsvOptions, - batch_size: Option, table_schema: TableSchema, projection: SplitProjection, metrics: ExecutionPlanMetricsSet, @@ -99,7 +97,6 @@ impl CsvSource { options: CsvOptions::default(), projection: SplitProjection::unprojected(&table_schema), table_schema, - batch_size: None, metrics: ExecutionPlanMetricsSet::new(), } } @@ -179,18 +176,15 @@ impl CsvSource { } impl CsvSource { - fn open(&self, reader: R) -> Result> { - Ok(self.builder().build(reader)?) + fn open(&self, reader: R, batch_size: usize) -> Result> { + Ok(self.builder(batch_size).build(reader)?) } - fn builder(&self) -> csv::ReaderBuilder { + fn builder(&self, batch_size: usize) -> csv::ReaderBuilder { let mut builder = csv::ReaderBuilder::new(Arc::clone(self.table_schema.file_schema())) .with_delimiter(self.delimiter()) - .with_batch_size( - self.batch_size - .expect("Batch size must be set before initializing builder"), - ) + .with_batch_size(batch_size) .with_header(self.has_header()) .with_quote(self.quote()) .with_truncated_rows(self.truncate_rows()); @@ -214,6 +208,7 @@ pub struct CsvOpener { config: Arc, file_compression_type: FileCompressionType, object_store: Arc, + batch_size: usize, partition_index: usize, } @@ -223,11 +218,13 @@ impl CsvOpener { config: Arc, file_compression_type: FileCompressionType, object_store: Arc, + batch_size: usize, ) -> Self { Self { config, file_compression_type, object_store, + batch_size, partition_index: 0, } } @@ -242,14 +239,14 @@ impl From for Arc { impl FileSource for CsvSource { fn create_file_opener( &self, - object_store: Arc, - base_config: &FileScanConfig, + args: &FileSourceArgs, partition_index: usize, ) -> Result> { let mut opener = Arc::new(CsvOpener { config: Arc::new(self.clone()), - file_compression_type: base_config.file_compression_type, - object_store, + file_compression_type: args.file_compression_type, + object_store: Arc::clone(&args.object_store), + batch_size: args.batch_size, partition_index, }) as Arc; opener = ProjectionOpener::try_new( @@ -264,12 +261,6 @@ impl FileSource for CsvSource { &self.table_schema } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - fn try_pushdown_projection( &self, projection: &ProjectionExprs, @@ -359,6 +350,7 @@ impl FileOpener for CsvOpener { } let store = Arc::clone(&self.object_store); + let batch_size = self.batch_size; let terminator = self.config.terminator(); let baseline_metrics = @@ -394,7 +386,7 @@ impl FileOpener for CsvOpener { .await? .map_err(DataFusionError::from); - let decoder = config.builder().build_decoder(); + let decoder = config.builder(batch_size).build_decoder(); let input = file_compression_type .convert_stream(aligned_stream.boxed())? .fuse(); @@ -413,7 +405,7 @@ impl FileOpener for CsvOpener { #[cfg(not(target_arch = "wasm32"))] GetResultPayload::File(file, _) => { let decoder = file_compression_type.convert_read(file)?; - let mut reader = config.open(decoder)?; + let mut reader = config.open(decoder, batch_size)?; // Use std::iter::from_fn to wrap execution of iterator's next() method. let iterator = std::iter::from_fn(move || { @@ -428,7 +420,7 @@ impl FileOpener for CsvOpener { .boxed()) } GetResultPayload::Stream(s) => { - let decoder = config.builder().build_decoder(); + let decoder = config.builder(batch_size).build_decoder(); let s = s.map_err(DataFusionError::from); let input = file_compression_type.convert_stream(s.boxed())?.fuse(); diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index 8632d6b942bc1..e79cfce7955e8 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -40,8 +40,7 @@ use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use arrow::array::RecordBatch; use arrow::json::ReaderBuilder; use arrow::{datatypes::SchemaRef, json}; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file::{FileSource, FileSourceArgs}; use datafusion_execution::TaskContext; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; @@ -129,7 +128,6 @@ impl JsonOpener { #[derive(Clone)] pub struct JsonSource { table_schema: datafusion_datasource::TableSchema, - batch_size: Option, metrics: ExecutionPlanMetricsSet, projection: SplitProjection, /// When `true` (default), expects newline-delimited JSON (NDJSON). @@ -144,7 +142,6 @@ impl JsonSource { Self { projection: SplitProjection::unprojected(&table_schema), table_schema, - batch_size: None, metrics: ExecutionPlanMetricsSet::new(), newline_delimited: true, } @@ -169,8 +166,7 @@ impl From for Arc { impl FileSource for JsonSource { fn create_file_opener( &self, - object_store: Arc, - base_config: &FileScanConfig, + args: &FileSourceArgs, _partition: usize, ) -> Result> { // Get the projected file schema for JsonOpener @@ -179,12 +175,10 @@ impl FileSource for JsonSource { Arc::new(file_schema.project(&self.projection.file_indices)?); let mut opener = Arc::new(JsonOpener { - batch_size: self - .batch_size - .expect("Batch size must set before creating opener"), + batch_size: args.batch_size, projected_schema, - file_compression_type: base_config.file_compression_type, - object_store, + file_compression_type: args.file_compression_type, + object_store: Arc::clone(&args.object_store), newline_delimited: self.newline_delimited, }) as Arc; @@ -202,12 +196,6 @@ impl FileSource for JsonSource { &self.table_schema } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - fn try_pushdown_projection( &self, projection: &ProjectionExprs, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3443b08475e0d..628e68576dbcd 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -41,8 +41,7 @@ use arrow::datatypes::TimeUnit; use datafusion_common::DataFusionError; use datafusion_common::config::TableParquetOptions; use datafusion_datasource::TableSchema; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file::{FileSource, FileSourceArgs}; use datafusion_functions::core::file_row_index::FileRowIndexFunc; use datafusion_physical_expr::expressions::{Column, DynamicFilterTracking}; use datafusion_physical_expr::projection::ProjectionExprs; @@ -67,7 +66,6 @@ use log::warn; use datafusion_execution::parquet_encryption::EncryptionFactory; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use itertools::Itertools; -use object_store::ObjectStore; use parquet::arrow::RowNumber; #[cfg(feature = "parquet_encryption")] use parquet::encryption::decrypt::FileDecryptionProperties; @@ -298,8 +296,6 @@ pub struct ParquetSource { pub(crate) predicate: Option>, /// Optional user defined parquet file reader factory pub(crate) parquet_file_reader_factory: Option>, - /// Batch size configuration - pub(crate) batch_size: Option, /// Optional hint for the size of the parquet metadata pub(crate) metadata_size_hint: Option, /// Projection to apply to the output. @@ -333,7 +329,6 @@ impl ParquetSource { metrics: ExecutionPlanMetricsSet::new(), predicate: None, parquet_file_reader_factory: None, - batch_size: None, metadata_size_hint: None, #[cfg(feature = "parquet_encryption")] encryption_factory: None, @@ -550,8 +545,7 @@ impl From for Arc { impl FileSource for ParquetSource { fn create_file_opener( &self, - _object_store: Arc, - _base_config: &FileScanConfig, + _args: &FileSourceArgs, _partition: usize, ) -> datafusion_common::Result> { datafusion_common::internal_err!( @@ -561,18 +555,19 @@ impl FileSource for ParquetSource { fn create_morselizer( &self, - object_store: Arc, - base_config: &FileScanConfig, + args: &FileSourceArgs, partition: usize, ) -> datafusion_common::Result> { - let expr_adapter_factory = base_config + let expr_adapter_factory = args .expr_adapter_factory .clone() .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); let parquet_file_reader_factory = self.parquet_file_reader_factory.clone().unwrap_or_else(|| { - Arc::new(DefaultParquetFileReaderFactory::new(object_store)) as _ + Arc::new(DefaultParquetFileReaderFactory::new(Arc::clone( + &args.object_store, + ))) as _ }); #[cfg(feature = "parquet_encryption")] @@ -623,11 +618,9 @@ impl FileSource for ParquetSource { Ok(Box::new(ParquetMorselizer { partition_index: partition, projection: self.projection.clone(), - batch_size: self - .batch_size - .expect("Batch size must set before creating ParquetMorselizer"), - limit: base_config.limit, - preserve_order: base_config.preserve_order, + batch_size: args.batch_size, + limit: args.limit, + preserve_order: args.preserve_order, predicate: self.predicate.clone(), table_schema: self.table_schema.clone(), metadata_size_hint: self.metadata_size_hint, @@ -673,12 +666,6 @@ impl FileSource for ParquetSource { self.predicate.clone() } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - fn try_pushdown_projection( &self, projection: &ProjectionExprs, diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 07460b23694b7..d1a289e966109 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -22,6 +22,7 @@ use std::fmt; use std::fmt::Formatter; use std::sync::Arc; +use crate::file_compression_type::FileCompressionType; use crate::file_groups::FileGroupPartitioner; use crate::file_scan_config::FileScanConfig; use crate::file_stream::FileOpener; @@ -32,6 +33,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::{Result, not_impl_err}; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; +use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_plan::DisplayFormatType; use datafusion_physical_plan::SortOrderPushdownResult; use datafusion_physical_plan::filter_pushdown::{FilterPushdownPropagation, PushedDown}; @@ -45,6 +47,71 @@ pub fn as_file_source(source: T) -> Arc Arc::new(source) } +/// Arguments a [`DataSource`] passes to its [`FileSource`] when creating an +/// opener or morselizer. +/// +/// [`DataSource`]: crate::source::DataSource +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct FileSourceArgs { + /// Store to read file bytes from. + pub object_store: Arc, + /// Target number of rows per output batch. + pub batch_size: usize, + /// Optional row limit for the scan. + pub limit: Option, + /// Whether file/row order must be preserved. + pub preserve_order: bool, + /// Compression of the files being read. + pub file_compression_type: FileCompressionType, + /// Adapter factory for rewriting expressions to the physical file schema. + pub expr_adapter_factory: Option>, +} + +impl FileSourceArgs { + pub fn new(object_store: Arc) -> Self { + Self { + object_store, + batch_size: 8192, + limit: None, + preserve_order: false, + file_compression_type: FileCompressionType::UNCOMPRESSED, + expr_adapter_factory: None, + } + } + + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + pub fn with_preserve_order(mut self, preserve_order: bool) -> Self { + self.preserve_order = preserve_order; + self + } + + pub fn with_file_compression_type( + mut self, + file_compression_type: FileCompressionType, + ) -> Self { + self.file_compression_type = file_compression_type; + self + } + + pub fn with_expr_adapter_factory( + mut self, + expr_adapter_factory: Option>, + ) -> Self { + self.expr_adapter_factory = expr_adapter_factory; + self + } +} + /// File format specific behaviors for [`DataSource`] /// /// # Schema information @@ -69,8 +136,7 @@ pub trait FileSource: Any + Send + Sync { /// error from this method and implementing [`Self::create_morselizer`] instead. fn create_file_opener( &self, - object_store: Arc, - base_config: &FileScanConfig, + args: &FileSourceArgs, partition: usize, ) -> Result>; @@ -83,11 +149,10 @@ pub trait FileSource: Any + Send + Sync { /// implementing this method. fn create_morselizer( &self, - object_store: Arc, - base_config: &FileScanConfig, + args: &FileSourceArgs, partition: usize, ) -> Result> { - let opener = self.create_file_opener(object_store, base_config, partition)?; + let opener = self.create_file_opener(args, partition)?; Ok(Box::new(FileOpenerMorselizer::new(opener))) } @@ -103,9 +168,6 @@ pub trait FileSource: Any + Send + Sync { /// after applying the projection. fn table_schema(&self) -> &crate::table_schema::TableSchema; - /// Initialize new type with batch size configuration - fn with_batch_size(&self, batch_size: usize) -> Arc; - /// Returns the filter expression that will be applied *during* the file scan. /// /// These expressions are in terms of the unprojected [`Self::table_schema`]. diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 1336ee69cd6dd..fca66bad495a3 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -22,7 +22,7 @@ pub(crate) mod sort_pushdown; use crate::file_groups::FileGroup; use crate::{ - PartitionedFile, display::FileGroupsDisplay, file::FileSource, + PartitionedFile, display::FileGroupsDisplay, file::FileSource, file::FileSourceArgs, file_compression_type::FileCompressionType, file_stream::FileStreamBuilder, file_stream::work_source::SharedWorkSource, source::DataSource, statistics::MinMaxStatistics, @@ -79,12 +79,11 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// ``` /// # use std::sync::Arc; /// # use arrow::datatypes::{Field, Fields, DataType, Schema, SchemaRef}; -/// # use object_store::ObjectStore; /// # use datafusion_common::Result; -/// # use datafusion_datasource::file::FileSource; +/// # use datafusion_datasource::file::{FileSource, FileSourceArgs}; /// # use datafusion_datasource::file_groups::FileGroup; /// # use datafusion_datasource::PartitionedFile; -/// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +/// # use datafusion_datasource::file_scan_config::FileScanConfigBuilder; /// # use datafusion_datasource::file_stream::FileOpener; /// # use datafusion_datasource::source::DataSourceExec; /// # use datafusion_datasource::table_schema::TableSchema; @@ -104,9 +103,8 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # table_schema: TableSchema, /// # }; /// # impl FileSource for ParquetSource { -/// # fn create_file_opener(&self, _: Arc, _: &FileScanConfig, _: usize) -> Result> { unimplemented!() } +/// # fn create_file_opener(&self, _: &FileSourceArgs, _: usize) -> Result> { unimplemented!() } /// # fn table_schema(&self) -> &TableSchema { &self.table_schema } -/// # fn with_batch_size(&self, _: usize) -> Arc { unimplemented!() } /// # fn metrics(&self) -> &ExecutionPlanMetricsSet { unimplemented!() } /// # fn file_type(&self) -> &str { "parquet" } /// # // Note that this implementation drops the projection on the floor, it is not complete! @@ -690,9 +688,16 @@ impl DataSource for FileScanConfig { .batch_size .unwrap_or_else(|| context.session_config().batch_size()); - let source = self.file_source.with_batch_size(batch_size); + let source_args = FileSourceArgs::new(object_store) + .with_batch_size(batch_size) + .with_limit(self.limit) + .with_preserve_order(self.preserve_order) + .with_file_compression_type(self.file_compression_type) + .with_expr_adapter_factory(self.expr_adapter_factory.clone()); - let morselizer = source.create_morselizer(object_store, self, partition)?; + let morselizer = self + .file_source + .create_morselizer(&source_args, partition)?; // Extract the shared work source from the sibling state if it exists. // This allows multiple sibling streams to steal work from a single @@ -706,7 +711,7 @@ impl DataSource for FileScanConfig { .with_partition(partition) .with_shared_work_source(shared_work_source) .with_morselizer(morselizer) - .with_metrics(source.metrics()) + .with_metrics(self.file_source.metrics()) .build()?; Ok(Box::pin(cooperative(stream))) } @@ -1567,7 +1572,6 @@ mod tests { use futures::FutureExt as _; use futures::StreamExt as _; use futures::stream; - use object_store::ObjectStore; use std::fmt::Debug; #[derive(Clone)] @@ -1588,8 +1592,7 @@ mod tests { impl FileSource for InexactSortPushdownSource { fn create_file_opener( &self, - _object_store: Arc, - _base_config: &FileScanConfig, + _args: &FileSourceArgs, _partition: usize, ) -> Result> { unimplemented!() @@ -1599,10 +1602,6 @@ mod tests { &self.table_schema } - fn with_batch_size(&self, _batch_size: usize) -> Arc { - Arc::new(self.clone()) - } - fn metrics(&self) -> &ExecutionPlanMetricsSet { &self.metrics } @@ -2801,8 +2800,7 @@ mod tests { impl FileSource for ExactSortPushdownSource { fn create_file_opener( &self, - _object_store: Arc, - _base_config: &FileScanConfig, + _args: &FileSourceArgs, _partition: usize, ) -> Result> { unimplemented!() @@ -2812,10 +2810,6 @@ mod tests { &self.table_schema } - fn with_batch_size(&self, _batch_size: usize) -> Arc { - Arc::new(self.clone()) - } - fn metrics(&self) -> &ExecutionPlanMetricsSet { &self.metrics } diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index 5bef6d44e1408..8968be0d98a99 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -16,7 +16,8 @@ // under the License. use crate::{ - file::FileSource, file_scan_config::FileScanConfig, file_stream::FileOpener, + file::{FileSource, FileSourceArgs}, + file_stream::FileOpener, }; use std::sync::Arc; @@ -77,8 +78,7 @@ impl MockSource { impl FileSource for MockSource { fn create_file_opener( &self, - _object_store: Arc, - _base_config: &FileScanConfig, + _args: &FileSourceArgs, _partition: usize, ) -> Result> { self.file_opener.clone().ok_or_else(|| { @@ -90,10 +90,6 @@ impl FileSource for MockSource { self.filter.clone() } - fn with_batch_size(&self, _batch_size: usize) -> Arc { - Arc::new(Self { ..self.clone() }) - } - fn metrics(&self) -> &ExecutionPlanMetricsSet { &self.metrics } diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 181e0e0b7f266..1e5609cf3225c 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -583,3 +583,83 @@ metadata and schema fingerprint match. - Pass the current schema fingerprint to `CachedFileMetadata::is_valid_for`. See [PR #23201](https://github.com/apache/datafusion/pull/23201) for details. + +### `FileSource` reader methods take `FileSourceArgs` instead of `FileScanConfig` + +`datafusion_datasource::file::FileSource` no longer depends on `FileScanConfig` +in its read path. `create_file_opener` and `create_morselizer` now take a single +`&FileSourceArgs` in place of the separate `object_store` and +`base_config: &FileScanConfig` arguments: + +```diff + fn create_file_opener( + &self, +- object_store: Arc, +- base_config: &FileScanConfig, ++ args: &FileSourceArgs, + partition: usize, + ) -> Result>; +``` + +`FileSourceArgs` carries the scan-level inputs a source needs to open a file: +`object_store`, `batch_size`, `limit`, `preserve_order`, +`file_compression_type`, and `expr_adapter_factory`. It is built with +`FileSourceArgs::new(object_store)` (the only field without a default) plus +`with_*` setters. Because these values are now passed in explicitly, a +`FileSource` can be driven by a custom `DataSource` without constructing a +`FileScanConfig`. + +`FileSource::with_batch_size` has also been removed. Batch size is supplied +through `FileSourceArgs::batch_size`, so sources no longer store it themselves. + +**Who is affected:** + +- Anyone with a custom `FileSource` implementation. +- Anyone calling `create_file_opener` / `create_morselizer` directly. +- Anyone calling `FileSource::with_batch_size`. + +**Migration guide:** + +Update your method signatures, read scan-level values from `args` instead of +`base_config`, and drop `with_batch_size`: + +```diff + fn create_file_opener( + &self, +- object_store: Arc, +- base_config: &FileScanConfig, ++ args: &FileSourceArgs, + partition: usize, + ) -> Result> { + let opener = MyOpener { +- object_store, +- file_compression_type: base_config.file_compression_type, +- batch_size: self.batch_size.expect("batch size set"), ++ object_store: Arc::clone(&args.object_store), ++ file_compression_type: args.file_compression_type, ++ batch_size: args.batch_size, + // ... + }; + // ... + } + +- fn with_batch_size(&self, batch_size: usize) -> Arc { +- let mut conf = self.clone(); +- conf.batch_size = Some(batch_size); +- Arc::new(conf) +- } +``` + +Callers that previously passed the object store and config separately now build +`FileSourceArgs`: + +```diff +- let opener = source.create_file_opener(object_store, &scan_config, partition)?; ++ let args = FileSourceArgs::new(object_store) ++ .with_batch_size(batch_size) ++ .with_limit(scan_config.limit) ++ .with_preserve_order(scan_config.preserve_order) ++ .with_file_compression_type(scan_config.file_compression_type) ++ .with_expr_adapter_factory(scan_config.expr_adapter_factory.clone()); ++ let opener = source.create_file_opener(&args, partition)?; +```