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/core/src/cache/policies/cache/three_queue.rs b/src/core/src/cache/policies/cache/three_queue.rs index 655c1984..e630c413 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,24 @@ 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) { + // 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) { let mut inner = self.inner.lock().unwrap(); @@ -249,6 +322,66 @@ 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_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 + // 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(); diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 2bb2485f..ebb43bb2 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -69,8 +69,10 @@ 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, 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, } @@ -86,7 +88,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 +150,22 @@ 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`) + /// 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). `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 } @@ -202,8 +215,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, tolerance, strict)) = self.admission { + optimizer = optimizer.with_admission_gate(expansion, safety, tolerance, strict); } let state = datafusion::execution::SessionStateBuilder::new() 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 804a5222..6b6eb638 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -2,14 +2,26 @@ mod squeeze_hint; +use std::collections::HashSet; 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 +30,39 @@ 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 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 (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). Inflates the + /// estimate (conservative direction). + pub expansion: f64, + /// 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, + /// 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. /// /// Rewrites `DataSourceExec` parquet scans to use [`LiquidParquetSource`], and @@ -26,10 +71,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 +82,7 @@ impl LocalModeOptimizer { pub fn new(cache: LiquidCacheParquetRef) -> Self { Self { cache, - max_scan_bytes: None, + admission: None, } } @@ -45,15 +90,41 @@ 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`) stays within `budget × tolerance`; otherwise it is read directly + /// from the parquet source, bypassing the cache. See [`AdmissionGate`]. + /// + /// 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. `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) + } else { + 3.0 + }; + self.admission = Some(AdmissionGate { + expansion: estimate_factor(expansion), + safety: estimate_factor(safety), + tolerance, + strict, + }); self } } @@ -66,11 +137,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) + && should_bypass_guarded(cfg, src, gate, budget) { return None; } @@ -125,19 +199,314 @@ 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::() +} + +/// 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, + /// 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 + /// 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 +/// 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 `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) -> 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 + // use it; byte accounting still charges only file columns. + let table_schema = cfg.file_source.table_schema().table_schema().clone(); + + // 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) { + 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(); + 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 FootprintEstimate { + required_cols, + ..FootprintEstimate::default() + }; + } + + let surviving = surviving_files(src, &table_schema, &files); + + // 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 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(), f.object_meta.size)) { + continue; + } + charged_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, + charged_files, + partitioned_files: files.len(), + fallback_files, + } +} + +/// 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 est = estimate_required_bytes(cfg, src); + // Multipliers are already sanitized to finite, >= 1.0 in `with_admission_gate`. + 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} \ + 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, + charged = est.charged_files, + total = est.partitioned_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. +/// +/// 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 — 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 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 +/// 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 { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + should_bypass(cfg, src, gate, budget) + })); + 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 (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. Configure the admission gate with strict=true to fail \ + loud instead." + ); + 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 +/// keeps files. `table_schema` (file + partition columns) is used so predicates +/// on partition columns resolve, matching `PartitionedFile::statistics`. +fn surviving_files( + src: &ParquetSource, + table_schema: &SchemaRef, + files: &[&PartitionedFile], +) -> Vec { + let Some(pred) = src.filter() else { + return vec![true; files.len()]; + }; + 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| 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, table_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 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. +/// +/// 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, true); + }; + let mut sum: u64 = 0; + for &c in required { + match stats.column_statistics.get(c).map(|cs| &cs.byte_size) { + Some(Precision::Exact(n) | Precision::Inexact(n)) => { + sum = sum.saturating_add(*n as u64) + } + _ => return (object_size, true), + } + } + (sum, false) } /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent @@ -209,6 +578,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 est = estimate_required_bytes(&cfg, &src); + // Whole-file fallback (no per-column Exact stats here) → non-zero, finite. + assert!(est.raw_bytes > 0, "estimate should be a real byte count"); + } + #[tokio::test] async fn test_plan_rewrite() { let ctx = SessionContext::new(); @@ -228,6 +685,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 +696,7 @@ mod tests { Arc::new( LiquidCacheParquet::new( 8192, - 1000000, + max_memory_bytes, usize::MAX, store, Box::new(LiquidPolicy::new()), @@ -266,11 +727,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 +748,105 @@ 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, 1.0, false) .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, 1.0, false) .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 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; + 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, true)); + } + + #[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. No fallback. + assert_eq!(file_required_bytes(Some(&s), 9999, &[0, 2]), (500, false)); + } + + #[test] + 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 (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] + 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), + ]); + // 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 (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, true)); + } +}