feat: update to datafusion 52 FileSource APIs (#147)

* feat: upgrade to datafusion 52 apis

* feat: update FileOpener implementation

Author: sunng87

Cargo.toml

@@ -34,8 +34,7 @@ all-features = true
 [dependencies]
 async-trait = { version = "0.1.77" }
 bytes = "1.4"
-datafusion = "51.0"
-datafusion-datasource = "51.0"
+datafusion = "52.0"
 futures = { version = "0.3", default-features = false, features = ["std"] }
 futures-util = { version = "0.3" }
 object_store = { version = "0.12" }

src/file_format.rs

@@ -25,6 +25,7 @@ use datafusion::common::Statistics;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
 use datafusion::datasource::file_format::FileFormat;
 use datafusion::datasource::physical_plan::{FileScanConfig, FileSource};
+use datafusion::datasource::table_schema::TableSchema;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::physical_plan::ExecutionPlan;
 use futures::TryStreamExt;
@@ -122,7 +123,7 @@ impl FileFormat for OrcFormat {
         Ok(DataSourceExec::from_data_source(conf))
     }
 
-    fn file_source(&self) -> Arc<dyn FileSource> {
-        Arc::new(OrcSource::default())
+    fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
+        Arc::new(OrcSource::new(table_schema))
     }
 }

src/file_source.rs

@@ -16,27 +16,35 @@
 // under the License.
 
 use crate::physical_exec::OrcOpener;
-use datafusion::common::Statistics;
+use datafusion::common::DataFusionError;
 use datafusion::datasource::physical_plan::{FileOpener, FileScanConfig, FileSource};
+use datafusion::datasource::table_schema::TableSchema;
 use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
-use datafusion_datasource::TableSchema;
+use datafusion::physical_plan::projection::ProjectionExprs;
 use object_store::ObjectStore;
 use std::any::Any;
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
 pub struct OrcSource {
     metrics: ExecutionPlanMetricsSet,
-    statistics: Statistics,
     batch_size: usize,
+    table_schema: TableSchema,
+    projection: ProjectionExprs,
 }
 
-impl Default for OrcSource {
-    fn default() -> Self {
+impl OrcSource {
+    pub fn new(table_schema: TableSchema) -> Self {
+        let table_schema_ref = table_schema.table_schema();
+        let projection = ProjectionExprs::from_indices(
+            &(0..table_schema_ref.fields().len()).collect::<Vec<_>>(),
+            table_schema_ref,
+        );
         Self {
             metrics: ExecutionPlanMetricsSet::default(),
-            statistics: Statistics::default(),
             batch_size: 1024,
+            table_schema,
+            projection,
         }
     }
 }
@@ -47,45 +55,49 @@ impl FileSource for OrcSource {
         object_store: Arc<dyn ObjectStore>,
         config: &FileScanConfig,
         _partition: usize,
-    ) -> Arc<dyn FileOpener> {
-        Arc::new(OrcOpener::new(object_store, config, self.batch_size))
+    ) -> Result<Arc<dyn FileOpener>, DataFusionError> {
+        OrcOpener::try_new(
+            object_store,
+            self.table_schema.table_schema().clone(),
+            config.batch_size.unwrap_or(self.batch_size),
+            self.projection.clone(),
+        )
+        .map(|f| Arc::new(f) as Arc<dyn FileOpener>)
     }
 
     fn as_any(&self) -> &dyn Any {
         self
     }
 
+    fn table_schema(&self) -> &TableSchema {
+        &self.table_schema
+    }
+
     fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
         Arc::new(Self {
             batch_size,
             ..self.clone()
         })
     }
 
-    fn with_schema(&self, _schema: TableSchema) -> Arc<dyn FileSource> {
-        Arc::new(self.clone())
-    }
-
-    fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
-        Arc::new(self.clone())
-    }
-
-    fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {
-        Arc::new(Self {
-            statistics,
-            ..self.clone()
-        })
+    fn projection(&self) -> Option<&ProjectionExprs> {
+        Some(&self.projection)
     }
 
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }
 
-    fn statistics(&self) -> datafusion::common::Result<Statistics> {
-        Ok(self.statistics.clone())
-    }
-
     fn file_type(&self) -> &str {
         "orc"
     }
+
+    fn try_pushdown_projection(
+        &self,
+        projection: &ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>, DataFusionError> {
+        let mut source = self.clone();
+        source.projection = self.projection.try_merge(projection)?;
+        Ok(Some(Arc::new(source)))
+    }
 }

src/physical_exec.rs

@@ -19,9 +19,10 @@ use std::sync::Arc;
 
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::arrow::error::ArrowError;
-use datafusion::datasource::physical_plan::{FileOpenFuture, FileOpener, FileScanConfig};
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileOpenFuture, FileOpener};
 use datafusion::error::Result;
-use datafusion_datasource::PartitionedFile;
+use datafusion::physical_plan::projection::ProjectionExprs;
 use orc_rust::projection::ProjectionMask;
 use orc_rust::ArrowReaderBuilder;
 
@@ -31,28 +32,25 @@ use object_store::ObjectStore;
 use super::object_store_reader::ObjectStoreReader;
 
 pub(crate) struct OrcOpener {
-    projection: Vec<usize>,
+    projection: ProjectionExprs,
     batch_size: usize,
     table_schema: SchemaRef,
     object_store: Arc<dyn ObjectStore>,
 }
 
 impl OrcOpener {
-    pub(crate) fn new(
+    pub(crate) fn try_new(
         object_store: Arc<dyn ObjectStore>,
-        config: &FileScanConfig,
+        table_schema: SchemaRef,
         batch_size: usize,
-    ) -> Self {
-        let projection = config
-            .file_column_projection_indices()
-            .unwrap_or_else(|| (0..config.file_schema().fields().len()).collect());
-
-        Self {
+        projection: ProjectionExprs,
+    ) -> Result<Self> {
+        Ok(Self {
             projection,
-            batch_size: config.batch_size.unwrap_or(batch_size),
-            table_schema: config.file_schema().clone(),
+            batch_size,
+            table_schema,
             object_store,
-        }
+        })
     }
 }
 
@@ -61,7 +59,7 @@ impl FileOpener for OrcOpener {
         let object_meta = &file.object_meta;
         let reader = ObjectStoreReader::new(self.object_store.clone(), object_meta.clone());
         let batch_size = self.batch_size;
-        let projected_schema = SchemaRef::from(self.table_schema.project(&self.projection)?);
+        let projected_schema = self.projection.project_schema(&self.table_schema)?;
 
         Ok(Box::pin(async move {
             let mut builder = ArrowReaderBuilder::try_new_async(reader)