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)?; +```