Add VirtualObjectStore to support routing paths to multiple ObjectStores

datafusion/datasource/src/file_scan_config.rs

@@ -64,6 +64,7 @@ use datafusion_physical_plan::{
     projection::{all_alias_free_columns, new_projections_for_columns, ProjectionExec},
     DisplayAs, DisplayFormatType, ExecutionPlan,
 };
+use object_store::ObjectStore;
 
 use datafusion_physical_plan::coop::cooperative;
 use datafusion_physical_plan::execution_plan::SchedulingType;
@@ -186,6 +187,8 @@ pub struct FileScanConfig {
     pub new_lines_in_values: bool,
     /// File source such as `ParquetSource`, `CsvSource`, `JsonSource`, etc.
     pub file_source: Arc<dyn FileSource>,
+    /// Optional virtual object store for routing paths to underlying stores
+    pub virtual_store: Option<Arc<dyn ObjectStore>>,
     /// Batch size while creating new batches
     /// Defaults to [`datafusion_common::config::ExecutionOptions`] batch_size.
     pub batch_size: Option<usize>,
@@ -258,6 +261,7 @@ pub struct FileScanConfigBuilder {
     /// This probably would be better named `table_schema`
     file_schema: SchemaRef,
     file_source: Arc<dyn FileSource>,
+    virtual_store: Option<Arc<dyn ObjectStore>>,
 
     limit: Option<usize>,
     projection: Option<Vec<usize>>,
@@ -288,6 +292,7 @@ impl FileScanConfigBuilder {
             object_store_url,
             file_schema,
             file_source,
+            virtual_store: None,
             file_groups: vec![],
             statistics: None,
             output_ordering: vec![],
@@ -318,6 +323,12 @@ impl FileScanConfigBuilder {
         self
     }
 
+    /// Set the virtual object store mapping
+    pub fn with_virtual_store(mut self, store: Arc<dyn ObjectStore>) -> Self {
+        self.virtual_store = Some(store);
+        self
+    }
+
     /// Set the columns on which to project the data. Indexes that are higher than the
     /// number of columns of `file_schema` refer to `table_partition_cols`.
     pub fn with_projection(mut self, projection: Option<Vec<usize>>) -> Self {
@@ -430,6 +441,7 @@ impl FileScanConfigBuilder {
             object_store_url,
             file_schema,
             file_source,
+            virtual_store,
             limit,
             projection,
             table_partition_cols,
@@ -458,6 +470,7 @@ impl FileScanConfigBuilder {
             object_store_url,
             file_schema,
             file_source,
+            virtual_store,
             limit,
             projection,
             table_partition_cols,
@@ -478,6 +491,7 @@ impl From<FileScanConfig> for FileScanConfigBuilder {
             object_store_url: config.object_store_url,
             file_schema: config.file_schema,
             file_source: Arc::<dyn FileSource>::clone(&config.file_source),
+            virtual_store: config.virtual_store,
             file_groups: config.file_groups,
             statistics: config.file_source.statistics().ok(),
             output_ordering: config.output_ordering,
@@ -499,7 +513,20 @@ impl DataSource for FileScanConfig {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        let object_store = context.runtime_env().object_store(&self.object_store_url)?;
+        let object_store: Arc<dyn ObjectStore> = if let Some(store) = &self.virtual_store
+        {
+            Arc::clone(store)
+        } else {
+            context
+                .runtime_env()
+                .object_store(&self.object_store_url)
+                .map_err(|e| {
+                    e.context(format!(
+                        "get object store for URL {}",
+                        self.object_store_url
+                    ))
+                })?
+        };
         let batch_size = self
             .batch_size
             .unwrap_or_else(|| context.session_config().batch_size());

datafusion/execution/Cargo.toml

@@ -44,6 +44,8 @@ parquet_encryption = [
 
 [dependencies]
 arrow = { workspace = true }
+async-trait = { workspace = true }
+chrono = { workspace = true }
 dashmap = { workspace = true }
 datafusion-common = { workspace = true, default-features = true }
 datafusion-expr = { workspace = true }
@@ -57,5 +59,5 @@ tempfile = { workspace = true }
 url = { workspace = true }
 
 [dev-dependencies]
-chrono = { workspace = true }
 insta = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt", "sync"] }

datafusion/execution/src/lib.rs

@@ -36,6 +36,7 @@ pub mod parquet_encryption;
 pub mod runtime_env;
 mod stream;
 mod task;
+pub mod virtual_object_store;
 
 pub mod registry {
     pub use datafusion_expr::registry::{