1use 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#[derive(Debug, Clone, Eq, PartialEq, Hash)]
39pub struct ListingTableUrl {
40 url: Url,
42 prefix: Path,
44 glob: Option<Pattern>,
46 table_ref: Option<TableReference>,
48}
49
50impl ListingTableUrl {
51 pub fn parse(s: impl AsRef<str>) -> Result<Self> {
107 let s = s.as_ref();
108
109 #[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 #[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 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 pub fn scheme(&self) -> &str {
162 self.url.scheme()
163 }
164
165 pub fn prefix(&self) -> &Path {
170 &self.prefix
171 }
172
173 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 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 None if ignore_subdirectory => segments.count() <= 1,
198 None => true,
200 }
201 }
202
203 pub fn is_collection(&self) -> bool {
205 self.url.path().ends_with(DELIMITER)
206 }
207
208 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 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 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 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 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 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 pub fn as_str(&self) -> &str {
319 self.as_ref()
320 }
321
322 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 pub fn is_folder(&self) -> bool {
330 self.url.scheme() == "file" && self.is_collection()
331 }
332
333 pub fn get_url(&self) -> &Url {
335 &self.url
336 }
337
338 pub fn get_glob(&self) -> &Option<Pattern> {
340 &self.glob
341 }
342
343 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 pub fn with_table_ref(mut self, table_ref: TableReference) -> Self {
352 self.table_ref = Some(table_ref);
353 self
354 }
355
356 pub fn get_table_ref(&self) -> &Option<TableReference> {
358 &self.table_ref
359 }
360}
361
362async 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 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 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 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 let mut vec = store
418 .list(Some(table_base_path))
419 .try_collect::<Vec<ObjectMeta>>()
420 .await?;
421 vec.shrink_to_fit(); 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#[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 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 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#[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 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 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 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_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 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 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 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 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 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 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 #[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_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 let session_no_cache = MockSession::new();
865
866 let runtime_with_cache = RuntimeEnvBuilder::new()
868 .with_object_list_cache_limit(1024 * 1024) .build_arc()?;
870 let session_with_cache = MockSession::with_runtime_env(runtime_with_cache);
871
872 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 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 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 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 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 #[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_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 let runtime = RuntimeEnvBuilder::new()
969 .with_object_list_cache_limit(1024 * 1024) .build_arc()?;
971 let session = MockSession::with_runtime_env(runtime);
972
973 let url = ListingTableUrl::parse("/sales/")?;
974
975 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 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 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 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 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 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 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 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}