diff --git a/src/core/src/cache/budget.rs b/src/core/src/cache/budget.rs index 2b7bc5d0..89b1c7a4 100644 --- a/src/core/src/cache/budget.rs +++ b/src/core/src/cache/budget.rs @@ -6,8 +6,8 @@ use crate::sync::{ #[derive(Debug)] pub struct BudgetAccounting { - max_memory_bytes: usize, - max_disk_bytes: usize, + max_memory_bytes: AtomicUsize, + max_disk_bytes: AtomicUsize, used_memory_bytes: AtomicUsize, used_disk_bytes: AtomicUsize, observer: Arc, @@ -20,8 +20,8 @@ impl BudgetAccounting { observer: Arc, ) -> Self { Self { - max_memory_bytes, - max_disk_bytes, + max_memory_bytes: AtomicUsize::new(max_memory_bytes), + max_disk_bytes: AtomicUsize::new(max_disk_bytes), used_memory_bytes: AtomicUsize::new(0), used_disk_bytes: AtomicUsize::new(0), observer, @@ -33,11 +33,29 @@ impl BudgetAccounting { self.used_disk_bytes.store(0, Ordering::Relaxed); } + /// Dynamically update the max memory limit. Takes effect for new reservations. + pub fn set_max_memory_bytes(&self, new_limit: usize) { + self.max_memory_bytes.store(new_limit, Ordering::Relaxed); + } + + /// Dynamically update the max disk limit. Takes effect for new reservations. + pub fn set_max_disk_bytes(&self, new_limit: usize) { + self.max_disk_bytes.store(new_limit, Ordering::Relaxed); + } + + pub fn max_memory_bytes(&self) -> usize { + self.max_memory_bytes.load(Ordering::Relaxed) + } + + pub fn max_disk_bytes(&self) -> usize { + self.max_disk_bytes.load(Ordering::Relaxed) + } + /// Try to reserve memory in the cache. /// Returns ok if the memory was reserved, err if the memory budget is full. pub(super) fn try_reserve_memory(&self, request_bytes: usize) -> Result<(), ()> { let used = self.used_memory_bytes.load(Ordering::Relaxed); - if used + request_bytes > self.max_memory_bytes { + if used + request_bytes > self.max_memory_bytes.load(Ordering::Relaxed) { return Err(()); } @@ -80,7 +98,7 @@ impl BudgetAccounting { pub(super) fn try_reserve_disk(&self, request_bytes: usize) -> Result<(), ()> { let used = self.used_disk_bytes.load(Ordering::Relaxed); - if used + request_bytes > self.max_disk_bytes { + if used + request_bytes > self.max_disk_bytes.load(Ordering::Relaxed) { self.observer.on_disk_reservation_failure(); return Err(()); } diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 84e0a795..3a702c73 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -112,8 +112,8 @@ impl LiquidCache { memory_squeezed_liquid_bytes, memory_usage_bytes, disk_usage_bytes, - max_memory_bytes: self.config.max_memory_bytes(), - max_disk_bytes: self.config.max_disk_bytes(), + max_memory_bytes: self.budget.max_memory_bytes(), + max_disk_bytes: self.budget.max_disk_bytes(), runtime, } } diff --git a/src/core/src/cache/observer/mod.rs b/src/core/src/cache/observer/mod.rs index 38de5fc4..5b848b07 100644 --- a/src/core/src/cache/observer/mod.rs +++ b/src/core/src/cache/observer/mod.rs @@ -131,7 +131,7 @@ impl Observer { } } - pub(crate) fn runtime_stats(&self) -> &RuntimeStats { + pub fn runtime_stats(&self) -> &RuntimeStats { &self.runtime } } diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index fa0c3d9a..23963248 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -97,6 +97,8 @@ define_runtime_stats! { (get, "Number of `get` calls issued via `CachedData`.", incr_get), (get_with_selection, "Number of `get_with_selection` calls issued via `CachedData`.", incr_get_with_selection), (eval_predicate, "Number of `eval_predicate` calls issued via `CachedData`.", incr_eval_predicate), + (cache_hit, "Number of cache hits (data found in cache).", incr_cache_hit), + (cache_miss, "Number of cache misses (data not in cache, fell back to Parquet).", incr_cache_miss), (get_squeezed_success, "Number of Squeezed-Liquid full evaluations finished without IO.", incr_get_squeezed_success), (get_squeezed_needs_io, "Number of Squeezed-Liquid full paths that required IO.", incr_get_squeezed_needs_io), (try_read_liquid_calls, "Number of `try_read_liquid` calls issued via `CachedData`.", incr_try_read_liquid), diff --git a/src/core/src/cache/policies/cache/lru.rs b/src/core/src/cache/policies/cache/lru.rs new file mode 100644 index 00000000..b8968c47 --- /dev/null +++ b/src/core/src/cache/policies/cache/lru.rs @@ -0,0 +1,291 @@ +//! LRU (Least Recently Used) cache eviction policy. +//! +//! Uses 4 queues by entry type (Arrow, Liquid, Squeezed, Disk), same as LiquidPolicy. +//! Within each queue, entries are ordered by recency (moved to back on access). +//! Eviction priority: Arrow (largest) first, then Liquid, then Squeezed. +//! Within each queue, evicts the LRU entry (front of the queue). + +use crate::cache::cached_batch::CachedBatchType; +use crate::cache::utils::EntryID; +use crate::sync::Mutex; +use ahash::AHashMap; +use std::ptr::NonNull; + +use super::CachePolicy; +use super::doubly_linked_list::{DoublyLinkedList, DoublyLinkedNode, drop_boxed_node}; + +/// Which queue an entry belongs to. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LruQueueKind { + Arrow, + Liquid, + Squeezed, + Disk, +} + +/// LRU cache eviction policy with type-aware queues. +/// +/// On insert: entry goes to the back of its type queue (most recently used). +/// On access: entry moves to the back of its type queue. +/// On eviction: picks the LRU entry from Arrow queue first, then Liquid, then Squeezed. +#[derive(Debug)] +pub struct LruPolicy { + inner: Mutex, +} + +#[derive(Debug)] +struct LruInner { + arrow: DoublyLinkedList, + liquid: DoublyLinkedList, + squeezed: DoublyLinkedList, + disk: DoublyLinkedList, + /// Maps entry_id → (node_ptr, which queue it's in) + map: AHashMap>, LruQueueKind)>, +} + +// Safety: We control access via Mutex and never expose raw pointers outside. +unsafe impl Send for LruInner {} +unsafe impl Sync for LruInner {} + +impl LruInner { + fn queue_mut(&mut self, kind: LruQueueKind) -> &mut DoublyLinkedList { + match kind { + LruQueueKind::Arrow => &mut self.arrow, + LruQueueKind::Liquid => &mut self.liquid, + LruQueueKind::Squeezed => &mut self.squeezed, + LruQueueKind::Disk => &mut self.disk, + } + } + + fn queue_kind_for(batch_type: CachedBatchType) -> LruQueueKind { + match batch_type { + CachedBatchType::MemoryArrow => LruQueueKind::Arrow, + CachedBatchType::MemoryLiquid => LruQueueKind::Liquid, + CachedBatchType::MemorySqueezedLiquid => LruQueueKind::Squeezed, + CachedBatchType::DiskLiquid | CachedBatchType::DiskArrow => LruQueueKind::Disk, + } + } + + /// Pop the LRU entry (front) from a specific queue. + fn pop_lru(&mut self, kind: LruQueueKind) -> Option { + let list = self.queue_mut(kind); + let node_ptr = list.head()?; + let entry_id = unsafe { node_ptr.as_ref().data }; + unsafe { list.unlink(node_ptr) }; + self.map.remove(&entry_id); + unsafe { drop_boxed_node(node_ptr) }; + Some(entry_id) + } +} + +impl LruPolicy { + /// Create a new LRU policy. + pub fn new() -> Self { + Self { + inner: Mutex::new(LruInner { + arrow: DoublyLinkedList::new(), + liquid: DoublyLinkedList::new(), + squeezed: DoublyLinkedList::new(), + disk: DoublyLinkedList::new(), + map: AHashMap::new(), + }), + } + } +} + +impl Default for LruPolicy { + fn default() -> Self { + Self::new() + } +} + +impl CachePolicy for LruPolicy { + fn find_memory_victim(&self, cnt: usize) -> Vec { + let mut inner = self.inner.lock().unwrap(); + let mut victims = Vec::with_capacity(cnt); + + while victims.len() < cnt { + // Priority: evict Arrow (largest) first, then Liquid, then Squeezed + if let Some(entry) = inner.pop_lru(LruQueueKind::Arrow) { + victims.push(entry); + continue; + } + if let Some(entry) = inner.pop_lru(LruQueueKind::Liquid) { + victims.push(entry); + continue; + } + if let Some(entry) = inner.pop_lru(LruQueueKind::Squeezed) { + victims.push(entry); + continue; + } + break; + } + + victims + } + + fn find_disk_victim(&self, cnt: usize) -> Vec { + if cnt == 0 { + return vec![]; + } + let mut inner = self.inner.lock().unwrap(); + let mut victims = Vec::with_capacity(cnt); + while victims.len() < cnt { + if let Some(entry) = inner.pop_lru(LruQueueKind::Disk) { + victims.push(entry); + } else { + break; + } + } + victims + } + + fn notify_insert(&self, entry_id: &EntryID, batch_type: CachedBatchType) { + let mut inner = self.inner.lock().unwrap(); + let target = LruInner::queue_kind_for(batch_type); + + // If already present, remove from old position/queue + if let Some((node_ptr, old_kind)) = inner.map.remove(entry_id) { + let old_list = inner.queue_mut(old_kind); + unsafe { old_list.unlink(node_ptr) }; + unsafe { drop_boxed_node(node_ptr) }; + } + + // Insert at the back of the target queue (most recently used) + let node = DoublyLinkedNode::new(*entry_id); + let node_ptr = NonNull::new(Box::into_raw(node)).unwrap(); + let list = inner.queue_mut(target); + unsafe { list.push_back(node_ptr) }; + inner.map.insert(*entry_id, (node_ptr, target)); + } + + fn notify_access(&self, entry_id: &EntryID, _batch_type: CachedBatchType) { + let mut inner = self.inner.lock().unwrap(); + + let Some(&(node_ptr, kind)) = inner.map.get(entry_id) else { + return; + }; + + // Move to back of its current queue (most recently used) + let list = inner.queue_mut(kind); + unsafe { + list.unlink(node_ptr); + list.push_back(node_ptr); + } + } + + fn notify_remove(&self, entry_id: &EntryID) { + let mut inner = self.inner.lock().unwrap(); + + if let Some((node_ptr, kind)) = inner.map.remove(entry_id) { + let list = inner.queue_mut(kind); + unsafe { list.unlink(node_ptr) }; + unsafe { drop_boxed_node(node_ptr) }; + } + } +} + +impl Drop for LruPolicy { + fn drop(&mut self) { + let mut inner = self.inner.lock().unwrap(); + unsafe { + inner.arrow.drop_all(); + inner.liquid.drop_all(); + inner.squeezed.drop_all(); + inner.disk.drop_all(); + } + inner.map.clear(); + } +} + + +#[cfg(test)] +mod tests { + use super::*; + + fn entry(id: usize) -> EntryID { + EntryID::from(id) + } + + #[test] + fn test_lru_evicts_arrow_first() { + let policy = LruPolicy::new(); + + // Insert entries: some Arrow, some Liquid + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(2), CachedBatchType::MemoryLiquid); + policy.notify_insert(&entry(3), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(4), CachedBatchType::MemoryLiquid); + + // Evict 1 — should pick from Arrow queue first (LRU = entry 1) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + + // Evict 1 more — next Arrow LRU = entry 3 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(3)]); + + // Evict 1 more — Arrow empty, now Liquid LRU = entry 2 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + } + + #[test] + fn test_lru_access_moves_to_back() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(2), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(3), CachedBatchType::MemoryArrow); + + // Access entry 1 — moves it to back + policy.notify_access(&entry(1), CachedBatchType::MemoryArrow); + + // Evict 1 — LRU is now entry 2 (entry 1 was moved to back) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + } + + #[test] + fn test_lru_insert_moves_between_queues() { + let policy = LruPolicy::new(); + + // Insert as Arrow + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + + // Re-insert same entry as Liquid (simulates transcode) + policy.notify_insert(&entry(1), CachedBatchType::MemoryLiquid); + + // Arrow queue should be empty now + let victims = policy.find_memory_victim(1); + // Should come from Liquid queue + assert_eq!(victims, vec![entry(1)]); + } + + #[test] + fn test_lru_disk_victims() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(1), CachedBatchType::DiskLiquid); + policy.notify_insert(&entry(2), CachedBatchType::DiskLiquid); + + // find_disk_victim should return LRU disk entries + let victims = policy.find_disk_victim(1); + assert_eq!(victims, vec![entry(1)]); + } + + #[test] + fn test_lru_remove() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(2), CachedBatchType::MemoryArrow); + + // Remove entry 1 + policy.notify_remove(&entry(1)); + + // Evict — should get entry 2 (entry 1 was removed) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + } +} diff --git a/src/core/src/cache/policies/cache/mod.rs b/src/core/src/cache/policies/cache/mod.rs index 352f83a0..7399a3d4 100644 --- a/src/core/src/cache/policies/cache/mod.rs +++ b/src/core/src/cache/policies/cache/mod.rs @@ -4,8 +4,10 @@ use crate::cache::cached_batch::CachedBatchType; use crate::cache::utils::EntryID; mod doubly_linked_list; +mod lru; mod three_queue; +pub use lru::LruPolicy; pub use three_queue::LiquidPolicy; /// The cache policy that guides the replacement of LiquidCache @@ -81,6 +83,7 @@ mod tests { fn run_concurrent_invariant_tests() { concurrent_invariant_advice_once(Arc::new(LiquidPolicy::new())); + concurrent_invariant_advice_once(Arc::new(super::LruPolicy::new())); } #[test] @@ -93,4 +96,104 @@ mod tests { fn shuttle_concurrent_invariant_advice_once() { crate::utils::shuttle_test(run_concurrent_invariant_tests); } + + #[test] + fn lru_evicts_least_recently_used() { + let policy = LruPolicy::new(); + + // Insert entries 0, 1, 2, 3 in order + for i in 0..4 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Evict 1 → should be entry 0 (oldest) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(0)]); + + // Evict 1 more → should be entry 1 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + } + + #[test] + fn lru_access_moves_to_back() { + let policy = LruPolicy::new(); + + // Insert 0, 1, 2 + for i in 0..3 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Access entry 0 → moves it to back + policy.notify_access(&entry(0), CachedBatchType::MemoryArrow); + + // Evict → should be entry 1 now (0 was moved to back) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + + // Evict → entry 2 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + + // Evict → entry 0 (was moved to back by access) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(0)]); + } + + #[test] + fn lru_remove_removes_entry() { + let policy = LruPolicy::new(); + + for i in 0..3 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Remove entry 1 + policy.notify_remove(&entry(1)); + + // Evict → entry 0, then entry 2 (entry 1 is gone) + let victims = policy.find_memory_victim(2); + assert_eq!(victims, vec![entry(0), entry(2)]); + } + + #[test] + fn lru_evict_more_than_available() { + let policy = LruPolicy::new(); + + policy.notify_insert(&entry(0), CachedBatchType::MemoryArrow); + policy.notify_insert(&entry(1), CachedBatchType::MemoryArrow); + + // Ask for 5 but only 2 exist + let victims = policy.find_memory_victim(5); + assert_eq!(victims.len(), 2); + assert_eq!(victims, vec![entry(0), entry(1)]); + + // Empty now + let victims = policy.find_memory_victim(1); + assert!(victims.is_empty()); + } + + #[test] + fn lru_reinsert_updates_position() { + let policy = LruPolicy::new(); + + for i in 0..3 { + policy.notify_insert(&entry(i), CachedBatchType::MemoryArrow); + } + + // Re-insert entry 0 → should move to back + policy.notify_insert(&entry(0), CachedBatchType::MemoryLiquid); + + // Evict → entry 1 (oldest now) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(1)]); + + // Evict → entry 2 + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(2)]); + + // Evict → entry 0 (re-inserted last) + let victims = policy.find_memory_victim(1); + assert_eq!(victims, vec![entry(0)]); + } } diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 81a259aa..cfe56bb6 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -15,11 +15,12 @@ use liquid_cache::cache::{AlwaysHydrate, HydrationPolicy}; use liquid_cache::cache_policies::{CachePolicy, LiquidPolicy}; use liquid_cache_datafusion::optimizers::{LineageOptimizer, LocalModeOptimizer}; use liquid_cache_datafusion::{ - LiquidCacheParquet, LiquidCacheParquetRef, VariantGetUdf, VariantPretty, VariantToJsonUdf, + LiquidCacheParquet, VariantGetUdf, VariantPretty, VariantToJsonUdf, }; pub use liquid_cache as storage; pub use liquid_cache_common as common; +pub use liquid_cache_datafusion::{LiquidCacheParquetRef, LiquidParquetSource}; /// Builder for in-process liquid cache session context /// @@ -59,6 +60,8 @@ pub struct LiquidCacheLocalBuilder { batch_size: usize, /// Maximum memory size in bytes max_memory_bytes: usize, + /// Maximum disk size in bytes + max_disk_bytes: usize, /// Directory for disk cache cache_dir: PathBuf, /// Cache policy @@ -75,6 +78,7 @@ impl Default for LiquidCacheLocalBuilder { Self { batch_size: 8192, max_memory_bytes: 1024 * 1024 * 1024, // 1GB + max_disk_bytes: usize::MAX, cache_dir: std::env::temp_dir(), cache_policy: Box::new(LiquidPolicy::new()), squeeze_policy: Box::new(TranscodeSqueezeEvict), @@ -102,6 +106,12 @@ impl LiquidCacheLocalBuilder { self } + /// Set maximum disk size in bytes + pub fn with_max_disk_bytes(mut self, max_disk_bytes: usize) -> Self { + self.max_disk_bytes = max_disk_bytes; + self + } + /// Set cache directory pub fn with_cache_dir(mut self, cache_dir: PathBuf) -> Self { self.cache_dir = cache_dir; @@ -155,7 +165,7 @@ impl LiquidCacheLocalBuilder { let cache = LiquidCacheParquet::new( self.batch_size, self.max_memory_bytes, - usize::MAX, + self.max_disk_bytes, store, self.cache_policy, self.squeeze_policy, @@ -167,7 +177,7 @@ impl LiquidCacheLocalBuilder { let cache = LiquidCacheParquet::new_with_squeeze_victim_concurrency( self.batch_size, self.max_memory_bytes, - usize::MAX, + self.max_disk_bytes, store, self.cache_policy, self.squeeze_policy, diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index f92b9852..d297c836 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -479,3 +479,206 @@ async fn test_provide_schema_with_filter() { } assert_eq!(formatted_results, reference); } + +/// Test that only predicate (WHERE clause) columns are cached. +/// Projection-only columns should NOT be cached — they read from Parquet directly. +#[tokio::test] +async fn test_predicate_only_caching() { + let cache_dir = TempDir::new().unwrap(); + + // Query: WHERE on "OS" (predicate column), SELECT "URL" (projection-only column) + let sql = r#"SELECT "URL" FROM hits WHERE "OS" > 0 ORDER BY "URL" LIMIT 5"#; + + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = 2; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_max_memory_bytes(usize::MAX) + .with_cache_dir(cache_dir.path().to_path_buf()) + .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .build(config) + .await + .unwrap(); + + ctx.register_parquet("hits", TEST_FILE, ParquetReadOptions::default()) + .await + .unwrap(); + + // First run: fills cache + let plan = get_physical_plan(sql, &ctx).await; + let batches_1 = collect(plan, ctx.task_ctx()).await.unwrap(); + let result_1 = pretty_format_batches(&batches_1).unwrap().to_string(); + + let stats_after_first = cache.storage().stats(); + let entries_after_first = stats_after_first.total_entries; + + // Should have cached ONLY OS (predicate column) batches, NOT URL + // With nano_hits.parquet (1 row group), we expect a small number of entries + // corresponding to OS batches only. + assert!( + entries_after_first > 0, + "Cache should have entries for predicate column OS" + ); + + // Second run: should reuse cache for predicate evaluation + let plan = get_physical_plan(sql, &ctx).await; + let batches_2 = collect(plan, ctx.task_ctx()).await.unwrap(); + let result_2 = pretty_format_batches(&batches_2).unwrap().to_string(); + + // Results must be identical + assert_eq!(result_1, result_2, "Results should be the same on hot run"); + + let stats_after_second = cache.storage().stats(); + + // No new entries on second run (all predicate columns already cached) + assert_eq!( + stats_after_second.total_entries, entries_after_first, + "No new cache entries on hot run" + ); + + // eval_predicate should be > 0 (OS predicate evaluated from cache on 2nd run) + assert!( + stats_after_second.runtime.eval_predicate > 0, + "Predicate should be evaluated from cache on hot run, got eval_predicate={}", + stats_after_second.runtime.eval_predicate + ); + + // get (projection read) should be 0 — URL is not cached + assert_eq!( + stats_after_second.runtime.get, 0, + "Projection column URL should NOT be read from cache (get={})", + stats_after_second.runtime.get + ); +} + +/// Test that when a column is used in BOTH predicate AND projection, +/// it IS cached and served from cache for both purposes. +#[tokio::test] +async fn test_predicate_column_in_projection_is_cached() { + let cache_dir = TempDir::new().unwrap(); + + // Query: OS is in both WHERE and SELECT + let sql = r#"SELECT "OS", COUNT(*) as cnt FROM hits WHERE "OS" > 0 GROUP BY "OS" ORDER BY cnt DESC LIMIT 5"#; + + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = 2; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_max_memory_bytes(usize::MAX) + .with_cache_dir(cache_dir.path().to_path_buf()) + .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .build(config) + .await + .unwrap(); + + ctx.register_parquet("hits", TEST_FILE, ParquetReadOptions::default()) + .await + .unwrap(); + + // First run: fills cache + let plan = get_physical_plan(sql, &ctx).await; + let batches_1 = collect(plan, ctx.task_ctx()).await.unwrap(); + let result_1 = pretty_format_batches(&batches_1).unwrap().to_string(); + + // Clear counters + cache.storage().stats(); + + // Second run: should hit cache for BOTH predicate eval AND projection read + let plan = get_physical_plan(sql, &ctx).await; + let batches_2 = collect(plan, ctx.task_ctx()).await.unwrap(); + let result_2 = pretty_format_batches(&batches_2).unwrap().to_string(); + + assert_eq!(result_1, result_2, "Results should be the same on hot run"); + + let stats = cache.storage().stats(); + + // OS is a predicate column → should be cached + // eval_predicate > 0 (WHERE clause evaluated from cache) + assert!( + stats.runtime.eval_predicate > 0, + "Predicate column OS should be evaluated from cache, got eval_predicate={}", + stats.runtime.eval_predicate + ); + + // OS is also in projection → get_with_selection > 0 or get > 0 + // Since OS is a predicate column, it IS cached and should be read from cache + let projection_reads = stats.runtime.get + stats.runtime.get_with_selection; + assert!( + projection_reads > 0, + "Predicate column OS in projection should be read from cache, got get={}, get_with_selection={}", + stats.runtime.get, stats.runtime.get_with_selection + ); +} + +/// Test that string predicate columns are NOT cached (only numeric predicates are cached). +/// Query uses a string predicate (URL LIKE) and a numeric predicate (OS > 0). +/// Only the numeric predicate column should be cached. +#[tokio::test] +async fn test_only_numeric_predicate_columns_cached() { + let cache_dir = TempDir::new().unwrap(); + + // Query: WHERE on "URL" (string predicate) AND "OS" (numeric predicate), SELECT "OS" + // Only OS should be cached; URL is a string predicate and should NOT be cached. + let sql = r#"SELECT "OS" FROM hits WHERE "URL" LIKE '%tours%' AND "OS" > 0 ORDER BY "OS" LIMIT 5"#; + + let mut config = SessionConfig::new(); + config.options_mut().execution.target_partitions = 2; + let (ctx, cache) = LiquidCacheLocalBuilder::new() + .with_max_memory_bytes(usize::MAX) + .with_cache_dir(cache_dir.path().to_path_buf()) + .with_squeeze_policy(Box::new(TranscodeSqueezeEvict)) + .with_cache_policy(Box::new(LiquidPolicy::new())) + .build(config) + .await + .unwrap(); + + ctx.register_parquet("hits", TEST_FILE, ParquetReadOptions::default()) + .await + .unwrap(); + + // First run: fills cache (only numeric predicate columns) + let plan = get_physical_plan(sql, &ctx).await; + let batches_1 = collect(plan, ctx.task_ctx()).await.unwrap(); + let result_1 = pretty_format_batches(&batches_1).unwrap().to_string(); + + let stats_after_first = cache.storage().stats(); + let entries_after_first = stats_after_first.total_entries; + + // Should have cached only OS (numeric predicate), NOT URL (string predicate) + assert!( + entries_after_first > 0, + "Cache should have entries for numeric predicate column OS" + ); + + // The memory should be small — OS is Int16, not a large string column. + // If URL were cached, memory would be much larger (URL is a huge string column). + // OS column for nano_hits: ~1000 rows × 2 bytes = ~2KB per batch. + // URL column would be megabytes. So we check memory is modest. + assert!( + stats_after_first.memory_usage_bytes < 1024 * 1024, // less than 1MB + "Cache should be small (only OS cached, not URL). Got {} bytes", + stats_after_first.memory_usage_bytes + ); + + // Second run: should reuse cache + let plan = get_physical_plan(sql, &ctx).await; + let batches_2 = collect(plan, ctx.task_ctx()).await.unwrap(); + let result_2 = pretty_format_batches(&batches_2).unwrap().to_string(); + + assert_eq!(result_1, result_2, "Results should be identical on hot run"); + + let stats_after_second = cache.storage().stats(); + + // No new entries on second run + assert_eq!( + stats_after_second.total_entries, entries_after_first, + "No new cache entries on hot run" + ); + + // eval_predicate should be > 0 (OS numeric predicate evaluated from cache) + assert!( + stats_after_second.runtime.eval_predicate > 0, + "Numeric predicate OS should be evaluated from cache, got eval_predicate={}", + stats_after_second.runtime.eval_predicate + ); +} diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index 65870a78..a6f893b6 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -22,6 +22,9 @@ pub struct CachedColumn { field: Arc, column_path: ColumnAccessPath, expression: Option>, + /// Whether this column is used in a predicate (WHERE clause). + /// In predicate-only mode, only predicate columns are cached. + is_predicate_column: bool, } /// A reference to a cached column. @@ -83,6 +86,7 @@ impl CachedColumn { cache_store, column_path: column_access_path, expression, + is_predicate_column, } } @@ -105,6 +109,11 @@ impl CachedColumn { self.expression.clone() } + /// Returns whether this column is a predicate column (used in WHERE clause). + pub fn is_predicate_column(&self) -> bool { + self.is_predicate_column + } + fn array_to_record_batch(&self, array: ArrayRef) -> RecordBatch { let schema = Arc::new(Schema::new(vec![self.field.clone()])); RecordBatch::try_new(schema, vec![array]).unwrap() @@ -166,18 +175,30 @@ impl CachedColumn { } /// Get an arrow array with a filter applied. + /// Returns None for non-predicate columns or string predicate columns + /// (only numeric predicate columns are cached and served from cache). pub async fn get_arrow_array_with_filter( &self, batch_id: BatchID, filter: &BooleanBuffer, ) -> Option { + if !self.is_predicate_column || is_string_type(self.field.data_type()) { + return None; + } let entry_id = self.entry_id(batch_id).into(); - self.cache_store + let result = self + .cache_store .get(&entry_id) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() - .await + .await; + if result.is_some() { + self.cache_store.observer().runtime_stats().incr_cache_hit(); + } else { + self.cache_store.observer().runtime_stats().incr_cache_miss(); + } + result } #[cfg(test)] @@ -187,11 +208,17 @@ impl CachedColumn { } /// Insert an array into the cache. + /// Only numeric predicate columns are cached; string predicates and + /// non-predicate columns return CacheFull. pub async fn insert( self: &Arc, batch_id: BatchID, array: ArrayRef, ) -> Result<(), InsertArrowArrayError> { + if !self.is_predicate_column || is_string_type(self.field.data_type()) { + return Err(InsertArrowArrayError::CacheFull); + } + if self.is_cached(batch_id) { return Err(InsertArrowArrayError::AlreadyCached); } @@ -206,6 +233,7 @@ impl CachedColumn { fn is_string_type(data_type: &DataType) -> bool { match data_type { DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 => true, + DataType::Binary | DataType::BinaryView | DataType::LargeBinary => true, DataType::Dictionary(_, value_type) => is_string_type(value_type.as_ref()), _ => false, } diff --git a/src/datafusion/src/lib.rs b/src/datafusion/src/lib.rs index 2c631f69..e1fdb0a3 100644 --- a/src/datafusion/src/lib.rs +++ b/src/datafusion/src/lib.rs @@ -13,4 +13,10 @@ pub use liquid_cache as storage; pub use liquid_cache_common as common; pub use reader::variant_udf::{VariantGetUdf, VariantPretty, VariantToJsonUdf}; pub use reader::{FilterCandidateBuilder, LiquidParquetSource, LiquidPredicate, LiquidRowFilter}; +pub use reader::plantime::engagement_policy::{ + AlwaysEngagePolicy, CacheEngagementPolicy, DEFAULT_SELECTIVITY_THRESHOLD, + EngagementContext, EngagementDecision, NeverEngagePolicy, SelectivityThresholdPolicy, + default_engagement_policy, +}; +pub use reader::plantime::source::pre_seed_metadata_cache; pub use utils::{boolean_buffer_and_then, extract_execution_metrics}; diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index b2490580..08975181 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -80,6 +80,20 @@ pub fn rewrite_data_source_plan( rewritten.data } +/// Returns true if a data type is uncacheable by LC (string/binary). +fn is_uncacheable_type(dt: &arrow_schema::DataType) -> bool { + use arrow_schema::DataType; + matches!( + dt, + DataType::Utf8 + | DataType::Utf8View + | DataType::LargeUtf8 + | DataType::Binary + | DataType::BinaryView + | DataType::LargeBinary + ) || matches!(dt, DataType::Dictionary(_, v) if is_uncacheable_type(v)) +} + fn try_optimize_parquet_source( plan: Arc, cache: &LiquidCacheParquetRef, @@ -89,6 +103,59 @@ fn try_optimize_parquet_source( && let Some((file_scan_config, parquet_source)) = data_source_exec.downcast_to_file_source::() { + // Skip LC wrapping if: + // - Output has zero columns (COUNT(*) — just needs row count from metadata) + // - ANY output column is string/binary (LC can't cache, fallback negates hits) + // - Predicate references a string column + let output_schema = plan.schema(); + if output_schema.fields().is_empty() { + log::debug!("[LC-Optimizer] SKIP: empty projection (COUNT(*))"); + return Ok(Transformed::no(plan)); + } + + // Skip LC when too many output columns — per-column cache overhead + // exceeds decode savings for wide projections. + const MAX_LC_COLUMNS: usize = 4; + if output_schema.fields().len() > MAX_LC_COLUMNS { + log::debug!( + "[LC-Optimizer] SKIP: too many columns ({} > {})", + output_schema.fields().len(), MAX_LC_COLUMNS + ); + return Ok(Transformed::no(plan)); + } + + let has_string_output = output_schema + .fields() + .iter() + .any(|f| is_uncacheable_type(f.data_type())); + + let predicate_has_string = parquet_source.filter().map_or(false, |pred| { + use datafusion::physical_expr::utils::collect_columns; + let file_schema = file_scan_config.file_schema(); + let cols = collect_columns(&pred); + cols.iter().any(|col| { + file_schema + .fields() + .get(col.index()) + .map_or(false, |f| is_uncacheable_type(f.data_type())) + }) + }); + + if has_string_output || predicate_has_string { + log::debug!( + "[LC-Optimizer] SKIP: string_in_output={}, string_in_predicate={}, output_cols={}", + has_string_output, predicate_has_string, output_schema.fields().len() + ); + return Ok(Transformed::no(plan)); + } + + let num_fields = output_schema.fields().len(); + let has_predicate = parquet_source.filter().is_some(); + log::debug!( + "[LC-Optimizer] WRAP: all {} output columns cacheable, predicate={}", + num_fields, has_predicate + ); + let mut new_config = file_scan_config.clone(); let mut new_source = diff --git a/src/datafusion/src/reader/mod.rs b/src/datafusion/src/reader/mod.rs index db64c969..0436b5e6 100644 --- a/src/datafusion/src/reader/mod.rs +++ b/src/datafusion/src/reader/mod.rs @@ -1,4 +1,4 @@ -mod plantime; +pub mod plantime; mod runtime; mod utils; pub(crate) mod variant_udf; diff --git a/src/datafusion/src/reader/plantime/engagement_policy.rs b/src/datafusion/src/reader/plantime/engagement_policy.rs new file mode 100644 index 00000000..f509d095 --- /dev/null +++ b/src/datafusion/src/reader/plantime/engagement_policy.rs @@ -0,0 +1,90 @@ +//! Cache engagement policy — decides per-file whether LC should serve +//! decoded batches from cache or delegate to plain parquet. + +use std::fmt::Debug; +use std::sync::Arc; + +/// Context passed to the policy at file-open time. +#[derive(Debug, Clone)] +pub struct EngagementContext { + /// Fraction of rows surviving RG + page pruning (0.0–1.0). + pub estimated_selectivity: f64, + /// Whether the query has a pushdown predicate on this file. + pub has_predicate: bool, + /// Total rows across selected row groups. + pub total_rows: usize, + /// Rows that survived pruning. + pub selected_rows: usize, + /// File path for logging. + pub file_path: String, +} + +/// Decision returned by the policy. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EngagementDecision { + /// Serve decoded batches from cache, fill on miss. + UseLiquidCache, + /// Bypass LC, delegate to plain parquet reader. + DelegateToParquet, +} + +/// Decides per-file whether Liquid Cache should engage or delegate. +pub trait CacheEngagementPolicy: Send + Sync + Debug { + fn decide(&self, ctx: &EngagementContext) -> EngagementDecision; +} + +/// Default selectivity threshold below which LC delegates to parquet. +pub const DEFAULT_SELECTIVITY_THRESHOLD: f64 = 0.5; + +/// Delegate when selectivity < threshold AND a predicate exists. +#[derive(Debug, Clone)] +pub struct SelectivityThresholdPolicy { + pub threshold: f64, +} + +impl Default for SelectivityThresholdPolicy { + fn default() -> Self { + Self { threshold: DEFAULT_SELECTIVITY_THRESHOLD } + } +} + +impl SelectivityThresholdPolicy { + pub fn new(threshold: f64) -> Self { + Self { threshold } + } +} + +impl CacheEngagementPolicy for SelectivityThresholdPolicy { + fn decide(&self, ctx: &EngagementContext) -> EngagementDecision { + if ctx.estimated_selectivity < self.threshold && ctx.has_predicate { + EngagementDecision::DelegateToParquet + } else { + EngagementDecision::UseLiquidCache + } + } +} + +/// Always use LC regardless of selectivity. +#[derive(Debug, Clone)] +pub struct AlwaysEngagePolicy; + +impl CacheEngagementPolicy for AlwaysEngagePolicy { + fn decide(&self, _ctx: &EngagementContext) -> EngagementDecision { + EngagementDecision::UseLiquidCache + } +} + +/// Never use LC — always delegate to plain parquet. +#[derive(Debug, Clone)] +pub struct NeverEngagePolicy; + +impl CacheEngagementPolicy for NeverEngagePolicy { + fn decide(&self, _ctx: &EngagementContext) -> EngagementDecision { + EngagementDecision::DelegateToParquet + } +} + +/// Returns the default engagement policy. +pub fn default_engagement_policy() -> Arc { + Arc::new(SelectivityThresholdPolicy::default()) +} diff --git a/src/datafusion/src/reader/plantime/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index b52a9bcf..36576d66 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -3,9 +3,15 @@ pub(crate) use source::CachedMetaReaderFactory; pub use source::LiquidParquetSource; pub(crate) use source::ParquetMetadataCacheReader; +pub mod engagement_policy; mod opener; mod row_filter; mod row_group_filter; -mod source; +pub mod source; +pub use engagement_policy::{ + AlwaysEngagePolicy, CacheEngagementPolicy, DEFAULT_SELECTIVITY_THRESHOLD, + EngagementContext, EngagementDecision, NeverEngagePolicy, SelectivityThresholdPolicy, + default_engagement_policy, +}; pub use row_filter::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}; diff --git a/src/datafusion/src/reader/plantime/opener.rs b/src/datafusion/src/reader/plantime/opener.rs index 7abe216e..059797a3 100644 --- a/src/datafusion/src/reader/plantime/opener.rs +++ b/src/datafusion/src/reader/plantime/opener.rs @@ -1,10 +1,17 @@ use std::sync::Arc; -use crate::{ - cache::LiquidCacheParquetRef, - reader::{ - plantime::{row_filter::build_row_filter, row_group_filter::RowGroupAccessPlanFilter}, - runtime::LiquidStreamBuilder, +use crate::{ + cache::LiquidCacheParquetRef, + reader::{ + plantime::{ + engagement_policy::{ + CacheEngagementPolicy, EngagementContext, EngagementDecision, + default_engagement_policy, + }, + row_filter::build_row_filter, + row_group_filter::RowGroupAccessPlanFilter, + }, + runtime::LiquidStreamBuilder, }, }; use ahash::AHashMap; @@ -15,7 +22,7 @@ use datafusion::{ datasource::{ listing::PartitionedFile, physical_plan::{ - FileOpenFuture, FileOpener, ParquetFileMetrics, + FileOpenFuture, FileOpener, ParquetFileMetrics, ParquetFileReaderFactory, parquet::{PagePruningAccessPlanFilter, ParquetAccessPlan}, }, table_schema::TableSchema, @@ -56,6 +63,10 @@ pub struct LiquidParquetOpener { liquid_cache: LiquidCacheParquetRef, expr_adapter_factory: Arc, span: Option>, + /// Optional caller-provided reader factory with pre-loaded metadata. + caller_reader_factory: Option>, + /// Policy that decides whether to use LC stream or delegate to parquet. + engagement_policy: Arc, } impl LiquidParquetOpener { @@ -73,6 +84,8 @@ impl LiquidParquetOpener { reorder_filters: bool, expr_adapter_factory: Arc, span: Option>, + caller_reader_factory: Option>, + engagement_policy: Arc, ) -> Self { Self { partition_index, @@ -87,6 +100,8 @@ impl LiquidParquetOpener { reorder_filters, expr_adapter_factory, span, + caller_reader_factory, + engagement_policy, } } } @@ -125,10 +140,26 @@ impl FileOpener for LiquidParquetOpener { let file_metrics = ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics); let metadata_size_hint = partitioned_file.metadata_size_hint; + let has_predicate = self.predicate.is_some(); + log::debug!( + "[LC-Opener] open file={}, predicate={}, batch_size={}, limit={:?}", + file_name, has_predicate, self.batch_size, self.limit + ); let lc = self.liquid_cache.clone(); let file_loc = partitioned_file.object_meta.location.to_string(); + // If caller provided a reader factory with pre-loaded metadata, create + // a reader from it. We'll use this for get_metadata() to avoid refetching. + let caller_metadata_reader = self.caller_reader_factory.as_ref().and_then(|factory| { + factory.create_reader( + self.partition_index, + partitioned_file.clone(), + metadata_size_hint, + &self.metrics, + ).ok() + }); + let mut async_file_reader = self.parquet_file_reader_factory.create_liquid_reader( self.partition_index, partitioned_file.clone(), @@ -169,6 +200,7 @@ impl FileOpener for LiquidParquetOpener { let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); let span = self.span.clone(); + let engagement_policy = Arc::clone(&self.engagement_policy); Ok(Box::pin(async move { // Prune this file using the file level statistics and partition values. // Since dynamic filters may have been updated since planning it is possible that we are able @@ -203,12 +235,22 @@ impl FileOpener for LiquidParquetOpener { .with_page_index_policy(parquet::file::metadata::PageIndexPolicy::Required); let mut metadata_timer = file_metrics.metadata_load_time.timer(); - // Begin by loading the metadata from the underlying reader (note - // the returned metadata may actually include page indexes as some - // readers may return page indexes even when not requested -- for - // example when they are cached) - let mut reader_metadata = - ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()).await?; + // Try to get metadata from caller's pre-loaded reader first (instant). + // Fall back to loading from the object store if not available. + let mut reader_metadata = if let Some(mut caller_reader) = caller_metadata_reader { + match caller_reader.get_metadata(Some(&options)).await { + Ok(meta) => { + log::debug!("[LC-Meta] REUSE from caller factory: {}", file_name); + ArrowReaderMetadata::try_new(meta, options.clone())? + } + Err(_) => { + log::debug!("[LC-Meta] caller factory failed, loading from store: {}", file_name); + ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()).await? + } + } + } else { + ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()).await? + }; // Note about schemas: we are actually dealing with **3 different schemas** here: // - The table schema as defined by the TableProvider. @@ -221,7 +263,7 @@ impl FileOpener for LiquidParquetOpener { Arc::clone(&logical_file_schema), Arc::clone(&physical_file_schema), )); - let cache_full_schema = Arc::clone(&physical_file_schema); + let cache_full_schema = Arc::clone(&physical_file_schema); options = options.with_schema(Arc::clone(&physical_file_schema)); reader_metadata = ArrowReaderMetadata::try_new(Arc::clone(reader_metadata.metadata()), options)?; @@ -327,8 +369,77 @@ impl FileOpener for LiquidParquetOpener { } let row_group_indexes = access_plan.row_group_indexes(); + + // Early exit: if all row groups were pruned, return empty stream + if row_group_indexes.is_empty() { + log::debug!("[LC-Opener] EMPTY: all RGs pruned, file={}", file_name); + return Ok(futures::stream::empty().boxed()); + } + let row_selection = access_plan.into_overall_row_selection(rg_metadata)?; + // Estimate selectivity from row_selection: how many rows survived + // RG pruning + page index pruning vs total rows in selected RGs. + let total_rows: usize = row_group_indexes + .iter() + .map(|&idx| rg_metadata[idx].num_rows() as usize) + .sum(); + let selected_rows = row_selection.as_ref() + .map(|sel| sel.row_count()) + .unwrap_or(total_rows); + let estimated_selectivity = if total_rows > 0 { + selected_rows as f64 / total_rows as f64 + } else { + 1.0 + }; + + // If selectivity is low (few rows match), the decode cost is already + // minimal — LC cache overhead would dominate for negligible savings. + // Delegate to plain parquet for fast pass-through. + // For high selectivity (most rows match = lots of decode), LC cache + // saves significant decode work on warm iterations. + let engagement_ctx = EngagementContext { + estimated_selectivity, + has_predicate: predicate.is_some(), + total_rows, + selected_rows, + file_path: file_name.clone(), + }; + + match engagement_policy.decide(&engagement_ctx) { + EngagementDecision::DelegateToParquet => { + log::debug!( + "[LC-Opener] DELEGATE to plain parquet: selectivity={:.3}, file={}", + estimated_selectivity, file_name + ); + let mut plain_builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + async_file_reader, + reader_metadata, + ) + .with_batch_size(batch_size) + .with_projection(mask) + .with_row_groups(row_group_indexes); + + if let Some(sel) = row_selection { + plain_builder = plain_builder.with_row_selection(sel); + } + if let Some(lim) = limit { + plain_builder = plain_builder.with_limit(lim); + } + + let stream = plain_builder.build()?; + let adapted = stream + .map_err(|e| DataFusionError::External(Box::new(e))); + return Ok(adapted.boxed()); + } + EngagementDecision::UseLiquidCache => { + log::debug!( + "[LC-Opener] LC STREAM: selectivity={:.3}, predicate={}, file={}", + estimated_selectivity, predicate.is_some(), file_name + ); + } + } + let mut liquid_builder = LiquidStreamBuilder::new(async_file_reader, Arc::clone(reader_metadata.metadata())) .with_batch_size(batch_size) diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 6c193c3f..626deaf8 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -1,5 +1,6 @@ use super::opener::LiquidParquetOpener; use crate::cache::LiquidCacheParquetRef; +use crate::reader::plantime::engagement_policy::{CacheEngagementPolicy, default_engagement_policy}; use ahash::{HashMap, HashMapExt}; use arrow_schema::Schema; use bytes::Bytes; @@ -40,6 +41,15 @@ use tokio::sync::RwLock; static META_CACHE: LazyLock = LazyLock::new(MetadataCache::new); +/// Pre-seed the metadata cache with already-loaded metadata. +/// Callers that have pre-loaded parquet metadata (e.g., from a custom +/// ParquetFileReaderFactory) can inject it here so that LC's opener +/// skips the expensive `ArrowReaderMetadata::load_async()` call. +pub async fn pre_seed_metadata_cache(path: &Path, metadata: Arc) { + let mut cache = META_CACHE.val.write().await; + cache.entry(path.clone()).or_insert(metadata); +} + #[derive(Debug)] pub(crate) struct CachedMetaReaderFactory { store: Arc, @@ -138,6 +148,7 @@ impl AsyncFileReader for ParquetMetadataCacheReader { { let cache = META_CACHE.val.read().await; if let Some(meta) = cache.get(&path) { + log::debug!("[LC-Meta] HIT path={}", path); return Ok(meta.clone()); } } @@ -145,8 +156,12 @@ impl AsyncFileReader for ParquetMetadataCacheReader { // Upgrade to write lock and double-check let mut cache = META_CACHE.val.write().await; match cache.entry(path.clone()) { - std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()), + std::collections::hash_map::Entry::Occupied(entry) => { + log::debug!("[LC-Meta] HIT (race) path={}", path); + Ok(entry.get().clone()) + } std::collections::hash_map::Entry::Vacant(entry) => { + log::debug!("[LC-Meta] MISS (loading from store) path={}", path); let meta = self.inner.get_metadata(options.as_ref()).await?; let meta = Arc::try_unwrap(meta).unwrap_or_else(|e| e.as_ref().clone()); let mut reader = ParquetMetaDataReader::new_with_metadata(meta.clone()) @@ -175,6 +190,12 @@ pub struct LiquidParquetSource { projection: ProjectionExprs, table_schema: TableSchema, span: Option>, + /// Optional caller-provided reader factory with pre-loaded metadata. + /// When set, LC's opener uses this to get metadata instantly instead of + /// fetching from the object store. + parquet_file_reader_factory: Option>, + /// Policy that decides per-file whether to use LC stream or delegate to parquet. + engagement_policy: Arc, } impl LiquidParquetSource { @@ -232,9 +253,41 @@ impl LiquidParquetSource { self } + /// Set predicate for row_filter only — no page-index pruning. + /// Used by the indexed path where the BoolNode's RowSelection is authoritative + /// and page-level statistics must not override it. + pub fn with_predicate_no_page_pruning( + mut self, + file_schema: Arc, + predicate: Arc, + ) -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let predicate_creation_errors = + MetricBuilder::new(&metrics).global_counter("num_predicate_creation_errors"); + self.metrics = metrics; + self.predicate = Some(Arc::clone(&predicate)); + + match PruningPredicate::try_new(Arc::clone(&predicate), Arc::clone(&file_schema)) { + Ok(pruning_predicate) => { + if !pruning_predicate.always_true() { + self.pruning_predicate = Some(Arc::new(pruning_predicate)); + } + } + Err(e) => { + log::debug!("Could not create pruning predicate for: {e}"); + predicate_creation_errors.add(1); + } + }; + + // Deliberately skip page_pruning_predicate — page-index stats must not + // prune pages that the caller's RowSelection already validated. + self + } + /// Create a new LiquidParquetSource from a ParquetSource pub fn from_parquet_source(source: ParquetSource, liquid_cache: LiquidCacheParquetRef) -> Self { let predicate = source.filter(); + let reader_factory = source.parquet_file_reader_factory().cloned(); let table_schema = source.table_schema().clone(); let file_schema = table_schema.file_schema().clone(); @@ -256,6 +309,8 @@ impl LiquidParquetSource { pruning_predicate: None, page_pruning_predicate: None, span: None, + parquet_file_reader_factory: reader_factory, + engagement_policy: default_engagement_policy(), }; if let Some(predicate) = predicate { @@ -265,6 +320,12 @@ impl LiquidParquetSource { v } + /// Set a custom cache engagement policy. + pub fn with_engagement_policy(mut self, policy: Arc) -> Self { + self.engagement_policy = policy; + self + } + /// Get the predicate for the LiquidParquetSource pub fn predicate(&self) -> Option> { self.predicate.clone() @@ -289,6 +350,14 @@ impl FileSource for LiquidParquetSource { let reader_factory = Arc::new(CachedMetaReaderFactory::new(object_store)); + // If the caller provided a ParquetFileReaderFactory (with pre-loaded + // metadata), pre-seed LC's metadata cache for all files in this partition. + // This avoids expensive ArrowReaderMetadata::load_async() calls in the + // opener for files whose metadata we already have. + // Pass caller's reader factory to the opener so it can use pre-loaded + // metadata (avoids re-fetching from object store). + let caller_reader_factory = self.parquet_file_reader_factory.clone(); + let execution_span = self .span .clone() @@ -307,6 +376,8 @@ impl FileSource for LiquidParquetSource { self.reorder_filters(), expr_adapter_factory, execution_span.map(Arc::new), + caller_reader_factory, + Arc::clone(&self.engagement_policy), ); Ok(Arc::new(opener)) diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 4a5a6b27..549d6b6e 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -356,8 +356,14 @@ impl LiquidCacheReaderInner { return Ok(Some(batch)); } - let mut arrays = Vec::with_capacity(self.projection_columns.len()); - for column_idx in self.projection_columns.clone() { + // Phase 1: Try cache for all columns, collect hits and track misses. + let mut arrays: Vec> = + vec![None; self.projection_columns.len()]; + let mut has_miss = false; + let mut hit_count = 0usize; + let mut miss_count = 0usize; + + for (i, &column_idx) in self.projection_columns.iter().enumerate() { let column = self .cached_row_group .get_column(column_idx as u64) @@ -371,22 +377,44 @@ impl LiquidCacheReaderInner { .get_arrow_array_with_filter(self.current_batch_id, selection) .await; - let array = match array { - Some(array) => array, - None => { - let record_batch = self - .read_parquet_batch_and_fill_cache(self.current_batch_id) - .await?; + if let Some(arr) = array { + arrays[i] = Some(arr); + hit_count += 1; + } else { + has_miss = true; + miss_count += 1; + } + } + + if *self.current_batch_id == 0 { + log::debug!( + "[LC-Reader] batch_id={}, selected_rows={}, cols={}, hits={}, misses={}, fallback={}", + *self.current_batch_id, + selected_rows, + self.projection_columns.len(), + hit_count, + miss_count, + has_miss, + ); + } + + // Phase 2: If any columns missed, read from parquet ONCE for all misses. + if has_miss { + let record_batch = self + .read_parquet_batch_and_fill_cache(self.current_batch_id) + .await?; + + for (i, &column_idx) in self.projection_columns.iter().enumerate() { + if arrays[i].is_none() { let array = self.parquet_array(&record_batch, column_idx)?; - filter_array(array, selection)? + arrays[i] = Some(filter_array(array, selection)?); } - }; - - arrays.push(array); + } } + let final_arrays: Vec = arrays.into_iter().map(|a| a.unwrap()).collect(); Ok(Some( - RecordBatch::try_new(self.schema.clone(), arrays).unwrap(), + RecordBatch::try_new(self.schema.clone(), final_arrays).unwrap(), )) } @@ -406,28 +434,21 @@ impl LiquidCacheReaderInner { .await .map_err(|e| ArrowError::ComputeError(format!("parquet fallback read failed: {e}")))?; - for (col_idx, file_column_id) in self - .parquet_fallback - .cache_column_ids - .iter() - .copied() - .enumerate() - { - let column = self - .cached_row_group - .get_column(file_column_id as u64) - .ok_or_else(|| { - ArrowError::ComputeError(format!( - "column {file_column_id} not present in liquid cache" - )) - })?; - let array = Arc::clone(record_batch.column(col_idx)); - - match column.insert(batch_id, array).await { - Ok(()) | Err(InsertArrowArrayError::AlreadyCached) => {} - Err(InsertArrowArrayError::CacheFull) => {} + // Spawn cache fill asynchronously — transcoding Arrow→Liquid should not + // block the query hot path. The batch is returned immediately; cache + // population happens in the background. + let cached_row_group = self.cached_row_group.clone(); + let cache_column_ids = self.parquet_fallback.cache_column_ids.clone(); + let batch_for_cache = record_batch.clone(); + tokio::spawn(async move { + for (col_idx, file_column_id) in cache_column_ids.iter().copied().enumerate() { + let Some(column) = cached_row_group.get_column(file_column_id as u64) else { + continue; + }; + let array = Arc::clone(batch_for_cache.column(col_idx)); + let _ = column.insert(batch_id, array).await; } - } + }); self.last_pull = Some((batch_id, record_batch.clone())); Ok(record_batch) diff --git a/src/datafusion/src/reader/runtime/liquid_stream.rs b/src/datafusion/src/reader/runtime/liquid_stream.rs index 2933d5d2..34ec325f 100644 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ b/src/datafusion/src/reader/runtime/liquid_stream.rs @@ -106,14 +106,30 @@ impl ReaderFactory { let schema_descr = self.metadata.file_metadata().schema_descr(); let cache_column_ids = get_root_column_ids(schema_descr, &cache_projection); + // When no predicate is present, all projected columns are cacheable. + // Without this, `is_predicate_column` is false for every column and + // the cache is effectively disabled (get returns None, insert rejects). let predicate_column_ids = if let Some(ref predicate_projection) = predicate_projection { get_root_column_ids(schema_descr, predicate_projection) } else { - Vec::new() + cache_column_ids.clone() }; + + if row_group_idx == 0 { + log::debug!( + "[LC-Stream] plan_row_group: rg={}, cache_cols={:?}, predicate_cols={:?}, \ + has_filter={}, rows={}", + row_group_idx, + &cache_column_ids, + &predicate_column_ids, + self.filter.is_some(), + self.metadata.row_group(row_group_idx).num_rows(), + ); + } + let cached_row_group = self .cached_file - .create_row_group(row_group_idx as u64, predicate_column_ids); + .create_row_group(row_group_idx as u64, predicate_column_ids.clone()); let projection_column_ids = get_root_column_ids(schema_descr, &projection);