Skip to main content

datafusion_datasource/
url.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use datafusion_common::{DataFusionError, Result, TableReference};
21use datafusion_execution::cache::cache_manager::CachedFileList;
22use datafusion_execution::cache::cache_manager::TableScopedPath;
23use datafusion_execution::object_store::ObjectStoreUrl;
24use datafusion_session::Session;
25
26use futures::stream::BoxStream;
27use futures::{StreamExt, TryStreamExt};
28use glob::Pattern;
29use itertools::Itertools;
30use log::debug;
31use object_store::path::DELIMITER;
32use object_store::path::Path;
33use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt};
34use url::Url;
35
36/// A parsed URL identifying files for a listing table, see [`ListingTableUrl::parse`]
37/// for more information on the supported expressions
38#[derive(Debug, Clone, Eq, PartialEq, Hash)]
39pub struct ListingTableUrl {
40    /// A URL that identifies a file or directory to list files from
41    url: Url,
42    /// The path prefix
43    prefix: Path,
44    /// An optional glob expression used to filter files
45    glob: Option<Pattern>,
46    /// Optional table reference for the table this url belongs to
47    table_ref: Option<TableReference>,
48}
49
50impl ListingTableUrl {
51    /// Parse a provided string as a `ListingTableUrl`
52    ///
53    /// A URL can either refer to a single object, or a collection of objects with a
54    /// common prefix, with the presence of a trailing `/` indicating a collection.
55    ///
56    /// For example, `file:///foo.txt` refers to the file at `/foo.txt`, whereas
57    /// `file:///foo/` refers to all the files under the directory `/foo` and its
58    /// subdirectories.
59    ///
60    /// Similarly `s3://BUCKET/blob.csv` refers to `blob.csv` in the S3 bucket `BUCKET`,
61    /// whereas `s3://BUCKET/foo/` refers to all objects with the prefix `foo/` in the
62    /// S3 bucket `BUCKET`
63    ///
64    /// # URL Encoding
65    ///
66    /// URL paths are expected to be URL-encoded. That is, the URL for a file named `bar%2Efoo`
67    /// would be `file:///bar%252Efoo`, as per the [URL] specification.
68    ///
69    /// It should be noted that some tools, such as the AWS CLI, take a different approach and
70    /// instead interpret the URL path verbatim. For example the object `bar%2Efoo` would be
71    /// addressed as `s3://BUCKET/bar%252Efoo` using [`ListingTableUrl`] but `s3://BUCKET/bar%2Efoo`
72    /// when using the aws-cli.
73    ///
74    /// # Paths without a Scheme
75    ///
76    /// If no scheme is provided, or the string is an absolute filesystem path
77    /// as determined by [`std::path::Path::is_absolute`], the string will be
78    /// interpreted as a path on the local filesystem using the operating
79    /// system's standard path delimiter, i.e. `\` on Windows, `/` on Unix.
80    ///
81    /// If the path contains any of `'?', '*', '['`, it will be considered
82    /// a glob expression and resolved as described in the section below.
83    ///
84    /// Otherwise, the path will be resolved to an absolute path based on the current
85    /// working directory, and converted to a [file URI].
86    ///
87    /// If the path already exists in the local filesystem this will be used to determine if this
88    /// [`ListingTableUrl`] refers to a collection or a single object, otherwise the presence
89    /// of a trailing path delimiter will be used to indicate a directory. For the avoidance
90    /// of ambiguity it is recommended users always include trailing `/` when intending to
91    /// refer to a directory.
92    ///
93    /// ## Glob File Paths
94    ///
95    /// If no scheme is provided, and the path contains a glob expression, it will
96    /// be resolved as follows.
97    ///
98    /// The string up to the first path segment containing a glob expression will be extracted,
99    /// and resolved in the same manner as a normal scheme-less path above.
100    ///
101    /// The remaining string will be interpreted as a [`glob::Pattern`] and used as a
102    /// filter when listing files from object storage
103    ///
104    /// [file URI]: https://en.wikipedia.org/wiki/File_URI_scheme
105    /// [URL]: https://url.spec.whatwg.org/
106    pub fn parse(s: impl AsRef<str>) -> Result<Self> {
107        let s = s.as_ref();
108
109        // This is necessary to handle the case of a path starting with a drive letter
110        #[cfg(not(target_arch = "wasm32"))]
111        if std::path::Path::new(s).is_absolute() {
112            return Self::parse_path(s);
113        }
114
115        match Url::parse(s) {
116            Ok(url) => Self::try_new(url, None),
117            #[cfg(not(target_arch = "wasm32"))]
118            Err(url::ParseError::RelativeUrlWithoutBase) => Self::parse_path(s),
119            Err(e) => Err(DataFusionError::External(Box::new(e))),
120        }
121    }
122
123    /// Creates a new [`ListingTableUrl`] interpreting `s` as a filesystem path
124    #[cfg(not(target_arch = "wasm32"))]
125    fn parse_path(s: &str) -> Result<Self> {
126        let (path, glob) = match split_glob_expression(s) {
127            Some((prefix, glob)) => {
128                let glob = Pattern::new(glob)
129                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
130                (prefix, Some(glob))
131            }
132            None => (s, None),
133        };
134
135        let url = url_from_filesystem_path(path).ok_or_else(|| {
136            DataFusionError::External(
137                format!("Failed to convert path to URL: {path}").into(),
138            )
139        })?;
140
141        Self::try_new(url, glob)
142    }
143
144    /// Creates a new [`ListingTableUrl`] from a url and optional glob expression
145    ///
146    /// [`Self::parse`] supports glob expression only for file system paths.
147    /// However, some applications may want to support glob expression for URLs with a scheme.
148    /// The application can split the URL into a base URL and a glob expression and use this method
149    /// to create a [`ListingTableUrl`].
150    pub fn try_new(url: Url, glob: Option<Pattern>) -> Result<Self> {
151        let prefix = Path::from_url_path(url.path())?;
152        Ok(Self {
153            url,
154            prefix,
155            glob,
156            table_ref: None,
157        })
158    }
159
160    /// Returns the URL scheme
161    pub fn scheme(&self) -> &str {
162        self.url.scheme()
163    }
164
165    /// Return the URL path not excluding any glob expression
166    ///
167    /// If [`Self::is_collection`], this is the listing prefix
168    /// Otherwise, this is the path to the object
169    pub fn prefix(&self) -> &Path {
170        &self.prefix
171    }
172
173    /// Returns `true` if `path` matches this [`ListingTableUrl`]
174    pub fn contains(&self, path: &Path, ignore_subdirectory: bool) -> bool {
175        let Some(all_segments) = self.strip_prefix(path) else {
176            return false;
177        };
178
179        // remove any segments that contain `=` as they are allowed even
180        // when ignore subdirectories is `true`.
181        let mut segments = all_segments.filter(|s| !s.contains('='));
182
183        match &self.glob {
184            Some(glob) => {
185                if ignore_subdirectory {
186                    segments
187                        .next()
188                        .is_some_and(|file_name| glob.matches(file_name))
189                } else {
190                    let stripped = segments.join(DELIMITER);
191                    glob.matches(&stripped)
192                }
193            }
194            // where we are ignoring subdirectories, we require
195            // the path to be either empty, or contain just the
196            // final file name segment.
197            None if ignore_subdirectory => segments.count() <= 1,
198            // in this case, any valid path at or below the url is allowed
199            None => true,
200        }
201    }
202
203    /// Returns `true` if `path` refers to a collection of objects
204    pub fn is_collection(&self) -> bool {
205        self.url.path().ends_with(DELIMITER)
206    }
207
208    /// Returns the file extension of the last path segment if it exists
209    ///
210    /// Examples:
211    /// ```rust
212    /// use datafusion_datasource::ListingTableUrl;
213    /// let url = ListingTableUrl::parse("file:///foo/bar.csv").unwrap();
214    /// assert_eq!(url.file_extension(), Some("csv"));
215    /// let url = ListingTableUrl::parse("file:///foo/bar").unwrap();
216    /// assert_eq!(url.file_extension(), None);
217    /// let url = ListingTableUrl::parse("file:///foo/bar.").unwrap();
218    /// assert_eq!(url.file_extension(), None);
219    /// ```
220    pub fn file_extension(&self) -> Option<&str> {
221        if let Some(mut segments) = self.url.path_segments()
222            && let Some(last_segment) = segments.next_back()
223            && last_segment.contains(".")
224            && !last_segment.ends_with(".")
225        {
226            return last_segment.split('.').next_back();
227        }
228
229        None
230    }
231
232    /// Strips the prefix of this [`ListingTableUrl`] from the provided path, returning
233    /// an iterator of the remaining path segments
234    pub fn strip_prefix<'a, 'b: 'a>(
235        &'a self,
236        path: &'b Path,
237    ) -> Option<impl Iterator<Item = &'b str> + 'a> {
238        let mut stripped = path.as_ref().strip_prefix(self.prefix.as_ref())?;
239        if !stripped.is_empty() && !self.prefix.as_ref().is_empty() {
240            stripped = stripped.strip_prefix(DELIMITER)?;
241        }
242        Some(stripped.split_terminator(DELIMITER))
243    }
244
245    /// List all files identified by this [`ListingTableUrl`] for the provided `file_extension`,
246    /// optionally filtering by a path prefix
247    pub async fn list_prefixed_files<'a>(
248        &'a self,
249        ctx: &'a dyn Session,
250        store: &'a dyn ObjectStore,
251        prefix: Option<Path>,
252        file_extension: &'a str,
253    ) -> Result<BoxStream<'a, Result<ObjectMeta>>> {
254        let exec_options = &ctx.config_options().execution;
255        let ignore_subdirectory = exec_options.listing_table_ignore_subdirectory;
256
257        // Build full_prefix for non-cached path and head() calls
258        let full_prefix = if let Some(ref p) = prefix {
259            let mut parts = self.prefix.parts().collect::<Vec<_>>();
260            parts.extend(p.parts());
261            Path::from_iter(parts)
262        } else {
263            self.prefix.clone()
264        };
265
266        let list: BoxStream<'a, Result<ObjectMeta>> = if self.is_collection() {
267            list_with_cache(
268                ctx,
269                store,
270                self.table_ref.as_ref(),
271                &self.prefix,
272                prefix.as_ref(),
273            )
274            .await?
275        } else {
276            match store.head(&full_prefix).await {
277                Ok(meta) => futures::stream::once(async { Ok(meta) })
278                    .map_err(|e| DataFusionError::ObjectStore(Box::new(e)))
279                    .boxed(),
280                // If the head command fails, it is likely that object doesn't exist.
281                // Retry as though it were a prefix (aka a collection)
282                Err(object_store::Error::NotFound { .. }) => {
283                    list_with_cache(
284                        ctx,
285                        store,
286                        self.table_ref.as_ref(),
287                        &self.prefix,
288                        prefix.as_ref(),
289                    )
290                    .await?
291                }
292                Err(e) => return Err(e.into()),
293            }
294        };
295
296        Ok(list
297            .try_filter(move |meta| {
298                let path = &meta.location;
299                let extension_match = path.as_ref().ends_with(file_extension);
300                let glob_match = self.contains(path, ignore_subdirectory);
301                futures::future::ready(extension_match && glob_match)
302            })
303            .boxed())
304    }
305
306    /// List all files identified by this [`ListingTableUrl`] for the provided `file_extension`
307    pub async fn list_all_files<'a>(
308        &'a self,
309        ctx: &'a dyn Session,
310        store: &'a dyn ObjectStore,
311        file_extension: &'a str,
312    ) -> Result<BoxStream<'a, Result<ObjectMeta>>> {
313        self.list_prefixed_files(ctx, store, None, file_extension)
314            .await
315    }
316
317    /// Returns this [`ListingTableUrl`] as a string
318    pub fn as_str(&self) -> &str {
319        self.as_ref()
320    }
321
322    /// Return the [`ObjectStoreUrl`] for this [`ListingTableUrl`]
323    pub fn object_store(&self) -> ObjectStoreUrl {
324        let url = &self.url[url::Position::BeforeScheme..url::Position::BeforePath];
325        ObjectStoreUrl::parse(url).unwrap()
326    }
327
328    /// Returns true if the [`ListingTableUrl`] points to the folder
329    pub fn is_folder(&self) -> bool {
330        self.url.scheme() == "file" && self.is_collection()
331    }
332
333    /// Return the `url` for [`ListingTableUrl`]
334    pub fn get_url(&self) -> &Url {
335        &self.url
336    }
337
338    /// Return the `glob` for [`ListingTableUrl`]
339    pub fn get_glob(&self) -> &Option<Pattern> {
340        &self.glob
341    }
342
343    /// Returns a copy of current [`ListingTableUrl`] with a specified `glob`
344    pub fn with_glob(mut self, glob: &str) -> Result<Self> {
345        self.glob =
346            Some(Pattern::new(glob).map_err(|e| DataFusionError::External(Box::new(e)))?);
347        Ok(self)
348    }
349
350    /// Set the table reference for this [`ListingTableUrl`]
351    pub fn with_table_ref(mut self, table_ref: TableReference) -> Self {
352        self.table_ref = Some(table_ref);
353        self
354    }
355
356    /// Return the table reference for this [`ListingTableUrl`]
357    pub fn get_table_ref(&self) -> &Option<TableReference> {
358        &self.table_ref
359    }
360}
361
362/// Lists files with cache support, using prefix-aware lookups.
363///
364/// # Arguments
365/// * `ctx` - The session context
366/// * `store` - The object store to list from
367/// * `table_base_path` - The table's base path (the stable cache key)
368/// * `prefix` - Optional prefix relative to table base for filtering results
369///
370/// # Cache Behavior:
371/// The cache key is always `table_base_path`. When a prefix-filtered listing
372/// is requested via `prefix`, the cache:
373/// - Looks up `table_base_path` in the cache
374/// - Filters results to match `table_base_path/prefix`
375/// - Returns filtered results without a storage call
376///
377/// On cache miss, the full table is always listed and cached, ensuring
378/// subsequent prefix queries can be served from cache.
379async fn list_with_cache<'b>(
380    ctx: &'b dyn Session,
381    store: &'b dyn ObjectStore,
382    table_ref: Option<&TableReference>,
383    table_base_path: &Path,
384    prefix: Option<&Path>,
385) -> Result<BoxStream<'b, Result<ObjectMeta>>> {
386    // Build the full listing path (table_base + prefix)
387    let full_prefix = match prefix {
388        Some(p) => {
389            let mut parts: Vec<_> = table_base_path.parts().collect();
390            parts.extend(p.parts());
391            Path::from_iter(parts)
392        }
393        None => table_base_path.clone(),
394    };
395
396    match ctx.runtime_env().cache_manager.get_list_files_cache() {
397        None => Ok(store
398            .list(Some(&full_prefix))
399            .map(|res| res.map_err(|e| DataFusionError::ObjectStore(Box::new(e))))
400            .boxed()),
401        Some(cache) => {
402            // Build the filter prefix (only Some if prefix was requested)
403            let filter_prefix = prefix.is_some().then(|| full_prefix.clone());
404
405            let table_scoped_base_path = TableScopedPath {
406                table: table_ref.cloned(),
407                path: table_base_path.clone(),
408            };
409
410            // Try cache lookup - get returns CachedFileList
411            let vec = if let Some(cached) = cache.get(&table_scoped_base_path) {
412                debug!("Hit list files cache");
413                cached.files_matching_prefix(&filter_prefix)
414            } else {
415                // Cache miss - always list and cache the full table
416                // This ensures we have complete data for future prefix queries
417                let mut vec = store
418                    .list(Some(table_base_path))
419                    .try_collect::<Vec<ObjectMeta>>()
420                    .await?;
421                vec.shrink_to_fit(); // Right-size before caching
422                let cached: CachedFileList = vec.into();
423                let result = cached.files_matching_prefix(&filter_prefix);
424                cache.put(&table_scoped_base_path, cached);
425                result
426            };
427            Ok(
428                futures::stream::iter(Arc::unwrap_or_clone(vec).into_iter().map(Ok))
429                    .boxed(),
430            )
431        }
432    }
433}
434
435/// Creates a file URL from a potentially relative filesystem path
436#[cfg(not(target_arch = "wasm32"))]
437fn url_from_filesystem_path(s: &str) -> Option<Url> {
438    let path = std::path::Path::new(s);
439    let is_dir = match path.exists() {
440        true => path.is_dir(),
441        // Fallback to inferring from trailing separator
442        false => std::path::is_separator(s.chars().last()?),
443    };
444
445    let from_absolute_path = |p| {
446        let first = match is_dir {
447            true => Url::from_directory_path(p).ok(),
448            false => Url::from_file_path(p).ok(),
449        }?;
450
451        // By default from_*_path preserve relative path segments
452        // We therefore parse the URL again to resolve these
453        Url::parse(first.as_str()).ok()
454    };
455
456    if path.is_absolute() {
457        return from_absolute_path(path);
458    }
459
460    let absolute = std::env::current_dir().ok()?.join(path);
461    from_absolute_path(&absolute)
462}
463
464impl AsRef<str> for ListingTableUrl {
465    fn as_ref(&self) -> &str {
466        self.url.as_ref()
467    }
468}
469
470impl AsRef<Url> for ListingTableUrl {
471    fn as_ref(&self) -> &Url {
472        &self.url
473    }
474}
475
476impl std::fmt::Display for ListingTableUrl {
477    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
478        self.as_str().fmt(f)
479    }
480}
481
482#[cfg(not(target_arch = "wasm32"))]
483const GLOB_START_CHARS: [char; 3] = ['?', '*', '['];
484
485/// Splits `path` at the first path segment containing a glob expression, returning
486/// `None` if no glob expression found.
487///
488/// Path delimiters are determined using [`std::path::is_separator`] which
489/// permits `/` as a path delimiter even on Windows platforms.
490#[cfg(not(target_arch = "wasm32"))]
491fn split_glob_expression(path: &str) -> Option<(&str, &str)> {
492    let mut last_separator = 0;
493
494    for (byte_idx, char) in path.char_indices() {
495        if GLOB_START_CHARS.contains(&char) {
496            if last_separator == 0 {
497                return Some((".", path));
498            }
499            return Some(path.split_at(last_separator));
500        }
501
502        if std::path::is_separator(char) {
503            last_separator = byte_idx + char.len_utf8();
504        }
505    }
506    None
507}
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512    use async_trait::async_trait;
513    use bytes::Bytes;
514    use datafusion_common::DFSchema;
515    use datafusion_common::config::TableOptions;
516    use datafusion_execution::TaskContext;
517    use datafusion_execution::config::SessionConfig;
518    use datafusion_execution::runtime_env::RuntimeEnv;
519    use datafusion_expr::execution_props::ExecutionProps;
520    use datafusion_expr::registry::ExtensionTypeRegistryRef;
521    use datafusion_expr::{
522        AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
523    };
524    use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
525    use datafusion_physical_plan::ExecutionPlan;
526    use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList};
527    use object_store::{
528        CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload,
529        PutMultipartOptions, PutPayload,
530    };
531    use std::any::Any;
532    use std::collections::HashMap;
533    use std::ops::Range;
534    use tempfile::tempdir;
535
536    #[test]
537    fn test_prefix_path() {
538        let root = std::env::current_dir().unwrap();
539        let root = root.to_string_lossy();
540
541        let url = ListingTableUrl::parse(root).unwrap();
542        let child = url.prefix.clone().join("partition").join("file");
543
544        let prefix: Vec<_> = url.strip_prefix(&child).unwrap().collect();
545        assert_eq!(prefix, vec!["partition", "file"]);
546
547        let url = ListingTableUrl::parse("file:///").unwrap();
548        let child = Path::parse("/foo/bar").unwrap();
549        let prefix: Vec<_> = url.strip_prefix(&child).unwrap().collect();
550        assert_eq!(prefix, vec!["foo", "bar"]);
551
552        let url = ListingTableUrl::parse("file:///foo").unwrap();
553        let child = Path::parse("/foob/bar").unwrap();
554        assert!(url.strip_prefix(&child).is_none());
555
556        let url = ListingTableUrl::parse("file:///foo/file").unwrap();
557        let child = Path::parse("/foo/file").unwrap();
558        assert_eq!(url.strip_prefix(&child).unwrap().count(), 0);
559
560        let url = ListingTableUrl::parse("file:///foo/ bar").unwrap();
561        assert_eq!(url.prefix.as_ref(), "foo/ bar");
562
563        let url = ListingTableUrl::parse("file:///foo/bar?").unwrap();
564        assert_eq!(url.prefix.as_ref(), "foo/bar");
565
566        let url = ListingTableUrl::parse("file:///foo/😺").unwrap();
567        assert_eq!(url.prefix.as_ref(), "foo/😺");
568
569        let url = ListingTableUrl::parse("file:///foo/bar%2Efoo").unwrap();
570        assert_eq!(url.prefix.as_ref(), "foo/bar.foo");
571
572        let url = ListingTableUrl::parse("file:///foo/bar%2Efoo").unwrap();
573        assert_eq!(url.prefix.as_ref(), "foo/bar.foo");
574
575        let url = ListingTableUrl::parse("file:///foo/bar%252Ffoo").unwrap();
576        assert_eq!(url.prefix.as_ref(), "foo/bar%2Ffoo");
577
578        let url = ListingTableUrl::parse("file:///foo/a%252Fb.txt").unwrap();
579        assert_eq!(url.prefix.as_ref(), "foo/a%2Fb.txt");
580
581        let dir = tempdir().unwrap();
582        let path = dir.path().join("bar%2Ffoo");
583        std::fs::File::create(&path).unwrap();
584
585        let url = ListingTableUrl::parse(path.to_str().unwrap()).unwrap();
586        assert!(url.prefix.as_ref().ends_with("bar%2Ffoo"), "{}", url.prefix);
587
588        let url = ListingTableUrl::parse("file:///foo/../a%252Fb.txt").unwrap();
589        assert_eq!(url.prefix.as_ref(), "a%2Fb.txt");
590
591        let url =
592            ListingTableUrl::parse("file:///foo/./bar/../../baz/./test.txt").unwrap();
593        assert_eq!(url.prefix.as_ref(), "baz/test.txt");
594
595        let workdir = std::env::current_dir().unwrap();
596        let t = workdir.join("non-existent");
597        let a = ListingTableUrl::parse(t.to_str().unwrap()).unwrap();
598        let b = ListingTableUrl::parse("non-existent").unwrap();
599        assert_eq!(a, b);
600        assert!(a.prefix.as_ref().ends_with("non-existent"));
601
602        let t = workdir.parent().unwrap();
603        let a = ListingTableUrl::parse(t.to_str().unwrap()).unwrap();
604        let b = ListingTableUrl::parse("..").unwrap();
605        assert_eq!(a, b);
606
607        let t = t.join("bar");
608        let a = ListingTableUrl::parse(t.to_str().unwrap()).unwrap();
609        let b = ListingTableUrl::parse("../bar").unwrap();
610        assert_eq!(a, b);
611        assert!(a.prefix.as_ref().ends_with("bar"));
612
613        let t = t.join(".").join("foo").join("..").join("baz");
614        let a = ListingTableUrl::parse(t.to_str().unwrap()).unwrap();
615        let b = ListingTableUrl::parse("../bar/./foo/../baz").unwrap();
616        assert_eq!(a, b);
617        assert!(a.prefix.as_ref().ends_with("bar/baz"));
618    }
619
620    #[test]
621    fn test_prefix_s3() {
622        let url = ListingTableUrl::parse("s3://bucket/foo/bar").unwrap();
623        assert_eq!(url.prefix.as_ref(), "foo/bar");
624
625        let path = Path::from("foo/bar/partition/foo.parquet");
626        let prefix: Vec<_> = url.strip_prefix(&path).unwrap().collect();
627        assert_eq!(prefix, vec!["partition", "foo.parquet"]);
628
629        let path = Path::from("other/bar/partition/foo.parquet");
630        assert!(url.strip_prefix(&path).is_none());
631    }
632
633    #[test]
634    fn test_split_glob() {
635        fn test(input: &str, expected: Option<(&str, &str)>) {
636            assert_eq!(
637                split_glob_expression(input),
638                expected,
639                "testing split_glob_expression with {input}"
640            );
641        }
642
643        // no glob patterns
644        test("/", None);
645        test("/a.txt", None);
646        test("/a", None);
647        test("/a/", None);
648        test("/a/b", None);
649        test("/a/b/", None);
650        test("/a/b.txt", None);
651        test("/a/b/c.txt", None);
652        // glob patterns, thus we build the longest path (os-specific)
653        test("*.txt", Some((".", "*.txt")));
654        test("/*.txt", Some(("/", "*.txt")));
655        test("/a/*b.txt", Some(("/a/", "*b.txt")));
656        test("/a/*/b.txt", Some(("/a/", "*/b.txt")));
657        test("/a/b/[123]/file*.txt", Some(("/a/b/", "[123]/file*.txt")));
658        test("/a/b*.txt", Some(("/a/", "b*.txt")));
659        test("/a/b/**/c*.txt", Some(("/a/b/", "**/c*.txt")));
660
661        // https://github.com/apache/datafusion/issues/2465
662        test(
663            "/a/b/c//alltypes_plain*.parquet",
664            Some(("/a/b/c//", "alltypes_plain*.parquet")),
665        );
666    }
667
668    #[test]
669    fn test_is_collection() {
670        fn test(input: &str, expected: bool, message: &str) {
671            let url = ListingTableUrl::parse(input).unwrap();
672            assert_eq!(url.is_collection(), expected, "{message}");
673        }
674
675        test("https://a.b.c/path/", true, "path ends with / - collection");
676        test(
677            "https://a.b.c/path/?a=b",
678            true,
679            "path ends with / - with query args - collection",
680        );
681        test(
682            "https://a.b.c/path?a=b/",
683            false,
684            "path not ends with / - query ends with / - not collection",
685        );
686        test(
687            "https://a.b.c/path/#a=b",
688            true,
689            "path ends with / - with fragment - collection",
690        );
691        test(
692            "https://a.b.c/path#a=b/",
693            false,
694            "path not ends with / - fragment ends with / - not collection",
695        );
696    }
697
698    #[test]
699    fn test_file_extension() {
700        fn test(input: &str, expected: Option<&str>, message: &str) {
701            let url = ListingTableUrl::parse(input).unwrap();
702            assert_eq!(url.file_extension(), expected, "{message}");
703        }
704
705        test("https://a.b.c/path/", None, "path ends with / - not a file");
706        test(
707            "https://a.b.c/path/?a=b",
708            None,
709            "path ends with / - with query args - not a file",
710        );
711        test(
712            "https://a.b.c/path?a=b/",
713            None,
714            "path not ends with / - query ends with / but no file extension",
715        );
716        test(
717            "https://a.b.c/path/#a=b",
718            None,
719            "path ends with / - with fragment - not a file",
720        );
721        test(
722            "https://a.b.c/path#a=b/",
723            None,
724            "path not ends with / - fragment ends with / but no file extension",
725        );
726        test(
727            "file///some/path/",
728            None,
729            "file path ends with / - not a file",
730        );
731        test(
732            "file///some/path/file",
733            None,
734            "file path does not end with - no extension",
735        );
736        test(
737            "file///some/path/file.",
738            None,
739            "file path ends with . - no value after .",
740        );
741        test(
742            "file///some/path/file.ext",
743            Some("ext"),
744            "file path ends with .ext - extension is ext",
745        );
746    }
747
748    #[tokio::test]
749    async fn test_list_files() -> Result<()> {
750        let store = MockObjectStore {
751            in_mem: object_store::memory::InMemory::new(),
752            forbidden_paths: vec!["forbidden/e.parquet".into()],
753        };
754
755        // Create some files:
756        create_file(&store, "a.parquet").await;
757        create_file(&store, "/t/b.parquet").await;
758        create_file(&store, "/t/c.csv").await;
759        create_file(&store, "/t/d.csv").await;
760
761        // This file returns a permission error.
762        create_file(&store, "/forbidden/e.parquet").await;
763
764        assert_eq!(
765            list_all_files("/", &store, "parquet").await?,
766            vec!["a.parquet"],
767        );
768
769        // test with and without trailing slash
770        assert_eq!(
771            list_all_files("/t/", &store, "parquet").await?,
772            vec!["t/b.parquet"],
773        );
774        assert_eq!(
775            list_all_files("/t", &store, "parquet").await?,
776            vec!["t/b.parquet"],
777        );
778
779        // test with and without trailing slash
780        assert_eq!(
781            list_all_files("/t", &store, "csv").await?,
782            vec!["t/c.csv", "t/d.csv"],
783        );
784        assert_eq!(
785            list_all_files("/t/", &store, "csv").await?,
786            vec!["t/c.csv", "t/d.csv"],
787        );
788
789        // Test a non existing prefix
790        assert_eq!(
791            list_all_files("/NonExisting", &store, "csv").await?,
792            vec![] as Vec<String>
793        );
794        assert_eq!(
795            list_all_files("/NonExisting/", &store, "csv").await?,
796            vec![] as Vec<String>
797        );
798
799        // Including forbidden.parquet generates an error.
800        let Err(DataFusionError::ObjectStore(err)) =
801            list_all_files("/forbidden/e.parquet", &store, "parquet").await
802        else {
803            panic!("Expected ObjectStore error");
804        };
805
806        let object_store::Error::PermissionDenied { .. } = &*err else {
807            panic!("Expected PermissionDenied error");
808        };
809
810        // Test prefix filtering with partition-style paths
811        create_file(&store, "/data/a=1/file1.parquet").await;
812        create_file(&store, "/data/a=1/b=100/file2.parquet").await;
813        create_file(&store, "/data/a=2/b=200/file3.parquet").await;
814        create_file(&store, "/data/a=2/b=200/file4.csv").await;
815
816        assert_eq!(
817            list_prefixed_files("/data/", &store, Some(Path::from("a=1")), "parquet")
818                .await?,
819            vec!["data/a=1/b=100/file2.parquet", "data/a=1/file1.parquet"],
820        );
821
822        assert_eq!(
823            list_prefixed_files(
824                "/data/",
825                &store,
826                Some(Path::from("a=1/b=100")),
827                "parquet"
828            )
829            .await?,
830            vec!["data/a=1/b=100/file2.parquet"],
831        );
832
833        assert_eq!(
834            list_prefixed_files("/data/", &store, Some(Path::from("a=2")), "parquet")
835                .await?,
836            vec!["data/a=2/b=200/file3.parquet"],
837        );
838
839        Ok(())
840    }
841
842    /// Tests that the cached code path produces identical results to the non-cached path.
843    ///
844    /// This is critical: the cache is a transparent optimization, so both paths
845    /// MUST return the same files. Note: order is not guaranteed by ObjectStore::list,
846    /// so we sort results before comparison.
847    #[tokio::test]
848    async fn test_cache_path_equivalence() -> Result<()> {
849        use datafusion_execution::runtime_env::RuntimeEnvBuilder;
850
851        let store = MockObjectStore {
852            in_mem: object_store::memory::InMemory::new(),
853            forbidden_paths: vec![],
854        };
855
856        // Create test files with partition-style paths
857        create_file(&store, "/table/year=2023/data1.parquet").await;
858        create_file(&store, "/table/year=2023/month=01/data2.parquet").await;
859        create_file(&store, "/table/year=2024/data3.parquet").await;
860        create_file(&store, "/table/year=2024/month=06/data4.parquet").await;
861        create_file(&store, "/table/year=2024/month=12/data5.parquet").await;
862
863        // Session WITHOUT cache
864        let session_no_cache = MockSession::new();
865
866        // Session WITH cache - use RuntimeEnvBuilder with cache limit (no TTL needed for this test)
867        let runtime_with_cache = RuntimeEnvBuilder::new()
868            .with_object_list_cache_limit(1024 * 1024) // 1MB limit
869            .build_arc()?;
870        let session_with_cache = MockSession::with_runtime_env(runtime_with_cache);
871
872        // Test cases: (url, prefix, description)
873        let test_cases = vec![
874            ("/table/", None, "full table listing"),
875            (
876                "/table/",
877                Some(Path::from("year=2023")),
878                "single partition filter",
879            ),
880            (
881                "/table/",
882                Some(Path::from("year=2024")),
883                "different partition filter",
884            ),
885            (
886                "/table/",
887                Some(Path::from("year=2024/month=06")),
888                "nested partition filter",
889            ),
890            (
891                "/table/",
892                Some(Path::from("year=2025")),
893                "non-existent partition",
894            ),
895        ];
896
897        for (url_str, prefix, description) in test_cases {
898            let url = ListingTableUrl::parse(url_str)?;
899
900            // Get results WITHOUT cache (sorted for comparison)
901            let mut results_no_cache: Vec<String> = url
902                .list_prefixed_files(&session_no_cache, &store, prefix.clone(), "parquet")
903                .await?
904                .try_collect::<Vec<_>>()
905                .await?
906                .into_iter()
907                .map(|m| m.location.to_string())
908                .collect();
909            results_no_cache.sort();
910
911            // Get results WITH cache (first call - cache miss, sorted for comparison)
912            let mut results_with_cache_miss: Vec<String> = url
913                .list_prefixed_files(
914                    &session_with_cache,
915                    &store,
916                    prefix.clone(),
917                    "parquet",
918                )
919                .await?
920                .try_collect::<Vec<_>>()
921                .await?
922                .into_iter()
923                .map(|m| m.location.to_string())
924                .collect();
925            results_with_cache_miss.sort();
926
927            // Get results WITH cache (second call - cache hit, sorted for comparison)
928            let mut results_with_cache_hit: Vec<String> = url
929                .list_prefixed_files(&session_with_cache, &store, prefix, "parquet")
930                .await?
931                .try_collect::<Vec<_>>()
932                .await?
933                .into_iter()
934                .map(|m| m.location.to_string())
935                .collect();
936            results_with_cache_hit.sort();
937
938            // All three should contain the same files
939            assert_eq!(
940                results_no_cache, results_with_cache_miss,
941                "Cache miss path should match non-cached path for: {description}"
942            );
943            assert_eq!(
944                results_no_cache, results_with_cache_hit,
945                "Cache hit path should match non-cached path for: {description}"
946            );
947        }
948
949        Ok(())
950    }
951
952    /// Tests that prefix queries can be served from a cached full-table listing
953    #[tokio::test]
954    async fn test_cache_serves_partition_from_full_listing() -> Result<()> {
955        use datafusion_execution::runtime_env::RuntimeEnvBuilder;
956
957        let store = MockObjectStore {
958            in_mem: object_store::memory::InMemory::new(),
959            forbidden_paths: vec![],
960        };
961
962        // Create test files
963        create_file(&store, "/sales/region=US/q1.parquet").await;
964        create_file(&store, "/sales/region=US/q2.parquet").await;
965        create_file(&store, "/sales/region=EU/q1.parquet").await;
966
967        // Create session with cache (no TTL needed for this test)
968        let runtime = RuntimeEnvBuilder::new()
969            .with_object_list_cache_limit(1024 * 1024) // 1MB limit
970            .build_arc()?;
971        let session = MockSession::with_runtime_env(runtime);
972
973        let url = ListingTableUrl::parse("/sales/")?;
974
975        // First: query full table (populates cache)
976        let full_results: Vec<String> = url
977            .list_prefixed_files(&session, &store, None, "parquet")
978            .await?
979            .try_collect::<Vec<_>>()
980            .await?
981            .into_iter()
982            .map(|m| m.location.to_string())
983            .collect();
984        assert_eq!(full_results.len(), 3);
985
986        // Second: query with prefix (should be served from cache)
987        let mut us_results: Vec<String> = url
988            .list_prefixed_files(
989                &session,
990                &store,
991                Some(Path::from("region=US")),
992                "parquet",
993            )
994            .await?
995            .try_collect::<Vec<_>>()
996            .await?
997            .into_iter()
998            .map(|m| m.location.to_string())
999            .collect();
1000        us_results.sort();
1001
1002        assert_eq!(
1003            us_results,
1004            vec!["sales/region=US/q1.parquet", "sales/region=US/q2.parquet"]
1005        );
1006
1007        // Third: different prefix (also from cache)
1008        let eu_results: Vec<String> = url
1009            .list_prefixed_files(
1010                &session,
1011                &store,
1012                Some(Path::from("region=EU")),
1013                "parquet",
1014            )
1015            .await?
1016            .try_collect::<Vec<_>>()
1017            .await?
1018            .into_iter()
1019            .map(|m| m.location.to_string())
1020            .collect();
1021
1022        assert_eq!(eu_results, vec!["sales/region=EU/q1.parquet"]);
1023
1024        Ok(())
1025    }
1026
1027    /// Creates a file with "hello world" content at the specified path
1028    async fn create_file(object_store: &dyn ObjectStore, path: &str) {
1029        object_store
1030            .put(&Path::from(path), PutPayload::from_static(b"hello world"))
1031            .await
1032            .expect("failed to create test file");
1033    }
1034
1035    /// Runs "list_prefixed_files"  with no prefix to list all files and returns their paths
1036    ///
1037    /// Panic's on error
1038    async fn list_all_files(
1039        url: &str,
1040        store: &dyn ObjectStore,
1041        file_extension: &str,
1042    ) -> Result<Vec<String>> {
1043        try_list_prefixed_files(url, store, None, file_extension).await
1044    }
1045
1046    /// Runs "list_prefixed_files" and returns their paths
1047    ///
1048    /// Panic's on error
1049    async fn list_prefixed_files(
1050        url: &str,
1051        store: &dyn ObjectStore,
1052        prefix: Option<Path>,
1053        file_extension: &str,
1054    ) -> Result<Vec<String>> {
1055        try_list_prefixed_files(url, store, prefix, file_extension).await
1056    }
1057
1058    /// Runs "list_prefixed_files" and returns their paths
1059    async fn try_list_prefixed_files(
1060        url: &str,
1061        store: &dyn ObjectStore,
1062        prefix: Option<Path>,
1063        file_extension: &str,
1064    ) -> Result<Vec<String>> {
1065        let session = MockSession::new();
1066        let url = ListingTableUrl::parse(url)?;
1067        let files = url
1068            .list_prefixed_files(&session, store, prefix, file_extension)
1069            .await?
1070            .try_collect::<Vec<_>>()
1071            .await?
1072            .into_iter()
1073            .map(|meta| meta.location.as_ref().to_string())
1074            .collect();
1075        Ok(files)
1076    }
1077
1078    #[derive(Debug)]
1079    struct MockObjectStore {
1080        in_mem: object_store::memory::InMemory,
1081        forbidden_paths: Vec<Path>,
1082    }
1083
1084    impl std::fmt::Display for MockObjectStore {
1085        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1086            self.in_mem.fmt(f)
1087        }
1088    }
1089
1090    #[async_trait]
1091    impl ObjectStore for MockObjectStore {
1092        async fn put_opts(
1093            &self,
1094            location: &Path,
1095            payload: PutPayload,
1096            opts: object_store::PutOptions,
1097        ) -> object_store::Result<object_store::PutResult> {
1098            self.in_mem.put_opts(location, payload, opts).await
1099        }
1100
1101        async fn put_multipart_opts(
1102            &self,
1103            location: &Path,
1104            opts: PutMultipartOptions,
1105        ) -> object_store::Result<Box<dyn MultipartUpload>> {
1106            self.in_mem.put_multipart_opts(location, opts).await
1107        }
1108
1109        async fn get_opts(
1110            &self,
1111            location: &Path,
1112            options: GetOptions,
1113        ) -> object_store::Result<GetResult> {
1114            if options.head && self.forbidden_paths.contains(location) {
1115                Err(object_store::Error::PermissionDenied {
1116                    path: location.to_string(),
1117                    source: "forbidden".into(),
1118                })
1119            } else {
1120                self.in_mem.get_opts(location, options).await
1121            }
1122        }
1123
1124        async fn get_ranges(
1125            &self,
1126            location: &Path,
1127            ranges: &[Range<u64>],
1128        ) -> object_store::Result<Vec<Bytes>> {
1129            self.in_mem.get_ranges(location, ranges).await
1130        }
1131
1132        fn delete_stream(
1133            &self,
1134            locations: BoxStream<'static, object_store::Result<Path>>,
1135        ) -> BoxStream<'static, object_store::Result<Path>> {
1136            self.in_mem.delete_stream(locations)
1137        }
1138
1139        fn list(
1140            &self,
1141            prefix: Option<&Path>,
1142        ) -> BoxStream<'static, object_store::Result<ObjectMeta>> {
1143            self.in_mem.list(prefix)
1144        }
1145
1146        async fn list_with_delimiter(
1147            &self,
1148            prefix: Option<&Path>,
1149        ) -> object_store::Result<ListResult> {
1150            self.in_mem.list_with_delimiter(prefix).await
1151        }
1152
1153        async fn copy_opts(
1154            &self,
1155            from: &Path,
1156            to: &Path,
1157            options: CopyOptions,
1158        ) -> object_store::Result<()> {
1159            self.in_mem.copy_opts(from, to, options).await
1160        }
1161    }
1162
1163    struct MockSession {
1164        config: SessionConfig,
1165        runtime_env: Arc<RuntimeEnv>,
1166    }
1167
1168    impl MockSession {
1169        fn new() -> Self {
1170            Self {
1171                config: SessionConfig::new(),
1172                runtime_env: Arc::new(RuntimeEnv::default()),
1173            }
1174        }
1175
1176        /// Create a MockSession with a custom RuntimeEnv (for cache testing)
1177        fn with_runtime_env(runtime_env: Arc<RuntimeEnv>) -> Self {
1178            Self {
1179                config: SessionConfig::new(),
1180                runtime_env,
1181            }
1182        }
1183    }
1184
1185    #[async_trait::async_trait]
1186    impl Session for MockSession {
1187        fn session_id(&self) -> &str {
1188            unimplemented!()
1189        }
1190
1191        fn config(&self) -> &SessionConfig {
1192            &self.config
1193        }
1194
1195        fn catalog_list(&self) -> Arc<dyn CatalogProviderList> {
1196            Arc::new(EmptyCatalogProviderList)
1197        }
1198
1199        async fn create_physical_plan(
1200            &self,
1201            _logical_plan: &LogicalPlan,
1202        ) -> Result<Arc<dyn ExecutionPlan>> {
1203            unimplemented!()
1204        }
1205
1206        fn create_physical_expr(
1207            &self,
1208            _expr: Expr,
1209            _df_schema: &DFSchema,
1210        ) -> Result<Arc<dyn PhysicalExpr>> {
1211            unimplemented!()
1212        }
1213
1214        fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> {
1215            unimplemented!()
1216        }
1217
1218        fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>> {
1219            unimplemented!()
1220        }
1221
1222        fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
1223            unimplemented!()
1224        }
1225
1226        fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
1227            unimplemented!()
1228        }
1229
1230        fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef {
1231            unimplemented!()
1232        }
1233
1234        fn runtime_env(&self) -> &Arc<RuntimeEnv> {
1235            &self.runtime_env
1236        }
1237
1238        fn execution_props(&self) -> &ExecutionProps {
1239            unimplemented!()
1240        }
1241
1242        fn as_any(&self) -> &dyn Any {
1243            unimplemented!()
1244        }
1245
1246        fn table_options(&self) -> &TableOptions {
1247            unimplemented!()
1248        }
1249
1250        fn table_options_mut(&mut self) -> &mut TableOptions {
1251            unimplemented!()
1252        }
1253
1254        fn task_ctx(&self) -> Arc<TaskContext> {
1255            unimplemented!()
1256        }
1257    }
1258}