[fix] get_by_uri also returns path

datafusion/src/datasource/listing.rs

@@ -74,9 +74,9 @@ impl ListingOptions {
         object_store_registry: Arc<ObjectStoreRegistry>,
         uri: &str,
     ) -> Result<SchemaRef> {
-        let object_store = object_store_registry.get_by_uri(uri)?;
+        let (object_store, path) = object_store_registry.get_by_uri(uri)?;
         let file_stream = object_store
-            .list_file_with_suffix(uri, &self.file_extension)
+            .list_file_with_suffix(path, &self.file_extension)
             .await?
             .map(move |file_meta| object_store.file_reader(file_meta?.sized_file));
         let file_schema = self.format.infer_schema(Box::pin(file_stream)).await?;
@@ -137,9 +137,9 @@ impl TableProvider for ListingTable {
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         // TODO object_store_registry should be provided as param here
-        let object_store = self.object_store_registry.get_by_uri(&self.uri)?;
+        let (object_store, path) = self.object_store_registry.get_by_uri(&self.uri)?;
         let (partitioned_file_lists, statistics) = self
-            .list_files_for_scan(Arc::clone(&object_store), filters, limit)
+            .list_files_for_scan(Arc::clone(&object_store), path, filters, limit)
             .await?;
         // create the execution plan
         self.options
@@ -169,13 +169,14 @@ impl ListingTable {
     async fn list_files_for_scan<'a>(
         &'a self,
         object_store: Arc<dyn ObjectStore>,
+        path: &'a str,
         filters: &'a [Expr],
         limit: Option<usize>,
     ) -> Result<(Vec<Vec<PartitionedFile>>, Statistics)> {
         // list files (with partitions)
         let file_list = pruned_partition_list(
             object_store.as_ref(),
-            &self.uri,
+            path,
             filters,
             &self.options.file_extension,
             &self.options.partitions,
@@ -387,7 +388,9 @@ mod tests {
             opt,
         );
 
-        let (file_list, _) = table.list_files_for_scan(mock_store, &[], None).await?;
+        let (file_list, _) = table
+            .list_files_for_scan(mock_store, "bucket/key-prefix", &[], None)
+            .await?;
 
         assert_eq!(file_list.len(), output_partitioning);
 

datafusion/src/datasource/object_store/mod.rs

@@ -216,22 +216,28 @@ impl ObjectStoreRegistry {
     }
 
     /// Get a suitable store for the URI based on it's scheme. For example:
-    /// URI with scheme file or no schema will return the default LocalFS store,
-    /// URI with scheme s3 will return the S3 store if it's registered.
-    pub fn get_by_uri(&self, uri: &str) -> Result<Arc<dyn ObjectStore>> {
-        if let Some((scheme, _)) = uri.split_once(':') {
+    /// - URI with scheme `file://` or no schema will return the default LocalFS store
+    /// - URI with scheme `s3://` will return the S3 store if it's registered
+    /// Returns a tuple with the store and the path of the file in that store
+    /// (URI=scheme://path).
+    pub fn get_by_uri<'a>(
+        &self,
+        uri: &'a str,
+    ) -> Result<(Arc<dyn ObjectStore>, &'a str)> {
+        if let Some((scheme, path)) = uri.split_once("://") {
             let stores = self.object_stores.read().unwrap();
-            stores
+            let store = stores
                 .get(&*scheme.to_lowercase())
                 .map(Clone::clone)
                 .ok_or_else(|| {
                     DataFusionError::Internal(format!(
                         "No suitable object store found for {}",
                         scheme
                     ))
-                })
+                })?;
+            Ok((store, path))
         } else {
-            Ok(Arc::new(LocalFileSystem))
+            Ok((Arc::new(LocalFileSystem), uri))
         }
     }
 }