diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 4641df671e448..155d6efe462c1 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -245,25 +245,26 @@ impl ListingTableUrl { let exec_options = &ctx.config_options().execution; let ignore_subdirectory = exec_options.listing_table_ignore_subdirectory; - let prefix = if let Some(prefix) = prefix { - let mut p = self.prefix.parts().collect::>(); - p.extend(prefix.parts()); - Path::from_iter(p.into_iter()) + // Build full_prefix for non-cached path and head() calls + let full_prefix = if let Some(ref p) = prefix { + let mut parts = self.prefix.parts().collect::>(); + parts.extend(p.parts()); + Path::from_iter(parts.into_iter()) } else { self.prefix.clone() }; let list: BoxStream<'a, Result> = if self.is_collection() { - list_with_cache(ctx, store, &prefix).await? + list_with_cache(ctx, store, &self.prefix, prefix.as_ref()).await? } else { - match store.head(&prefix).await { + match store.head(&full_prefix).await { Ok(meta) => futures::stream::once(async { Ok(meta) }) .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) .boxed(), // If the head command fails, it is likely that object doesn't exist. // Retry as though it were a prefix (aka a collection) Err(object_store::Error::NotFound { .. }) => { - list_with_cache(ctx, store, &prefix).await? + list_with_cache(ctx, store, &self.prefix, prefix.as_ref()).await? } Err(e) => return Err(e.into()), } @@ -324,27 +325,74 @@ impl ListingTableUrl { } } +/// Lists files with cache support, using prefix-aware lookups. +/// +/// # Arguments +/// * `ctx` - The session context +/// * `store` - The object store to list from +/// * `table_base_path` - The table's base path (the stable cache key) +/// * `prefix` - Optional prefix relative to table base for filtering results +/// +/// # Cache Behavior: +/// The cache key is always `table_base_path`. When a prefix-filtered listing +/// is requested via `prefix`, the cache: +/// - Looks up `table_base_path` in the cache +/// - Filters results to match `table_base_path/prefix` +/// - Returns filtered results without a storage call +/// +/// On cache miss, the full table is always listed and cached, ensuring +/// subsequent prefix queries can be served from cache. async fn list_with_cache<'b>( ctx: &'b dyn Session, store: &'b dyn ObjectStore, - prefix: &Path, + table_base_path: &Path, + prefix: Option<&Path>, ) -> Result>> { + // Build the full listing path (table_base + prefix) + let full_prefix = match prefix { + Some(p) => { + let mut parts: Vec<_> = table_base_path.parts().collect(); + parts.extend(p.parts()); + Path::from_iter(parts) + } + None => table_base_path.clone(), + }; + match ctx.runtime_env().cache_manager.get_list_files_cache() { None => Ok(store - .list(Some(prefix)) + .list(Some(&full_prefix)) .map(|res| res.map_err(|e| DataFusionError::ObjectStore(Box::new(e)))) .boxed()), Some(cache) => { - let vec = if let Some(res) = cache.get(prefix) { - debug!("Hit list all files cache"); + // Convert prefix to Option for cache lookup + let prefix_filter = prefix.cloned(); + + // Try cache lookup with optional prefix filter + let vec = if let Some(res) = + cache.get_with_extra(table_base_path, &prefix_filter) + { + debug!("Hit list files cache"); res.as_ref().clone() } else { + // Cache miss - always list and cache the full table + // This ensures we have complete data for future prefix queries let vec = store - .list(Some(prefix)) + .list(Some(table_base_path)) .try_collect::>() .await?; - cache.put(prefix, Arc::new(vec.clone())); - vec + cache.put(table_base_path, Arc::new(vec.clone())); + + // If a prefix filter was requested, apply it to the results + if prefix.is_some() { + let full_prefix_str = full_prefix.as_ref(); + vec.into_iter() + .filter(|meta| { + meta.location.as_ref().starts_with(full_prefix_str) + }) + .collect() + } else { + vec + } }; Ok(futures::stream::iter(vec.into_iter().map(Ok)).boxed()) } @@ -754,6 +802,191 @@ mod tests { Ok(()) } + /// Tests that the cached code path produces identical results to the non-cached path. + /// + /// This is critical: the cache is a transparent optimization, so both paths + /// MUST return the same files. Note: order is not guaranteed by ObjectStore::list, + /// so we sort results before comparison. + #[tokio::test] + async fn test_cache_path_equivalence() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let store = MockObjectStore { + in_mem: object_store::memory::InMemory::new(), + forbidden_paths: vec![], + }; + + // Create test files with partition-style paths + create_file(&store, "/table/year=2023/data1.parquet").await; + create_file(&store, "/table/year=2023/month=01/data2.parquet").await; + create_file(&store, "/table/year=2024/data3.parquet").await; + create_file(&store, "/table/year=2024/month=06/data4.parquet").await; + create_file(&store, "/table/year=2024/month=12/data5.parquet").await; + + // Session WITHOUT cache + let session_no_cache = MockSession::new(); + + // Session WITH cache - use RuntimeEnvBuilder with cache limit (no TTL needed for this test) + let runtime_with_cache = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(1024 * 1024) // 1MB limit + .build_arc()?; + let session_with_cache = MockSession::with_runtime_env(runtime_with_cache); + + // Test cases: (url, prefix, description) + let test_cases = vec![ + ("/table/", None, "full table listing"), + ( + "/table/", + Some(Path::from("year=2023")), + "single partition filter", + ), + ( + "/table/", + Some(Path::from("year=2024")), + "different partition filter", + ), + ( + "/table/", + Some(Path::from("year=2024/month=06")), + "nested partition filter", + ), + ( + "/table/", + Some(Path::from("year=2025")), + "non-existent partition", + ), + ]; + + for (url_str, prefix, description) in test_cases { + let url = ListingTableUrl::parse(url_str)?; + + // Get results WITHOUT cache (sorted for comparison) + let mut results_no_cache: Vec = url + .list_prefixed_files(&session_no_cache, &store, prefix.clone(), "parquet") + .await? + .try_collect::>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + results_no_cache.sort(); + + // Get results WITH cache (first call - cache miss, sorted for comparison) + let mut results_with_cache_miss: Vec = url + .list_prefixed_files( + &session_with_cache, + &store, + prefix.clone(), + "parquet", + ) + .await? + .try_collect::>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + results_with_cache_miss.sort(); + + // Get results WITH cache (second call - cache hit, sorted for comparison) + let mut results_with_cache_hit: Vec = url + .list_prefixed_files(&session_with_cache, &store, prefix, "parquet") + .await? + .try_collect::>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + results_with_cache_hit.sort(); + + // All three should contain the same files + assert_eq!( + results_no_cache, results_with_cache_miss, + "Cache miss path should match non-cached path for: {description}" + ); + assert_eq!( + results_no_cache, results_with_cache_hit, + "Cache hit path should match non-cached path for: {description}" + ); + } + + Ok(()) + } + + /// Tests that prefix queries can be served from a cached full-table listing + #[tokio::test] + async fn test_cache_serves_partition_from_full_listing() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let store = MockObjectStore { + in_mem: object_store::memory::InMemory::new(), + forbidden_paths: vec![], + }; + + // Create test files + create_file(&store, "/sales/region=US/q1.parquet").await; + create_file(&store, "/sales/region=US/q2.parquet").await; + create_file(&store, "/sales/region=EU/q1.parquet").await; + + // Create session with cache (no TTL needed for this test) + let runtime = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(1024 * 1024) // 1MB limit + .build_arc()?; + let session = MockSession::with_runtime_env(runtime); + + let url = ListingTableUrl::parse("/sales/")?; + + // First: query full table (populates cache) + let full_results: Vec = url + .list_prefixed_files(&session, &store, None, "parquet") + .await? + .try_collect::>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + assert_eq!(full_results.len(), 3); + + // Second: query with prefix (should be served from cache) + let mut us_results: Vec = url + .list_prefixed_files( + &session, + &store, + Some(Path::from("region=US")), + "parquet", + ) + .await? + .try_collect::>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + us_results.sort(); + + assert_eq!( + us_results, + vec!["sales/region=US/q1.parquet", "sales/region=US/q2.parquet"] + ); + + // Third: different prefix (also from cache) + let eu_results: Vec = url + .list_prefixed_files( + &session, + &store, + Some(Path::from("region=EU")), + "parquet", + ) + .await? + .try_collect::>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + + assert_eq!(eu_results, vec!["sales/region=EU/q1.parquet"]); + + Ok(()) + } + /// Creates a file with "hello world" content at the specified path async fn create_file(object_store: &dyn ObjectStore, path: &str) { object_store @@ -906,6 +1139,14 @@ mod tests { runtime_env: Arc::new(RuntimeEnv::default()), } } + + /// Create a MockSession with a custom RuntimeEnv (for cache testing) + fn with_runtime_env(runtime_env: Arc) -> Self { + Self { + config: SessionConfig::new(), + runtime_env, + } + } } #[async_trait::async_trait] diff --git a/datafusion/execution/src/cache/cache_manager.rs b/datafusion/execution/src/cache/cache_manager.rs index 2df5ef1b4458c..89b48e9629acb 100644 --- a/datafusion/execution/src/cache/cache_manager.rs +++ b/datafusion/execution/src/cache/cache_manager.rs @@ -66,9 +66,21 @@ pub struct FileStatisticsCacheEntry { /// command on the local filesystem. This operation can be expensive, /// especially when done over remote object stores. /// +/// The cache key is always the table's base path, ensuring a stable cache key. +/// The `Extra` type is `Option`, representing an optional prefix filter +/// (relative to the table base path) for partition-aware lookups. +/// +/// When `get_with_extra(key, Some(prefix))` is called: +/// - The cache entry for `key` (table base path) is fetched +/// - Results are filtered to only include files matching `key/prefix` +/// - Filtered results are returned without making a storage call +/// +/// This enables efficient partition pruning: a single cached listing of the +/// full table can serve queries for any partition subset. +/// /// See [`crate::runtime_env::RuntimeEnv`] for more details. pub trait ListFilesCache: - CacheAccessor>, Extra = ObjectMeta> + CacheAccessor>, Extra = Option> { /// Returns the cache's memory limit in bytes. fn cache_limit(&self) -> usize; diff --git a/datafusion/execution/src/cache/list_files_cache.rs b/datafusion/execution/src/cache/list_files_cache.rs index 285abea2e66ae..661bc47b5468a 100644 --- a/datafusion/execution/src/cache/list_files_cache.rs +++ b/datafusion/execution/src/cache/list_files_cache.rs @@ -175,20 +175,64 @@ impl DefaultListFilesCacheState { } } - /// Returns the respective entry from the cache, if it exists and the entry - /// has not expired by `now`. + /// Performs a prefix-aware cache lookup. /// - /// If the entry exists it becomes the most recently used. If the entry has expired it is - /// removed from the cache - fn get(&mut self, key: &Path, now: Instant) -> Option>> { - let entry = self.lru_queue.get(key)?; + /// # Arguments + /// * `table_base` - The table's base path (the cache key) + /// * `prefix` - Optional prefix filter relative to the table base path + /// * `now` - Current time for expiration checking + /// + /// # Behavior + /// - Fetches the cache entry for `table_base` + /// - If `prefix` is `Some`, filters results to only files matching `table_base/prefix` + /// - Returns the (potentially filtered) results + /// + /// # Example + /// ```text + /// get_with_prefix("my_table", Some("a=1"), now) + /// → Fetch cache entry for "my_table" + /// → Filter to files matching "my_table/a=1/*" + /// → Return filtered results + /// ``` + fn get_with_prefix( + &mut self, + table_base: &Path, + prefix: Option<&Path>, + now: Instant, + ) -> Option>> { + let entry = self.lru_queue.get(table_base)?; - match entry.expires { - Some(exp) if now > exp => { - self.remove(key); - None - } - _ => Some(Arc::clone(&entry.metas)), + // Check expiration + if let Some(exp) = entry.expires + && now > exp + { + self.remove(table_base); + return None; + } + + // Early return if no prefix filter - return all files + let Some(prefix) = prefix else { + return Some(Arc::clone(&entry.metas)); + }; + + // Build the full prefix path: table_base/prefix + let mut parts: Vec<_> = table_base.parts().collect(); + parts.extend(prefix.parts()); + let full_prefix = Path::from_iter(parts); + let full_prefix_str = full_prefix.as_ref(); + + // Filter files to only those matching the prefix + let filtered: Vec = entry + .metas + .iter() + .filter(|meta| meta.location.as_ref().starts_with(full_prefix_str)) + .cloned() + .collect(); + + if filtered.is_empty() { + None + } else { + Some(Arc::new(filtered)) } } @@ -306,16 +350,36 @@ impl ListFilesCache for DefaultListFilesCache { } impl CacheAccessor>> for DefaultListFilesCache { - type Extra = ObjectMeta; + type Extra = Option; + /// Gets all files for the given table base path. + /// + /// This is equivalent to calling `get_with_extra(k, &None)`. fn get(&self, k: &Path) -> Option>> { - let mut state = self.state.lock().unwrap(); - let now = self.time_provider.now(); - state.get(k, now) + self.get_with_extra(k, &None) } - fn get_with_extra(&self, k: &Path, _e: &Self::Extra) -> Option>> { - self.get(k) + /// Performs a prefix-aware cache lookup. + /// + /// # Arguments + /// * `table_base` - The table's base path (the cache key) + /// * `prefix` - Optional prefix filter (relative to table base) for partition filtering + /// + /// # Behavior + /// - Fetches the cache entry for `table_base` + /// - If `prefix` is `Some`, filters results to only files matching `table_base/prefix` + /// - Returns the (potentially filtered) results + /// + /// This enables efficient partition pruning - a single cached listing of the full table + /// can serve queries for any partition subset without additional storage calls. + fn get_with_extra( + &self, + table_base: &Path, + prefix: &Self::Extra, + ) -> Option>> { + let mut state = self.state.lock().unwrap(); + let now = self.time_provider.now(); + state.get_with_prefix(table_base, prefix.as_ref(), now) } fn put( @@ -807,4 +871,202 @@ mod tests { assert_eq!(state.memory_used, 0); } } + + // Prefix-aware cache tests + + /// Helper function to create ObjectMeta with a specific location path + fn create_object_meta_with_path(location: &str) -> ObjectMeta { + ObjectMeta { + location: Path::from(location), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size: 1024, + e_tag: None, + version: None, + } + } + + #[test] + fn test_prefix_aware_cache_hit() { + // Scenario: Cache has full table listing, query for partition returns filtered results + let cache = DefaultListFilesCache::new(100000, None); + + // Create files for a partitioned table + let table_base = Path::from("my_table"); + let files = Arc::new(vec![ + create_object_meta_with_path("my_table/a=1/file1.parquet"), + create_object_meta_with_path("my_table/a=1/file2.parquet"), + create_object_meta_with_path("my_table/a=2/file3.parquet"), + create_object_meta_with_path("my_table/a=2/file4.parquet"), + ]); + + // Cache the full table listing + cache.put(&table_base, files); + + // Query for partition a=1 using get_with_extra + // New API: get_with_extra(table_base, Some(relative_prefix)) + let prefix_a1 = Some(Path::from("a=1")); + let result = cache.get_with_extra(&table_base, &prefix_a1); + + // Should return filtered results (only files from a=1) + assert!(result.is_some()); + let filtered = result.unwrap(); + assert_eq!(filtered.len(), 2); + assert!( + filtered + .iter() + .all(|m| m.location.as_ref().starts_with("my_table/a=1")) + ); + + // Query for partition a=2 + let prefix_a2 = Some(Path::from("a=2")); + let result_2 = cache.get_with_extra(&table_base, &prefix_a2); + + assert!(result_2.is_some()); + let filtered_2 = result_2.unwrap(); + assert_eq!(filtered_2.len(), 2); + assert!( + filtered_2 + .iter() + .all(|m| m.location.as_ref().starts_with("my_table/a=2")) + ); + } + + #[test] + fn test_prefix_aware_cache_no_filter_returns_all() { + // Scenario: Query with no prefix filter should return all files + let cache = DefaultListFilesCache::new(100000, None); + + let table_base = Path::from("my_table"); + + // Cache full table listing with 4 files + let full_files = Arc::new(vec![ + create_object_meta_with_path("my_table/a=1/file1.parquet"), + create_object_meta_with_path("my_table/a=1/file2.parquet"), + create_object_meta_with_path("my_table/a=2/file3.parquet"), + create_object_meta_with_path("my_table/a=2/file4.parquet"), + ]); + cache.put(&table_base, full_files); + + // Query with no prefix filter (None) should return all 4 files + let result = cache.get_with_extra(&table_base, &None); + assert!(result.is_some()); + let files = result.unwrap(); + assert_eq!(files.len(), 4); + + // Also test using get() which delegates to get_with_extra(&None) + let result_get = cache.get(&table_base); + assert!(result_get.is_some()); + assert_eq!(result_get.unwrap().len(), 4); + } + + #[test] + fn test_prefix_aware_cache_miss_no_entry() { + // Scenario: Table not cached, query should miss + let cache = DefaultListFilesCache::new(100000, None); + + let table_base = Path::from("my_table"); + + // Query for full table should miss (nothing cached) + let result = cache.get_with_extra(&table_base, &None); + assert!(result.is_none()); + + // Query with prefix should also miss + let prefix = Some(Path::from("a=1")); + let result_2 = cache.get_with_extra(&table_base, &prefix); + assert!(result_2.is_none()); + } + + #[test] + fn test_prefix_aware_cache_no_matching_files() { + // Scenario: Cache has table listing but no files match the requested partition + let cache = DefaultListFilesCache::new(100000, None); + + let table_base = Path::from("my_table"); + let files = Arc::new(vec![ + create_object_meta_with_path("my_table/a=1/file1.parquet"), + create_object_meta_with_path("my_table/a=2/file2.parquet"), + ]); + cache.put(&table_base, files); + + // Query for partition a=3 which doesn't exist + let prefix_a3 = Some(Path::from("a=3")); + let result = cache.get_with_extra(&table_base, &prefix_a3); + + // Should return None since no files match + assert!(result.is_none()); + } + + #[test] + fn test_prefix_aware_nested_partitions() { + // Scenario: Table with multiple partition levels (e.g., year/month/day) + let cache = DefaultListFilesCache::new(100000, None); + + let table_base = Path::from("events"); + let files = Arc::new(vec![ + create_object_meta_with_path( + "events/year=2024/month=01/day=01/file1.parquet", + ), + create_object_meta_with_path( + "events/year=2024/month=01/day=02/file2.parquet", + ), + create_object_meta_with_path( + "events/year=2024/month=02/day=01/file3.parquet", + ), + create_object_meta_with_path( + "events/year=2025/month=01/day=01/file4.parquet", + ), + ]); + cache.put(&table_base, files); + + // Query for year=2024/month=01 (should get 2 files) + let prefix_month = Some(Path::from("year=2024/month=01")); + let result = cache.get_with_extra(&table_base, &prefix_month); + assert!(result.is_some()); + assert_eq!(result.unwrap().len(), 2); + + // Query for year=2024 (should get 3 files) + let prefix_year = Some(Path::from("year=2024")); + let result_year = cache.get_with_extra(&table_base, &prefix_year); + assert!(result_year.is_some()); + assert_eq!(result_year.unwrap().len(), 3); + + // Query for specific day (should get 1 file) + let prefix_day = Some(Path::from("year=2024/month=01/day=01")); + let result_day = cache.get_with_extra(&table_base, &prefix_day); + assert!(result_day.is_some()); + assert_eq!(result_day.unwrap().len(), 1); + } + + #[test] + fn test_prefix_aware_different_tables() { + // Scenario: Multiple tables cached, queries should not cross-contaminate + let cache = DefaultListFilesCache::new(100000, None); + + let table_a = Path::from("table_a"); + let table_b = Path::from("table_b"); + + let files_a = Arc::new(vec![create_object_meta_with_path( + "table_a/part=1/file1.parquet", + )]); + let files_b = Arc::new(vec![ + create_object_meta_with_path("table_b/part=1/file1.parquet"), + create_object_meta_with_path("table_b/part=2/file2.parquet"), + ]); + + cache.put(&table_a, files_a); + cache.put(&table_b, files_b); + + // Query table_a should only return table_a files + let result_a = cache.get(&table_a); + assert!(result_a.is_some()); + assert_eq!(result_a.unwrap().len(), 1); + + // Query table_b with prefix should only return matching table_b files + let prefix = Some(Path::from("part=1")); + let result_b = cache.get_with_extra(&table_b, &prefix); + assert!(result_b.is_some()); + assert_eq!(result_b.unwrap().len(), 1); + } }