Move schema and batch_size into file scan config

Author: friendlymatthew

datafusion-examples/examples/csv_json_opener.rs

@@ -61,14 +61,13 @@ async fn csv_opener() -> Result<()> {
         Arc::new(CsvSource::default()),
     )
     .with_projection(Some(vec![12, 0]))
+    .with_batch_size(Some(8192))
     .with_limit(Some(5))
     .with_file(PartitionedFile::new(path.display().to_string(), 10))
     .build();
 
     let config = CsvSource::new(true, b',', b'"')
         .with_comment(Some(b'#'))
-        .with_schema(schema)
-        .with_batch_size(8192)
         .with_projection(&scan_config);
 
     let opener = config.create_file_opener(object_store, &scan_config, 0);

datafusion/core/src/datasource/physical_plan/arrow_file.rs

@@ -24,7 +24,6 @@ use datafusion_datasource::as_file_source;
 use datafusion_datasource::schema_adapter::SchemaAdapterFactory;
 
 use arrow::buffer::Buffer;
-use arrow::datatypes::SchemaRef;
 use arrow_ipc::reader::FileDecoder;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
@@ -66,14 +65,6 @@ impl FileSource for ArrowSource {
         self
     }
 
-    fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
-    fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
         Arc::new(Self { ..self.clone() })
     }

datafusion/core/src/datasource/physical_plan/parquet.rs

@@ -185,7 +185,9 @@ mod tests {
                 source = source.with_bloom_filter_on_read(false);
             }
 
-            source.with_schema(Arc::clone(&table_schema))
+            // source.with_schema(Arc::clone(&table_schema))
+
+            Arc::new(source) as Arc<dyn FileSource>
         }
 
         fn build_parquet_exec(

datafusion/core/src/test_util/parquet.rs

@@ -37,7 +37,6 @@ use crate::physical_plan::metrics::MetricsSet;
 use crate::physical_plan::ExecutionPlan;
 use crate::prelude::{Expr, SessionConfig, SessionContext};
 
-use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_datasource::source::DataSourceExec;
 use object_store::path::Path;
@@ -185,8 +184,8 @@ impl TestParquetFile {
             let source = Arc::new(
                 ParquetSource::new(parquet_options)
                     .with_predicate(Arc::clone(&physical_filter_expr)),
-            )
-            .with_schema(Arc::clone(&self.schema));
+            );
+
             let config = scan_config_builder.with_source(source).build();
             let parquet_exec = DataSourceExec::from_data_source(config);
 

datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs

@@ -19,7 +19,7 @@ use arrow::datatypes::SchemaRef;
 use arrow::error::ArrowError;
 use arrow::{array::RecordBatch, compute::concat_batches};
 use datafusion::{datasource::object_store::ObjectStoreUrl, physical_plan::PhysicalExpr};
-use datafusion_common::{config::ConfigOptions, internal_err, Result, Statistics};
+use datafusion_common::{config::ConfigOptions, internal_err, Result};
 use datafusion_datasource::{
     file::FileSource, file_meta::FileMeta, file_scan_config::FileScanConfig,
     file_scan_config::FileScanConfigBuilder, file_stream::FileOpenFuture,
@@ -107,10 +107,7 @@ impl FileOpener for TestOpener {
 pub struct TestSource {
     support: bool,
     predicate: Option<Arc<dyn PhysicalExpr>>,
-    statistics: Option<Statistics>,
-    batch_size: Option<usize>,
     batches: Vec<RecordBatch>,
-    schema: Option<SchemaRef>,
     metrics: ExecutionPlanMetricsSet,
     projection: Option<Vec<usize>>,
     schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
@@ -131,13 +128,13 @@ impl FileSource for TestSource {
     fn create_file_opener(
         &self,
         _object_store: Arc<dyn ObjectStore>,
-        _base_config: &FileScanConfig,
+        base_config: &FileScanConfig,
         _partition: usize,
     ) -> Arc<dyn FileOpener> {
         Arc::new(TestOpener {
             batches: self.batches.clone(),
-            batch_size: self.batch_size,
-            schema: self.schema.clone(),
+            batch_size: base_config.batch_size,
+            schema: Some(Arc::clone(&base_config.file_schema)),
             projection: self.projection.clone(),
         })
     }
@@ -146,20 +143,6 @@ impl FileSource for TestSource {
         todo!("should not be called")
     }
 
-    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
-        Arc::new(TestSource {
-            batch_size: Some(batch_size),
-            ..self.clone()
-        })
-    }
-
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource> {
-        Arc::new(TestSource {
-            schema: Some(schema),
-            ..self.clone()
-        })
-    }
-
     fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource> {
         Arc::new(TestSource {
             projection: config.projection.clone(),
@@ -175,7 +158,12 @@ impl FileSource for TestSource {
         "test"
     }
 
-    fn fmt_extra(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+    fn fmt_extra(
+        &self,
+        t: DisplayFormatType,
+        f: &mut Formatter,
+        _config: &FileScanConfig,
+    ) -> std::fmt::Result {
         match t {
             DisplayFormatType::Default | DisplayFormatType::Verbose => {
                 let support = format!(", pushdown_supported={}", self.support);
@@ -202,6 +190,7 @@ impl FileSource for TestSource {
         &self,
         mut filters: Vec<Arc<dyn PhysicalExpr>>,
         config: &ConfigOptions,
+        _fs_config: &FileScanConfig,
     ) -> Result<FilterPushdownPropagation<Arc<dyn FileSource>>> {
         if self.support && config.execution.parquet.pushdown_filters {
             if let Some(internal) = self.predicate.as_ref() {

datafusion/datasource-avro/src/source.rs

@@ -36,8 +36,6 @@ use object_store::ObjectStore;
 /// AvroSource holds the extra configuration that is necessary for opening avro files
 #[derive(Clone, Default)]
 pub struct AvroSource {
-    schema: Option<SchemaRef>,
-    batch_size: Option<usize>,
     projection: Option<Vec<String>>,
     metrics: ExecutionPlanMetricsSet,
     schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
@@ -49,11 +47,16 @@ impl AvroSource {
         Self::default()
     }
 
-    fn open<R: std::io::Read>(&self, reader: R) -> Result<AvroReader<'static, R>> {
+    fn open<R: std::io::Read>(
+        &self,
+        reader: R,
+        file_schema: SchemaRef,
+        batch_size: Option<usize>,
+    ) -> Result<AvroReader<'static, R>> {
         AvroReader::try_new(
             reader,
-            Arc::clone(self.schema.as_ref().expect("Schema must set before open")),
-            self.batch_size.expect("Batch size must set before open"),
+            file_schema,
+            batch_size.expect("Batch size must set before open"),
             self.projection.clone(),
         )
     }
@@ -63,31 +66,21 @@ impl FileSource for AvroSource {
     fn create_file_opener(
         &self,
         object_store: Arc<dyn ObjectStore>,
-        _base_config: &FileScanConfig,
+        base_config: &FileScanConfig,
         _partition: usize,
     ) -> Arc<dyn FileOpener> {
         Arc::new(private::AvroOpener {
-            config: Arc::new(self.clone()),
+            source: Arc::new(self.clone()),
             object_store,
+            file_schema: base_config.file_schema.clone(),
+            batch_size: base_config.batch_size,
         })
     }
 
     fn as_any(&self) -> &dyn Any {
         self
     }
 
-    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
-        let mut conf = self.clone();
-        conf.batch_size = Some(batch_size);
-        Arc::new(conf)
-    }
-
-    fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource> {
-        let mut conf = self.clone();
-        conf.schema = Some(schema);
-        Arc::new(conf)
-    }
-
     fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource> {
         let mut conf = self.clone();
         conf.projection = config.projected_file_column_names();
@@ -138,8 +131,11 @@ mod private {
     use object_store::{GetResultPayload, ObjectStore};
 
     pub struct AvroOpener {
-        pub config: Arc<AvroSource>,
+        pub source: Arc<AvroSource>,
         pub object_store: Arc<dyn ObjectStore>,
+
+        pub file_schema: SchemaRef,
+        pub batch_size: Option<usize>,
     }
 
     impl FileOpener for AvroOpener {
@@ -148,18 +144,21 @@ mod private {
             file_meta: FileMeta,
             _file: PartitionedFile,
         ) -> Result<FileOpenFuture> {
-            let config = Arc::clone(&self.config);
+            let source = Arc::clone(&self.source);
             let object_store = Arc::clone(&self.object_store);
+            let file_schema = Arc::clone(&self.file_schema);
+            let batch_size = self.batch_size;
             Ok(Box::pin(async move {
                 let r = object_store.get(file_meta.location()).await?;
                 match r.payload {
                     GetResultPayload::File(file, _) => {
-                        let reader = config.open(file)?;
+                        let reader = source.open(file, file_schema, batch_size)?;
                         Ok(futures::stream::iter(reader).boxed())
                     }
                     GetResultPayload::Stream(_) => {
                         let bytes = r.bytes().await?;
-                        let reader = config.open(bytes.reader())?;
+                        let reader =
+                            source.open(bytes.reader(), file_schema, batch_size)?;
                         Ok(futures::stream::iter(reader).boxed())
                     }
                 }