This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23402-21a3215b66a8620cb932899de419e6b43d1848e6 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7a4dcc4e2134ab297fcb47e061c6f2d59e8017bd Author: Kumar Ujjawal <[email protected]> AuthorDate: Thu Sep 24 12:59:26 2026 +0000 fix: avoid full listings for cached pruned partitions (#23402) ## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #23341. ## Rationale for this change The list files cache currently performs a full table listing on a cold cache miss, even when partition pruning has already narrowed the scan to a specific path prefix. This makes selective partition queries pay the cost of listing the full table unless users disable the cache. <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> ## What changes are included in this PR? - List only the requested prefix on cold list-cache misses. - Keep using full-table cache entries to serve later prefix queries. - Ensure prefix-scoped cache entries do not satisfy full-table lookups. - Invalidate matching prefix-scoped list-cache entries on table writes. - Add regression tests that check the actual object store list prefix. <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? Yes <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? Partition-pruned listing table queries with the list files cache enabled avoid full table listings on cold cache misses. No API Change <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: Adam Gutglick <[email protected]> --- datafusion/catalog-listing/src/table.rs | 22 ++- .../core/tests/datasource/object_store_access.rs | 99 +++++++++++ datafusion/datasource/src/url.rs | 193 ++++++++++++++++++--- 3 files changed, 282 insertions(+), 32 deletions(-) diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index b04c267d1e..71802fa253 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -823,11 +823,23 @@ impl ListingTable { // Invalidate cache entries for this table if they exist if let Some(lfc) = state.runtime_env().cache_manager.get_list_files_cache() { - let key = TableScopedPath { - table: table_path.get_table_ref().clone(), - path: table_path.prefix().clone(), - }; - let _ = lfc.remove(&key); + if let Some(table_ref) = table_path.get_table_ref() { + lfc.drop_table_entries(table_ref)?; + } else { + let table_prefix = table_path.prefix(); + let keys: Vec<_> = lfc + .list_entries() + .into_keys() + .filter(|key| { + key.table.is_none() + && (key.path.prefix_matches(table_prefix) + || table_prefix.prefix_matches(&key.path)) + }) + .collect(); + for key in keys { + let _ = lfc.remove(&key); + } + } } // Sink related option, apart from format diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 16d894bde1..c494f0ee8e 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -25,15 +25,18 @@ //! [`ListingTable`]: datafusion::datasource::listing::ListingTable use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; use async_trait::async_trait; use bytes::Bytes; use datafusion::prelude::{ CsvReadOptions, JsonReadOptions, ParquetReadOptions, SessionContext, }; use datafusion_catalog_listing::{ListingOptions, ListingTable, ListingTableConfig}; +use datafusion_common::assert_batches_eq; use datafusion_datasource::ListingTableUrl; use datafusion_datasource_csv::CsvFormat; use datafusion_datasource_json::JsonFormat; +use datafusion_execution::cache::TableScopedPath; use futures::stream::BoxStream; use insta::assert_snapshot; use object_store::memory::InMemory; @@ -204,6 +207,102 @@ async fn multi_query_multi_file_csv_file() { ); } +#[tokio::test] +async fn insert_invalidates_overlapping_unscoped_listings() { + let store = Arc::new(InMemory::new()); + for (path, data) in [ + ("table/region=US/q1/data.csv", "1\n"), + ("table/region=EU/data.csv", "2\n"), + ] { + store.put(&Path::from(path), data.into()).await.unwrap(); + } + let ctx = SessionContext::new(); + ctx.runtime_env() + .register_object_store(&Url::parse("mem://").unwrap(), store); + ctx.sql("SET datafusion.execution.listing_table_ignore_subdirectory = false") + .await + .unwrap() + .collect() + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let cache = ctx + .runtime_env() + .cache_manager + .get_list_files_cache() + .unwrap(); + let tables = [ + ("root_table", "table"), + ("child_table", "table/region=US"), + ("descendant_table", "table/region=US/q1"), + ("sibling_table", "table/region=EU"), + ]; + + for (name, path) in tables { + // Keep the URLs unscoped so overlapping tables share path-based cache entries. + let url = ListingTableUrl::parse(format!("mem:///{path}/")).unwrap(); + let options = + ListingOptions::new(Arc::new(CsvFormat::default().with_has_header(false))) + .with_file_extension(".csv"); + let config = ListingTableConfig::new(url) + .with_listing_options(options) + .with_schema(Arc::clone(&schema)); + ctx.register_table(name, Arc::new(ListingTable::try_new(config).unwrap())) + .unwrap(); + ctx.sql(&format!("SELECT * FROM {name}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + let key = TableScopedPath { + table: None, + path: Path::from(path), + }; + assert!(cache.get(&key).is_some()); + } + + ctx.sql("INSERT INTO child_table VALUES (3)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + for (name, path) in tables { + let key = TableScopedPath { + table: None, + path: Path::from(path), + }; + let cached = cache.get(&key); + assert_eq!(cached.is_some(), name == "sibling_table", "{name}"); + } + + let batches = ctx + .sql("SELECT * FROM root_table ORDER BY value") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_batches_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 1 |", + "| 2 |", + "| 3 |", + "+-------+" + ], + &batches + ); +} + #[tokio::test] async fn query_multi_csv_file() { let test = Test::new().with_multi_file_csv().await; diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index cfb6608ca0..50d277803e 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -368,14 +368,8 @@ impl ListingTableUrl { /// * `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. +/// A full table listing can satisfy a prefix-filtered request by filtering the +/// cached files. On cache miss, only the requested path is listed and cached. async fn list_with_cache<'b>( ctx: &'b dyn Session, store: &'b dyn ObjectStore, @@ -399,29 +393,32 @@ async fn list_with_cache<'b>( .map(|res| res.map_err(|e| DataFusionError::ObjectStore(Box::new(e)))) .boxed()), Some(cache) => { - // Build the filter prefix (only Some if prefix was requested) let filter_prefix = prefix.is_some().then(|| full_prefix.clone()); let table_scoped_base_path = TableScopedPath { table: table_ref.cloned(), path: table_base_path.clone(), }; + let table_scoped_list_path = TableScopedPath { + table: table_ref.cloned(), + path: full_prefix.clone(), + }; - // Try cache lookup - get returns CachedFileList let vec = if let Some(cached) = cache.get(&table_scoped_base_path) { debug!("Hit list files cache"); cached.files_matching_prefix(&filter_prefix) + } else if let Some(cached) = cache.get(&table_scoped_list_path) { + debug!("Hit list files cache for requested path"); + cached.files_matching_prefix(&None) } else { - // Cache miss - always list and cache the full table - // This ensures we have complete data for future prefix queries let mut vec = store - .list(Some(table_base_path)) + .list(Some(&full_prefix)) .try_collect::<Vec<ObjectMeta>>() .await?; vec.shrink_to_fit(); // Right-size before caching let cached: CachedFileList = vec.into(); - let result = cached.files_matching_prefix(&filter_prefix); - cache.put(&table_scoped_base_path, cached); + let result = cached.files_matching_prefix(&None); + cache.put(&table_scoped_list_path, cached); result }; Ok( @@ -531,6 +528,7 @@ mod tests { use std::any::Any; use std::collections::HashMap; use std::ops::Range; + use std::sync::Mutex; use tempfile::tempdir; #[test] @@ -747,10 +745,8 @@ mod tests { #[tokio::test] async fn test_list_files() -> Result<()> { - let store = MockObjectStore { - in_mem: object_store::memory::InMemory::new(), - forbidden_paths: vec!["forbidden/e.parquet".into()], - }; + let store = + MockObjectStore::with_forbidden_paths(vec!["forbidden/e.parquet".into()]); // Create some files: create_file(&store, "a.parquet").await; @@ -848,10 +844,7 @@ mod tests { 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![], - }; + let store = MockObjectStore::new(); // Create test files with partition-style paths create_file(&store, "/table/year=2023/data1.parquet").await; @@ -949,15 +942,138 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_cache_miss_with_prefix_lists_prefixed_path() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let store = MockObjectStore::new(); + create_file(&store, "/table/year=2023/data1.parquet").await; + create_file(&store, "/table/year=2024/month=06/data2.parquet").await; + create_file(&store, "/table/year=2024/month=12/data3.parquet").await; + + let runtime = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(1024 * 1024) + .build_arc()?; + let session = MockSession::with_runtime_env(runtime); + let url = ListingTableUrl::parse("/table/")?; + let prefix = Path::from("year=2024/month=06"); + + let results: Vec<String> = url + .list_prefixed_files(&session, &store, Some(prefix.clone()), "parquet") + .await? + .try_collect::<Vec<_>>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + + assert_eq!(results, vec!["table/year=2024/month=06/data2.parquet"]); + assert_eq!( + store.list_prefixes(), + vec![Some(Path::from("table/year=2024/month=06"))] + ); + + let results: Vec<String> = url + .list_prefixed_files(&session, &store, Some(prefix), "parquet") + .await? + .try_collect::<Vec<_>>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + + assert_eq!(results, vec!["table/year=2024/month=06/data2.parquet"]); + assert_eq!( + store.list_prefixes(), + vec![Some(Path::from("table/year=2024/month=06"))] + ); + + Ok(()) + } + + #[tokio::test] + async fn test_prefix_cache_does_not_satisfy_broader_listing() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let store = MockObjectStore::new(); + create_file(&store, "/sales/region=US/q1/data.parquet").await; + create_file(&store, "/sales/region=US/q2/data.parquet").await; + create_file(&store, "/sales/region=EU/q1.parquet").await; + + let runtime = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(1024 * 1024) + .build_arc()?; + let mut session = MockSession::with_runtime_env(runtime); + session + .config + .options_mut() + .execution + .listing_table_ignore_subdirectory = false; + let url = ListingTableUrl::parse("/sales/")?; + + let q1_results: Vec<String> = url + .list_prefixed_files( + &session, + &store, + Some(Path::from("region=US/q1")), + "parquet", + ) + .await? + .try_collect::<Vec<_>>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + assert_eq!(q1_results, vec!["sales/region=US/q1/data.parquet"]); + + let us_results: Vec<String> = url + .list_prefixed_files( + &session, + &store, + Some(Path::from("region=US")), + "parquet", + ) + .await? + .try_collect::<Vec<_>>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + assert_eq!( + us_results, + vec![ + "sales/region=US/q1/data.parquet", + "sales/region=US/q2/data.parquet" + ] + ); + + let full_results: Vec<String> = url + .list_prefixed_files(&session, &store, None, "parquet") + .await? + .try_collect::<Vec<_>>() + .await? + .into_iter() + .map(|m| m.location.to_string()) + .collect(); + assert_eq!(full_results.len(), 3); + assert_eq!( + store.list_prefixes(), + vec![ + Some(Path::from("sales/region=US/q1")), + Some(Path::from("sales/region=US")), + Some(Path::from("sales")) + ] + ); + + 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![], - }; + let store = MockObjectStore::new(); // Create test files create_file(&store, "/sales/region=US/q1.parquet").await; @@ -982,6 +1098,7 @@ mod tests { .map(|m| m.location.to_string()) .collect(); assert_eq!(full_results.len(), 3); + assert_eq!(store.list_prefixes(), vec![Some(Path::from("sales"))]); // Second: query with prefix (should be served from cache) let mut us_results: Vec<String> = url @@ -1003,6 +1120,7 @@ mod tests { us_results, vec!["sales/region=US/q1.parquet", "sales/region=US/q2.parquet"] ); + assert_eq!(store.list_prefixes(), vec![Some(Path::from("sales"))]); // Third: different prefix (also from cache) let eu_results: Vec<String> = url @@ -1020,6 +1138,7 @@ mod tests { .collect(); assert_eq!(eu_results, vec!["sales/region=EU/q1.parquet"]); + assert_eq!(store.list_prefixes(), vec![Some(Path::from("sales"))]); Ok(()) } @@ -1079,6 +1198,25 @@ mod tests { struct MockObjectStore { in_mem: object_store::memory::InMemory, forbidden_paths: Vec<Path>, + list_prefixes: Arc<Mutex<Vec<Option<Path>>>>, + } + + impl MockObjectStore { + fn new() -> Self { + Self::with_forbidden_paths(vec![]) + } + + fn with_forbidden_paths(forbidden_paths: Vec<Path>) -> Self { + Self { + in_mem: object_store::memory::InMemory::new(), + forbidden_paths, + list_prefixes: Default::default(), + } + } + + fn list_prefixes(&self) -> Vec<Option<Path>> { + self.list_prefixes.lock().unwrap().clone() + } } impl std::fmt::Display for MockObjectStore { @@ -1140,6 +1278,7 @@ mod tests { &self, prefix: Option<&Path>, ) -> BoxStream<'static, object_store::Result<ObjectMeta>> { + self.list_prefixes.lock().unwrap().push(prefix.cloned()); self.in_mem.list(prefix) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
