From 9bdb8421c856c72b601386a911579ac7c6f9bb0b Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 15 Jul 2026 20:03:05 +0530 Subject: [PATCH 01/13] feat(cache): make LiquidPolicy eviction scan-resistant (CLOCK) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LiquidPolicy was FIFO-per-type with a no-op notify_access, so a large one-pass scan floods the cache and evicts the reused working set (and, once full, read-once batches still displace hot ones). Add a CLOCK/second-chance reference bit: notify_access sets it on a cache hit, and find_memory_victim passes over a referenced entry once (clearing the bit, moving it to the back) before evicting it. A batch read exactly once starts and stays cold, so it is evicted before any reused batch — the cache converges on the working set regardless of how large the surrounding column/scan is. FIFO behaviour is unchanged when nothing is reused. --- .../src/cache/policies/cache/three_queue.rs | 135 +++++++++++++++--- 1 file changed, 118 insertions(+), 17 deletions(-) diff --git a/src/core/src/cache/policies/cache/three_queue.rs b/src/core/src/cache/policies/cache/three_queue.rs index 655c1984..a958e55f 100644 --- a/src/core/src/cache/policies/cache/three_queue.rs +++ b/src/core/src/cache/policies/cache/three_queue.rs @@ -1,4 +1,7 @@ -use std::{collections::HashMap, ptr::NonNull}; +use std::{ + collections::{HashMap, HashSet}, + ptr::NonNull, +}; use crate::{ cache::{CachePolicy, EntryID, cached_batch::CachedBatchType}, @@ -19,6 +22,12 @@ enum QueueKind { struct QueueNode { entry_id: EntryID, queue: QueueKind, + /// Second-chance (CLOCK) reference bit. Set when the entry is accessed + /// (a cache hit); cleared when it is passed over during victim search. + /// A fresh insert starts cold (`false`), so a batch read exactly once — + /// e.g. every batch of a large one-pass scan — is evicted before any + /// batch that has been reused, keeping the working set scan-resistant. + referenced: bool, } type NodePtr = NonNull>; @@ -68,6 +77,7 @@ impl LiquidQueueInternalState { let node = DoublyLinkedNode::new(QueueNode { entry_id, queue: target, + referenced: false, }); let node_ptr = NonNull::from(Box::leak(node)); @@ -107,6 +117,46 @@ impl LiquidQueueInternalState { } Some(removed) } + + /// Entry id and reference bit of the queue's front node, if any. + fn head_info(&self, queue: QueueKind) -> Option<(EntryID, bool)> { + let list = match queue { + QueueKind::Arrow => &self.arrow, + QueueKind::Liquid => &self.liquid, + QueueKind::Squeezed => &self.squeezed, + QueueKind::Disk => &self.disk, + }; + let head_ptr = list.head()?; + let data = unsafe { &head_ptr.as_ref().data }; + Some((data.entry_id, data.referenced)) + } + + /// Set the reference bit of an entry, if present. + fn set_referenced(&mut self, entry_id: &EntryID, value: bool) { + if let Some(node_ptr) = self.map.get(entry_id).copied() { + unsafe { + (*node_ptr.as_ptr()).data.referenced = value; + } + } + } + + /// Give the queue's front node a second chance: clear its reference bit and + /// move it to the back. Caller has already confirmed the queue is non-empty. + fn second_chance_head(&mut self, queue: QueueKind) { + let head_ptr = match queue { + QueueKind::Arrow => self.arrow.head(), + QueueKind::Liquid => self.liquid.head(), + QueueKind::Squeezed => self.squeezed.head(), + QueueKind::Disk => self.disk.head(), + }; + if let Some(head_ptr) = head_ptr { + unsafe { + (*head_ptr.as_ptr()).data.referenced = false; + self.detach(head_ptr); + self.push_back(queue, head_ptr); + } + } + } } impl Drop for LiquidQueueInternalState { @@ -173,23 +223,29 @@ impl CachePolicy for LiquidPolicy { let mut inner = self.inner.lock().unwrap(); let mut victims = Vec::with_capacity(cnt); - while victims.len() < cnt { - if let Some(entry) = inner.pop_front(QueueKind::Arrow) { - victims.push(entry); - continue; - } - - if let Some(entry) = inner.pop_front(QueueKind::Liquid) { - victims.push(entry); - continue; + // Evict decoded (Arrow) first, then Liquid, then Squeezed — cheapest to + // reconstruct first. Within each queue apply CLOCK/second-chance: a + // referenced entry is passed over once (bit cleared, moved to the back) + // before it can be evicted, so a reused batch outlives a read-once one. + // `chanced` bounds each entry to a single second chance, so the loop + // always terminates (at most one pass of reprieves, then eviction). + for queue in [QueueKind::Arrow, QueueKind::Liquid, QueueKind::Squeezed] { + let mut chanced: HashSet = HashSet::new(); + while victims.len() < cnt { + let Some((entry_id, referenced)) = inner.head_info(queue) else { + break; + }; + if referenced && chanced.insert(entry_id) { + inner.second_chance_head(queue); + } else { + let popped = inner.pop_front(queue); + debug_assert_eq!(popped, Some(entry_id)); + victims.push(entry_id); + } } - - if let Some(entry) = inner.pop_front(QueueKind::Squeezed) { - victims.push(entry); - continue; + if victims.len() >= cnt { + break; } - - break; } victims @@ -213,7 +269,12 @@ impl CachePolicy for LiquidPolicy { victims } - fn notify_access(&self, _entry_id: &EntryID, _batch_type: CachedBatchType) {} + fn notify_access(&self, entry_id: &EntryID, _batch_type: CachedBatchType) { + // A cache hit marks the entry as reused, protecting it from the next + // victim search (see `find_memory_victim`). + let mut inner = self.inner.lock().unwrap(); + inner.set_referenced(entry_id, true); + } fn notify_remove(&self, entry_id: &EntryID) { let mut inner = self.inner.lock().unwrap(); @@ -249,6 +310,46 @@ mod tests { assert_eq!(policy.find_memory_victim(1), vec![liquid_b]); } + #[test] + fn test_accessed_entry_gets_second_chance() { + let policy = LiquidPolicy::new(); + let a = entry(1); + let b = entry(2); + let c = entry(3); + policy.notify_insert(&a, CachedBatchType::MemoryArrow); + policy.notify_insert(&b, CachedBatchType::MemoryArrow); + policy.notify_insert(&c, CachedBatchType::MemoryArrow); + + // `a` is reused, so it should be spared on the next victim search even + // though it was inserted first (plain FIFO would evict it). + policy.notify_access(&a, CachedBatchType::MemoryArrow); + + assert_eq!(policy.find_memory_victim(1), vec![b]); + assert_eq!(policy.find_memory_victim(1), vec![c]); + // `a` survived the longest and is evicted last (its bit was cleared when + // it was passed over, so it gets no further reprieve). + assert_eq!(policy.find_memory_victim(1), vec![a]); + } + + #[test] + fn test_read_once_entries_evicted_before_reused() { + // Simulates a large one-pass scan (read-once batches) alongside a single + // reused batch: the scan batches must all be evicted before the reused + // one, so the scan cannot displace the working set. + let policy = LiquidPolicy::new(); + let ids: Vec<_> = (1..=5).map(entry).collect(); + for id in &ids { + policy.notify_insert(id, CachedBatchType::MemoryArrow); + } + // ids[2] (entry 3) is reused. + policy.notify_access(&ids[2], CachedBatchType::MemoryArrow); + + let victims = policy.find_memory_victim(4); + assert_eq!(victims, vec![ids[0], ids[1], ids[3], ids[4]]); + // The reused entry is the sole survivor. + assert_eq!(policy.find_memory_victim(4), vec![ids[2]]); + } + #[test] fn test_queue_priority_order() { let policy = LiquidPolicy::new(); From bf9b2160fde6dd015372f355626053233d1cbcc4 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Tue, 21 Jul 2026 11:13:40 +0530 Subject: [PATCH 02/13] fix(cache): only memory hits set the CLOCK reference bit A disk-served hit called notify_access before hydration, and the following Disk->Memory notify_insert preserved the bit, so a batch read once from disk arrived in memory looking reused and could evict a genuinely reused batch during a large one-pass scan (the exact overflow case scan-resistance targets). Gate the reference bit on memory batch types only. (codex review) --- .../src/cache/policies/cache/three_queue.rs | 42 ++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/src/core/src/cache/policies/cache/three_queue.rs b/src/core/src/cache/policies/cache/three_queue.rs index a958e55f..e630c413 100644 --- a/src/core/src/cache/policies/cache/three_queue.rs +++ b/src/core/src/cache/policies/cache/three_queue.rs @@ -269,11 +269,23 @@ impl CachePolicy for LiquidPolicy { victims } - fn notify_access(&self, entry_id: &EntryID, _batch_type: CachedBatchType) { - // A cache hit marks the entry as reused, protecting it from the next - // victim search (see `find_memory_victim`). - let mut inner = self.inner.lock().unwrap(); - inner.set_referenced(entry_id, true); + fn notify_access(&self, entry_id: &EntryID, batch_type: CachedBatchType) { + // Only a memory-resident hit counts as reuse for CLOCK. A disk hit must + // not set the bit: the disk read path calls `notify_access` before + // hydration, and the subsequent Disk->Memory `notify_insert` preserves + // the bit — so a batch read once from disk would arrive in memory + // looking hot and could evict a genuinely reused memory batch during a + // large one-pass scan. + let is_memory = matches!( + batch_type, + CachedBatchType::MemoryArrow + | CachedBatchType::MemoryLiquid + | CachedBatchType::MemorySqueezedLiquid + ); + if is_memory { + let mut inner = self.inner.lock().unwrap(); + inner.set_referenced(entry_id, true); + } } fn notify_remove(&self, entry_id: &EntryID) { @@ -331,6 +343,26 @@ mod tests { assert_eq!(policy.find_memory_victim(1), vec![a]); } + #[test] + fn test_disk_hit_does_not_protect_only_memory_hit_does() { + // A hit served from disk must NOT protect the entry — otherwise a + // one-pass disk-backed scan arrives hot on hydration and evicts the + // real working set. Only a memory hit grants a second chance. + let policy = LiquidPolicy::new(); + let a = entry(1); + let b = entry(2); + let c = entry(3); + for id in [&a, &b, &c] { + policy.notify_insert(id, CachedBatchType::MemoryArrow); + } + policy.notify_access(&a, CachedBatchType::DiskArrow); // must not protect + policy.notify_access(&b, CachedBatchType::MemoryArrow); // protects + + // `a` (disk hit) and `c` (untouched) evicted; `b` (memory hit) survives. + assert_eq!(policy.find_memory_victim(2), vec![a, c]); + assert_eq!(policy.find_memory_victim(1), vec![b]); + } + #[test] fn test_read_once_entries_evicted_before_reused() { // Simulates a large one-pass scan (read-once batches) alongside a single From 242556d1d094bb924c8223fbf78c82fbf71dc2d1 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 00:37:58 +0530 Subject: [PATCH 03/13] feat(admission): footprint-based cache admission gate Replace LocalModeOptimizer's raw-file-size cap with a byte-accurate, filter-aware footprint gate. For a parquet scan it estimates the in-memory liquid footprint from DuckLake per-file stats: required columns (output projection UNION predicate columns), summed over the files that survive the predicate (PruningPredicate + PrunableStatistics), using only Exact per-column byte sizes and falling back to the whole-file size otherwise. A scan whose estimate (raw x expansion x safety) exceeds the cache memory budget is left as a plain parquet read (bypass -> mount) instead of thrashing the cache. Conservative by construction: pruning only drops proven-non-matching files; non-exact/missing byte sizes over-count; multipliers inflate the estimate, never the budget -- so the gate never wrongly admits an oversized scan. Complements the CLOCK eviction (admission vs. eviction). Pure byte-math unit-tested. Builder: with_max_scan_bytes -> with_admission_gate(expansion, safety). --- src/datafusion-local/src/lib.rs | 22 ++- src/datafusion/src/optimizers/mod.rs | 279 ++++++++++++++++++++++----- 2 files changed, 248 insertions(+), 53 deletions(-) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 2bb2485f..aa8a0c2e 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -69,8 +69,9 @@ pub struct LiquidCacheLocalBuilder { squeeze_policy: Box, /// Hydration policy hydration_policy: Box, - /// Maximum total file size for a scan to be routed through LiquidCache. - max_scan_bytes: Option, + /// Footprint-based admission gate `(expansion, safety)`. When set, a scan is + /// cached only if its estimated liquid footprint fits the memory budget. + admission: Option<(f64, f64)>, span: fastrace::Span, } @@ -86,7 +87,7 @@ impl Default for LiquidCacheLocalBuilder { cache_policy: Box::new(LiquidPolicy::new()), squeeze_policy: Box::new(TranscodeSqueezeEvict), hydration_policy: Box::new(AlwaysHydrate::new()), - max_scan_bytes: None, + admission: None, span: fastrace::Span::enter_with_local_parent("liquid_cache_datafusion_local_builder"), } } @@ -148,11 +149,12 @@ impl LiquidCacheLocalBuilder { self } - /// Set the maximum total file size (in bytes) for a scan to be routed - /// through LiquidCache. Scans exceeding this threshold are read directly - /// from the parquet source, bypassing the cache entirely. - pub fn with_max_scan_bytes(mut self, max_bytes: u64) -> Self { - self.max_scan_bytes = Some(max_bytes); + /// Enable the footprint-based admission gate. A scan is cached only when its + /// estimated liquid footprint (raw required bytes x `expansion` x `safety`) + /// fits the cache's memory budget; oversized scans are read directly from the + /// parquet source, bypassing the cache. `expansion` and `safety` are `>= 1.0`. + pub fn with_admission_gate(mut self, expansion: f64, safety: f64) -> Self { + self.admission = Some((expansion, safety)); self } @@ -202,8 +204,8 @@ impl LiquidCacheLocalBuilder { let cache_ref = Arc::new(cache); let mut optimizer = LocalModeOptimizer::new(cache_ref.clone()); - if let Some(max_bytes) = self.max_scan_bytes { - optimizer = optimizer.with_max_scan_bytes(max_bytes); + if let Some((expansion, safety)) = self.admission { + optimizer = optimizer.with_admission_gate(expansion, safety); } let state = datafusion::execution::SessionStateBuilder::new() diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 804a5222..78369030 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -5,11 +5,22 @@ mod squeeze_hint; use std::sync::Arc; use datafusion::{ + arrow::datatypes::SchemaRef, catalog::memory::DataSourceExec, - common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + common::{ + Statistics, + pruning::PrunableStatistics, + stats::Precision, + tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + }, config::ConfigOptions, - datasource::{physical_plan::ParquetSource, source::DataSource}, - physical_optimizer::PhysicalOptimizerRule, + datasource::{ + listing::PartitionedFile, + physical_plan::{FileScanConfig, FileSource, ParquetSource}, + source::DataSource, + }, + physical_expr::utils::collect_columns, + physical_optimizer::{PhysicalOptimizerRule, pruning::PruningPredicate}, physical_plan::ExecutionPlan, }; @@ -18,6 +29,25 @@ pub use squeeze_hint::SqueezeHintMap; use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnSqueezeHints}; +/// Parameters for the footprint-based admission gate. +/// +/// A parquet scan is routed through LiquidCache only if its estimated in-memory +/// liquid footprint fits the cache budget; oversized scans are left as vanilla +/// parquet reads (which, if the object store is a cached mount, read from it). +/// +/// The estimate multiplies the raw required parquet bytes by `expansion` +/// (parquet -> liquid in-memory blow-up) and `safety` (extra margin); both are +/// `>= 1.0` so the estimate is conservative. The multipliers are applied to the +/// *estimate*, never to the budget, so raising them only makes admission +/// stricter. +#[derive(Debug, Clone, Copy)] +pub struct AdmissionGate { + /// Parquet-bytes -> liquid-in-memory-bytes multiplier (>= 1.0). + pub expansion: f64, + /// Extra safety margin on the estimate (>= 1.0). + pub safety: f64, +} + /// Physical optimizer rule for local mode liquid cache. /// /// Rewrites `DataSourceExec` parquet scans to use [`LiquidParquetSource`], and @@ -26,10 +56,10 @@ use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnSqueezeHint #[derive(Debug)] pub struct LocalModeOptimizer { cache: LiquidCacheParquetRef, - /// When set, parquet scans whose total file size exceeds this threshold - /// are left as vanilla DataFusion reads instead of being wrapped by - /// LiquidCache. `None` means cache every scan. - max_scan_bytes: Option, + /// When set, a scan whose estimated footprint exceeds the cache budget is + /// left as a vanilla parquet read instead of being wrapped by LiquidCache. + /// `None` means cache every scan. + admission: Option, } impl LocalModeOptimizer { @@ -37,7 +67,7 @@ impl LocalModeOptimizer { pub fn new(cache: LiquidCacheParquetRef) -> Self { Self { cache, - max_scan_bytes: None, + admission: None, } } @@ -45,15 +75,16 @@ impl LocalModeOptimizer { pub fn with_cache(cache: LiquidCacheParquetRef) -> Self { Self { cache, - max_scan_bytes: None, + admission: None, } } - /// Set the maximum total file size (in bytes) for a parquet scan to be - /// routed through LiquidCache. Scans exceeding this are read directly - /// from the underlying parquet source, bypassing the cache. - pub fn with_max_scan_bytes(mut self, max_bytes: u64) -> Self { - self.max_scan_bytes = Some(max_bytes); + /// Enable the footprint-based admission gate. A parquet scan is cached only + /// when its estimated liquid footprint (raw required bytes x `expansion` x + /// `safety`) fits the cache's memory budget; otherwise it is read directly + /// from the parquet source, bypassing the cache. See [`AdmissionGate`]. + pub fn with_admission_gate(mut self, expansion: f64, safety: f64) -> Self { + self.admission = Some(AdmissionGate { expansion, safety }); self } } @@ -66,11 +97,14 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { ) -> Result, datafusion::error::DataFusionError> { let analysis = HintAnalyzer::analyze(&plan); let cache = self.cache.clone(); - let max_scan_bytes = self.max_scan_bytes; + let admission = self.admission; + let budget = self.cache.max_memory_bytes() as u64; let mut convert = |node: &Arc, hints: ColumnSqueezeHints| { - // Leave oversized scans as vanilla parquet reads, bypassing the cache. - if let Some(max_bytes) = max_scan_bytes - && parquet_scan_total_bytes(node).is_some_and(|total| total > max_bytes) + // Leave scans whose estimated liquid footprint exceeds the budget as + // vanilla parquet reads, so oversized scans don't thrash the cache. + if let Some(gate) = admission + && let Some((cfg, src)) = parquet_scan_parts(node) + && estimate_liquid_footprint(cfg, src, gate) > budget { return None; } @@ -125,19 +159,110 @@ pub fn rewrite_data_source_plan( rewrite_data_source_plan_with_hints(plan, cache, &ColumnSqueezeHints::default()) } -/// Total size in bytes of all files scanned by a parquet `DataSourceExec`, or -/// `None` if `node` is not such a scan. -fn parquet_scan_total_bytes(node: &Arc) -> Option { - let data_source_exec = node.downcast_ref::()?; - let (file_scan_config, _) = data_source_exec.downcast_to_file_source::()?; - Some( - file_scan_config - .file_groups - .iter() - .flat_map(|g| g.files()) - .map(|f| f.object_meta.size) - .sum(), - ) +/// If `node` is a parquet `DataSourceExec`, return its `FileScanConfig` and +/// `ParquetSource`. +fn parquet_scan_parts( + node: &Arc, +) -> Option<(&FileScanConfig, &ParquetSource)> { + let dse = node.downcast_ref::()?; + dse.downcast_to_file_source::() +} + +/// Estimate the in-memory liquid footprint (bytes) a parquet scan would occupy +/// if cached: the required-column bytes of the files that survive the scan's +/// predicate, scaled by the admission gate's expansion + safety factors. +/// +/// "Required columns" is the output projection **unioned with the predicate +/// columns**, since LiquidCache materializes both. Byte sizing uses only +/// `Exact` per-column sizes; anything else falls back to the whole file size, +/// so the estimate never *under*-counts (which could wrongly admit an oversized +/// scan and thrash the cache). +fn estimate_liquid_footprint(cfg: &FileScanConfig, src: &ParquetSource, gate: AdmissionGate) -> u64 { + let file_schema = cfg.file_schema(); + + let mut required: Vec = { + // Removed in df58; we are pinned to df54. + #[allow(deprecated)] + match cfg.file_column_projection_indices() { + Some(cols) => cols, + None => (0..file_schema.fields().len()).collect(), + } + }; + if let Some(pred) = src.filter() { + for col in collect_columns(&pred) { + required.push(col.index()); + } + } + required.sort_unstable(); + required.dedup(); + + let files: Vec<&PartitionedFile> = cfg.file_groups.iter().flat_map(|g| g.files()).collect(); + if files.is_empty() { + return 0; + } + + let surviving = surviving_files(src, file_schema, &files); + + let raw: u64 = files + .iter() + .zip(surviving.iter()) + .filter(|(_, keep)| **keep) + .map(|(f, _)| { + file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required) + }) + .sum(); + + ((raw as f64) * gate.expansion.max(1.0) * gate.safety.max(1.0)) as u64 +} + +/// Boolean per file: `true` if the file may match the predicate (keep it), +/// `false` if the predicate's stats prove it cannot match (prune it). +/// Conservative: no predicate, missing stats, or any pruning error keeps files. +fn surviving_files( + src: &ParquetSource, + file_schema: &SchemaRef, + files: &[&PartitionedFile], +) -> Vec { + let Some(pred) = src.filter() else { + return vec![true; files.len()]; + }; + let pruning = match PruningPredicate::try_new(pred, file_schema.clone()) { + Ok(p) => p, + Err(_) => return vec![true; files.len()], + }; + let stats: Vec> = files + .iter() + .map(|f| { + f.statistics + .clone() + .unwrap_or_else(|| Arc::new(Statistics::new_unknown(file_schema))) + }) + .collect(); + let prunable = PrunableStatistics::new(stats, file_schema.clone()); + match pruning.prune(&prunable) { + Ok(mask) if mask.len() == files.len() => mask, + _ => vec![true; files.len()], + } +} + +/// Bytes the `required` columns of one file contribute to the footprint. +/// +/// Uses only `Exact` per-column byte sizes. If the file has no stats, or any +/// required column lacks an exact size, fall back to the whole-file size — a +/// deliberate over-estimate, since under-counting could wrongly admit an +/// oversized scan. +fn file_required_bytes(stats: Option<&Statistics>, object_size: u64, required: &[usize]) -> u64 { + let Some(stats) = stats else { + return object_size; + }; + let mut sum: u64 = 0; + for &c in required { + match stats.column_statistics.get(c).map(|cs| &cs.byte_size) { + Some(Precision::Exact(n)) => sum += *n as u64, + _ => return object_size, + } + } + sum } /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent @@ -228,6 +353,10 @@ mod tests { } async fn build_cache() -> LiquidCacheParquetRef { + build_cache_with_budget(1_000_000).await + } + + async fn build_cache_with_budget(max_memory_bytes: usize) -> LiquidCacheParquetRef { let tmp_dir = tempfile::tempdir().unwrap(); let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) .await @@ -235,7 +364,7 @@ mod tests { Arc::new( LiquidCacheParquet::new( 8192, - 1000000, + max_memory_bytes, usize::MAX, store, Box::new(LiquidPolicy::new()), @@ -266,11 +395,10 @@ mod tests { found } - /// A parquet scan whose total file size exceeds `max_scan_bytes` is left as - /// a plain `ParquetSource`; one under the threshold is still wrapped in - /// `LiquidParquetSource`. + /// The admission gate bypasses a scan whose estimated footprint exceeds the + /// budget (large expansion here forces that), and caches one that fits. #[tokio::test] - async fn test_max_scan_bytes_pass_through() { + async fn test_admission_gate_pass_through() { let ctx = SessionContext::new(); ctx.register_parquet( "nano_hits", @@ -288,24 +416,89 @@ mod tests { .unwrap(); let config = ConfigOptions::default(); - // Threshold below the file size → scan bypasses the cache. + // A huge expansion inflates the estimate past the 1 MB budget → bypass. let capped = LocalModeOptimizer::new(build_cache().await) - .with_max_scan_bytes(1) + .with_admission_gate(1e9, 1.0) .optimize(plan.clone(), &config) .unwrap(); assert!( !has_liquid_source(&capped), - "oversized scan should stay a plain ParquetSource" + "oversized estimate should stay a plain ParquetSource" ); - // Threshold above the file size → scan is cached as usual. - let uncapped = LocalModeOptimizer::new(build_cache().await) - .with_max_scan_bytes(u64::MAX) + // With a large budget the footprint fits (expansion 1.0) → cached. + let uncapped = LocalModeOptimizer::new(build_cache_with_budget(usize::MAX).await) + .with_admission_gate(1.0, 1.0) .optimize(plan, &config) .unwrap(); assert!( has_liquid_source(&uncapped), - "under-threshold scan should be wrapped in LiquidParquetSource" + "fitting scan should be wrapped in LiquidParquetSource" ); } } + +/// Pure unit tests for the footprint byte-math (no cache / no t4 mount, so they +/// run everywhere). These cover the correctness-critical fallback direction: +/// under-counting could wrongly admit an oversized scan, so anything not backed +/// by an `Exact` per-column byte size must fall back to the whole-file size. +#[cfg(test)] +mod footprint_tests { + use super::file_required_bytes; + use datafusion::common::{ColumnStatistics, Statistics, stats::Precision}; + + fn col(byte_size: Precision) -> ColumnStatistics { + let mut c = ColumnStatistics::new_unknown(); + c.byte_size = byte_size; + c + } + + fn stats(cols: Vec) -> Statistics { + Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: cols, + } + } + + #[test] + fn no_stats_falls_back_to_full_file() { + assert_eq!(file_required_bytes(None, 5000, &[0, 1]), 5000); + } + + #[test] + fn sums_only_required_exact_columns() { + let s = stats(vec![ + col(Precision::Exact(100)), + col(Precision::Exact(200)), + col(Precision::Exact(400)), + ]); + // required = cols 0 and 2 → 100 + 400, ignoring col 1. + assert_eq!(file_required_bytes(Some(&s), 9999, &[0, 2]), 500); + } + + #[test] + fn inexact_required_column_falls_back_to_full_file() { + let s = stats(vec![ + col(Precision::Exact(100)), + col(Precision::Inexact(200)), + ]); + // col 1 inexact → not a safe upper bound → whole file. + assert_eq!(file_required_bytes(Some(&s), 7000, &[0, 1]), 7000); + // col 0 alone is exact → its bytes only. + assert_eq!(file_required_bytes(Some(&s), 7000, &[0]), 100); + } + + #[test] + fn missing_required_column_falls_back_to_full_file() { + let s = stats(vec![col(Precision::Exact(100))]); + // required col 5 doesn't exist → conservative whole file. + assert_eq!(file_required_bytes(Some(&s), 3000, &[0, 5]), 3000); + } + + #[test] + fn absent_byte_size_falls_back_to_full_file() { + let s = stats(vec![col(Precision::Absent)]); + assert_eq!(file_required_bytes(Some(&s), 2000, &[0]), 2000); + } +} From 2f067176dce49c922c1936a96b6ce10f902f8b4b Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 00:52:48 +0530 Subject: [PATCH 04/13] fix(admission): handle partition columns in footprint gate (codex review) - Prune with the full table schema (file + partition columns) so predicates on partition columns resolve; PartitionedFile stats are table-schema too. - Exclude partition columns from byte accounting -- they are literals, never materialized in LiquidCache; only file columns contribute to the footprint. - Guard per-file stats width against the schema to avoid a mismatch; a mismatched/absent file is treated as unknown (kept, not pruned). --- src/datafusion/src/optimizers/mod.rs | 40 +++++++++++++++++++--------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 78369030..82363f27 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -178,19 +178,29 @@ fn parquet_scan_parts( /// so the estimate never *under*-counts (which could wrongly admit an oversized /// scan and thrash the cache). fn estimate_liquid_footprint(cfg: &FileScanConfig, src: &ParquetSource, gate: AdmissionGate) -> u64 { - let file_schema = cfg.file_schema(); + let num_file_cols = cfg.file_schema().fields().len(); + // Full table schema (file + partition columns). The pushed-down predicate + // and `PartitionedFile` stats are expressed against it, so file pruning must + // use it; byte accounting still charges only file columns. + let table_schema = cfg.file_source.table_schema().table_schema().clone(); let mut required: Vec = { - // Removed in df58; we are pinned to df54. + // Removed in df58; we are pinned to df54. Returns file-column indices. #[allow(deprecated)] match cfg.file_column_projection_indices() { Some(cols) => cols, - None => (0..file_schema.fields().len()).collect(), + None => (0..num_file_cols).collect(), } }; if let Some(pred) = src.filter() { for col in collect_columns(&pred) { - required.push(col.index()); + let idx = col.index(); + // Partition columns are appended after file columns in the table + // schema and are literals (never materialized in LiquidCache), so + // only file columns contribute to the footprint. + if idx < num_file_cols { + required.push(idx); + } } } required.sort_unstable(); @@ -201,7 +211,7 @@ fn estimate_liquid_footprint(cfg: &FileScanConfig, src: &ParquetSource, gate: Ad return 0; } - let surviving = surviving_files(src, file_schema, &files); + let surviving = surviving_files(src, &table_schema, &files); let raw: u64 = files .iter() @@ -217,28 +227,32 @@ fn estimate_liquid_footprint(cfg: &FileScanConfig, src: &ParquetSource, gate: Ad /// Boolean per file: `true` if the file may match the predicate (keep it), /// `false` if the predicate's stats prove it cannot match (prune it). -/// Conservative: no predicate, missing stats, or any pruning error keeps files. +/// Conservative: no predicate, missing/mismatched stats, or any pruning error +/// keeps files. `table_schema` (file + partition columns) is used so predicates +/// on partition columns resolve, matching `PartitionedFile::statistics`. fn surviving_files( src: &ParquetSource, - file_schema: &SchemaRef, + table_schema: &SchemaRef, files: &[&PartitionedFile], ) -> Vec { let Some(pred) = src.filter() else { return vec![true; files.len()]; }; - let pruning = match PruningPredicate::try_new(pred, file_schema.clone()) { + let pruning = match PruningPredicate::try_new(pred, table_schema.clone()) { Ok(p) => p, Err(_) => return vec![true; files.len()], }; + let expected = table_schema.fields().len(); let stats: Vec> = files .iter() - .map(|f| { - f.statistics - .clone() - .unwrap_or_else(|| Arc::new(Statistics::new_unknown(file_schema))) + .map(|f| match &f.statistics { + // Only trust stats whose width matches the table schema; otherwise + // treat as unknown (kept, not pruned) to avoid a schema mismatch. + Some(s) if s.column_statistics.len() == expected => s.clone(), + _ => Arc::new(Statistics::new_unknown(table_schema)), }) .collect(); - let prunable = PrunableStatistics::new(stats, file_schema.clone()); + let prunable = PrunableStatistics::new(stats, table_schema.clone()); match pruning.prune(&prunable) { Ok(mask) if mask.len() == files.len() => mask, _ => vec![true; files.len()], From 920663a2d7ef17377d06bf6b6fb6faceb37e9bff Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 09:13:57 +0530 Subject: [PATCH 05/13] feat(admission): add overcommit tolerance to footprint gate Deny only when estimated liquid footprint exceeds budget x tolerance, not budget x1. LiquidCache compacts in RAM and beats the fallback mount until ~5x over budget, so gating at 1x over-bypasses a whole band of scans that would cache and win. Compare as a finite f64 pressure ratio (footprint/budget > tolerance) with a saturating raw-byte sum and an explicit budget==0 case. tolerance is the one relaxing knob: clamped to [1.0, 5.0], defaulting to 3.0, non-finite rejected. --- src/datafusion-local/src/lib.rs | 21 +++--- src/datafusion/src/optimizers/mod.rs | 103 ++++++++++++++++++++------- 2 files changed, 88 insertions(+), 36 deletions(-) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index aa8a0c2e..7342340a 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -69,9 +69,10 @@ pub struct LiquidCacheLocalBuilder { squeeze_policy: Box, /// Hydration policy hydration_policy: Box, - /// Footprint-based admission gate `(expansion, safety)`. When set, a scan is - /// cached only if its estimated liquid footprint fits the memory budget. - admission: Option<(f64, f64)>, + /// Footprint-based admission gate `(expansion, safety, tolerance)`. When set, + /// a scan is cached only if its estimated liquid footprint stays within + /// `budget × tolerance`. + admission: Option<(f64, f64, f64)>, span: fastrace::Span, } @@ -151,10 +152,12 @@ impl LiquidCacheLocalBuilder { /// Enable the footprint-based admission gate. A scan is cached only when its /// estimated liquid footprint (raw required bytes x `expansion` x `safety`) - /// fits the cache's memory budget; oversized scans are read directly from the - /// parquet source, bypassing the cache. `expansion` and `safety` are `>= 1.0`. - pub fn with_admission_gate(mut self, expansion: f64, safety: f64) -> Self { - self.admission = Some((expansion, safety)); + /// stays within `budget × tolerance`; larger scans are read directly from the + /// parquet source, bypassing the cache. `expansion`/`safety` are `>= 1.0` + /// (inflate the estimate); `tolerance` is `>= 1.0` (overcommit the budget, + /// clamped to the measured ~5x compaction crossover). + pub fn with_admission_gate(mut self, expansion: f64, safety: f64, tolerance: f64) -> Self { + self.admission = Some((expansion, safety, tolerance)); self } @@ -204,8 +207,8 @@ impl LiquidCacheLocalBuilder { let cache_ref = Arc::new(cache); let mut optimizer = LocalModeOptimizer::new(cache_ref.clone()); - if let Some((expansion, safety)) = self.admission { - optimizer = optimizer.with_admission_gate(expansion, safety); + if let Some((expansion, safety, tolerance)) = self.admission { + optimizer = optimizer.with_admission_gate(expansion, safety, tolerance); } let state = datafusion::execution::SessionStateBuilder::new() diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 82363f27..9e15069e 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -32,20 +32,29 @@ use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnSqueezeHint /// Parameters for the footprint-based admission gate. /// /// A parquet scan is routed through LiquidCache only if its estimated in-memory -/// liquid footprint fits the cache budget; oversized scans are left as vanilla -/// parquet reads (which, if the object store is a cached mount, read from it). +/// liquid footprint stays within `budget × tolerance`; larger scans are left as +/// vanilla parquet reads (which, if the object store is a cached mount, read +/// from it). /// /// The estimate multiplies the raw required parquet bytes by `expansion` /// (parquet -> liquid in-memory blow-up) and `safety` (extra margin); both are -/// `>= 1.0` so the estimate is conservative. The multipliers are applied to the -/// *estimate*, never to the budget, so raising them only makes admission -/// stricter. +/// `>= 1.0` so the estimate is conservative (over-counts). `tolerance` (`>= 1.0`) +/// is the one knob applied to the *budget*: it encodes that LiquidCache compacts +/// in RAM and still beats the fallback mount until footprint reaches ~5x the +/// budget, so the gate tolerates that overcommit before bypassing. #[derive(Debug, Clone, Copy)] pub struct AdmissionGate { - /// Parquet-bytes -> liquid-in-memory-bytes multiplier (>= 1.0). + /// Parquet-bytes -> liquid-in-memory-bytes multiplier (>= 1.0). Inflates the + /// estimate (conservative direction). pub expansion: f64, - /// Extra safety margin on the estimate (>= 1.0). + /// Extra safety margin on the estimate (>= 1.0). Conservative direction. pub safety: f64, + /// Overcommit tolerance: how far the estimated liquid footprint may exceed + /// the budget before the scan is denied, in multiples of the budget. + /// LiquidCache compacts in RAM up to ~5x over budget before it starts + /// thrashing, so caching still wins in that band. This is the one *relaxing* + /// knob, so it is clamped to `[1.0, 5.0]` (5.0 = the measured crossover). + pub tolerance: f64, } /// Physical optimizer rule for local mode liquid cache. @@ -81,10 +90,27 @@ impl LocalModeOptimizer { /// Enable the footprint-based admission gate. A parquet scan is cached only /// when its estimated liquid footprint (raw required bytes x `expansion` x - /// `safety`) fits the cache's memory budget; otherwise it is read directly + /// `safety`) stays within `budget × tolerance`; otherwise it is read directly /// from the parquet source, bypassing the cache. See [`AdmissionGate`]. - pub fn with_admission_gate(mut self, expansion: f64, safety: f64) -> Self { - self.admission = Some(AdmissionGate { expansion, safety }); + /// + /// Inputs are sanitized so a misconfigured value can never make the gate + /// unsound: `expansion`/`safety` are forced finite and `>= 1.0` (their only + /// effect is to inflate the estimate, the conservative direction), and + /// `tolerance` — the one *relaxing* knob — is forced finite and clamped to + /// `[1.0, 5.0]` (5.0 = the measured compaction crossover), defaulting to 3.0 + /// if non-finite. + pub fn with_admission_gate(mut self, expansion: f64, safety: f64, tolerance: f64) -> Self { + let estimate_factor = |v: f64| if v.is_finite() && v >= 1.0 { v } else { 1.0 }; + let tolerance = if tolerance.is_finite() { + tolerance.clamp(1.0, 5.0) + } else { + 3.0 + }; + self.admission = Some(AdmissionGate { + expansion: estimate_factor(expansion), + safety: estimate_factor(safety), + tolerance, + }); self } } @@ -104,7 +130,7 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { // vanilla parquet reads, so oversized scans don't thrash the cache. if let Some(gate) = admission && let Some((cfg, src)) = parquet_scan_parts(node) - && estimate_liquid_footprint(cfg, src, gate) > budget + && should_bypass(cfg, src, gate, budget) { return None; } @@ -161,23 +187,24 @@ pub fn rewrite_data_source_plan( /// If `node` is a parquet `DataSourceExec`, return its `FileScanConfig` and /// `ParquetSource`. -fn parquet_scan_parts( - node: &Arc, -) -> Option<(&FileScanConfig, &ParquetSource)> { +fn parquet_scan_parts(node: &Arc) -> Option<(&FileScanConfig, &ParquetSource)> { let dse = node.downcast_ref::()?; dse.downcast_to_file_source::() } -/// Estimate the in-memory liquid footprint (bytes) a parquet scan would occupy -/// if cached: the required-column bytes of the files that survive the scan's -/// predicate, scaled by the admission gate's expansion + safety factors. +/// Estimate the raw required-column parquet bytes a scan reads: the sum, over +/// the files that survive the scan's predicate, of the byte sizes of the columns +/// it materializes. This is the byte-accurate, filter-aware size the admission +/// decision is built on (see [`should_bypass`]); the expansion/safety/tolerance +/// factors are applied there, not here. /// /// "Required columns" is the output projection **unioned with the predicate /// columns**, since LiquidCache materializes both. Byte sizing uses only /// `Exact` per-column sizes; anything else falls back to the whole file size, /// so the estimate never *under*-counts (which could wrongly admit an oversized -/// scan and thrash the cache). -fn estimate_liquid_footprint(cfg: &FileScanConfig, src: &ParquetSource, gate: AdmissionGate) -> u64 { +/// scan and thrash the cache). The sum saturates rather than overflowing on +/// pathological file lists. +fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> u64 { let num_file_cols = cfg.file_schema().fields().len(); // Full table schema (file + partition columns). The pushed-down predicate // and `PartitionedFile` stats are expressed against it, so file pruning must @@ -213,16 +240,38 @@ fn estimate_liquid_footprint(cfg: &FileScanConfig, src: &ParquetSource, gate: Ad let surviving = surviving_files(src, &table_schema, &files); - let raw: u64 = files + files .iter() .zip(surviving.iter()) .filter(|(_, keep)| **keep) - .map(|(f, _)| { - file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required) - }) - .sum(); + .map(|(f, _)| file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required)) + .fold(0u64, u64::saturating_add) +} - ((raw as f64) * gate.expansion.max(1.0) * gate.safety.max(1.0)) as u64 +/// Decide whether a parquet scan should bypass the cache (read as vanilla +/// parquet) rather than be transcoded into LiquidCache. +/// +/// The scan's estimated liquid footprint is `raw × expansion × safety`, where +/// `raw` is [`estimate_required_bytes`]. It bypasses when that footprint exceeds +/// `budget × tolerance` — equivalently, when the pressure ratio +/// `footprint / budget` exceeds `tolerance`. The comparison is done in finite +/// `f64` space to avoid the overflow/truncation of integer `budget × tolerance`. +/// +/// `budget == 0` (cache disabled / unsized) bypasses any non-empty scan and +/// admits only zero-footprint ones. +fn should_bypass( + cfg: &FileScanConfig, + src: &ParquetSource, + gate: AdmissionGate, + budget: u64, +) -> bool { + let raw = estimate_required_bytes(cfg, src); + // Multipliers are already sanitized to finite, >= 1.0 in `with_admission_gate`. + let footprint = (raw as f64) * gate.expansion * gate.safety; + if budget == 0 { + return footprint > 0.0; + } + footprint / (budget as f64) > gate.tolerance } /// Boolean per file: `true` if the file may match the predicate (keep it), @@ -432,7 +481,7 @@ mod tests { // A huge expansion inflates the estimate past the 1 MB budget → bypass. let capped = LocalModeOptimizer::new(build_cache().await) - .with_admission_gate(1e9, 1.0) + .with_admission_gate(1e9, 1.0, 1.0) .optimize(plan.clone(), &config) .unwrap(); assert!( @@ -442,7 +491,7 @@ mod tests { // With a large budget the footprint fits (expansion 1.0) → cached. let uncapped = LocalModeOptimizer::new(build_cache_with_budget(usize::MAX).await) - .with_admission_gate(1.0, 1.0) + .with_admission_gate(1.0, 1.0, 1.0) .optimize(plan, &config) .unwrap(); assert!( From 8f3f98a8bb8a96c60f48d68d293bff927f8ca635 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 09:34:58 +0530 Subject: [PATCH 06/13] fix(admission): use non-panicking column extraction in footprint gate file_column_projection_indices() delegates to the deprecated ProjectionExprs::ordered_column_indices(), which .expect()-panics on any non-Column projection expression. DataFusion pushes non-column exprs (get_field on struct columns, via enable_leaf_expression_pushdown) into the scan projection, so the gate panicked and killed the query. Use src.projection().column_indices() (collect_columns-based, handles arbitrary exprs) instead, matching the reader's own opener. Add a struct-column get_field regression test that runs without a cache mount. --- src/datafusion/src/optimizers/mod.rs | 109 +++++++++++++++++++++++++-- 1 file changed, 102 insertions(+), 7 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 9e15069e..16c681c1 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -211,13 +211,20 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> u64 { // use it; byte accounting still charges only file columns. let table_schema = cfg.file_source.table_schema().table_schema().clone(); - let mut required: Vec = { - // Removed in df58; we are pinned to df54. Returns file-column indices. - #[allow(deprecated)] - match cfg.file_column_projection_indices() { - Some(cols) => cols, - None => (0..num_file_cols).collect(), - } + // Columns the scan projects. `column_indices()` collects the source columns + // referenced by each projection expression, so it is correct for compound + // projections (e.g. `a * b` reads a and b) and — unlike the deprecated + // `FileScanConfig::file_column_projection_indices` / + // `ProjectionExprs::ordered_column_indices` — does not panic on a + // non-column projection expression. Indices are table-schema-relative; keep + // only file columns (partition columns are literals, never materialized). + let mut required: Vec = match src.projection() { + Some(p) => p + .column_indices() + .into_iter() + .filter(|&i| i < num_file_cols) + .collect(), + None => (0..num_file_cols).collect(), }; if let Some(pred) = src.filter() { for col in collect_columns(&pred) { @@ -397,6 +404,94 @@ mod tests { .unwrap(); } + /// Regression: a `get_field` on a struct column is pushed into the scan + /// projection as a non-`Column` expression (via + /// `enable_leaf_expression_pushdown`). The footprint gate must estimate such + /// a scan without panicking (the old code called the deprecated + /// `ordered_column_indices`, which `.expect`s a bare column and killed the + /// query). Runs everywhere: it builds a physical plan and calls the estimate + /// directly, so it needs no cache / t4 mount. + #[tokio::test] + async fn estimate_survives_non_column_scan_projection() { + use arrow::array::{ArrayRef, Int64Array, RecordBatch, StructArray}; + use arrow_schema::{DataType, Field, Fields}; + use datafusion::physical_expr::expressions::Column; + use parquet::arrow::ArrowWriter; + + // A struct column `s {a, b}` plus a flat column `p`. + let struct_fields = Fields::from(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ]); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("s", DataType::Struct(struct_fields.clone()), false), + Field::new("p", DataType::Int64, false), + ])); + let s = StructArray::new( + struct_fields, + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef, + Arc::new(Int64Array::from(vec![4, 5, 6])) as ArrayRef, + ], + None, + ); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(s) as ArrayRef, + Arc::new(Int64Array::from(vec![7, 8, 9])), + ], + ) + .unwrap(); + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("structs.parquet"); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let ctx = SessionContext::new(); + ctx.register_parquet("t", path.to_str().unwrap(), Default::default()) + .await + .unwrap(); + let plan = ctx + .sql("SELECT s.a FROM t") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + + // Locate the parquet scan and assert its projection really does carry a + // non-column expression — otherwise the test wouldn't exercise the bug. + let mut parts = None; + plan.apply(|node| { + if let Some((cfg, src)) = parquet_scan_parts(node) { + let has_expr = src + .projection() + .map(|p| p.iter().any(|e| e.expr.downcast_ref::().is_none())) + .unwrap_or(false); + assert!( + has_expr, + "expected a non-column scan projection to exercise the gate; \ + if this fails, DataFusion changed leaf-expression pushdown" + ); + parts = Some((cfg.clone(), src.clone())); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + let (cfg, src) = parts.expect("no parquet scan in plan"); + + // The call that used to panic. It must return a finite estimate; `s` + // (the struct column read by `s.a`) is the one required file column. + let bytes = estimate_required_bytes(&cfg, &src); + // Whole-file fallback (no per-column Exact stats here) → non-zero, finite. + assert!(bytes > 0, "estimate should be a real byte count"); + } + #[tokio::test] async fn test_plan_rewrite() { let ctx = SessionContext::new(); From cfa05a894b45c2d62c70aeb2eb974b7727a8c50d Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 10:37:34 +0530 Subject: [PATCH 07/13] fix(admission): count Inexact per-column byte sizes in footprint gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gate counted only Precision::Exact per-column byte_size and fell back to the whole-file size otherwise. DuckLake (the target backend) always labels byte_size Inexact — it is the real recorded compressed column size, marked Inexact because catalog stats can go stale after deletes/compaction — so the fallback fired on every scan, inflating the estimate ~50-100x (charging all columns for a single-column read) and leaving the cache near-empty. Count Inexact sizes too; only Absent (or missing file stats) still falls back to whole-file. expansion/safety remain the margin against stale-low drift. Sum with saturating_add. --- src/datafusion/src/optimizers/mod.rs | 59 +++++++++++++++++++--------- 1 file changed, 41 insertions(+), 18 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 16c681c1..546a7b16 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -199,11 +199,11 @@ fn parquet_scan_parts(node: &Arc) -> Option<(&FileScanConfig, /// factors are applied there, not here. /// /// "Required columns" is the output projection **unioned with the predicate -/// columns**, since LiquidCache materializes both. Byte sizing uses only -/// `Exact` per-column sizes; anything else falls back to the whole file size, -/// so the estimate never *under*-counts (which could wrongly admit an oversized -/// scan and thrash the cache). The sum saturates rather than overflowing on -/// pathological file lists. +/// columns**, since LiquidCache materializes both. Byte sizing uses `Exact` or +/// `Inexact` per-column sizes (DuckLake records real column sizes but labels +/// them `Inexact`); a column with an `Absent` size or a file with no stats +/// falls back to the whole file size. The sum saturates rather than overflowing +/// on pathological file lists. fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> u64 { let num_file_cols = cfg.file_schema().fields().len(); // Full table schema (file + partition columns). The pushed-down predicate @@ -317,10 +317,18 @@ fn surviving_files( /// Bytes the `required` columns of one file contribute to the footprint. /// -/// Uses only `Exact` per-column byte sizes. If the file has no stats, or any -/// required column lacks an exact size, fall back to the whole-file size — a -/// deliberate over-estimate, since under-counting could wrongly admit an -/// oversized scan. +/// Uses per-column byte sizes that are either `Exact` or `Inexact`. DuckLake +/// records the real compressed on-disk column size but always labels it +/// `Inexact` (catalog stats can go stale after deletes/compaction), so +/// rejecting `Inexact` would make the gate fall back to the whole-file size on +/// *every* DuckLake scan — charging all columns for a single-column read and +/// bypassing everything. `Inexact` is a real measurement, not a guess; the +/// caller's `expansion`/`safety` margin absorbs modest drift, and even a large +/// stale-low under-count only risks a too-eager admit (perf), never wrong +/// results. +/// +/// If the file has no stats, or any required column's size is `Absent`, fall +/// back to the whole-file size — a deliberate over-estimate. fn file_required_bytes(stats: Option<&Statistics>, object_size: u64, required: &[usize]) -> u64 { let Some(stats) = stats else { return object_size; @@ -328,7 +336,9 @@ fn file_required_bytes(stats: Option<&Statistics>, object_size: u64, required: & let mut sum: u64 = 0; for &c in required { match stats.column_statistics.get(c).map(|cs| &cs.byte_size) { - Some(Precision::Exact(n)) => sum += *n as u64, + Some(Precision::Exact(n) | Precision::Inexact(n)) => { + sum = sum.saturating_add(*n as u64) + } _ => return object_size, } } @@ -597,9 +607,9 @@ mod tests { } /// Pure unit tests for the footprint byte-math (no cache / no t4 mount, so they -/// run everywhere). These cover the correctness-critical fallback direction: -/// under-counting could wrongly admit an oversized scan, so anything not backed -/// by an `Exact` per-column byte size must fall back to the whole-file size. +/// run everywhere). These cover the fallback direction: a column with an +/// `Absent` size (or a file with no stats) falls back to the whole-file size, +/// while `Exact`/`Inexact` per-column sizes are counted. #[cfg(test)] mod footprint_tests { use super::file_required_bytes; @@ -636,15 +646,28 @@ mod footprint_tests { } #[test] - fn inexact_required_column_falls_back_to_full_file() { + fn inexact_required_column_is_counted() { + let s = stats(vec![ + col(Precision::Exact(100)), + col(Precision::Inexact(200)), + ]); + // Inexact is a real (possibly-stale) size, so it counts. DuckLake always + // labels byte_size Inexact; rejecting it would bypass every scan. + assert_eq!(file_required_bytes(Some(&s), 7000, &[0, 1]), 300); + assert_eq!(file_required_bytes(Some(&s), 7000, &[1]), 200); + } + + #[test] + fn absent_among_sized_columns_falls_back_to_full_file() { let s = stats(vec![ col(Precision::Exact(100)), col(Precision::Inexact(200)), + col(Precision::Absent), ]); - // col 1 inexact → not a safe upper bound → whole file. - assert_eq!(file_required_bytes(Some(&s), 7000, &[0, 1]), 7000); - // col 0 alone is exact → its bytes only. - assert_eq!(file_required_bytes(Some(&s), 7000, &[0]), 100); + // Any Absent required column → whole file, even mixed with sized ones. + assert_eq!(file_required_bytes(Some(&s), 8000, &[0, 1, 2]), 8000); + // Dropping the absent column, Exact + Inexact are summed. + assert_eq!(file_required_bytes(Some(&s), 8000, &[0, 1]), 300); } #[test] From 39c69861ff8890d64fb83a29c2f26ce079c2e85d Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 11:06:52 +0530 Subject: [PATCH 08/13] feat(admission): fail-safe the footprint gate against panics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gate runs in the physical optimizer on every query, so a panic in footprint estimation (e.g. a DataFusion API that panics on an unusual plan shape) aborts the whole query. Since the gate is a pure performance optimization — caching a scan or not yields identical results — guard the decision with catch_unwind: on panic, log and cache the scan normally instead of taking down the query. --- src/datafusion/src/optimizers/mod.rs | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 546a7b16..09768aa9 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -130,7 +130,7 @@ impl PhysicalOptimizerRule for LocalModeOptimizer { // vanilla parquet reads, so oversized scans don't thrash the cache. if let Some(gate) = admission && let Some((cfg, src)) = parquet_scan_parts(node) - && should_bypass(cfg, src, gate, budget) + && should_bypass_guarded(cfg, src, gate, budget) { return None; } @@ -281,6 +281,32 @@ fn should_bypass( footprint / (budget as f64) > gate.tolerance } +/// [`should_bypass`], guarded against panics. +/// +/// The gate is a pure performance optimization: caching a scan or reading it +/// as vanilla parquet yields identical results. So if footprint estimation ever +/// panics — e.g. a DataFusion API that panics on an unusual plan shape, the +/// class of bug that `ordered_column_indices` was — we must not let it abort the +/// query. Catch it, log it (so it is observable, not silently swallowed), and +/// fall back to caching the scan normally, exactly as if the gate were off. +fn should_bypass_guarded( + cfg: &FileScanConfig, + src: &ParquetSource, + gate: AdmissionGate, + budget: u64, +) -> bool { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + should_bypass(cfg, src, gate, budget) + })) + .unwrap_or_else(|_| { + log::warn!( + "liquid-cache admission gate panicked during footprint estimation; \ + caching scan normally" + ); + false + }) +} + /// Boolean per file: `true` if the file may match the predicate (keep it), /// `false` if the predicate's stats prove it cannot match (prune it). /// Conservative: no predicate, missing/mismatched stats, or any pruning error From 28c2fe5e57b78ebfe5020aca41f474f1abf719d0 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 11:18:13 +0530 Subject: [PATCH 09/13] feat(admission): strict fail-loud toggle for the gate panic guard Add a strict flag to the admission gate. Either way the panic is caught and logged at ERROR with its message (never silently swallowed): - strict: log, advise disabling the flag to keep queries running, then re-raise so the query aborts and the bug surfaces immediately. - non-strict: log, advise enabling the flag to fail loud, then cache the scan normally so the query survives. The caller (runtimedb) defaults strict on. --- src/datafusion-local/src/lib.rs | 26 ++++++---- src/datafusion/src/optimizers/mod.rs | 78 +++++++++++++++++++++------- 2 files changed, 76 insertions(+), 28 deletions(-) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 7342340a..ebb43bb2 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -69,10 +69,10 @@ pub struct LiquidCacheLocalBuilder { squeeze_policy: Box, /// Hydration policy hydration_policy: Box, - /// Footprint-based admission gate `(expansion, safety, tolerance)`. When set, - /// a scan is cached only if its estimated liquid footprint stays within - /// `budget × tolerance`. - admission: Option<(f64, f64, f64)>, + /// Footprint-based admission gate `(expansion, safety, tolerance, strict)`. + /// When set, a scan is cached only if its estimated liquid footprint stays + /// within `budget × tolerance`; `strict` toggles fail-loud panic handling. + admission: Option<(f64, f64, f64, bool)>, span: fastrace::Span, } @@ -155,9 +155,17 @@ impl LiquidCacheLocalBuilder { /// stays within `budget × tolerance`; larger scans are read directly from the /// parquet source, bypassing the cache. `expansion`/`safety` are `>= 1.0` /// (inflate the estimate); `tolerance` is `>= 1.0` (overcommit the budget, - /// clamped to the measured ~5x compaction crossover). - pub fn with_admission_gate(mut self, expansion: f64, safety: f64, tolerance: f64) -> Self { - self.admission = Some((expansion, safety, tolerance)); + /// clamped to the measured ~5x compaction crossover). `strict == true` lets a + /// footprint-estimation panic abort the query (fail loud); `false` catches it + /// and caches the scan normally. + pub fn with_admission_gate( + mut self, + expansion: f64, + safety: f64, + tolerance: f64, + strict: bool, + ) -> Self { + self.admission = Some((expansion, safety, tolerance, strict)); self } @@ -207,8 +215,8 @@ impl LiquidCacheLocalBuilder { let cache_ref = Arc::new(cache); let mut optimizer = LocalModeOptimizer::new(cache_ref.clone()); - if let Some((expansion, safety, tolerance)) = self.admission { - optimizer = optimizer.with_admission_gate(expansion, safety, tolerance); + if let Some((expansion, safety, tolerance, strict)) = self.admission { + optimizer = optimizer.with_admission_gate(expansion, safety, tolerance, strict); } let state = datafusion::execution::SessionStateBuilder::new() diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 09768aa9..56460369 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -55,6 +55,11 @@ pub struct AdmissionGate { /// thrashing, so caching still wins in that band. This is the one *relaxing* /// knob, so it is clamped to `[1.0, 5.0]` (5.0 = the measured crossover). pub tolerance: f64, + /// Fail-loud mode. When `true`, a panic during footprint estimation is *not* + /// caught — it aborts the query — so estimation bugs surface immediately. + /// When `false`, the panic is caught, logged at ERROR, and the scan is cached + /// normally so the query survives. The caller chooses the default. + pub strict: bool, } /// Physical optimizer rule for local mode liquid cache. @@ -98,8 +103,15 @@ impl LocalModeOptimizer { /// effect is to inflate the estimate, the conservative direction), and /// `tolerance` — the one *relaxing* knob — is forced finite and clamped to /// `[1.0, 5.0]` (5.0 = the measured compaction crossover), defaulting to 3.0 - /// if non-finite. - pub fn with_admission_gate(mut self, expansion: f64, safety: f64, tolerance: f64) -> Self { + /// if non-finite. `strict` (see [`AdmissionGate::strict`]) turns off the + /// panic guard so estimation bugs fail loud. + pub fn with_admission_gate( + mut self, + expansion: f64, + safety: f64, + tolerance: f64, + strict: bool, + ) -> Self { let estimate_factor = |v: f64| if v.is_finite() && v >= 1.0 { v } else { 1.0 }; let tolerance = if tolerance.is_finite() { tolerance.clamp(1.0, 5.0) @@ -110,6 +122,7 @@ impl LocalModeOptimizer { expansion: estimate_factor(expansion), safety: estimate_factor(safety), tolerance, + strict, }); self } @@ -281,30 +294,57 @@ fn should_bypass( footprint / (budget as f64) > gate.tolerance } -/// [`should_bypass`], guarded against panics. +/// [`should_bypass`], with an optional panic guard. /// -/// The gate is a pure performance optimization: caching a scan or reading it -/// as vanilla parquet yields identical results. So if footprint estimation ever +/// The gate is a pure performance optimization: caching a scan or reading it as +/// vanilla parquet yields identical results. So if footprint estimation ever /// panics — e.g. a DataFusion API that panics on an unusual plan shape, the -/// class of bug that `ordered_column_indices` was — we must not let it abort the -/// query. Catch it, log it (so it is observable, not silently swallowed), and -/// fall back to caching the scan normally, exactly as if the gate were off. +/// class of bug that `ordered_column_indices` was — a non-strict gate must not +/// let it abort the query. +/// +/// Either way the panic is caught and logged at ERROR with its message (never +/// silently swallowed), and the log advises the `LIQUID_CACHE_ADMISSION_STRICT` +/// flag as the alternative behavior: +/// +/// - `gate.strict == true`: log, then re-raise so the panic aborts the query and +/// the bug surfaces immediately. The log advises turning the flag *off* to +/// keep queries running while it is fixed. +/// - `gate.strict == false`: log, then cache the scan normally so the query +/// survives. The log advises turning the flag *on* to fail loud instead. fn should_bypass_guarded( cfg: &FileScanConfig, src: &ParquetSource, gate: AdmissionGate, budget: u64, ) -> bool { - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { should_bypass(cfg, src, gate, budget) - })) - .unwrap_or_else(|_| { - log::warn!( - "liquid-cache admission gate panicked during footprint estimation; \ - caching scan normally" - ); - false - }) + })); + match result { + Ok(bypass) => bypass, + Err(payload) => { + let msg = payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or(""); + if gate.strict { + log::error!( + "liquid-cache admission gate panicked during footprint estimation: \ + {msg}. Aborting query (strict mode). Set \ + LIQUID_CACHE_ADMISSION_STRICT=off to fall back to caching and keep \ + queries running while this is fixed." + ); + std::panic::resume_unwind(payload); + } + log::error!( + "liquid-cache admission gate panicked during footprint estimation: {msg}; \ + caching scan normally. Set LIQUID_CACHE_ADMISSION_STRICT=on to fail loud \ + instead." + ); + false + } + } } /// Boolean per file: `true` if the file may match the predicate (keep it), @@ -612,7 +652,7 @@ mod tests { // A huge expansion inflates the estimate past the 1 MB budget → bypass. let capped = LocalModeOptimizer::new(build_cache().await) - .with_admission_gate(1e9, 1.0, 1.0) + .with_admission_gate(1e9, 1.0, 1.0, false) .optimize(plan.clone(), &config) .unwrap(); assert!( @@ -622,7 +662,7 @@ mod tests { // With a large budget the footprint fits (expansion 1.0) → cached. let uncapped = LocalModeOptimizer::new(build_cache_with_budget(usize::MAX).await) - .with_admission_gate(1.0, 1.0, 1.0) + .with_admission_gate(1.0, 1.0, 1.0, false) .optimize(plan, &config) .unwrap(); assert!( From aaf0ed95a3be6dc88480b0df5c1f6ae505c471f6 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 12:59:55 +0530 Subject: [PATCH 10/13] feat(admission): log every gate decision with its breakdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gate was a black box — no way to see why a scan was admitted or bypassed without a full benchmark run. Log one INFO line per decision (target liquid_cache::admission): file, projected cols, surviving/total files, raw and footprint bytes, budget, threshold, and the multipliers. Crucially, surface fallback_files: when a file lacks per-column byte sizes the estimate charges the whole file, so a non-zero count is the direct signal that the catalog has no column_size_bytes for the table (the DuckLake write-side stat). estimate_required_bytes now returns a FootprintEstimate breakdown and file_required_bytes reports whether it fell back. --- src/datafusion/src/optimizers/mod.rs | 153 +++++++++++++++++++++------ 1 file changed, 119 insertions(+), 34 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 56460369..6e3f40d4 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -205,6 +205,26 @@ fn parquet_scan_parts(node: &Arc) -> Option<(&FileScanConfig, dse.downcast_to_file_source::() } +/// A scan's footprint estimate plus the breakdown behind it, for the admission +/// decision ([`should_bypass`]) and its diagnostic log line. +#[derive(Debug, Default, Clone, Copy)] +struct FootprintEstimate { + /// Raw required-column parquet bytes (the decision input). + raw_bytes: u64, + /// Number of file columns the scan materializes (projection ∪ predicate). + required_cols: usize, + /// Files kept after predicate pruning (only these are charged). + surviving_files: usize, + /// Total files in the scan before pruning. + total_files: usize, + /// Surviving files charged the *whole file* size because they lacked + /// per-column byte sizes. A non-zero count means the catalog has no + /// `column_size_bytes` for this table, so the estimate is coarse and + /// over-counts — the signal that the DuckLake write-side size stat is + /// missing for this table. + fallback_files: usize, +} + /// Estimate the raw required-column parquet bytes a scan reads: the sum, over /// the files that survive the scan's predicate, of the byte sizes of the columns /// it materializes. This is the byte-accurate, filter-aware size the admission @@ -217,7 +237,7 @@ fn parquet_scan_parts(node: &Arc) -> Option<(&FileScanConfig, /// them `Inexact`); a column with an `Absent` size or a file with no stats /// falls back to the whole file size. The sum saturates rather than overflowing /// on pathological file lists. -fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> u64 { +fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> FootprintEstimate { let num_file_cols = cfg.file_schema().fields().len(); // Full table schema (file + partition columns). The pushed-down predicate // and `PartitionedFile` stats are expressed against it, so file pruning must @@ -254,18 +274,39 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> u64 { required.dedup(); let files: Vec<&PartitionedFile> = cfg.file_groups.iter().flat_map(|g| g.files()).collect(); + let required_cols = required.len(); if files.is_empty() { - return 0; + return FootprintEstimate { + required_cols, + ..FootprintEstimate::default() + }; } let surviving = surviving_files(src, &table_schema, &files); - files - .iter() - .zip(surviving.iter()) - .filter(|(_, keep)| **keep) - .map(|(f, _)| file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required)) - .fold(0u64, u64::saturating_add) + let mut raw_bytes = 0u64; + let mut surviving_files = 0usize; + let mut fallback_files = 0usize; + for (f, keep) in files.iter().zip(surviving.iter()) { + if !*keep { + continue; + } + surviving_files += 1; + let (bytes, fell_back) = + file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required); + if fell_back { + fallback_files += 1; + } + raw_bytes = raw_bytes.saturating_add(bytes); + } + + FootprintEstimate { + raw_bytes, + required_cols, + surviving_files, + total_files: files.len(), + fallback_files, + } } /// Decide whether a parquet scan should bypass the cache (read as vanilla @@ -285,13 +326,46 @@ fn should_bypass( gate: AdmissionGate, budget: u64, ) -> bool { - let raw = estimate_required_bytes(cfg, src); + let est = estimate_required_bytes(cfg, src); // Multipliers are already sanitized to finite, >= 1.0 in `with_admission_gate`. - let footprint = (raw as f64) * gate.expansion * gate.safety; - if budget == 0 { - return footprint > 0.0; - } - footprint / (budget as f64) > gate.tolerance + let footprint = (est.raw_bytes as f64) * gate.expansion * gate.safety; + let bypass = if budget == 0 { + footprint > 0.0 + } else { + footprint / (budget as f64) > gate.tolerance + }; + + // One line per admission decision. Without this the gate is a black box and + // every tuning cycle costs a benchmark run to infer decisions from side + // effects. `fallback_files > 0` is the flag that the catalog has no + // per-column sizes for this table (estimate is coarse / over-counts). + let threshold = (budget as f64) * gate.tolerance; + let path = cfg + .file_groups + .iter() + .flat_map(|g| g.files()) + .next() + .map(|f| f.object_meta.location.as_ref()) + .unwrap_or(""); + log::info!( + target: "liquid_cache::admission", + "admission {verdict}: file={path} projected_cols={cols} files={surv}/{total} \ + fallback_files={fb} raw_bytes={raw} footprint_bytes={fp} budget_bytes={budget} \ + threshold_bytes={thr} expansion={exp} safety={saf} tolerance={tol}", + verdict = if bypass { "BYPASS" } else { "ADMIT" }, + cols = est.required_cols, + surv = est.surviving_files, + total = est.total_files, + fb = est.fallback_files, + raw = est.raw_bytes, + fp = footprint as u64, + thr = threshold as u64, + exp = gate.expansion, + saf = gate.safety, + tol = gate.tolerance, + ); + + bypass } /// [`should_bypass`], with an optional panic guard. @@ -395,9 +469,17 @@ fn surviving_files( /// /// If the file has no stats, or any required column's size is `Absent`, fall /// back to the whole-file size — a deliberate over-estimate. -fn file_required_bytes(stats: Option<&Statistics>, object_size: u64, required: &[usize]) -> u64 { +/// +/// Returns `(bytes, fell_back)`, where `fell_back` is `true` when the whole-file +/// over-estimate was used (surfaced in the decision log as the "no per-column +/// sizes in the catalog" signal). +fn file_required_bytes( + stats: Option<&Statistics>, + object_size: u64, + required: &[usize], +) -> (u64, bool) { let Some(stats) = stats else { - return object_size; + return (object_size, true); }; let mut sum: u64 = 0; for &c in required { @@ -405,10 +487,10 @@ fn file_required_bytes(stats: Option<&Statistics>, object_size: u64, required: & Some(Precision::Exact(n) | Precision::Inexact(n)) => { sum = sum.saturating_add(*n as u64) } - _ => return object_size, + _ => return (object_size, true), } } - sum + (sum, false) } /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent @@ -563,9 +645,9 @@ mod tests { // The call that used to panic. It must return a finite estimate; `s` // (the struct column read by `s.a`) is the one required file column. - let bytes = estimate_required_bytes(&cfg, &src); + let est = estimate_required_bytes(&cfg, &src); // Whole-file fallback (no per-column Exact stats here) → non-zero, finite. - assert!(bytes > 0, "estimate should be a real byte count"); + assert!(est.raw_bytes > 0, "estimate should be a real byte count"); } #[tokio::test] @@ -697,7 +779,7 @@ mod footprint_tests { #[test] fn no_stats_falls_back_to_full_file() { - assert_eq!(file_required_bytes(None, 5000, &[0, 1]), 5000); + assert_eq!(file_required_bytes(None, 5000, &[0, 1]), (5000, true)); } #[test] @@ -707,8 +789,8 @@ mod footprint_tests { col(Precision::Exact(200)), col(Precision::Exact(400)), ]); - // required = cols 0 and 2 → 100 + 400, ignoring col 1. - assert_eq!(file_required_bytes(Some(&s), 9999, &[0, 2]), 500); + // required = cols 0 and 2 → 100 + 400, ignoring col 1. No fallback. + assert_eq!(file_required_bytes(Some(&s), 9999, &[0, 2]), (500, false)); } #[test] @@ -717,10 +799,10 @@ mod footprint_tests { col(Precision::Exact(100)), col(Precision::Inexact(200)), ]); - // Inexact is a real (possibly-stale) size, so it counts. DuckLake always - // labels byte_size Inexact; rejecting it would bypass every scan. - assert_eq!(file_required_bytes(Some(&s), 7000, &[0, 1]), 300); - assert_eq!(file_required_bytes(Some(&s), 7000, &[1]), 200); + // Inexact is a real (possibly-stale) size, so it counts (no fallback). + // DuckLake always labels byte_size Inexact; rejecting it would bypass all. + assert_eq!(file_required_bytes(Some(&s), 7000, &[0, 1]), (300, false)); + assert_eq!(file_required_bytes(Some(&s), 7000, &[1]), (200, false)); } #[test] @@ -730,22 +812,25 @@ mod footprint_tests { col(Precision::Inexact(200)), col(Precision::Absent), ]); - // Any Absent required column → whole file, even mixed with sized ones. - assert_eq!(file_required_bytes(Some(&s), 8000, &[0, 1, 2]), 8000); - // Dropping the absent column, Exact + Inexact are summed. - assert_eq!(file_required_bytes(Some(&s), 8000, &[0, 1]), 300); + // Any Absent required column → whole file (fallback), even mixed with sized ones. + assert_eq!( + file_required_bytes(Some(&s), 8000, &[0, 1, 2]), + (8000, true) + ); + // Dropping the absent column, Exact + Inexact are summed (no fallback). + assert_eq!(file_required_bytes(Some(&s), 8000, &[0, 1]), (300, false)); } #[test] fn missing_required_column_falls_back_to_full_file() { let s = stats(vec![col(Precision::Exact(100))]); - // required col 5 doesn't exist → conservative whole file. - assert_eq!(file_required_bytes(Some(&s), 3000, &[0, 5]), 3000); + // required col 5 doesn't exist → conservative whole file (fallback). + assert_eq!(file_required_bytes(Some(&s), 3000, &[0, 5]), (3000, true)); } #[test] fn absent_byte_size_falls_back_to_full_file() { let s = stats(vec![col(Precision::Absent)]); - assert_eq!(file_required_bytes(Some(&s), 2000, &[0]), 2000); + assert_eq!(file_required_bytes(Some(&s), 2000, &[0]), (2000, true)); } } From 7e3efb46e4c94253e34dc5af626454f9018ea4f1 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 14:55:57 +0530 Subject: [PATCH 11/13] fix(admission): dedupe split byte-range files in footprint estimate DataFusion splits one file into several byte-range PartitionedFiles for parallelism, each cloning the whole file's statistics and object size. The estimate summed every range, multiplying a single file's footprint by the split count (e.g. 7x14.8GB). Charge each distinct file path once; pruning is file-granular, so a surviving file means its whole required columns regardless of how it was split for scanning. --- src/datafusion/src/optimizers/mod.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 6e3f40d4..3d844b53 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -2,6 +2,7 @@ mod squeeze_hint; +use std::collections::HashSet; use std::sync::Arc; use datafusion::{ @@ -284,6 +285,13 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> Footpri let surviving = surviving_files(src, &table_schema, &files); + // Dedupe by physical file path. DataFusion splits one file into several + // byte-range `PartitionedFile`s for parallelism, each a *clone* of the whole + // file's statistics (and whole-file object size). Counting every range would + // multiply a single file's footprint by the split count, so charge each + // distinct file once — pruning is file-granular, so a surviving file means + // caching its whole required columns regardless of how it was split. + let mut seen: HashSet = HashSet::new(); let mut raw_bytes = 0u64; let mut surviving_files = 0usize; let mut fallback_files = 0usize; @@ -291,6 +299,9 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> Footpri if !*keep { continue; } + if !seen.insert(f.object_meta.location.clone()) { + continue; + } surviving_files += 1; let (bytes, fell_back) = file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required); From 8bfec46e4fe62a7e4b50f9bcb950a0c532a83cdd Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 17:05:48 +0530 Subject: [PATCH 12/13] refactor(admission): dedupe by (path,size); clearer decision-log fields Codex review follow-ups: key the file dedupe by (location, size) so two distinct objects sharing a path can't be collapsed, and rename the decision-log counters to charged_files / partitioned_files so a split-into-N file reads unambiguously (e.g. charged_files=1 partitioned_files=7). --- src/datafusion/src/optimizers/mod.rs | 43 ++++++++++++++++------------ 1 file changed, 24 insertions(+), 19 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 3d844b53..e1d606a0 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -214,10 +214,12 @@ struct FootprintEstimate { raw_bytes: u64, /// Number of file columns the scan materializes (projection ∪ predicate). required_cols: usize, - /// Files kept after predicate pruning (only these are charged). - surviving_files: usize, - /// Total files in the scan before pruning. - total_files: usize, + /// Distinct physical files charged (survived pruning, deduped across the + /// byte-range splits of the same file). + charged_files: usize, + /// Raw `PartitionedFile` count before dedupe (one file may be split into + /// several byte-range partitions for parallelism). + partitioned_files: usize, /// Surviving files charged the *whole file* size because they lacked /// per-column byte sizes. A non-zero count means the catalog has no /// `column_size_bytes` for this table, so the estimate is coarse and @@ -285,24 +287,26 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> Footpri let surviving = surviving_files(src, &table_schema, &files); - // Dedupe by physical file path. DataFusion splits one file into several - // byte-range `PartitionedFile`s for parallelism, each a *clone* of the whole - // file's statistics (and whole-file object size). Counting every range would - // multiply a single file's footprint by the split count, so charge each - // distinct file once — pruning is file-granular, so a surviving file means - // caching its whole required columns regardless of how it was split. - let mut seen: HashSet = HashSet::new(); + // Dedupe by physical file identity (path + size). DataFusion splits one file + // into several byte-range `PartitionedFile`s for parallelism, each a *clone* + // of the whole file's statistics (and whole-file object size), differing only + // in `range`. Counting every range would multiply a single file's footprint + // by the split count, so charge each distinct file once — pruning is + // file-granular, so a surviving file means caching its whole required columns + // regardless of how it was split. Size is part of the key so two genuinely + // distinct objects that share a path are never collapsed. + let mut seen: HashSet<(object_store::path::Path, u64)> = HashSet::new(); let mut raw_bytes = 0u64; - let mut surviving_files = 0usize; + let mut charged_files = 0usize; let mut fallback_files = 0usize; for (f, keep) in files.iter().zip(surviving.iter()) { if !*keep { continue; } - if !seen.insert(f.object_meta.location.clone()) { + if !seen.insert((f.object_meta.location.clone(), f.object_meta.size)) { continue; } - surviving_files += 1; + charged_files += 1; let (bytes, fell_back) = file_required_bytes(f.statistics.as_deref(), f.object_meta.size, &required); if fell_back { @@ -314,8 +318,8 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> Footpri FootprintEstimate { raw_bytes, required_cols, - surviving_files, - total_files: files.len(), + charged_files, + partitioned_files: files.len(), fallback_files, } } @@ -360,13 +364,14 @@ fn should_bypass( .unwrap_or(""); log::info!( target: "liquid_cache::admission", - "admission {verdict}: file={path} projected_cols={cols} files={surv}/{total} \ + "admission {verdict}: file={path} projected_cols={cols} \ + charged_files={charged} partitioned_files={total} \ fallback_files={fb} raw_bytes={raw} footprint_bytes={fp} budget_bytes={budget} \ threshold_bytes={thr} expansion={exp} safety={saf} tolerance={tol}", verdict = if bypass { "BYPASS" } else { "ADMIT" }, cols = est.required_cols, - surv = est.surviving_files, - total = est.total_files, + charged = est.charged_files, + total = est.partitioned_files, fb = est.fallback_files, raw = est.raw_bytes, fp = footprint as u64, From 4e0c361d3af5d623597e94efeaeff56c35497ff5 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Wed, 22 Jul 2026 17:26:06 +0530 Subject: [PATCH 13/13] test/chore: update squeeze snapshot for CLOCK order; CI + review nits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - squeeze_strings snapshot: CLOCK second-chance eviction reorders the squeeze victims ([851969,851968,917504] vs prior) — same victims, same outcomes, deterministic queue order. Accept the reordered trace. - admission log/doc: advise the gate's 'strict' builder flag rather than a runtimedb env var the crate doesn't read (review nit). - dev-tools cache_state_view: iterate map keys() (clippy for_kv_map, newly flagged by a clippy toolchain bump; pre-existing code unrelated to this change but blocks -D warnings CI). --- dev/dev-tools/src/components/cache_state_view.rs | 2 +- ...ion_local__tests__squeeze__squeeze_strings.snap | 8 ++++---- src/datafusion/src/optimizers/mod.rs | 14 +++++++------- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/dev/dev-tools/src/components/cache_state_view.rs b/dev/dev-tools/src/components/cache_state_view.rs index e909f175..fe8d4a6d 100644 --- a/dev/dev-tools/src/components/cache_state_view.rs +++ b/dev/dev-tools/src/components/cache_state_view.rs @@ -50,7 +50,7 @@ pub fn CacheStateView(simulator: Signal) -> Element { // Create a unified sorted list of all entry IDs (both actual and failed) let mut all_entry_ids: Vec = entries.iter().map(|e| e.entry_id).collect(); - for (entry_id, _) in state.failed_inserts.iter() { + for entry_id in state.failed_inserts.keys() { if !all_entry_ids.contains(entry_id) { all_entry_ids.push(*entry_id); } diff --git a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap index 2943782b..b54a6066 100644 --- a/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap +++ b/src/datafusion-local/src/tests/snapshots/liquid_cache_datafusion_local__tests__squeeze__squeeze_strings.snap @@ -22,15 +22,15 @@ event=hydrate entry=851968 cached=DiskLiquid new=MemoryLiquid event=insert_success entry=851968 kind=MemoryLiquid event=insert_success entry=851969 kind=MemoryArrow event=insert_failed entry=917505 kind=MemoryArrow -event=squeeze_begin victims=[851969,917504,851968] +event=squeeze_begin victims=[851969,851968,917504] event=squeeze_victim entry=851969 event=insert_success entry=851969 kind=MemoryLiquid -event=squeeze_victim entry=917504 -event=io_write entry=917504 kind=MemorySqueezedLiquid bytes=136440 -event=insert_success entry=917504 kind=MemorySqueezedLiquid event=squeeze_victim entry=851968 event=io_write entry=851968 kind=DiskLiquid bytes=139416 event=insert_success entry=851968 kind=DiskLiquid +event=squeeze_victim entry=917504 +event=io_write entry=917504 kind=MemorySqueezedLiquid bytes=136440 +event=insert_success entry=917504 kind=MemorySqueezedLiquid event=insert_success entry=917505 kind=MemoryArrow event=eval_predicate entry=917505 selection=true cached=MemoryArrow event=read entry=851969 selection=true expr=None cached=MemoryLiquid diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index e1d606a0..6b6eb638 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -393,8 +393,8 @@ fn should_bypass( /// let it abort the query. /// /// Either way the panic is caught and logged at ERROR with its message (never -/// silently swallowed), and the log advises the `LIQUID_CACHE_ADMISSION_STRICT` -/// flag as the alternative behavior: +/// silently swallowed), and the log advises flipping the admission gate's +/// `strict` flag for the alternative behavior: /// /// - `gate.strict == true`: log, then re-raise so the panic aborts the query and /// the bug surfaces immediately. The log advises turning the flag *off* to @@ -421,16 +421,16 @@ fn should_bypass_guarded( if gate.strict { log::error!( "liquid-cache admission gate panicked during footprint estimation: \ - {msg}. Aborting query (strict mode). Set \ - LIQUID_CACHE_ADMISSION_STRICT=off to fall back to caching and keep \ - queries running while this is fixed." + {msg}. Aborting query (admission gate is in strict mode). Configure the \ + admission gate with strict=false to fall back to caching and keep queries \ + running while this is fixed." ); std::panic::resume_unwind(payload); } log::error!( "liquid-cache admission gate panicked during footprint estimation: {msg}; \ - caching scan normally. Set LIQUID_CACHE_ADMISSION_STRICT=on to fail loud \ - instead." + caching scan normally. Configure the admission gate with strict=true to fail \ + loud instead." ); false }