polars_io/cloud/
glob.rs

1use std::borrow::Cow;
2
3use futures::TryStreamExt;
4use object_store::path::Path;
5use polars_error::{PolarsResult, polars_bail, polars_err};
6use polars_utils::pl_path::{CloudScheme, PlPath, PlRefPath};
7use polars_utils::pl_str::PlSmallStr;
8use regex::Regex;
9
10use super::CloudOptions;
11
12/// Converts a glob to regex form.
13///
14/// # Returns
15/// 1. the prefix part (all path components until the first one with '*')
16/// 2. a regular expression representation of the rest.
17pub(crate) fn extract_prefix_expansion(path: &str) -> PolarsResult<(Cow<'_, str>, Option<String>)> {
18    // (offset, len, replacement)
19    let mut replacements: Vec<(usize, usize, &[u8])> = vec![];
20
21    // The position after the last slash before glob characters begin.
22    // `a/b/c*/`
23    //      ^
24    let mut pos: usize = if let Some(after_last_slash) =
25        memchr::memchr2(b'*', b'[', path.as_bytes()).map(|i| {
26            path.as_bytes()[..i]
27                .iter()
28                .rposition(|x| *x == b'/')
29                .map_or(0, |x| 1 + x)
30        }) {
31        // First value is used as the starting point later.
32        replacements.push((after_last_slash, 0, &[]));
33        after_last_slash
34    } else {
35        usize::MAX
36    };
37
38    while pos < path.len() {
39        match memchr::memchr2(b'*', b'.', &path.as_bytes()[pos..]) {
40            None => break,
41            Some(i) => pos += i,
42        }
43
44        let (len, replace): (usize, &[u8]) = match &path[pos..] {
45            // Accept:
46            // - `**/`
47            // - `**` only if it is the end of the path
48            v if v.starts_with("**") && (v.len() == 2 || v.as_bytes()[2] == b'/') => {
49                // Wrapping in a capture group ensures we also match non-nested paths.
50                (3, b"(.*/)?" as _)
51            },
52            v if v.starts_with("**") => {
53                polars_bail!(ComputeError: "invalid ** glob pattern")
54            },
55            v if v.starts_with('*') => (1, b"[^/]*" as _),
56            // Dots need to be escaped in regex.
57            v if v.starts_with('.') => (1, b"\\." as _),
58            _ => {
59                pos += 1;
60                continue;
61            },
62        };
63
64        replacements.push((pos, len, replace));
65        pos += len;
66    }
67
68    if replacements.is_empty() {
69        return Ok((Cow::Borrowed(path), None));
70    }
71
72    let prefix = Cow::Borrowed(&path[..replacements[0].0]);
73
74    let mut pos = replacements[0].0;
75    let mut expansion = Vec::with_capacity(path.len() - pos);
76    expansion.push(b'^');
77
78    for (offset, len, replace) in replacements {
79        expansion.extend_from_slice(&path.as_bytes()[pos..offset]);
80        expansion.extend_from_slice(replace);
81        pos = offset + len;
82    }
83
84    if pos < path.len() {
85        expansion.extend_from_slice(&path.as_bytes()[pos..]);
86    }
87
88    expansion.push(b'$');
89
90    Ok((prefix, Some(String::from_utf8(expansion).unwrap())))
91}
92
93/// A location on cloud storage, may have wildcards.
94#[derive(PartialEq, Debug, Default)]
95pub struct CloudLocation {
96    /// The scheme (s3, ...).
97    pub scheme: &'static str,
98    /// The bucket name.
99    pub bucket: PlSmallStr,
100    /// The prefix inside the bucket, this will be the full key when wildcards are not used.
101    pub prefix: String,
102    /// The path components that need to be expanded.
103    pub expansion: Option<PlSmallStr>,
104}
105
106/// Splits an HTTP(S) URL into `scheme://authority` and the object path after it, if the object
107/// path round-trips through [`Path`] unchanged. Otherwise the URL gets its own store.
108pub(crate) fn split_http_url(path: &PlPath) -> Option<(&str, &str)> {
109    let (base_url, rest) = path.as_str().split_at(path.authority_end_position());
110    // `authority_end_position` only stops at '/'.
111    if base_url.contains(['?', '#', '\\']) {
112        return None;
113    }
114    let object_path = rest.strip_prefix('/').unwrap_or(rest);
115    Path::parse(object_path)
116        .is_ok_and(|p| p.as_ref() == object_path)
117        .then_some((base_url, object_path))
118}
119
120impl CloudLocation {
121    pub fn new(path: PlRefPath, glob: bool) -> PolarsResult<Self> {
122        if let Some(scheme @ CloudScheme::Http | scheme @ CloudScheme::Https) = path.scheme() {
123            // Http/s does not glob; the prefix is empty if the store is tied to the URL.
124            return Ok(CloudLocation {
125                scheme: scheme.as_str(),
126                prefix: split_http_url(&path).map_or_else(String::new, |(_, p)| p.to_string()),
127                ..Default::default()
128            });
129        }
130
131        let path_is_local = matches!(
132            path.scheme(),
133            None | Some(CloudScheme::File | CloudScheme::FileNoHostname)
134        );
135
136        let (bucket, key) = path
137            .strip_scheme_split_authority()
138            .ok_or(Cow::Borrowed(
139                "could not extract bucket/key (path did not contain '/')",
140            ))
141            .and_then(|x @ (bucket, _)| {
142                let bucket_is_empty = bucket.is_empty();
143
144                if path_is_local && !bucket_is_empty {
145                    Err(Cow::Owned(format!(
146                        "unsupported: non-empty hostname for 'file:' URI: '{bucket}'",
147                    )))
148                } else if bucket_is_empty && !path_is_local {
149                    Err(Cow::Borrowed("empty bucket name"))
150                } else {
151                    Ok(x)
152                }
153            })
154            .map_err(|failed_reason| {
155                polars_err!(
156                    ComputeError:
157                    "failed to create CloudLocation: {} (path: '{}')",
158                    failed_reason,
159                    path,
160                )
161            })?;
162
163        let key = if path_is_local {
164            key
165        } else {
166            key.strip_prefix('/').unwrap_or(key)
167        };
168
169        let (prefix, expansion) = if glob {
170            let (prefix, expansion) = extract_prefix_expansion(key)?;
171
172            assert_eq!(prefix.starts_with('/'), key.starts_with('/'));
173
174            (prefix, expansion.map(|x| x.into()))
175        } else {
176            (key.into(), None)
177        };
178
179        Ok(CloudLocation {
180            scheme: path.scheme().unwrap_or(CloudScheme::File).as_str(),
181            bucket: PlSmallStr::from_str(bucket),
182            prefix: prefix.into_owned(),
183            expansion,
184        })
185    }
186}
187
188/// Return a full url from a key relative to the given location.
189fn full_url(scheme: &str, bucket: &str, key: Path) -> String {
190    format!("{scheme}://{bucket}/{key}")
191}
192
193/// A simple matcher, if more is required consider depending on https://crates.io/crates/globset.
194/// The Cloud list api returns a list of all the file names under a prefix, there is no additional cost of `readdir`.
195pub(crate) struct Matcher {
196    prefix: String,
197    re: Option<Regex>,
198}
199
200impl Matcher {
201    /// Build a Matcher for the given prefix and expansion.
202    pub(crate) fn new(prefix: String, expansion: Option<&str>) -> PolarsResult<Matcher> {
203        // Cloud APIs accept a prefix without any expansion, extract it.
204        let re = expansion
205            .map(polars_utils::regex_cache::compile_regex)
206            .transpose()?;
207        Ok(Matcher { prefix, re })
208    }
209
210    pub(crate) fn is_matching(&self, key: &str) -> bool {
211        if !key.starts_with(self.prefix.as_str()) {
212            // Prefix does not match, should not happen.
213            return false;
214        }
215        if self.re.is_none() {
216            return true;
217        }
218        let last = &key[self.prefix.len()..];
219        self.re.as_ref().unwrap().is_match(last.as_ref())
220    }
221}
222
223/// List files with a prefix derived from the pattern.
224pub async fn glob(
225    url: PlRefPath,
226    cloud_options: Option<&CloudOptions>,
227) -> PolarsResult<Vec<(String, u64)>> {
228    // Find the fixed prefix, up to the first '*'.
229
230    let (
231        CloudLocation {
232            scheme,
233            bucket,
234            prefix,
235            expansion,
236        },
237        store,
238    ) = super::build_object_store(url, cloud_options, true).await?;
239    let matcher = &Matcher::new(
240        if scheme == "file" {
241            // For local paths the returned location has the leading slash stripped.
242            prefix[1..].into()
243        } else {
244            prefix.clone()
245        },
246        expansion.as_deref(),
247    )?;
248
249    let path = Path::from(prefix.as_str());
250    let path = Some(&path);
251
252    let mut locations = store
253        .exec_with_rebuild_retry_on_err(|store| async move {
254            store
255                .list(path)
256                .try_filter_map(|x| async move {
257                    // Keep the LIST-reported byte size alongside the path.
258                    let out = (x.size > 0 && matcher.is_matching(x.location.as_ref()))
259                        .then_some((x.location, x.size));
260                    Ok(out)
261                })
262                .try_collect::<Vec<_>>()
263                .await
264        })
265        .await?;
266
267    locations.sort_unstable_by(|a, b| a.0.cmp(&b.0));
268    Ok(locations
269        .into_iter()
270        .map(|(l, size)| (full_url(scheme, &bucket, l), size))
271        .collect::<Vec<_>>())
272}
273
274#[cfg(test)]
275mod test {
276    use super::*;
277
278    #[test]
279    fn test_cloud_location() {
280        assert_eq!(
281            CloudLocation::new(PlRefPath::new("s3://a/b"), true).unwrap(),
282            CloudLocation {
283                scheme: "s3",
284                bucket: "a".into(),
285                prefix: "b".into(),
286                expansion: None,
287            }
288        );
289        assert_eq!(
290            CloudLocation::new(PlRefPath::new("s3://a/b/*.c"), true).unwrap(),
291            CloudLocation {
292                scheme: "s3",
293                bucket: "a".into(),
294                prefix: "b/".into(),
295                expansion: Some("^[^/]*\\.c$".into()),
296            }
297        );
298        assert_eq!(
299            CloudLocation::new(PlRefPath::new("file:///a/b"), true).unwrap(),
300            CloudLocation {
301                scheme: "file",
302                bucket: "".into(),
303                prefix: "/a/b".into(),
304                expansion: None,
305            }
306        );
307        assert_eq!(
308            CloudLocation::new(PlRefPath::new("file:/a/b"), true).unwrap(),
309            CloudLocation {
310                scheme: "file",
311                bucket: "".into(),
312                prefix: "/a/b".into(),
313                expansion: None,
314            }
315        );
316    }
317
318    #[test]
319    fn test_extract_prefix_expansion() {
320        assert!(extract_prefix_expansion("**url").is_err());
321        assert_eq!(
322            extract_prefix_expansion("a/b.c").unwrap(),
323            ("a/b.c".into(), None)
324        );
325        assert_eq!(
326            extract_prefix_expansion("a/**").unwrap(),
327            ("a/".into(), Some("^(.*/)?$".into()))
328        );
329        assert_eq!(
330            extract_prefix_expansion("a/**/b").unwrap(),
331            ("a/".into(), Some("^(.*/)?b$".into()))
332        );
333        assert_eq!(
334            extract_prefix_expansion("a/**/*b").unwrap(),
335            ("a/".into(), Some("^(.*/)?[^/]*b$".into()))
336        );
337        assert_eq!(
338            extract_prefix_expansion("a/**/data/*b").unwrap(),
339            ("a/".into(), Some("^(.*/)?data/[^/]*b$".into()))
340        );
341        assert_eq!(
342            extract_prefix_expansion("a/*b").unwrap(),
343            ("a/".into(), Some("^[^/]*b$".into()))
344        );
345    }
346
347    #[test]
348    fn test_matcher_file_name() {
349        let cloud_location =
350            CloudLocation::new(PlRefPath::new("s3://bucket/folder/*.parquet"), true).unwrap();
351        let a = Matcher::new(cloud_location.prefix, cloud_location.expansion.as_deref()).unwrap();
352        // Regular match.
353        assert!(a.is_matching(Path::from("folder/1.parquet").as_ref()));
354        // Require . in the file name.
355        assert!(!a.is_matching(Path::from("folder/1parquet").as_ref()));
356        // Intermediary folders are not allowed.
357        assert!(!a.is_matching(Path::from("folder/other/1.parquet").as_ref()));
358    }
359
360    #[test]
361    fn test_matcher_folders() {
362        let cloud_location =
363            CloudLocation::new(PlRefPath::new("s3://bucket/folder/**/*.parquet"), true).unwrap();
364
365        let a = Matcher::new(cloud_location.prefix, cloud_location.expansion.as_deref()).unwrap();
366        // Intermediary folders are optional.
367        assert!(a.is_matching(Path::from("folder/1.parquet").as_ref()));
368        // Intermediary folders are allowed.
369        assert!(a.is_matching(Path::from("folder/other/1.parquet").as_ref()));
370
371        let cloud_location =
372            CloudLocation::new(PlRefPath::new("s3://bucket/folder/**/data/*.parquet"), true)
373                .unwrap();
374        let a = Matcher::new(cloud_location.prefix, cloud_location.expansion.as_deref()).unwrap();
375
376        // Required folder `data` is missing.
377        assert!(!a.is_matching(Path::from("folder/1.parquet").as_ref()));
378        // Required folder is present.
379        assert!(a.is_matching(Path::from("folder/data/1.parquet").as_ref()));
380        // Required folder is present and additional folders are allowed.
381        assert!(a.is_matching(Path::from("folder/other/data/1.parquet").as_ref()));
382    }
383
384    #[test]
385    fn test_cloud_location_no_glob() {
386        let cloud_location = CloudLocation::new(PlRefPath::new("s3://bucket/[*"), false).unwrap();
387        assert_eq!(
388            cloud_location,
389            CloudLocation {
390                scheme: "s3",
391                bucket: "bucket".into(),
392                prefix: "[*".into(),
393                expansion: None,
394            },
395        )
396    }
397
398    #[test]
399    fn test_cloud_location_percentages() {
400        use super::CloudLocation;
401
402        let path = "s3://bucket/%25";
403        let cloud_location = CloudLocation::new(PlRefPath::new(path), true).unwrap();
404
405        assert_eq!(
406            cloud_location,
407            CloudLocation {
408                scheme: "s3",
409                bucket: "bucket".into(),
410                prefix: "%25".into(),
411                expansion: None,
412            }
413        );
414
415        let path = "https://pola.rs/%25";
416        let cloud_location = CloudLocation::new(PlRefPath::new(path), true).unwrap();
417
418        assert_eq!(
419            cloud_location,
420            CloudLocation {
421                scheme: "https",
422                bucket: "".into(),
423                prefix: "%25".into(),
424                expansion: None,
425            }
426        );
427    }
428
429    #[test]
430    fn test_split_http_url() {
431        use super::split_http_url;
432
433        let check = |url: &str, expected: Option<(&str, &str)>| {
434            assert_eq!(split_http_url(&PlRefPath::new(url)), expected, "{url}")
435        };
436
437        check(
438            "https://pola.rs/a/b.parquet",
439            Some(("https://pola.rs", "a/b.parquet")),
440        );
441        check(
442            "http://u:[email protected]:9000/b/k.ipc?X-Amz-Signature=a%2Fb#f",
443            Some((
444                "http://u:[email protected]:9000",
445                "b/k.ipc?X-Amz-Signature=a%2Fb#f",
446            )),
447        );
448        check("https://pola.rs/%25", Some(("https://pola.rs", "%25")));
449        check("https://pola.rs/", Some(("https://pola.rs", "")));
450        check("https://pola.rs", Some(("https://pola.rs", "")));
451
452        // Not preserved by `Path`, or authority not terminated by '/'.
453        check("https://pola.rs/dir/", None);
454        check("https://pola.rs//a", None);
455        check("https://pola.rs/a//b", None);
456        check("https://pola.rs/a/../b", None);
457        check("https://pola.rs?a=/b", None);
458        check("https://pola.rs\\a/b", None);
459    }
460
461    #[test]
462    fn test_glob_wildcard_21736() {
463        let path = "s3://bucket/folder/**/data.parquet";
464        let cloud_location = CloudLocation::new(PlRefPath::new(path), true).unwrap();
465
466        let a = Matcher::new(cloud_location.prefix, cloud_location.expansion.as_deref()).unwrap();
467
468        assert!(!a.is_matching("folder/_data.parquet"));
469
470        assert!(a.is_matching("folder/data.parquet"));
471        assert!(a.is_matching("folder/abc/data.parquet"));
472        assert!(a.is_matching("folder/abc/def/data.parquet"));
473
474        let path = "s3://bucket/folder/data_*.parquet";
475        let cloud_location = CloudLocation::new(PlRefPath::new(path), true).unwrap();
476
477        let a = Matcher::new(cloud_location.prefix, cloud_location.expansion.as_deref()).unwrap();
478
479        assert!(!a.is_matching("folder/data_1.ipc"))
480    }
481}
polars_io/cloud/glob.rs

polars_io/cloud/
glob.rs
