Allow Overriding AsyncFileReader used by ParquetExec

datafusion/core/src/datasource/file_format/mod.rs

@@ -108,6 +108,7 @@ pub(crate) mod test_util {
             object_meta: meta,
             partition_values: vec![],
             range: None,
+            extensions: None,
         }]];
 
         let exec = format

datafusion/core/src/datasource/listing/helpers.rs

@@ -183,24 +183,27 @@ pub async fn pruned_partition_list<'a>(
         // Note: We might avoid parsing the partition values if they are not used in any projection,
         // but the cost of parsing will likely be far dominated by the time to fetch the listing from
         // the object store.
-        Ok(Box::pin(list.try_filter_map(move |file_meta| async move {
-            let parsed_path = parse_partitions_for_path(
-                table_path,
-                &file_meta.location,
-                table_partition_cols,
-            )
-            .map(|p| {
-                p.iter()
-                    .map(|&pn| ScalarValue::Utf8(Some(pn.to_owned())))
-                    .collect()
-            });
-
-            Ok(parsed_path.map(|partition_values| PartitionedFile {
-                partition_values,
-                object_meta: file_meta,
-                range: None,
-            }))
-        })))
+        Ok(Box::pin(list.try_filter_map(
+            move |object_meta| async move {
+                let parsed_path = parse_partitions_for_path(
+                    table_path,
+                    &object_meta.location,
+                    table_partition_cols,
+                )
+                .map(|p| {
+                    p.iter()
+                        .map(|&pn| ScalarValue::Utf8(Some(pn.to_owned())))
+                        .collect()
+                });
+
+                Ok(parsed_path.map(|partition_values| PartitionedFile {
+                    partition_values,
+                    object_meta,
+                    range: None,
+                    extensions: None,
+                }))
+            },
+        )))
     } else {
         // parse the partition values and serde them as a RecordBatch to filter them
         let metas: Vec<_> = list.try_collect().await?;
@@ -317,6 +320,7 @@ fn batches_to_paths(batches: &[RecordBatch]) -> Result<Vec<PartitionedFile>> {
                         })
                         .collect(),
                     range: None,
+                    extensions: None,
                 })
             })
         })

datafusion/core/src/datasource/listing/mod.rs

@@ -28,6 +28,7 @@ use datafusion_common::ScalarValue;
 use futures::Stream;
 use object_store::{path::Path, ObjectMeta};
 use std::pin::Pin;
+use std::sync::Arc;
 
 pub use self::url::ListingTableUrl;
 pub use table::{ListingOptions, ListingTable, ListingTableConfig};
@@ -58,6 +59,8 @@ pub struct PartitionedFile {
     pub partition_values: Vec<ScalarValue>,
     /// An optional file range for a more fine-grained parallel execution
     pub range: Option<FileRange>,
+    /// An optional field for user defined per object metadata  
+    pub extensions: Option<Arc<dyn std::any::Any + Send + Sync>>,
 }
 
 impl PartitionedFile {
@@ -71,6 +74,7 @@ impl PartitionedFile {
             },
             partition_values: vec![],
             range: None,
+            extensions: None,
         }
     }
 
@@ -84,6 +88,7 @@ impl PartitionedFile {
             },
             partition_values: vec![],
             range: Some(FileRange { start, end }),
+            extensions: None,
         }
     }
 }
@@ -94,6 +99,7 @@ impl From<ObjectMeta> for PartitionedFile {
             object_meta,
             partition_values: vec![],
             range: None,
+            extensions: None,
         }
     }
 }

datafusion/core/src/physical_plan/file_format/avro.rs

@@ -155,6 +155,7 @@ mod private {
     use super::*;
     use crate::datasource::listing::FileRange;
     use crate::physical_plan::file_format::file_stream::{FileOpenFuture, FileOpener};
+    use crate::physical_plan::file_format::FileMeta;
     use bytes::Buf;
     use futures::StreamExt;
     use object_store::{GetResult, ObjectMeta, ObjectStore};
@@ -187,12 +188,11 @@ mod private {
         fn open(
             &self,
             store: Arc<dyn ObjectStore>,
-            file: ObjectMeta,
-            _range: Option<FileRange>,
-        ) -> FileOpenFuture {
+            file_meta: FileMeta,
+        ) -> Result<FileOpenFuture> {
             let config = self.config.clone();
-            Box::pin(async move {
-                match store.get(&file.location).await? {
+            Ok(Box::pin(async move {
+                match store.get(file_meta.location()).await? {
                     GetResult::File(file, _) => {
                         let reader = config.open(file)?;
                         Ok(futures::stream::iter(reader).boxed())
@@ -203,7 +203,7 @@ mod private {
                         Ok(futures::stream::iter(reader).boxed())
                     }
                 }
-            })
+            }))
         }
     }
 }

datafusion/core/src/physical_plan/file_format/csv.rs

@@ -20,16 +20,15 @@
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::{SessionState, TaskContext};
 use crate::physical_plan::expressions::PhysicalSortExpr;
-use crate::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
-};
-
-use crate::datasource::listing::FileRange;
 use crate::physical_plan::file_format::delimited_stream::newline_delimited_stream;
 use crate::physical_plan::file_format::file_stream::{
     FileOpenFuture, FileOpener, FileStream,
 };
+use crate::physical_plan::file_format::FileMeta;
 use crate::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
+use crate::physical_plan::{
+    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
+};
 use arrow::csv;
 use arrow::datatypes::SchemaRef;
 use bytes::Buf;
@@ -201,12 +200,11 @@ impl FileOpener for CsvOpener {
     fn open(
         &self,
         store: Arc<dyn ObjectStore>,
-        file: ObjectMeta,
-        _range: Option<FileRange>,
-    ) -> FileOpenFuture {
+        file_meta: FileMeta,
+    ) -> Result<FileOpenFuture> {
         let config = self.config.clone();
-        Box::pin(async move {
-            match store.get(&file.location).await? {
+        Ok(Box::pin(async move {
+            match store.get(file_meta.location()).await? {
                 GetResult::File(file, _) => {
                     Ok(futures::stream::iter(config.open(file, true)).boxed())
                 }
@@ -222,7 +220,7 @@ impl FileOpener for CsvOpener {
                         .boxed())
                 }
             }
-        })
+        }))
     }
 }
 

datafusion/core/src/physical_plan/file_format/file_stream.rs

@@ -31,14 +31,16 @@ use arrow::{error::Result as ArrowResult, record_batch::RecordBatch};
 use futures::future::BoxFuture;
 use futures::stream::BoxStream;
 use futures::{ready, FutureExt, Stream, StreamExt};
-use object_store::{ObjectMeta, ObjectStore};
+use object_store::ObjectStore;
 
 use datafusion_common::ScalarValue;
 
-use crate::datasource::listing::{FileRange, PartitionedFile};
+use crate::datasource::listing::PartitionedFile;
 use crate::error::Result;
 use crate::execution::context::TaskContext;
-use crate::physical_plan::file_format::{FileScanConfig, PartitionColumnProjector};
+use crate::physical_plan::file_format::{
+    FileMeta, FileScanConfig, PartitionColumnProjector,
+};
 use crate::physical_plan::metrics::BaselineMetrics;
 use crate::physical_plan::RecordBatchStream;
 
@@ -50,9 +52,8 @@ pub trait FileOpener: Unpin {
     fn open(
         &self,
         store: Arc<dyn ObjectStore>,
-        file: ObjectMeta,
-        range: Option<FileRange>,
-    ) -> FileOpenFuture;
+        file_meta: FileMeta,
+    ) -> Result<FileOpenFuture>;
 }
 
 /// A stream that iterates record batch by record batch, file over file.
@@ -143,20 +144,28 @@ impl<F: FileOpener> FileStream<F> {
         loop {
             match &mut self.state {
                 FileStreamState::Idle => {
-                    let file = match self.file_iter.pop_front() {
+                    let part_file = match self.file_iter.pop_front() {
                         Some(file) => file,
                         None => return Poll::Ready(None),
                     };
 
-                    let future = self.file_reader.open(
-                        self.object_store.clone(),
-                        file.object_meta,
-                        file.range,
-                    );
+                    let file_meta = FileMeta {
+                        object_meta: part_file.object_meta,
+                        range: part_file.range,
+                        extensions: part_file.extensions,
+                    };
 
-                    self.state = FileStreamState::Open {
-                        future,
-                        partition_values: file.partition_values,
+                    match self.file_reader.open(self.object_store.clone(), file_meta) {
+                        Ok(future) => {
+                            self.state = FileStreamState::Open {
+                                future,
+                                partition_values: part_file.partition_values,
+                            }
+                        }
+                        Err(e) => {
+                            self.state = FileStreamState::Error;
+                            return Poll::Ready(Some(Err(e.into())));
+                        }
                     }
                 }
                 FileStreamState::Open {
@@ -254,12 +263,11 @@ mod tests {
         fn open(
             &self,
             _store: Arc<dyn ObjectStore>,
-            _file: ObjectMeta,
-            _range: Option<FileRange>,
-        ) -> FileOpenFuture {
+            _file_meta: FileMeta,
+        ) -> Result<FileOpenFuture> {
             let iterator = self.records.clone().into_iter().map(Ok);
             let stream = futures::stream::iter(iterator).boxed();
-            futures::future::ready(Ok(stream)).boxed()
+            Ok(futures::future::ready(Ok(stream)).boxed())
         }
     }
 

datafusion/core/src/physical_plan/file_format/json.rs

@@ -16,9 +16,6 @@
 // under the License.
 
 //! Execution plan for reading line-delimited JSON files
-use arrow::json::reader::DecoderOptions;
-
-use crate::datasource::listing::FileRange;
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::SessionState;
 use crate::execution::context::TaskContext;
@@ -27,14 +24,16 @@ use crate::physical_plan::file_format::delimited_stream::newline_delimited_strea
 use crate::physical_plan::file_format::file_stream::{
     FileOpenFuture, FileOpener, FileStream,
 };
+use crate::physical_plan::file_format::FileMeta;
 use crate::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
 use crate::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
+use arrow::json::reader::DecoderOptions;
 use arrow::{datatypes::SchemaRef, json};
 use bytes::Buf;
 use futures::{StreamExt, TryStreamExt};
-use object_store::{GetResult, ObjectMeta, ObjectStore};
+use object_store::{GetResult, ObjectStore};
 use std::any::Any;
 use std::fs;
 use std::path::Path;
@@ -163,13 +162,12 @@ impl FileOpener for JsonOpener {
     fn open(
         &self,
         store: Arc<dyn ObjectStore>,
-        file: ObjectMeta,
-        _range: Option<FileRange>,
-    ) -> FileOpenFuture {
+        file_meta: FileMeta,
+    ) -> Result<FileOpenFuture> {
         let options = self.options.clone();
         let schema = self.file_schema.clone();
-        Box::pin(async move {
-            match store.get(&file.location).await? {
+        Ok(Box::pin(async move {
+            match store.get(file_meta.location()).await? {
                 GetResult::File(file, _) => {
                     let reader = json::Reader::new(file, schema.clone(), options);
                     Ok(futures::stream::iter(reader).boxed())
@@ -188,7 +186,7 @@ impl FileOpener for JsonOpener {
                         .boxed())
                 }
             }
-        })
+        }))
     }
 }
 

datafusion/core/src/physical_plan/file_format/mod.rs

@@ -41,6 +41,7 @@ pub use avro::AvroExec;
 pub(crate) use json::plan_to_json;
 pub use json::NdJsonExec;
 
+use crate::datasource::listing::FileRange;
 use crate::datasource::{listing::PartitionedFile, object_store::ObjectStoreUrl};
 use crate::{
     error::{DataFusionError, Result},
@@ -50,6 +51,8 @@ use arrow::array::{new_null_array, UInt16BufferBuilder};
 use arrow::record_batch::RecordBatchOptions;
 use lazy_static::lazy_static;
 use log::info;
+use object_store::path::Path;
+use object_store::ObjectMeta;
 use std::{
     collections::HashMap,
     fmt::{Display, Formatter, Result as FmtResult},
@@ -401,6 +404,33 @@ fn create_dict_array(
     ))
 }
 
+/// A single file or part of a file that should be read, along with its schema, statistics
+pub struct FileMeta {
+    /// Path for the file (e.g. URL, filesystem path, etc)
+    pub object_meta: ObjectMeta,
+    /// An optional file range for a more fine-grained parallel execution
+    pub range: Option<FileRange>,
+    /// An optional field for user defined per object metadata  
+    pub extensions: Option<Arc<dyn std::any::Any + Send + Sync>>,
+}
+
+impl FileMeta {
+    /// The full path to the object
+    pub fn location(&self) -> &Path {
+        &self.object_meta.location
+    }
+}
+
+impl From<ObjectMeta> for FileMeta {
+    fn from(object_meta: ObjectMeta) -> Self {
+        Self {
+            object_meta,
+            range: None,
+            extensions: None,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{