diff --git a/datafusion_iceberg/AGENTS.md b/datafusion_iceberg/AGENTS.md new file mode 100644 index 00000000..059b73a5 --- /dev/null +++ b/datafusion_iceberg/AGENTS.md @@ -0,0 +1,22 @@ +# datafusion_iceberg + +This module owns DataFusion providers and physical scans over Iceberg snapshots. +Keep schema, delete-file and snapshot semantics intact when optimizing reads. + +`parquet_metadata_cache.rs` wraps every Iceberg Parquet reader with the shared +footer cache. `parquet_data_cache.rs` optionally caches immutable byte ranges; +the default capacity is zero. Both use store-qualified file identity and bounded +memory. Keep data-cache metrics attached to each scan. Preserve batched I/O on +misses, do not hold synchronous locks across awaits, and account for retained +buffer ownership rather than only a shared slice's visible length. + +Validate cache changes with focused unit tests for identities, memory limits, +partial hits, error recovery, and representative Iceberg query result checks. +Consumers must include configured cache capacity in their memory budgeting. + +Metadata readers must remain independently pollable. Never await a per-file +load gate owned by another query input: a prefetched probe can be paused while +the build input needs the same footer. Contended cold loads currently read +independently, expose a bypass counter and preserve the byte-capped warm cache. +Any future miss coalescing must pass the paused-loader liveness regression and +must not detach unbounded work or retain canceled readers in a global cache. diff --git a/datafusion_iceberg/README.md b/datafusion_iceberg/README.md index 885d0ec3..8dd32034 100644 --- a/datafusion_iceberg/README.md +++ b/datafusion_iceberg/README.md @@ -1,3 +1,40 @@ # Datafusion iceberg -Provides the functionality to use apache iceberg with datafusion including the `TableProvider`, `SchemaProvider` and `CatalogProvider` traits. \ No newline at end of file +Provides the functionality to use apache iceberg with datafusion including the `TableProvider`, `SchemaProvider` and `CatalogProvider` traits. + +## Parquet footer cache and pipeline progress + +The process-wide parsed-footer cache is byte-capped by +`ICEBERG_PARQUET_METADATA_CACHE_MB` (default 64 MiB, zero disables it) and preserves +store-qualified immutable-file identity plus a file-size sanity stamp. Warm +reads reuse parsed metadata without object-store I/O. + +Cold readers do not wait for a per-file gate owned by another reader. A lazy +query pipeline can stop polling a prefetched probe while a build input needs +the same footer; blocking single-flight would then deadlock. A contended cold +reader instead performs its own bounded read and records +`parquet_metadata_cache_contention_bypasses`. Duplicate cold fetches are an +intentional liveness trade, not a query-result cache or a background fill task. +The retained cache remains byte-capped; concurrent reads and query-owned +metadata require their own headroom. A deterministic paused-loader unit test +covers this scheduling condition independently of benchmark table names. + +## Immutable Parquet data cache + +`ICEBERG_PARQUET_DATA_CACHE_MB` enables a process-wide, byte-bounded LRU for exact +Parquet read ranges. Its default is `0` (disabled). Cache identity includes the +store-qualified file URI, file size, and range. Iceberg rewrites and time travel +use immutable data-file paths; query results, table snapshots and credentials +are not cached here. Metadata resolution and delete-file handling still run. + +Warm hits avoid object-store reads while retaining normal Parquet decoding and +DataFusion execution. Misses keep the reader's batched multi-range API. Retained +bytes own compact buffers, so a small slice cannot pin a large coalesced response. +Entries plus conservative key overhead count against the cap. The process's +other memory users and in-flight reads still need their own budget. + +Scan metrics expose `parquet_data_cache_hits`, `parquet_data_cache_misses`, +`parquet_data_cache_hit_bytes`, and `parquet_data_cache_miss_bytes` through +`EXPLAIN ANALYZE VERBOSE`. This initial implementation does not reuse overlapping +ranges or deduplicate concurrent misses. It is intended for immutable Iceberg +files, not files overwritten in place at the same URI and size. diff --git a/datafusion_iceberg/src/lib.rs b/datafusion_iceberg/src/lib.rs index 7f8360be..7da70c77 100644 --- a/datafusion_iceberg/src/lib.rs +++ b/datafusion_iceberg/src/lib.rs @@ -1,6 +1,7 @@ pub mod catalog; pub mod error; pub mod materialized_view; +mod parquet_data_cache; mod parquet_metadata_cache; pub mod planner; mod pruning_statistics; diff --git a/datafusion_iceberg/src/parquet_data_cache.rs b/datafusion_iceberg/src/parquet_data_cache.rs new file mode 100644 index 00000000..1e866171 --- /dev/null +++ b/datafusion_iceberg/src/parquet_data_cache.rs @@ -0,0 +1,362 @@ +//! Optional process-wide cache of immutable Iceberg Parquet byte ranges. +//! +//! `ICEBERG_PARQUET_DATA_CACHE_MB` bounds resident cached bytes (including a +//! conservative key/entry charge); zero, the default, bypasses this cache. Keys +//! include the full store-qualified file identity, file size, and exact range. +//! Iceberg rewrites create new paths, so snapshot changes require no invalidation. +//! Metadata and credentials are still resolved by the catalog in the usual way. +//! +//! Each retained range owns its buffer: an object-store multi-range response can +//! return slices backed by a much larger coalesced allocation. Retaining those +//! slices while charging only their lengths would evade the cache's byte cap. + +use std::ops::Range; +use std::sync::{Arc, LazyLock, Mutex}; + +use bytes::Bytes; +use datafusion::parquet::arrow::async_reader::AsyncFileReader; +use datafusion::parquet::errors::{ParquetError, Result}; +use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}; +use lru::LruCache; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct Key { + file: Arc, + size: u64, + start: u64, + end: u64, +} + +impl Key { + fn weight(&self, bytes: usize) -> usize { + // Includes the LRU/hash entry, shared-string allocation and allocator + // overhead conservatively, even though sibling keys share the string. + bytes.saturating_add(self.file.len()).saturating_add(192) + } +} + +struct State { + entries: LruCache, + used: usize, +} + +pub(crate) struct DataCache { + cap: usize, + state: Mutex, +} + +impl DataCache { + fn new(cap: usize) -> Self { + Self { + cap, + state: Mutex::new(State { + entries: LruCache::unbounded(), + used: 0, + }), + } + } + + pub(crate) fn enabled(&self) -> bool { + self.cap != 0 + } + + fn get(&self, key: &Key) -> Option { + self.state.lock().ok()?.entries.get(key).cloned() + } + + fn put(&self, key: Key, bytes: &Bytes) { + let weight = key.weight(bytes.len()); + if weight > self.cap || bytes.is_empty() { + return; + } + // Copy outside the shared lock; never retain a slice of a larger range. + let owned = Bytes::copy_from_slice(bytes); + let Ok(mut state) = self.state.lock() else { + return; + }; + if let Some(previous) = state.entries.put(key.clone(), owned) { + state.used = state.used.saturating_sub(key.weight(previous.len())); + } + state.used = state.used.saturating_add(weight); + while state.used > self.cap { + let Some((old_key, old_bytes)) = state.entries.pop_lru() else { + break; + }; + state.used = state.used.saturating_sub(old_key.weight(old_bytes.len())); + } + } + + pub(crate) async fn read( + &self, + inner: &mut (dyn AsyncFileReader + Send), + file: &Arc, + size: u64, + ranges: Vec>, + metrics: &DataMetrics, + ) -> Result> { + let mut output = Vec::with_capacity(ranges.len()); + let mut missing = Vec::new(); + let mut missing_keys = Vec::new(); + for (index, range) in ranges.into_iter().enumerate() { + if range.start > range.end || range.end > size { + return Err(ParquetError::General( + "Parquet range outside file bounds".into(), + )); + } + let key = Key { + file: Arc::clone(file), + size, + start: range.start, + end: range.end, + }; + if let Some(bytes) = self.get(&key) { + metrics.hits.add(1); + metrics.hit_bytes.add(bytes.len()); + output.push(Some(bytes)); + } else { + metrics.misses.add(1); + missing.push(range); + missing_keys.push((index, key)); + output.push(None); + } + } + if !missing.is_empty() { + let fetched = inner.get_byte_ranges(missing).await?; + if fetched.len() != missing_keys.len() { + return Err(ParquetError::General( + "Incomplete Parquet multi-range response".into(), + )); + } + for ((index, key), bytes) in missing_keys.into_iter().zip(fetched) { + if bytes.len() as u64 != key.end - key.start { + return Err(ParquetError::General("Truncated Parquet byte range".into())); + } + metrics.miss_bytes.add(bytes.len()); + self.put(key, &bytes); + output[index] = Some(bytes); + } + } + output + .into_iter() + .map(|bytes| { + bytes.ok_or_else(|| ParquetError::General("Missing Parquet byte range".into())) + }) + .collect() + } +} + +pub(crate) struct DataMetrics { + hits: Count, + misses: Count, + hit_bytes: Count, + miss_bytes: Count, +} + +impl DataMetrics { + pub(crate) fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + hits: MetricBuilder::new(metrics).counter("parquet_data_cache_hits", partition), + misses: MetricBuilder::new(metrics).counter("parquet_data_cache_misses", partition), + hit_bytes: MetricBuilder::new(metrics) + .counter("parquet_data_cache_hit_bytes", partition), + miss_bytes: MetricBuilder::new(metrics) + .counter("parquet_data_cache_miss_bytes", partition), + } + } +} + +pub(crate) static DATA_CACHE: LazyLock = LazyLock::new(|| { + let cap_mb = std::env::var("ICEBERG_PARQUET_DATA_CACHE_MB") + .ok() + .and_then(|s| s.trim().parse::().ok()) + .unwrap_or(0); + tracing::info!(cap_mb, "Iceberg Parquet data cache initialized"); + DataCache::new(cap_mb.saturating_mul(1024 * 1024)) +}); + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::parquet::arrow::arrow_reader::ArrowReaderOptions; + use datafusion::parquet::file::metadata::ParquetMetaData; + use futures::{future::BoxFuture, FutureExt}; + + struct Reader { + data: Bytes, + reads: Vec>, + truncate: bool, + } + + impl AsyncFileReader for Reader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + self.reads.push(range.clone()); + let end = range.end as usize - usize::from(self.truncate); + futures::future::ready(Ok(self.data.slice(range.start as usize..end))).boxed() + } + + fn get_metadata<'a>( + &'a mut self, + _: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + futures::future::ready(Err(ParquetError::General( + "unused in data cache tests".into(), + ))) + .boxed() + } + } + + #[tokio::test] + async fn mixed_hits_and_misses_preserve_requested_order() { + let cache = DataCache::new(4096); + let metrics = DataMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut reader = Reader { + data: Bytes::from_static(b"abcdefgh"), + reads: vec![], + truncate: false, + }; + let file = Arc::from("s3://bucket/file"); + let first = cache + .read(&mut reader, &file, 8, vec![0..2, 4..6], &metrics) + .await + .unwrap(); + assert_eq!( + first, + vec![Bytes::from_static(b"ab"), Bytes::from_static(b"ef")] + ); + reader.reads.clear(); + let mixed = cache + .read(&mut reader, &file, 8, vec![4..6, 2..4, 0..2], &metrics) + .await + .unwrap(); + assert_eq!( + mixed, + vec![ + Bytes::from_static(b"ef"), + Bytes::from_static(b"cd"), + Bytes::from_static(b"ab") + ] + ); + assert_eq!(reader.reads, vec![2..4]); + assert_eq!(metrics.hit_bytes.value(), 4); + } + + #[tokio::test] + async fn failed_read_is_not_cached_and_retry_can_succeed() { + let cache = DataCache::new(4096); + let metrics = DataMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut reader = Reader { + data: Bytes::from_static(b"abcd"), + reads: vec![], + truncate: true, + }; + let file = Arc::from("s3://bucket/file"); + assert!(cache + .read( + &mut reader, + &file, + 4, + std::iter::once(0..4).collect(), + &metrics + ) + .await + .is_err()); + reader.truncate = false; + let retry = cache + .read( + &mut reader, + &file, + 4, + std::iter::once(0..4).collect(), + &metrics, + ) + .await + .unwrap(); + assert_eq!(retry, vec![Bytes::from_static(b"abcd")]); + assert_eq!(reader.reads.len(), 2); + } + + #[tokio::test] + async fn empty_requests_and_invalid_bounds_do_not_read_storage() { + let cache = DataCache::new(4096); + let metrics = DataMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut reader = Reader { + data: Bytes::from_static(b"abcd"), + reads: vec![], + truncate: false, + }; + let file = Arc::from("s3://bucket/file"); + assert!(cache + .read(&mut reader, &file, 4, vec![], &metrics) + .await + .unwrap() + .is_empty()); + for range in [0..5, Range { start: 3, end: 2 }] { + assert!(cache + .read(&mut reader, &file, 4, vec![range], &metrics) + .await + .is_err()); + } + assert!(reader.reads.is_empty()); + } + + fn key(file: &str, size: u64, start: u64, end: u64) -> Key { + Key { + file: Arc::from(file), + size, + start, + end, + } + } + + #[test] + fn keys_isolate_stores_sizes_and_ranges() { + let cache = DataCache::new(4096); + let a = key("s3://a/file", 20, 0, 4); + cache.put(a.clone(), &Bytes::from_static(b"data")); + assert_eq!(cache.get(&a).unwrap(), "data"); + for different in [ + key("s3://b/file", 20, 0, 4), + key("s3://a/file", 21, 0, 4), + key("s3://a/file", 20, 1, 5), + ] { + assert!(cache.get(&different).is_none()); + } + } + + #[test] + fn lru_replacement_and_eviction_stay_bounded() { + let a = key("a", 20, 0, 4); + let b = key("b", 20, 0, 4); + let cache = DataCache::new(a.weight(4)); + cache.put(a.clone(), &Bytes::from_static(b"abcd")); + cache.put(a.clone(), &Bytes::from_static(b"efgh")); + assert_eq!(cache.state.lock().unwrap().used, a.weight(4)); + cache.put(b.clone(), &Bytes::from_static(b"ijkl")); + assert!(cache.get(&a).is_none()); + assert_eq!(cache.get(&b).unwrap(), "ijkl"); + assert!(cache.state.lock().unwrap().used <= cache.cap); + } + + #[test] + fn retained_slice_owns_only_its_range() { + let cache = DataCache::new(4096); + let large = Bytes::from(vec![7; 1024 * 1024]); + let slice = large.slice(100..104); + let a = key("a", 1024 * 1024, 100, 104); + cache.put(a.clone(), &slice); + let retained = cache.get(&a).unwrap(); + assert_eq!(retained, slice); + assert_ne!(retained.as_ptr(), slice.as_ptr()); + } + + #[test] + fn disabled_and_oversized_entries_are_not_retained() { + for cap in [0, 1] { + let cache = DataCache::new(cap); + let a = key("a", 4, 0, 4); + cache.put(a.clone(), &Bytes::from_static(b"data")); + assert!(cache.get(&a).is_none()); + assert_eq!(cache.state.lock().unwrap().used, 0); + } + } +} diff --git a/datafusion_iceberg/src/parquet_metadata_cache.rs b/datafusion_iceberg/src/parquet_metadata_cache.rs index 6bb8ec7d..d3795e76 100644 --- a/datafusion_iceberg/src/parquet_metadata_cache.rs +++ b/datafusion_iceberg/src/parquet_metadata_cache.rs @@ -82,24 +82,27 @@ indistinguishable — both just look like slow metadata loading. Interception happens through a custom [`ParquetFileReaderFactory`] ([`CachingParquetFileReaderFactory`]) installed on every [`ParquetSource`] the -scan builds, so both cold and warm queries flow through it. Only footer -metadata (`get_metadata`) is cached; row-group/page data reads pass straight -through to the underlying reader. - -Concurrent *cold* readers of the same file are collapsed by per-file -single-flight. Without it, N queries racing on the same not-yet-cached footer -each fetch and parse it independently (a dogpile) before the first `put` lands. -The first reader to miss takes the file's single-flight gate — a per-key async -mutex — and performs the sole fetch+parse+insert; the rest wait on the gate, -then re-check the cache after acquiring it and return the winner's `Arc`. So N -racing cold scans pay one round-trip, not N. The synchronous LRU lock is only -taken to look up or insert an entry and is never held across the fetch await; -only the per-file async gate is. The gate map is itself a small LRU capped at -[`INFLIGHT_GATES_CAP`] live gates, so it stays bounded no matter how many -distinct files the process reads; a gate evicted while a load is still in flight -merely lets that one file be fetched twice, never breaking correctness. When -caching is disabled (cap `0`) this path is skipped entirely — reads pass -straight through with no gate and no lock, exactly as before single-flight. +scan builds, so both cold and warm queries flow through it. Footer metadata +(`get_metadata`) uses this cache. Row-group/page reads optionally use the separate +byte-bounded `parquet_data_cache`; with its default zero capacity they pass +straight through to the underlying reader. + +Cold readers must remain independently pollable. A query may prefetch metadata +on a probe input and then stop polling that input while it builds a join whose +other input reads the same file. Waiting on a per-file single-flight gate owned +by the paused probe creates a cycle: build waits for metadata, metadata waits +for probe, probe waits for build. This was reproduced with cold TPC-H SF100 Q2. + +Per-file gates therefore never block a reader: a contender rechecks the cache +and, if still cold, independently fetches and parses the immutable footer. +`parquet_metadata_cache_contention_bypasses` exposes this trade. Concurrent cold +readers can issue duplicate I/O, but no task is detached and warm hits retain +the zero-I/O path. The synchronous LRU lock is never held across an await. +The gate map is bounded by [`INFLIGHT_GATES_CAP`]; eviction can only weaken +contention accounting, not correctness. With metadata cap `0`, the footer cache +and its gates are bypassed; the data cache has its own independent capacity. +Retained parsed metadata remains byte-capped; in-flight reads +and query-owned metadata still require separate memory headroom. [`ParquetSource`]: datafusion::datasource::physical_plan::parquet::source::ParquetSource */ @@ -109,6 +112,7 @@ use std::ops::Range; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, LazyLock, Mutex}; +use crate::parquet_data_cache::{DataMetrics, DATA_CACHE}; use bytes::Bytes; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::physical_plan::parquet::{ @@ -127,7 +131,7 @@ use tokio::sync::Mutex as AsyncMutex; const DEFAULT_CAP_MB: usize = 64; -/// Upper bound on the number of distinct per-file single-flight gates kept live +/// Upper bound on the number of distinct nonblocking per-file gates kept live /// at once. The gate map is an LRU capped here so it stays bounded no matter how /// many distinct files the process reads over its lifetime; a gate evicted while /// a load is still in flight only weakens de-duplication for that one file (an @@ -232,6 +236,8 @@ struct CacheMetrics { /// latched. Non-zero means this deployment's working set outgrew the cache and /// the optimization backed itself off. page_index_suppressed: Count, + /// Independent loads used instead of waiting for another input's loader. + contention_bypasses: Count, } impl CacheMetrics { @@ -243,12 +249,13 @@ impl CacheMetrics { .counter("parquet_metadata_cache_evictions", partition), page_index_suppressed: MetricBuilder::new(metrics) .counter("parquet_metadata_cache_page_index_suppressed", partition), + contention_bypasses: MetricBuilder::new(metrics) + .counter("parquet_metadata_cache_contention_bypasses", partition), } } } -/// The parsed-footer LRU together with the per-file single-flight gates that -/// collapse concurrent cold loads of the same file into one fetch+parse. +/// The parsed-footer LRU together with nonblocking per-file load gates. /// /// A single instance backs the whole process ([`CACHE`]); the tests construct /// isolated instances (including disabled ones) to drive [`load`] deterministically. @@ -256,10 +263,10 @@ impl CacheMetrics { /// [`load`]: MetadataCache::load struct MetadataCache { /// Parsed footers, byte-capped. `None` when caching is disabled (cap `0`), - /// in which case lookups, inserts, and the single-flight gate are all + /// in which case lookups, inserts, and load gates are all /// bypassed. store: Option>, - /// Per-file single-flight gates, keyed exactly like [`ByteCappedCache`] + /// Nonblocking per-file gates, keyed exactly like [`ByteCappedCache`] /// entries. Bounded to [`INFLIGHT_GATES_CAP`] live gates by its own LRU. inflight: Mutex>>>, /// Entries evicted over the process lifetime. Drives [`page_index_backed_off`]. @@ -296,7 +303,7 @@ impl MetadataCache { Some(guard.stats()) } - /// Whether caching (and therefore single-flight) is active. + /// Whether caching and load-gate accounting are active. fn enabled(&self) -> bool { self.store.is_some() } @@ -373,7 +380,7 @@ impl MetadataCache { self.backed_off.load(Ordering::Relaxed) } - /// Return the single-flight gate for `key`, creating it on first use. The + /// Return the nonblocking load gate for `key`, creating it on first use. The /// gate map is a small LRU, so live gates stay bounded to /// [`INFLIGHT_GATES_CAP`]; a gate evicted while still held survives through /// its holders' `Arc` clones. Takes the map lock only briefly and never @@ -392,15 +399,14 @@ impl MetadataCache { } /// Resolve the footer for `key`: serve it from cache, or fetch+parse it - /// through `reader` on a miss, collapsing concurrent cold misses of the same - /// file into a single fetch. + /// through `reader` on a miss, without depending on another reader's polling. /// /// When caching is disabled the fetch is issued straight through with no - /// gate and no lock. Otherwise the first caller to miss takes the file's - /// single-flight gate and performs the sole fetch+parse+insert; the rest - /// wait on the gate and, once it releases, re-check the cache and return the - /// winner's `Arc`. The synchronous LRU locks are never held across the - /// `reader.get_metadata` await — only the per-file async gate is. + /// gate and no lock. Contenders never wait for an in-flight loader: it may + /// be a prefetched input that the query parent has stopped polling. They + /// recheck the cache, then load independently when necessary. The trade is + /// duplicate cold I/O, not a cross-input dependency or detached background + /// task. Synchronous LRU locks are never held across the metadata await. async fn load( &self, reader: &mut (dyn AsyncFileReader + Send), @@ -419,12 +425,15 @@ impl MetadataCache { return reader.get_metadata(options).await; } let gate = self.inflight_gate(key); - let _permit = gate.lock().await; - // The winner may have populated the cache while we waited on the gate. + let permit = gate.try_lock().ok(); + // A racing reader may have populated the cache since our first lookup. if let Some(meta) = self.get(key, size) { metrics.hits.add(1); return Ok(meta); } + if permit.is_none() { + metrics.contention_bypasses.add(1); + } metrics.misses.add(1); // Fetch the page index in this same round trip when configured, so the // cached entry is complete and the opener's own `load_page_index` finds @@ -517,7 +526,7 @@ fn page_index_options( /// The process-wide metadata cache. Capacity comes from /// `ICEBERG_PARQUET_METADATA_CACHE_MB` (read once per process; default 64 MiB; -/// `0` disables caching and single-flight entirely). +/// `0` disables caching and load gates entirely). static CACHE: LazyLock = LazyLock::new(|| { let raw = std::env::var("ICEBERG_PARQUET_METADATA_CACHE_MB").ok(); let cap_mb = raw @@ -578,6 +587,7 @@ impl ParquetFileReaderFactory for CachingParquetFileReaderFactory { ); let size = partitioned_file.object_meta.size; let cache_metrics = CacheMetrics::new(metrics, partition_index); + let data_metrics = DataMetrics::new(metrics, partition_index); let inner = self.inner.create_reader( partition_index, partitioned_file, @@ -586,9 +596,10 @@ impl ParquetFileReaderFactory for CachingParquetFileReaderFactory { )?; Ok(Box::new(CachingMetadataReader { inner, - key, + key: Arc::from(key), size, metrics: cache_metrics, + data_metrics, })) } } @@ -597,21 +608,52 @@ impl ParquetFileReaderFactory for CachingParquetFileReaderFactory { /// process-wide cache and forwards data reads to `inner`. struct CachingMetadataReader { inner: Box, - key: String, + key: Arc, size: u64, metrics: CacheMetrics, + data_metrics: DataMetrics, } impl AsyncFileReader for CachingMetadataReader { fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, ParquetResult> { - self.inner.get_bytes(range) + if !DATA_CACHE.enabled() { + return self.inner.get_bytes(range); + } + async move { + let mut bytes = DATA_CACHE + .read( + &mut *self.inner, + &self.key, + self.size, + vec![range], + &self.data_metrics, + ) + .await?; + bytes.pop().ok_or_else(|| { + datafusion::parquet::errors::ParquetError::General( + "Missing Parquet byte range".into(), + ) + }) + } + .boxed() } fn get_byte_ranges( &mut self, ranges: Vec>, ) -> BoxFuture<'_, ParquetResult>> { - self.inner.get_byte_ranges(ranges) + if !DATA_CACHE.enabled() { + return self.inner.get_byte_ranges(ranges); + } + DATA_CACHE + .read( + &mut *self.inner, + &self.key, + self.size, + ranges, + &self.data_metrics, + ) + .boxed() } fn get_metadata<'a>( @@ -885,7 +927,7 @@ mod tests { /// An [`ObjectStore`] that counts `get_opts` calls (all footer/data reads /// funnel through it) and delegates everything else to an inner store. An /// optional per-`get` `delay` widens the window in which concurrent cold - /// readers overlap, making the single-flight tests deterministic. + /// readers overlap, exercising cold-read contention. #[derive(Debug)] struct CountingStore { inner: Arc, @@ -924,8 +966,8 @@ mod tests { options: GetOptions, ) -> OsResult { if let Some(delay) = self.delay { - // Yield while "fetching" so racing cold readers pile onto the - // single-flight gate before the winner inserts. + // Yield while "fetching" so cold readers overlap before any + // reader inserts its result. tokio::time::sleep(delay).await; } self.gets.fetch_add(1, Ordering::SeqCst); @@ -1076,10 +1118,51 @@ mod tests { gets.swap(0, Ordering::SeqCst) } - /// N concurrent cold readers of the same file collapse to exactly one - /// fetch+parse, and every reader receives the winner's `Arc`. + #[tokio::test] + async fn paused_cold_reader_cannot_block_an_independently_polled_reader() { + let path = "db/tbl/data/paused_leader.parquet"; + let (_inner, factory, size, _gets) = + counting_setup(path, Some(Duration::from_millis(1))).await; + let cache = MetadataCache::new(64 * 1024 * 1024); + let plan_metrics = ExecutionPlanMetricsSet::new(); + let key = "iceberg-rust://paused-leader/file"; + let mut first = factory + .create_reader(0, partitioned_file(path, size), None, &plan_metrics) + .unwrap(); + let mut second = factory + .create_reader(0, partitioned_file(path, size), None, &plan_metrics) + .unwrap(); + let first_metrics = discard_metrics(); + let second_metrics = discard_metrics(); + let mut paused = Box::pin(cache.load(&mut *first, None, key, size, &first_metrics)); + assert!(paused.as_mut().now_or_never().is_none()); + assert!( + cache.inflight_gate(key).try_lock().is_err(), + "the first load is paused while holding the per-file gate" + ); + + // A query parent may stop polling its prefetched probe input while it + // consumes a build input reading the same immutable file. The cache + // must not introduce a dependency on that paused probe future. + let completed = tokio::time::timeout( + Duration::from_millis(100), + cache.load(&mut *second, None, key, size, &second_metrics), + ) + .await + .expect("a parked metadata loader must not deadlock another query input") + .unwrap(); + assert_eq!(completed.file_metadata().num_rows(), 3); + assert_eq!(second_metrics.contention_bypasses.value(), 1); + let resumed = paused.await.unwrap(); + assert_eq!(resumed.as_ref(), completed.as_ref()); + assert_eq!(cache.stats().unwrap().0, 1); + } + + /// Racing cold readers may load independently, but values and the cache + /// byte bound remain exact. This deliberately replaces the old one-fetch + /// performance guarantee, whose blocking gate could deadlock a query. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn concurrent_cold_reads_single_flight_one_fetch() { + async fn concurrent_cold_reads_preserve_metadata_and_bounded_cache() { let path = "db/tbl/data/single_flight.parquet"; let (_inner, factory, size, gets) = counting_setup(path, Some(Duration::from_millis(25))).await; @@ -1111,17 +1194,14 @@ mod tests { .map(|joined| joined.unwrap()) .collect(); - assert_eq!( - gets.load(Ordering::SeqCst), - per_fetch, - "N concurrent cold readers must trigger exactly one underlying fetch" - ); + let actual_gets = gets.load(Ordering::SeqCst); + assert!(actual_gets >= per_fetch && actual_gets <= per_fetch * n); for meta in &metas[1..] { - assert!( - Arc::ptr_eq(&metas[0], meta), - "every reader shares the single-flight winner's Arc" - ); + assert_eq!(metas[0].as_ref(), meta.as_ref()); } + let (entries, bytes) = cache.stats().unwrap(); + assert_eq!(entries, 1); + assert!(bytes <= 64 * 1024 * 1024); } /// Concurrent cold reads of *different* files are not serialized against