From 65aecf26acc512cbb8fbcfaf44972df4576a17c5 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Wed, 16 Sep 2026 05:51:07 +0300 Subject: [PATCH] Support regex filters for listing table paths --- Cargo.lock | 1 + datafusion/datasource/Cargo.toml | 1 + datafusion/datasource/src/mod.rs | 19 +++++++ datafusion/datasource/src/url.rs | 92 +++++++++++++++++++++++++++++--- 4 files changed, 105 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4f7b8bb279c20..d0efd6d41a4e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1959,6 +1959,7 @@ dependencies = [ "object_store", "parking_lot", "rand 0.9.4", + "regex", "tempfile", "tokio", "tokio-util", diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index f09447b694f52..98bf9ac13f1ff 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -75,6 +75,7 @@ log = { workspace = true } object_store = { workspace = true } parking_lot = { workspace = true } rand = { workspace = true } +regex = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true } tokio-util = { version = "0.7.17", features = ["io"], optional = true } diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index e415b3e48a02a..d03ea6d3417e2 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -731,4 +731,23 @@ mod tests { // testing an empty path with `ignore_subdirectory` set to false assert!(url.contains(&Path::parse("/var/data/mytable/").unwrap(), false)); } + + #[test] + fn test_url_contains_full_path_regex() { + let url = ListingTableUrl::parse("s3://bucket/prefix/") + .unwrap() + .with_full_path_regex( + regex::Regex::new(r"\Aprefix/(?:nested/)?selected_[ab][.]csv\z").unwrap(), + ); + + assert!(url.get_glob().is_none()); + assert_eq!( + url.get_full_path_regex().map(regex::Regex::as_str), + Some(r"\Aprefix/(?:nested/)?selected_[ab][.]csv\z") + ); + assert!(url.contains(&Path::from("prefix/selected_a.csv"), true)); + assert!(url.contains(&Path::from("prefix/nested/selected_b.csv"), true)); + assert!(!url.contains(&Path::from("prefix/ignored.csv"), false)); + assert!(!url.contains(&Path::from("other/selected_a.csv"), false)); + } } diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 9ac4c5f50d1f7..e865e7ee6d519 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -31,8 +31,26 @@ use log::debug; use object_store::path::DELIMITER; use object_store::path::Path; use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; +use regex::Regex; use url::Url; +#[derive(Debug, Clone)] +struct FullPathRegex(Regex); + +impl PartialEq for FullPathRegex { + fn eq(&self, other: &Self) -> bool { + self.0.as_str() == other.0.as_str() + } +} + +impl Eq for FullPathRegex {} + +impl std::hash::Hash for FullPathRegex { + fn hash(&self, state: &mut H) { + std::hash::Hash::hash(self.0.as_str(), state); + } +} + /// A parsed URL identifying files for a listing table, see [`ListingTableUrl::parse`] /// for more information on the supported expressions #[derive(Debug, Clone, Eq, PartialEq, Hash)] @@ -43,6 +61,8 @@ pub struct ListingTableUrl { prefix: Path, /// An optional glob expression used to filter files glob: Option, + /// An optional regular expression applied to the complete object path + full_path_regex: Option, /// Optional table reference for the table this url belongs to table_ref: Option, } @@ -153,6 +173,7 @@ impl ListingTableUrl { url, prefix, glob, + full_path_regex: None, table_ref: None, }) } @@ -172,6 +193,14 @@ impl ListingTableUrl { /// Returns `true` if `path` matches this [`ListingTableUrl`] pub fn contains(&self, path: &Path, ignore_subdirectory: bool) -> bool { + if let Some(regex) = &self.full_path_regex { + return self.strip_prefix(path).is_some() && regex.0.is_match(path.as_ref()); + } + + self.contains_without_regex(path, ignore_subdirectory) + } + + fn contains_without_regex(&self, path: &Path, ignore_subdirectory: bool) -> bool { let Some(all_segments) = self.strip_prefix(path) else { return false; }; @@ -293,14 +322,27 @@ impl ListingTableUrl { } }; - Ok(list - .try_filter(move |meta| { - let path = &meta.location; - let extension_match = path.as_ref().ends_with(file_extension); - let glob_match = self.contains(path, ignore_subdirectory); - futures::future::ready(extension_match && glob_match) - }) - .boxed()) + if let Some(regex) = &self.full_path_regex { + Ok(list + .try_filter(move |meta| { + let path = &meta.location; + let extension_match = path.as_ref().ends_with(file_extension); + let regex_match = self.strip_prefix(path).is_some() + && regex.0.is_match(path.as_ref()); + futures::future::ready(extension_match && regex_match) + }) + .boxed()) + } else { + Ok(list + .try_filter(move |meta| { + let path = &meta.location; + let extension_match = path.as_ref().ends_with(file_extension); + let glob_match = + self.contains_without_regex(path, ignore_subdirectory); + futures::future::ready(extension_match && glob_match) + }) + .boxed()) + } } /// List all files identified by this [`ListingTableUrl`] for the provided `file_extension` @@ -344,9 +386,27 @@ impl ListingTableUrl { pub fn with_glob(mut self, glob: &str) -> Result { self.glob = Some(Pattern::new(glob).map_err(|e| DataFusionError::External(Box::new(e)))?); + self.full_path_regex = None; Ok(self) } + /// Returns the regular expression applied to the complete object path. + pub fn get_full_path_regex(&self) -> Option<&Regex> { + self.full_path_regex.as_ref().map(|regex| ®ex.0) + } + + /// Returns a copy of this [`ListingTableUrl`] with a precompiled regular expression that is + /// applied to the complete object path. + /// + /// The regular expression replaces any existing glob. Unlike the default path matching, + /// matching nested paths is controlled entirely by the regular expression and is therefore + /// unaffected by `listing_table_ignore_subdirectory`. + pub fn with_full_path_regex(mut self, regex: Regex) -> Self { + self.glob = None; + self.full_path_regex = Some(FullPathRegex(regex)); + self + } + /// Set the table reference for this [`ListingTableUrl`] pub fn with_table_ref(mut self, table_ref: TableReference) -> Self { self.table_ref = Some(table_ref); @@ -786,6 +846,22 @@ mod tests { vec!["t/c.csv", "t/d.csv"], ); + let session = MockSession::new(); + let regex_url = ListingTableUrl::parse("memory:///t/")?.with_full_path_regex( + Regex::new(r"\At/(?:c|nested/e)[.]csv\z") + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + create_file(&store, "/t/nested/e.csv").await; + let regex_files = regex_url + .list_all_files(&session, &store, "csv") + .await? + .try_collect::>() + .await? + .into_iter() + .map(|meta| meta.location.to_string()) + .collect::>(); + assert_eq!(regex_files, ["t/c.csv", "t/nested/e.csv"]); + // Test a non existing prefix assert_eq!( list_all_files("/NonExisting", &store, "csv").await?,