Adds prefix filtering for table URLs

datafusion/catalog-listing/src/helpers.rs

@@ -380,8 +380,14 @@ pub async fn pruned_partition_list<'a>(
     file_extension: &'a str,
     partition_cols: &'a [(String, DataType)],
 ) -> Result<BoxStream<'a, Result<PartitionedFile>>> {
+    let prefix = if !partition_cols.is_empty() {
+        evaluate_partition_prefix(partition_cols, filters)
+    } else {
+        None
+    };
+
     let objects = table_path
-        .list_all_files(ctx, store, file_extension)
+        .list_prefixed_files(ctx, store, prefix, file_extension)
         .await?
         .try_filter(|object_meta| futures::future::ready(object_meta.size > 0));
 

datafusion/core/tests/datasource/object_store_access.rs

@@ -166,7 +166,7 @@ async fn query_partitioned_csv_file() {
     ------- Object Store Request Summary -------
     RequestCountingObjectStore()
     Total Requests: 2
-    - LIST prefix=data
+    - LIST prefix=data/a=2
     - GET  (opts) path=data/a=2/b=20/c=200/file_2.csv
     "
     );
@@ -220,7 +220,7 @@ async fn query_partitioned_csv_file() {
     ------- Object Store Request Summary -------
     RequestCountingObjectStore()
     Total Requests: 2
-    - LIST prefix=data
+    - LIST prefix=data/a=2/b=20
     - GET  (opts) path=data/a=2/b=20/c=200/file_2.csv
     "
     );

datafusion/datasource/src/url.rs

@@ -233,27 +233,37 @@ impl ListingTableUrl {
         Some(stripped.split_terminator(DELIMITER))
     }
 
-    /// List all files identified by this [`ListingTableUrl`] for the provided `file_extension`
-    pub async fn list_all_files<'a>(
+    /// List all files identified by this [`ListingTableUrl`] for the provided `file_extension`,
+    /// optionally filtering by a path prefix
+    pub async fn list_prefixed_files<'a>(
         &'a self,
         ctx: &'a dyn Session,
         store: &'a dyn ObjectStore,
+        prefix: Option<Path>,
         file_extension: &'a str,
     ) -> Result<BoxStream<'a, Result<ObjectMeta>>> {
         let exec_options = &ctx.config_options().execution;
         let ignore_subdirectory = exec_options.listing_table_ignore_subdirectory;
 
+        let prefix = if let Some(prefix) = prefix {
+            let mut p = self.prefix.parts().collect::<Vec<_>>();
+            p.extend(prefix.parts());
+            Path::from_iter(p.into_iter())
+        } else {
+            self.prefix.clone()
+        };
+
         let list: BoxStream<'a, Result<ObjectMeta>> = if self.is_collection() {
-            list_with_cache(ctx, store, &self.prefix).await?
+            list_with_cache(ctx, store, &prefix).await?
         } else {
-            match store.head(&self.prefix).await {
+            match store.head(&prefix).await {
                 Ok(meta) => futures::stream::once(async { Ok(meta) })
                     .map_err(|e| DataFusionError::ObjectStore(Box::new(e)))
                     .boxed(),
                 // If the head command fails, it is likely that object doesn't exist.
                 // Retry as though it were a prefix (aka a collection)
                 Err(object_store::Error::NotFound { .. }) => {
-                    list_with_cache(ctx, store, &self.prefix).await?
+                    list_with_cache(ctx, store, &prefix).await?
                 }
                 Err(e) => return Err(e.into()),
             }
@@ -269,6 +279,17 @@ impl ListingTableUrl {
             .boxed())
     }
 
+    /// List all files identified by this [`ListingTableUrl`] for the provided `file_extension`
+    pub async fn list_all_files<'a>(
+        &'a self,
+        ctx: &'a dyn Session,
+        store: &'a dyn ObjectStore,
+        file_extension: &'a str,
+    ) -> Result<BoxStream<'a, Result<ObjectMeta>>> {
+        self.list_prefixed_files(ctx, store, None, file_extension)
+            .await
+    }
+
     /// Returns this [`ListingTableUrl`] as a string
     pub fn as_str(&self) -> &str {
         self.as_ref()
@@ -306,7 +327,7 @@ impl ListingTableUrl {
 async fn list_with_cache<'b>(
     ctx: &'b dyn Session,
     store: &'b dyn ObjectStore,
-    prefix: &'b Path,
+    prefix: &Path,
 ) -> Result<BoxStream<'b, Result<ObjectMeta>>> {
     match ctx.runtime_env().cache_manager.get_list_files_cache() {
         None => Ok(store
@@ -701,6 +722,35 @@ mod tests {
             panic!("Expected PermissionDenied error");
         };
 
+        // Test prefix filtering with partition-style paths
+        create_file(&store, "/data/a=1/file1.parquet").await;
+        create_file(&store, "/data/a=1/b=100/file2.parquet").await;
+        create_file(&store, "/data/a=2/b=200/file3.parquet").await;
+        create_file(&store, "/data/a=2/b=200/file4.csv").await;
+
+        assert_eq!(
+            list_prefixed_files("/data/", &store, Some(Path::from("a=1")), "parquet")
+                .await?,
+            vec!["data/a=1/b=100/file2.parquet", "data/a=1/file1.parquet"],
+        );
+
+        assert_eq!(
+            list_prefixed_files(
+                "/data/",
+                &store,
+                Some(Path::from("a=1/b=100")),
+                "parquet"
+            )
+            .await?,
+            vec!["data/a=1/b=100/file2.parquet"],
+        );
+
+        assert_eq!(
+            list_prefixed_files("/data/", &store, Some(Path::from("a=2")), "parquet")
+                .await?,
+            vec!["data/a=2/b=200/file3.parquet"],
+        );
+
         Ok(())
     }
 
@@ -712,27 +762,40 @@ mod tests {
             .expect("failed to create test file");
     }
 
-    /// Runs "list_all_files" and returns their paths
+    /// Runs "list_prefixed_files"  with no prefix to list all files and returns their paths
     ///
     /// Panic's on error
     async fn list_all_files(
         url: &str,
         store: &dyn ObjectStore,
         file_extension: &str,
     ) -> Result<Vec<String>> {
-        try_list_all_files(url, store, file_extension).await
+        try_list_prefixed_files(url, store, None, file_extension).await
+    }
+
+    /// Runs "list_prefixed_files" and returns their paths
+    ///
+    /// Panic's on error
+    async fn list_prefixed_files(
+        url: &str,
+        store: &dyn ObjectStore,
+        prefix: Option<Path>,
+        file_extension: &str,
+    ) -> Result<Vec<String>> {
+        try_list_prefixed_files(url, store, prefix, file_extension).await
     }
 
-    /// Runs "list_all_files" and returns their paths
-    async fn try_list_all_files(
+    /// Runs "list_prefixed_files" and returns their paths
+    async fn try_list_prefixed_files(
         url: &str,
         store: &dyn ObjectStore,
+        prefix: Option<Path>,
         file_extension: &str,
     ) -> Result<Vec<String>> {
         let session = MockSession::new();
         let url = ListingTableUrl::parse(url)?;
         let files = url
-            .list_all_files(&session, store, file_extension)
+            .list_prefixed_files(&session, store, prefix, file_extension)
             .await?
             .try_collect::<Vec<_>>()
             .await?