Move projection into file scan config

Author: friendlymatthew

datafusion-examples/examples/csv_json_opener.rs

@@ -66,9 +66,7 @@ async fn csv_opener() -> Result<()> {
     .with_file(PartitionedFile::new(path.display().to_string(), 10))
     .build();
 
-    let config = CsvSource::new(true, b',', b'"')
-        .with_comment(Some(b'#'))
-        .with_projection(&scan_config);
+    let config = CsvSource::new(true, b',', b'"').with_comment(Some(b'#'));
 
     let opener = config.create_file_opener(object_store, &scan_config, 0);
 

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

@@ -65,10 +65,6 @@ impl FileSource for ArrowSource {
         self
     }
 
-    fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

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

@@ -109,7 +109,6 @@ pub struct TestSource {
     predicate: Option<Arc<dyn PhysicalExpr>>,
     batches: Vec<RecordBatch>,
     metrics: ExecutionPlanMetricsSet,
-    projection: Option<Vec<usize>>,
     schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
 }
 
@@ -135,21 +134,14 @@ impl FileSource for TestSource {
             batches: self.batches.clone(),
             batch_size: base_config.batch_size,
             schema: Some(Arc::clone(&base_config.file_schema)),
-            projection: self.projection.clone(),
+            projection: base_config.projection.clone(),
         })
     }
 
     fn as_any(&self) -> &dyn Any {
         todo!("should not be called")
     }
 
-    fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource> {
-        Arc::new(TestSource {
-            projection: config.projection.clone(),
-            ..self.clone()
-        })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource-avro/src/source.rs

@@ -36,7 +36,6 @@ use object_store::ObjectStore;
 /// AvroSource holds the extra configuration that is necessary for opening avro files
 #[derive(Clone, Default)]
 pub struct AvroSource {
-    projection: Option<Vec<String>>,
     metrics: ExecutionPlanMetricsSet,
     schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
 }
@@ -52,12 +51,13 @@ impl AvroSource {
         reader: R,
         file_schema: SchemaRef,
         batch_size: Option<usize>,
+        projected_file_column_names: Option<Vec<String>>,
     ) -> Result<AvroReader<'static, R>> {
         AvroReader::try_new(
             reader,
             file_schema,
             batch_size.expect("Batch size must set before open"),
-            self.projection.clone(),
+            projected_file_column_names,
         )
     }
 }
@@ -74,19 +74,14 @@ impl FileSource for AvroSource {
             object_store,
             file_schema: base_config.file_schema.clone(),
             batch_size: base_config.batch_size,
+            projected_file_column_names: base_config.projected_file_column_names(),
         })
     }
 
     fn as_any(&self) -> &dyn Any {
         self
     }
 
-    fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource> {
-        let mut conf = self.clone();
-        conf.projection = config.projected_file_column_names();
-        Arc::new(conf)
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }
@@ -136,6 +131,7 @@ mod private {
 
         pub file_schema: SchemaRef,
         pub batch_size: Option<usize>,
+        pub projected_file_column_names: Option<Vec<String>>,
     }
 
     impl FileOpener for AvroOpener {
@@ -148,17 +144,28 @@ mod private {
             let object_store = Arc::clone(&self.object_store);
             let file_schema = Arc::clone(&self.file_schema);
             let batch_size = self.batch_size;
+            let projected_file_names = self.projected_file_column_names.clone();
+
             Ok(Box::pin(async move {
                 let r = object_store.get(file_meta.location()).await?;
                 match r.payload {
                     GetResultPayload::File(file, _) => {
-                        let reader = source.open(file, file_schema, batch_size)?;
+                        let reader = source.open(
+                            file,
+                            file_schema,
+                            batch_size,
+                            projected_file_names,
+                        )?;
                         Ok(futures::stream::iter(reader).boxed())
                     }
                     GetResultPayload::Stream(_) => {
                         let bytes = r.bytes().await?;
-                        let reader =
-                            source.open(bytes.reader(), file_schema, batch_size)?;
+                        let reader = source.open(
+                            bytes.reader(),
+                            file_schema,
+                            batch_size,
+                            projected_file_names,
+                        )?;
                         Ok(futures::stream::iter(reader).boxed())
                     }
                 }

datafusion/datasource-csv/src/source.rs

@@ -82,7 +82,6 @@ use tokio::io::AsyncWriteExt;
 /// ```
 #[derive(Debug, Clone, Default)]
 pub struct CsvSource {
-    file_projection: Option<Vec<usize>>,
     pub(crate) has_header: bool,
     delimiter: u8,
     quote: u8,
@@ -161,14 +160,18 @@ impl CsvSource {
         reader: R,
         file_schema: SchemaRef,
         batch_size: Option<usize>,
+        file_column_projection_indices: Option<Vec<usize>>,
     ) -> Result<csv::Reader<R>> {
-        Ok(self.builder(file_schema, batch_size).build(reader)?)
+        Ok(self
+            .builder(file_schema, batch_size, file_column_projection_indices)
+            .build(reader)?)
     }
 
     fn builder(
         &self,
         file_schema: SchemaRef,
         batch_size: Option<usize>,
+        file_column_projection_indices: Option<Vec<usize>>,
     ) -> csv::ReaderBuilder {
         let mut builder = csv::ReaderBuilder::new(Arc::clone(&file_schema))
             .with_delimiter(self.delimiter)
@@ -180,7 +183,7 @@ impl CsvSource {
         if let Some(terminator) = self.terminator {
             builder = builder.with_terminator(terminator);
         }
-        if let Some(proj) = &self.file_projection {
+        if let Some(proj) = &file_column_projection_indices {
             builder = builder.with_projection(proj.clone());
         }
         if let Some(escape) = self.escape {
@@ -199,8 +202,11 @@ pub struct CsvOpener {
     source: Arc<CsvSource>,
     file_compression_type: FileCompressionType,
     object_store: Arc<dyn ObjectStore>,
+
+    // matthew remove and use source.config
     file_schema: SchemaRef,
     batch_size: Option<usize>,
+    file_column_projection_indices: Option<Vec<usize>>,
 }
 
 impl CsvOpener {
@@ -213,13 +219,15 @@ impl CsvOpener {
         // matthew: delete this and use source.config
         file_schema: SchemaRef,
         batch_size: Option<usize>,
+        file_column_projection_indices: Option<Vec<usize>>,
     ) -> Self {
         Self {
             source,
             file_compression_type,
             object_store,
             file_schema,
             batch_size,
+            file_column_projection_indices,
         }
     }
 }
@@ -243,19 +251,14 @@ impl FileSource for CsvSource {
             object_store,
             file_schema: base_config.file_schema.clone(),
             batch_size: base_config.batch_size,
+            file_column_projection_indices: base_config.file_column_projection_indices(),
         })
     }
 
     fn as_any(&self) -> &dyn Any {
         self
     }
 
-    fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource> {
-        let mut conf = self.clone();
-        conf.file_projection = config.file_column_projection_indices();
-        Arc::new(conf)
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }
@@ -350,6 +353,7 @@ impl FileOpener for CsvOpener {
 
         let file_schema = Arc::clone(&self.file_schema);
         let batch_size = self.batch_size;
+        let file_column_projection_indices = self.file_column_projection_indices.clone();
 
         Ok(Box::pin(async move {
             // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries)
@@ -392,11 +396,14 @@ impl FileOpener for CsvOpener {
                         decoder,
                         file_schema,
                         batch_size,
+                        file_column_projection_indices,
                     )?)
                     .boxed())
                 }
                 GetResultPayload::Stream(s) => {
-                    let decoder = source.builder(file_schema, batch_size).build_decoder();
+                    let decoder = source
+                        .builder(file_schema, batch_size, file_column_projection_indices)
+                        .build_decoder();
                     let s = s.map_err(DataFusionError::from);
                     let input = file_compression_type.convert_stream(s.boxed())?.fuse();
 

datafusion/datasource-json/src/source.rs

@@ -114,10 +114,6 @@ impl FileSource for JsonSource {
         self
     }
 
-    fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource-parquet/src/source.rs

@@ -580,10 +580,6 @@ impl FileSource for ParquetSource {
         self
     }
 
-    fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource/src/file.rs

@@ -60,8 +60,6 @@ pub trait FileSource: Send + Sync {
     ) -> Arc<dyn FileOpener>;
     /// Any
     fn as_any(&self) -> &dyn Any;
-    /// Initialize new instance with projection information
-    fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource>;
     /// Return execution plan metrics
     fn metrics(&self) -> &ExecutionPlanMetricsSet;
     /// Return projected statistics

datafusion/datasource/src/file_scan_config.rs

@@ -516,11 +516,12 @@ impl DataSource for FileScanConfig {
             .with_batch_size(Some(batch_size))
             .build();
 
-        let source = self.file_source.with_projection(&config);
+        let opener =
+            self.file_source
+                .create_file_opener(object_store, &config, partition);
 
-        let opener = source.create_file_opener(object_store, &config, partition);
-
-        let stream = FileStream::new(&config, partition, opener, source.metrics())?;
+        let stream =
+            FileStream::new(&config, partition, opener, self.file_source.metrics())?;
         Ok(Box::pin(cooperative(stream)))
     }
 

datafusion/datasource/src/test_util.rs

@@ -49,10 +49,6 @@ impl FileSource for MockSource {
         self
     }
 
-    fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }