From 6194d58e987da1e0babb7505a3c0f3026c793faa Mon Sep 17 00:00:00 2001 From: Shehab Ali Date: Fri, 6 Nov 2026 13:33:07 -0500 Subject: [PATCH] perf: bound file-listing cache prefetch by its memory budget --- benchmarks/Cargo.toml | 4 + benchmarks/benches/list_files_cache.rs | 174 +++++++++++++++ datafusion/datasource/src/url.rs | 280 ++++++++++++++++++++++++- 3 files changed, 454 insertions(+), 4 deletions(-) create mode 100644 benchmarks/benches/list_files_cache.rs diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 2ddcc48a322f5..0e00508996401 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -63,3 +63,7 @@ name = "sql" [[bench]] harness = false name = "parquet_pruning_setup_cache" + +[[bench]] +harness = false +name = "list_files_cache" diff --git a/benchmarks/benches/list_files_cache.rs b/benchmarks/benches/list_files_cache.rs new file mode 100644 index 0000000000000..bfafafd1235c6 --- /dev/null +++ b/benchmarks/benches/list_files_cache.rs @@ -0,0 +1,174 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Benchmarks listing-cache misses through the public `ListingTableUrl` API. +//! +//! A lazy object store generates fresh metadata without retaining the listing. +//! There is no network or file-data I/O. Consumers count results without collecting +//! them, so large allocations come from listing/cache behavior, not the harness. +//! Cases vary whether a complete listing fits the 1 MiB cache, whether the consumer +//! stops after ten matches, and whether only one percent of extensions match. +//! +//! Run with `cargo bench -p datafusion-benchmarks --bench list_files_cache`. + +use std::fmt::{Display, Formatter}; + +use async_trait::async_trait; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::datasource::listing::ListingTableUrl; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::prelude::{SessionConfig, SessionContext}; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt, future, stream}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, +}; + +const CACHE_LIMIT: usize = 1024 * 1024; + +/// Generates metadata on demand; even a million-file listing has constant store +/// memory. Every `list` call generates its own paths rather than cloning a fixture. +#[derive(Debug)] +struct GeneratedObjectStore { + files: usize, + parquet_every: usize, +} + +impl Display for GeneratedObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "GeneratedObjectStore") + } +} + +fn unsupported(operation: &str) -> object_store::Error { + object_store::Error::NotImplemented { + operation: operation.to_owned(), + implementer: "GeneratedObjectStore".to_owned(), + } +} + +#[async_trait] +impl ObjectStore for GeneratedObjectStore { + async fn put_opts( + &self, + _location: &Path, + _payload: PutPayload, + _opts: PutOptions, + ) -> object_store::Result { + Err(unsupported("put_opts")) + } + + async fn put_multipart_opts( + &self, + _location: &Path, + _opts: PutMultipartOptions, + ) -> object_store::Result> { + Err(unsupported("put_multipart_opts")) + } + + async fn get_opts( + &self, + _location: &Path, + _opts: GetOptions, + ) -> object_store::Result { + Err(unsupported("get_opts")) + } + + fn delete_stream( + &self, + _locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + stream::once(future::ready(Err(unsupported("delete_stream")))).boxed() + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'static, object_store::Result> { + assert_eq!(prefix.map(Path::as_ref), Some("table")); + let parquet_every = self.parquet_every; + stream::iter(0..self.files) + .map(move |index| { + let extension = if index % parquet_every == 0 { + "parquet" + } else { + "csv" + }; + ObjectMeta { + location: Path::from(format!("table/part-{index:08}.{extension}")), + last_modified: Default::default(), + size: 1024, + e_tag: None, + version: None, + } + }) + .map(Ok) + .boxed() + } + + async fn list_with_delimiter( + &self, + _prefix: Option<&Path>, + ) -> object_store::Result { + Err(unsupported("list_with_delimiter")) + } + + async fn copy_opts( + &self, + _from: &Path, + _to: &Path, + _options: CopyOptions, + ) -> object_store::Result<()> { + Err(unsupported("copy_opts")) + } +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let runtime_env = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(CACHE_LIMIT) + .build_arc() + .unwrap(); + let cache = runtime_env.cache_manager.get_list_files_cache().unwrap(); + let context = SessionContext::new_with_config_rt(SessionConfig::new(), runtime_env); + let state = context.state(); + let url = ListingTableUrl::parse("test:///table/").unwrap(); + + let mut group = criterion.benchmark_group("list_files_cache_miss"); + for (name, files, take, parquet_every) in [ + ("cacheable/full", 1_000, usize::MAX, 1), + ("cacheable/first_10", 1_000, 10, 1), + ("oversized/full", 100_000, usize::MAX, 1), + ("oversized/first_10", 100_000, 10, 1), + ("oversized/selective_full", 100_000, usize::MAX, 100), + ("oversized/selective_first_10", 100_000, 10, 100), + ] { + let store = GeneratedObjectStore { + files, + parquet_every, + }; + let expected_files = files.div_ceil(parquet_every).min(take); + group.bench_function(name, |bencher| { + bencher.to_async(&runtime).iter(|| async { + // Include cache teardown, and ensure every iteration is a miss. + cache.clear(); + let count = url + .list_all_files(&state, &store, ".parquet") + .await + .unwrap() + .take(take) + .try_fold(0usize, |count, _| future::ready(Ok(count + 1))) + .await + .unwrap(); + assert_eq!(count, expected_files); + }); + }); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index bae80abac6498..a74d5faa91cb5 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -5,8 +5,9 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, TableReference}; -use datafusion_execution::cache::cache_manager::CachedFileList; +use datafusion_execution::cache::CacheKey; use datafusion_execution::cache::cache_manager::TableScopedPath; +use datafusion_execution::cache::cache_manager::{CachedFileList, meta_heap_bytes}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_session::Session; @@ -357,6 +358,8 @@ impl ListingTableUrl { /// # Cache Behavior: /// 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. +/// Once the listing exceeds the cache budget, the buffered prefix and remaining +/// objects are streamed without caching the incomplete listing. async fn list_with_cache<'b>( ctx: &'b dyn Session, store: &'b dyn ObjectStore, @@ -398,10 +401,32 @@ async fn list_with_cache<'b>( debug!("Hit list files cache for requested path"); cached.files_matching_prefix(&None) } else { - let mut vec = store + let mut list = store .list(Some(&full_prefix)) - .try_collect::>() - .await?; + .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) + .boxed(); + let limit = cache.cache_limit(); + let mut size = table_scoped_list_path.size(); + if size > limit { + return Ok(list); + } + let mut vec = Vec::new(); + while let Some(meta) = list.try_next().await? { + // Admission uses the final, right-sized vector, not spare capacity + // from growth. Count only bytes that shrink_to_fit cannot remove. + let next_size = size + .checked_add(size_of::()) + .and_then(|size| size.checked_add(meta_heap_bytes(&meta))); + vec.push(meta); + match next_size { + Some(next_size) if next_size <= limit => size = next_size, + _ => { + return Ok(futures::stream::iter(vec.into_iter().map(Ok)) + .chain(list) + .boxed()); + } + } + } vec.shrink_to_fit(); // Right-size before caching let cached: CachedFileList = vec.into(); let result = cached.files_matching_prefix(&None); @@ -498,6 +523,7 @@ mod tests { use datafusion_common::DFSchema; use datafusion_common::config::TableOptions; use datafusion_execution::TaskContext; + use datafusion_execution::cache::CacheValue; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::execution_props::ExecutionProps; @@ -516,6 +542,7 @@ mod tests { use std::collections::HashMap; use std::ops::Range; use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; use tempfile::tempdir; #[test] @@ -1130,6 +1157,189 @@ mod tests { Ok(()) } + fn listing_cache_key(url: &ListingTableUrl) -> TableScopedPath { + TableScopedPath { + table: url.get_table_ref().clone(), + path: url.prefix().clone(), + } + } + + fn generated_listing_weight( + url: &ListingTableUrl, + listing: &GeneratedListing, + count: usize, + ) -> usize { + let mut files = (0..count).map(|index| listing.object(index)).collect_vec(); + files.shrink_to_fit(); + listing_cache_key(url).size() + CachedFileList::new(files).size() + } + + fn session_with_listing_cache(limit: usize) -> Result { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + Ok(MockSession::with_runtime_env( + RuntimeEnvBuilder::new() + .with_object_list_cache_limit(limit) + .build_arc()?, + )) + } + + #[tokio::test] + async fn test_oversized_listing_only_prefetches_a_bounded_prefix() -> Result<()> { + let store = MockObjectStore::with_generated_listing(100, 128, None); + let listing = store.generated_listing.as_ref().unwrap(); + let url = ListingTableUrl::parse("/table/")?; + let session = + session_with_listing_cache(generated_listing_weight(&url, listing, 2))?; + let stream = url.list_all_files(&session, &store, "parquet").await?; + + // Two objects fit; the third proves this listing cannot be cached. + assert_eq!(listing.polled.load(Ordering::Relaxed), 3); + let files = stream.try_collect::>().await?; + assert_eq!( + files, + (0..100).map(|index| listing.object(index)).collect_vec() + ); + assert_eq!(listing.polled.load(Ordering::Relaxed), 100); + let cache = session + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap(); + assert!(cache.get(&listing_cache_key(&url)).is_none()); + Ok(()) + } + + #[tokio::test] + async fn test_dropping_oversized_listing_does_not_read_or_cache_its_tail() + -> Result<()> { + let store = MockObjectStore::with_generated_listing(100, 128, None); + let listing = store.generated_listing.as_ref().unwrap(); + let url = ListingTableUrl::parse("/table/")?; + let session = + session_with_listing_cache(generated_listing_weight(&url, listing, 1))?; + let mut stream = url.list_all_files(&session, &store, "parquet").await?; + assert_eq!(stream.try_next().await?, Some(listing.object(0))); + drop(stream); + + assert_eq!(listing.polled.load(Ordering::Relaxed), 2); + let cache = session + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap(); + assert!(cache.get(&listing_cache_key(&url)).is_none()); + Ok(()) + } + + #[tokio::test] + async fn test_listing_cache_admission_boundary_includes_key_and_metadata() + -> Result<()> { + let url = ListingTableUrl::parse("/table/")? + .with_table_ref(TableReference::bare("a_table_reference")); + for below_limit in [0, 1] { + let store = MockObjectStore::with_generated_listing(2, 128, None); + let listing = store.generated_listing.as_ref().unwrap(); + let limit = generated_listing_weight(&url, listing, 2) - below_limit; + let session = session_with_listing_cache(limit)?; + let cache = session + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap(); + let expected = (0..2).map(|index| listing.object(index)).collect_vec(); + + for _ in 0..2 { + let files = url + .list_all_files(&session, &store, "parquet") + .await? + .try_collect::>() + .await?; + assert_eq!(files, expected); + } + + assert_eq!( + cache.get(&listing_cache_key(&url)).is_some(), + below_limit == 0 + ); + assert_eq!( + listing.polled.load(Ordering::Relaxed), + if below_limit == 0 { 2 } else { 4 } + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_single_oversized_listing_object_is_returned_without_reading_tail() + -> Result<()> { + let store = MockObjectStore::with_generated_listing(100, 1024 * 1024, None); + let listing = store.generated_listing.as_ref().unwrap(); + let url = ListingTableUrl::parse("/table/")?; + let session = session_with_listing_cache(1024)?; + let mut stream = url.list_all_files(&session, &store, "parquet").await?; + + assert_eq!(listing.polled.load(Ordering::Relaxed), 1); + assert_eq!(stream.try_next().await?, Some(listing.object(0))); + drop(stream); + assert_eq!(listing.polled.load(Ordering::Relaxed), 1); + let cache = session + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap(); + assert!(cache.get(&listing_cache_key(&url)).is_none()); + Ok(()) + } + + #[tokio::test] + async fn test_listing_key_over_budget_skips_prefetch() -> Result<()> { + let store = MockObjectStore::with_generated_listing(3, 128, None); + let listing = store.generated_listing.as_ref().unwrap(); + let url = ListingTableUrl::parse("/table/")?; + let session = session_with_listing_cache(listing_cache_key(&url).size() - 1)?; + let stream = url.list_all_files(&session, &store, "parquet").await?; + + assert_eq!(listing.polled.load(Ordering::Relaxed), 0); + assert_eq!(stream.try_collect::>().await?.len(), 3); + let cache = session + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap(); + assert!(cache.get(&listing_cache_key(&url)).is_none()); + Ok(()) + } + + #[tokio::test] + async fn test_listing_errors_never_cache_a_partial_listing() -> Result<()> { + let url = ListingTableUrl::parse("/table/")?; + for fitting_objects in [1, 10] { + let store = MockObjectStore::with_generated_listing(10, 128, Some(2)); + let listing = store.generated_listing.as_ref().unwrap(); + let session = session_with_listing_cache(generated_listing_weight( + &url, + listing, + fitting_objects, + ))?; + let result = match url.list_all_files(&session, &store, "parquet").await { + Ok(stream) => stream.try_collect::>().await, + Err(error) => Err(error), + }; + + let error = result.unwrap_err(); + assert!(matches!(&error, DataFusionError::ObjectStore(_))); + assert!(error.to_string().contains("generated listing failure")); + let cache = session + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap(); + assert!(cache.get(&listing_cache_key(&url)).is_none()); + } + Ok(()) + } + /// Creates a file with "hello world" content at the specified path async fn create_file(object_store: &dyn ObjectStore, path: &str) { object_store @@ -1186,6 +1396,28 @@ mod tests { in_mem: object_store::memory::InMemory, forbidden_paths: Vec, list_prefixes: Arc>>>, + generated_listing: Option, + } + + /// Generate metadata only when polled, so the store cannot hide eager listing. + #[derive(Clone, Debug)] + struct GeneratedListing { + num_objects: usize, + metadata_bytes: usize, + failure_at: Option, + polled: Arc, + } + + impl GeneratedListing { + fn object(&self, index: usize) -> ObjectMeta { + ObjectMeta { + location: Path::from(format!("table/file-{index:08}.parquet")), + last_modified: chrono::DateTime::from_timestamp(0, 0).unwrap(), + size: 1, + e_tag: Some("e".repeat(self.metadata_bytes)), + version: Some("version".to_string()), + } + } } impl MockObjectStore { @@ -1198,6 +1430,23 @@ mod tests { in_mem: object_store::memory::InMemory::new(), forbidden_paths, list_prefixes: Default::default(), + generated_listing: None, + } + } + + fn with_generated_listing( + num_objects: usize, + metadata_bytes: usize, + failure_at: Option, + ) -> Self { + Self { + generated_listing: Some(GeneratedListing { + num_objects, + metadata_bytes, + failure_at, + polled: Default::default(), + }), + ..Self::new() } } @@ -1266,6 +1515,29 @@ mod tests { prefix: Option<&Path>, ) -> BoxStream<'static, object_store::Result> { self.list_prefixes.lock().unwrap().push(prefix.cloned()); + if let Some(listing) = self.generated_listing.clone() { + let prefix = prefix.cloned(); + return futures::stream::iter(0..listing.num_objects) + .then(move |index| { + listing.polled.fetch_add(1, Ordering::Relaxed); + futures::future::ready(if listing.failure_at == Some(index) { + Err(object_store::Error::Generic { + store: "GeneratedListing", + source: "generated listing failure".into(), + }) + } else { + Ok(listing.object(index)) + }) + }) + .try_filter(move |meta| { + futures::future::ready( + prefix.as_ref().is_none_or(|prefix| { + meta.location.prefix_matches(prefix) + }), + ) + }) + .boxed(); + } self.in_mem.list(prefix) }