From a59a5555456f3e4218afd857d5477370ddc78c29 Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Thu, 21 May 2026 16:11:55 +0530 Subject: [PATCH 01/12] Cache only numeric predicate columns, skip string/binary and projection-only columns Signed-off-by: Tanvir Alam --- src/core/src/cache/observer/mod.rs | 2 +- src/core/src/cache/observer/stats.rs | 2 + src/datafusion-local/src/tests/mod.rs | 203 ++++++++++++++++++++++++++ src/datafusion/src/cache/column.rs | 32 +++- 4 files changed, 236 insertions(+), 3 deletions(-) diff --git a/src/core/src/cache/observer/mod.rs b/src/core/src/cache/observer/mod.rs index 38de5fc45..5b848b071 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 fa0c3d9ae..23963248a 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/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index f92b9852d..d297c836a 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 65870a78f..a6f893b69 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, } From a2138d448230b6e01f69393d79d17be8093d08c3 Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Thu, 21 May 2026 16:12:31 +0530 Subject: [PATCH 02/12] Add LRU cache eviction policy with type-aware queues Signed-off-by: Tanvir Alam --- src/core/src/cache/policies/cache/lru.rs | 291 +++++++++++++++++++++++ src/core/src/cache/policies/cache/mod.rs | 103 ++++++++ 2 files changed, 394 insertions(+) create mode 100644 src/core/src/cache/policies/cache/lru.rs 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 000000000..b8968c479 --- /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 352f83a07..7399a3d46 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)]); + } } From 62f9ab105383d57f39861353f649b9c80b1c0b00 Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Thu, 21 May 2026 16:12:55 +0530 Subject: [PATCH 03/12] Add with_max_disk_bytes to LiquidCacheLocalBuilder Signed-off-by: Tanvir Alam --- src/datafusion-local/src/lib.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 81a259aa1..1e0cd4b19 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -59,6 +59,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 +77,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 +105,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 +164,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 +176,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, From e7484d3eb155909cb4f3755cad3ccbd7b028f35f Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Mon, 25 May 2026 03:10:24 +0530 Subject: [PATCH 04/12] Make cache budget limits dynamically resizable at runtime Signed-off-by: Tanvir Alam --- src/core/src/cache/budget.rs | 30 ++++++++++++++++++++++++------ src/core/src/cache/core.rs | 4 ++-- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/src/core/src/cache/budget.rs b/src/core/src/cache/budget.rs index 2b7bc5d0a..89b1c7a45 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 84e0a795f..3a702c735 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, } } From c9415406cff433547dd74d7a0ca8b16c52ea7223 Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Thu, 4 Jun 2026 23:52:39 +0530 Subject: [PATCH 05/12] Export LiquidParquetSource and LiquidCacheParquetRef from datafusion-local crate --- src/datafusion-local/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index 1e0cd4b19..c6cf23bdd 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -20,6 +20,8 @@ use liquid_cache_datafusion::{ pub use liquid_cache as storage; pub use liquid_cache_common as common; +pub use liquid_cache_datafusion::LiquidParquetSource; +pub use liquid_cache_datafusion::LiquidCacheParquetRef; /// Builder for in-process liquid cache session context /// From 5e113e9035ed613021d1214ec353e2efc82fb9a0 Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Fri, 5 Jun 2026 00:14:08 +0530 Subject: [PATCH 06/12] Fix re-export conflict: LiquidCacheParquetRef defined twice Remove LiquidCacheParquetRef from the private use statement (line 18) since it's already made available via pub use (line 24). Having it in both causes E0252 (name defined multiple times). Also consolidate both pub use statements into a single line. Signed-off-by: Bukhtawar Khan --- src/datafusion-local/src/lib.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/datafusion-local/src/lib.rs b/src/datafusion-local/src/lib.rs index c6cf23bdd..cfe56bb6b 100644 --- a/src/datafusion-local/src/lib.rs +++ b/src/datafusion-local/src/lib.rs @@ -15,13 +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::LiquidParquetSource; -pub use liquid_cache_datafusion::LiquidCacheParquetRef; +pub use liquid_cache_datafusion::{LiquidCacheParquetRef, LiquidParquetSource}; /// Builder for in-process liquid cache session context /// From cb96cceec7a302ce3f9168c01674bdd21254ca8d Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Fri, 5 Jun 2026 20:10:54 +0530 Subject: [PATCH 07/12] Fix cache disabled when no predicate and make insert async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes for the OpenSearch indexed-table integration path where predicate=None (filtering is handled externally by the BoolNode evaluator): 1. When no row filter is present, treat all projected columns as cacheable (predicate_column_ids = cache_column_ids). Previously, an empty predicate_column_ids meant is_predicate_column=false for all columns, causing get/insert to always bail out — the cache was effectively a no-op passthrough with pure overhead. 2. Move cache insert (Arrow→Liquid transcoding) into tokio::spawn so it runs asynchronously. The query batch is returned immediately without waiting for cache population. This ensures cache MISS has near-zero overhead vs the non-LC path. Together these ensure: repeated numeric queries get cache HITs (served from in-memory Arrow arrays), first execution has minimal overhead, and string columns are still correctly rejected by the is_string_type guard. Signed-off-by: Bukhtawar Khan --- .../src/reader/runtime/liquid_cache_reader.rs | 35 ++++++++----------- .../src/reader/runtime/liquid_stream.rs | 5 ++- 2 files changed, 18 insertions(+), 22 deletions(-) diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 4a5a6b27b..601761438 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -406,28 +406,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 2933d5d2a..21444d29c 100644 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ b/src/datafusion/src/reader/runtime/liquid_stream.rs @@ -106,10 +106,13 @@ 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() }; let cached_row_group = self .cached_file From 3dcc0e458b41ea96ee990ff6fa8524c12bde4361 Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Fri, 5 Jun 2026 20:10:54 +0530 Subject: [PATCH 08/12] Fix cache disabled when no predicate and make insert async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes for the OpenSearch indexed-table integration path where predicate=None (filtering is handled externally by the BoolNode evaluator): 1. When no row filter is present, treat all projected columns as cacheable (predicate_column_ids = cache_column_ids). Previously, an empty predicate_column_ids meant is_predicate_column=false for all columns, causing get/insert to always bail out — the cache was effectively a no-op passthrough with pure overhead. 2. Move cache insert (Arrow→Liquid transcoding) into tokio::spawn so it runs asynchronously. The query batch is returned immediately without waiting for cache population. This ensures cache MISS has near-zero overhead vs the non-LC path. Together these ensure: repeated numeric queries get cache HITs (served from in-memory Arrow arrays), first execution has minimal overhead, and string columns are still correctly rejected by the is_string_type guard. Signed-off-by: Bukhtawar Khan --- src/datafusion/src/lib.rs | 1 + src/datafusion/src/optimizers/mod.rs | 67 ++++++++++ src/datafusion/src/reader/mod.rs | 2 +- src/datafusion/src/reader/plantime/mod.rs | 2 +- src/datafusion/src/reader/plantime/opener.rs | 118 ++++++++++++++++-- src/datafusion/src/reader/plantime/source.rs | 31 ++++- .../src/reader/runtime/liquid_cache_reader.rs | 89 ++++++++----- .../src/reader/runtime/liquid_stream.rs | 20 ++- 8 files changed, 278 insertions(+), 52 deletions(-) diff --git a/src/datafusion/src/lib.rs b/src/datafusion/src/lib.rs index 2c631f698..9abc3099d 100644 --- a/src/datafusion/src/lib.rs +++ b/src/datafusion/src/lib.rs @@ -13,4 +13,5 @@ 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::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 b24905800..256ed9a66 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::info!("[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::info!( + "[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::info!( + "[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::info!( + "[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 db64c969e..0436b5e67 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/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index b52a9bcf3..15524da53 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -6,6 +6,6 @@ pub(crate) use source::ParquetMetadataCacheReader; mod opener; mod row_filter; mod row_group_filter; -mod source; +pub mod source; 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 7abe216e9..1d3714b5c 100644 --- a/src/datafusion/src/reader/plantime/opener.rs +++ b/src/datafusion/src/reader/plantime/opener.rs @@ -1,10 +1,10 @@ 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::{row_filter::build_row_filter, row_group_filter::RowGroupAccessPlanFilter}, + runtime::LiquidStreamBuilder, }, }; use ahash::AHashMap; @@ -15,7 +15,7 @@ use datafusion::{ datasource::{ listing::PartitionedFile, physical_plan::{ - FileOpenFuture, FileOpener, ParquetFileMetrics, + FileOpenFuture, FileOpener, ParquetFileMetrics, ParquetFileReaderFactory, parquet::{PagePruningAccessPlanFilter, ParquetAccessPlan}, }, table_schema::TableSchema, @@ -56,6 +56,8 @@ 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>, } impl LiquidParquetOpener { @@ -73,6 +75,7 @@ impl LiquidParquetOpener { reorder_filters: bool, expr_adapter_factory: Arc, span: Option>, + caller_reader_factory: Option>, ) -> Self { Self { partition_index, @@ -87,6 +90,7 @@ impl LiquidParquetOpener { reorder_filters, expr_adapter_factory, span, + caller_reader_factory, } } } @@ -125,10 +129,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::info!( + "[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(), @@ -203,12 +223,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::info!("[LC-Meta] REUSE from caller factory: {}", file_name); + ArrowReaderMetadata::try_new(meta, options.clone())? + } + Err(_) => { + log::info!("[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 +251,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 +357,70 @@ 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::info!("[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. + if estimated_selectivity < 0.5 && predicate.is_some() { + log::info!( + "[LC-Opener] DELEGATE to plain parquet: selectivity={:.3} (low, fast already), file={}", + estimated_selectivity, file_name + ); + // Low selectivity: few rows match, decode cost minimal. + // Plain parquet is faster than LC's per-batch cache overhead. + 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()); + } + + // High selectivity (many rows to decode) or no predicate → use LC + // Warm cache avoids repeated parquet decompress+decode. + log::info!( + "[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 6c193c3f0..c077be263 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -40,6 +40,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 +147,7 @@ impl AsyncFileReader for ParquetMetadataCacheReader { { let cache = META_CACHE.val.read().await; if let Some(meta) = cache.get(&path) { + log::info!("[LC-Meta] HIT path={}", path); return Ok(meta.clone()); } } @@ -145,8 +155,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::info!("[LC-Meta] HIT (race) path={}", path); + Ok(entry.get().clone()) + } std::collections::hash_map::Entry::Vacant(entry) => { + log::info!("[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 +189,10 @@ 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>, } impl LiquidParquetSource { @@ -235,6 +253,7 @@ impl LiquidParquetSource { /// 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 +275,7 @@ impl LiquidParquetSource { pruning_predicate: None, page_pruning_predicate: None, span: None, + parquet_file_reader_factory: reader_factory, }; if let Some(predicate) = predicate { @@ -289,6 +309,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 +335,7 @@ impl FileSource for LiquidParquetSource { self.reorder_filters(), expr_adapter_factory, execution_span.map(Arc::new), + caller_reader_factory, ); 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 4a5a6b27b..68dfc1ee5 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::info!( + "[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 2933d5d2a..a056a68d2 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::info!( + "[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); From 663af35a1501886729e016a11586ece950da1c83 Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Fri, 5 Jun 2026 20:10:54 +0530 Subject: [PATCH 09/12] Fix cache disabled when no predicate and make insert async MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes for the OpenSearch indexed-table integration path where predicate=None (filtering is handled externally by the BoolNode evaluator): 1. When no row filter is present, treat all projected columns as cacheable (predicate_column_ids = cache_column_ids). Previously, an empty predicate_column_ids meant is_predicate_column=false for all columns, causing get/insert to always bail out — the cache was effectively a no-op passthrough with pure overhead. 2. Move cache insert (Arrow→Liquid transcoding) into tokio::spawn so it runs asynchronously. The query batch is returned immediately without waiting for cache population. This ensures cache MISS has near-zero overhead vs the non-LC path. Together these ensure: repeated numeric queries get cache HITs (served from in-memory Arrow arrays), first execution has minimal overhead, and string columns are still correctly rejected by the is_string_type guard. Signed-off-by: Bukhtawar Khan --- src/datafusion/src/lib.rs | 1 + src/datafusion/src/optimizers/mod.rs | 67 ++++++++++ src/datafusion/src/reader/mod.rs | 2 +- src/datafusion/src/reader/plantime/mod.rs | 2 +- src/datafusion/src/reader/plantime/opener.rs | 116 ++++++++++++++++-- src/datafusion/src/reader/plantime/source.rs | 31 ++++- .../src/reader/runtime/liquid_cache_reader.rs | 104 +++++++++++----- .../src/reader/runtime/liquid_stream.rs | 25 +++- 8 files changed, 292 insertions(+), 56 deletions(-) diff --git a/src/datafusion/src/lib.rs b/src/datafusion/src/lib.rs index 2c631f698..9abc3099d 100644 --- a/src/datafusion/src/lib.rs +++ b/src/datafusion/src/lib.rs @@ -13,4 +13,5 @@ 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::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 b24905800..f5e011da8 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::info!("[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 = 10; + if output_schema.fields().len() > MAX_LC_COLUMNS { + log::info!( + "[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::info!( + "[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::info!( + "[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 db64c969e..0436b5e67 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/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index b52a9bcf3..15524da53 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -6,6 +6,6 @@ pub(crate) use source::ParquetMetadataCacheReader; mod opener; mod row_filter; mod row_group_filter; -mod source; +pub mod source; 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 7abe216e9..5a90ebe7d 100644 --- a/src/datafusion/src/reader/plantime/opener.rs +++ b/src/datafusion/src/reader/plantime/opener.rs @@ -1,10 +1,10 @@ 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::{row_filter::build_row_filter, row_group_filter::RowGroupAccessPlanFilter}, + runtime::LiquidStreamBuilder, }, }; use ahash::AHashMap; @@ -15,7 +15,7 @@ use datafusion::{ datasource::{ listing::PartitionedFile, physical_plan::{ - FileOpenFuture, FileOpener, ParquetFileMetrics, + FileOpenFuture, FileOpener, ParquetFileMetrics, ParquetFileReaderFactory, parquet::{PagePruningAccessPlanFilter, ParquetAccessPlan}, }, table_schema::TableSchema, @@ -56,6 +56,8 @@ 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>, } impl LiquidParquetOpener { @@ -73,6 +75,7 @@ impl LiquidParquetOpener { reorder_filters: bool, expr_adapter_factory: Arc, span: Option>, + caller_reader_factory: Option>, ) -> Self { Self { partition_index, @@ -87,6 +90,7 @@ impl LiquidParquetOpener { reorder_filters, expr_adapter_factory, span, + caller_reader_factory, } } } @@ -125,10 +129,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::info!( + "[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(), @@ -203,12 +223,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::info!("[LC-Meta] REUSE from caller factory: {}", file_name); + ArrowReaderMetadata::try_new(meta, options.clone())? + } + Err(_) => { + log::info!("[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 +251,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 +357,68 @@ 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::info!("[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 + }; + + // Only use LC STREAM when selectivity is high (>80% rows match) OR + // no predicate (pure cache for full scans). High selectivity means + // lots of decode work that LC cache can save on warm iterations. + // Low selectivity (<80%) means most rows already skipped by pruning, + // so decode cost is minimal and LC overhead isn't worth it. + const LC_STREAM_THRESHOLD: f64 = 0.8; + if predicate.is_some() && estimated_selectivity < LC_STREAM_THRESHOLD { + log::info!( + "[LC-Opener] DELEGATE: selectivity={:.3} (<{:.1}), rows={}/{}, file={}", + estimated_selectivity, LC_STREAM_THRESHOLD, selected_rows, total_rows, 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()); + } + + log::info!( + "[LC-Opener] LC STREAM: selectivity={:.3}, predicate={}, rows={}/{}, rgs={}, file={}", + estimated_selectivity, predicate.is_some(), selected_rows, total_rows, + row_group_indexes.len(), 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 6c193c3f0..c077be263 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -40,6 +40,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 +147,7 @@ impl AsyncFileReader for ParquetMetadataCacheReader { { let cache = META_CACHE.val.read().await; if let Some(meta) = cache.get(&path) { + log::info!("[LC-Meta] HIT path={}", path); return Ok(meta.clone()); } } @@ -145,8 +155,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::info!("[LC-Meta] HIT (race) path={}", path); + Ok(entry.get().clone()) + } std::collections::hash_map::Entry::Vacant(entry) => { + log::info!("[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 +189,10 @@ 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>, } impl LiquidParquetSource { @@ -235,6 +253,7 @@ impl LiquidParquetSource { /// 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 +275,7 @@ impl LiquidParquetSource { pruning_predicate: None, page_pruning_predicate: None, span: None, + parquet_file_reader_factory: reader_factory, }; if let Some(predicate) = predicate { @@ -289,6 +309,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 +335,7 @@ impl FileSource for LiquidParquetSource { self.reorder_filters(), expr_adapter_factory, execution_span.map(Arc::new), + caller_reader_factory, ); 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 4a5a6b27b..446abf45d 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::info!( + "[LC-Reader] batch_id={}, selected_rows={}, cols={:?}, hits={}, misses={}, fallback={}", + *self.current_batch_id, + selected_rows, + &self.projection_columns, + 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,36 @@ 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(); + let log_batch_id = *batch_id as usize; + 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 { + if log_batch_id == 0 { + log::info!("[LC-Insert] col {} not found in cached_row_group", file_column_id); + } + continue; + }; + let array = Arc::clone(batch_for_cache.column(col_idx)); + match column.insert(batch_id, array).await { + Ok(()) => { + if log_batch_id == 0 { + log::info!("[LC-Insert] OK batch={} col={}", log_batch_id, file_column_id); + } + } + Err(e) => { + if log_batch_id == 0 { + log::info!("[LC-Insert] FAIL batch={} col={} err={:?}", log_batch_id, file_column_id, e); + } + } + } } - } + }); 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 2933d5d2a..a12899d43 100644 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ b/src/datafusion/src/reader/runtime/liquid_stream.rs @@ -106,14 +106,27 @@ impl ReaderFactory { let schema_descr = self.metadata.file_metadata().schema_descr(); let cache_column_ids = get_root_column_ids(schema_descr, &cache_projection); - let predicate_column_ids = if let Some(ref predicate_projection) = predicate_projection { - get_root_column_ids(schema_descr, predicate_projection) - } else { - Vec::new() - }; + // All projected columns are cacheable regardless of whether a predicate + // exists. Without this, only predicate-referenced columns get + // is_predicate_column=true, and non-predicate projected columns are + // rejected by get/insert (returns None/CacheFull). + let predicate_column_ids = cache_column_ids.clone(); + + if row_group_idx == 0 { + log::info!( + "[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); From 8a8507748e8550ce9fd972f8a19cf5ac3f42658a Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Mon, 8 Jun 2026 17:58:59 +0530 Subject: [PATCH 10/12] Add CacheEngagementPolicy trait for pluggable selectivity-based delegation --- src/datafusion/src/lib.rs | 5 ++ .../src/reader/plantime/engagement_policy.rs | 90 +++++++++++++++++++ src/datafusion/src/reader/plantime/mod.rs | 6 ++ src/datafusion/src/reader/plantime/opener.rs | 83 ++++++++++------- src/datafusion/src/reader/plantime/source.rs | 11 +++ 5 files changed, 163 insertions(+), 32 deletions(-) create mode 100644 src/datafusion/src/reader/plantime/engagement_policy.rs diff --git a/src/datafusion/src/lib.rs b/src/datafusion/src/lib.rs index 9abc3099d..e1fdb0a37 100644 --- a/src/datafusion/src/lib.rs +++ b/src/datafusion/src/lib.rs @@ -13,5 +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/reader/plantime/engagement_policy.rs b/src/datafusion/src/reader/plantime/engagement_policy.rs new file mode 100644 index 000000000..f509d095d --- /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 15524da53..36576d666 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; 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 1d3714b5c..52d29dfde 100644 --- a/src/datafusion/src/reader/plantime/opener.rs +++ b/src/datafusion/src/reader/plantime/opener.rs @@ -3,7 +3,14 @@ use std::sync::Arc; use crate::{ cache::LiquidCacheParquetRef, reader::{ - plantime::{row_filter::build_row_filter, row_group_filter::RowGroupAccessPlanFilter}, + plantime::{ + engagement_policy::{ + CacheEngagementPolicy, EngagementContext, EngagementDecision, + default_engagement_policy, + }, + row_filter::build_row_filter, + row_group_filter::RowGroupAccessPlanFilter, + }, runtime::LiquidStreamBuilder, }, }; @@ -58,6 +65,8 @@ pub struct LiquidParquetOpener { 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 { @@ -76,6 +85,7 @@ impl LiquidParquetOpener { expr_adapter_factory: Arc, span: Option>, caller_reader_factory: Option>, + engagement_policy: Arc, ) -> Self { Self { partition_index, @@ -91,6 +101,7 @@ impl LiquidParquetOpener { expr_adapter_factory, span, caller_reader_factory, + engagement_policy, } } } @@ -189,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 @@ -386,41 +398,48 @@ impl FileOpener for LiquidParquetOpener { // 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. - if estimated_selectivity < 0.5 && predicate.is_some() { - log::info!( - "[LC-Opener] DELEGATE to plain parquet: selectivity={:.3} (low, fast already), file={}", - estimated_selectivity, file_name - ); - // Low selectivity: few rows match, decode cost minimal. - // Plain parquet is faster than LC's per-batch cache overhead. - 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); + 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::info!( + "[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()); } - if let Some(lim) = limit { - plain_builder = plain_builder.with_limit(lim); + EngagementDecision::UseLiquidCache => { + log::info!( + "[LC-Opener] LC STREAM: selectivity={:.3}, predicate={}, file={}", + estimated_selectivity, predicate.is_some(), file_name + ); } - - let stream = plain_builder.build()?; - let adapted = stream - .map_err(|e| DataFusionError::External(Box::new(e))); - return Ok(adapted.boxed()); } - // High selectivity (many rows to decode) or no predicate → use LC - // Warm cache avoids repeated parquet decompress+decode. - log::info!( - "[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 c077be263..0c6268492 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; @@ -193,6 +194,8 @@ pub struct LiquidParquetSource { /// 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 { @@ -276,6 +279,7 @@ impl LiquidParquetSource { page_pruning_predicate: None, span: None, parquet_file_reader_factory: reader_factory, + engagement_policy: default_engagement_policy(), }; if let Some(predicate) = predicate { @@ -285,6 +289,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() @@ -336,6 +346,7 @@ impl FileSource for LiquidParquetSource { expr_adapter_factory, execution_span.map(Arc::new), caller_reader_factory, + Arc::clone(&self.engagement_policy), ); Ok(Arc::new(opener)) From 1d5a6680034e0ab5fc09850b564a8566d50e323f Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Tue, 9 Jun 2026 04:04:34 +0530 Subject: [PATCH 11/12] Change LC log statements from info to debug level --- src/datafusion/src/optimizers/mod.rs | 8 ++++---- src/datafusion/src/reader/plantime/opener.rs | 12 ++++++------ src/datafusion/src/reader/plantime/source.rs | 6 +++--- .../src/reader/runtime/liquid_cache_reader.rs | 2 +- src/datafusion/src/reader/runtime/liquid_stream.rs | 2 +- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 256ed9a66..08975181c 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -109,7 +109,7 @@ fn try_optimize_parquet_source( // - Predicate references a string column let output_schema = plan.schema(); if output_schema.fields().is_empty() { - log::info!("[LC-Optimizer] SKIP: empty projection (COUNT(*))"); + log::debug!("[LC-Optimizer] SKIP: empty projection (COUNT(*))"); return Ok(Transformed::no(plan)); } @@ -117,7 +117,7 @@ fn try_optimize_parquet_source( // exceeds decode savings for wide projections. const MAX_LC_COLUMNS: usize = 4; if output_schema.fields().len() > MAX_LC_COLUMNS { - log::info!( + log::debug!( "[LC-Optimizer] SKIP: too many columns ({} > {})", output_schema.fields().len(), MAX_LC_COLUMNS ); @@ -142,7 +142,7 @@ fn try_optimize_parquet_source( }); if has_string_output || predicate_has_string { - log::info!( + log::debug!( "[LC-Optimizer] SKIP: string_in_output={}, string_in_predicate={}, output_cols={}", has_string_output, predicate_has_string, output_schema.fields().len() ); @@ -151,7 +151,7 @@ fn try_optimize_parquet_source( let num_fields = output_schema.fields().len(); let has_predicate = parquet_source.filter().is_some(); - log::info!( + log::debug!( "[LC-Optimizer] WRAP: all {} output columns cacheable, predicate={}", num_fields, has_predicate ); diff --git a/src/datafusion/src/reader/plantime/opener.rs b/src/datafusion/src/reader/plantime/opener.rs index 52d29dfde..059797a3b 100644 --- a/src/datafusion/src/reader/plantime/opener.rs +++ b/src/datafusion/src/reader/plantime/opener.rs @@ -141,7 +141,7 @@ impl FileOpener for LiquidParquetOpener { let metadata_size_hint = partitioned_file.metadata_size_hint; let has_predicate = self.predicate.is_some(); - log::info!( + log::debug!( "[LC-Opener] open file={}, predicate={}, batch_size={}, limit={:?}", file_name, has_predicate, self.batch_size, self.limit ); @@ -240,11 +240,11 @@ impl FileOpener for LiquidParquetOpener { let mut reader_metadata = if let Some(mut caller_reader) = caller_metadata_reader { match caller_reader.get_metadata(Some(&options)).await { Ok(meta) => { - log::info!("[LC-Meta] REUSE from caller factory: {}", file_name); + log::debug!("[LC-Meta] REUSE from caller factory: {}", file_name); ArrowReaderMetadata::try_new(meta, options.clone())? } Err(_) => { - log::info!("[LC-Meta] caller factory failed, loading from store: {}", file_name); + log::debug!("[LC-Meta] caller factory failed, loading from store: {}", file_name); ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone()).await? } } @@ -372,7 +372,7 @@ impl FileOpener for LiquidParquetOpener { // Early exit: if all row groups were pruned, return empty stream if row_group_indexes.is_empty() { - log::info!("[LC-Opener] EMPTY: all RGs pruned, file={}", file_name); + log::debug!("[LC-Opener] EMPTY: all RGs pruned, file={}", file_name); return Ok(futures::stream::empty().boxed()); } @@ -408,7 +408,7 @@ impl FileOpener for LiquidParquetOpener { match engagement_policy.decide(&engagement_ctx) { EngagementDecision::DelegateToParquet => { - log::info!( + log::debug!( "[LC-Opener] DELEGATE to plain parquet: selectivity={:.3}, file={}", estimated_selectivity, file_name ); @@ -433,7 +433,7 @@ impl FileOpener for LiquidParquetOpener { return Ok(adapted.boxed()); } EngagementDecision::UseLiquidCache => { - log::info!( + log::debug!( "[LC-Opener] LC STREAM: selectivity={:.3}, predicate={}, file={}", estimated_selectivity, predicate.is_some(), file_name ); diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 0c6268492..31c0823ee 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -148,7 +148,7 @@ impl AsyncFileReader for ParquetMetadataCacheReader { { let cache = META_CACHE.val.read().await; if let Some(meta) = cache.get(&path) { - log::info!("[LC-Meta] HIT path={}", path); + log::debug!("[LC-Meta] HIT path={}", path); return Ok(meta.clone()); } } @@ -157,11 +157,11 @@ impl AsyncFileReader for ParquetMetadataCacheReader { let mut cache = META_CACHE.val.write().await; match cache.entry(path.clone()) { std::collections::hash_map::Entry::Occupied(entry) => { - log::info!("[LC-Meta] HIT (race) path={}", path); + log::debug!("[LC-Meta] HIT (race) path={}", path); Ok(entry.get().clone()) } std::collections::hash_map::Entry::Vacant(entry) => { - log::info!("[LC-Meta] MISS (loading from store) path={}", path); + 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()) diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 68dfc1ee5..549d6b6e8 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -387,7 +387,7 @@ impl LiquidCacheReaderInner { } if *self.current_batch_id == 0 { - log::info!( + log::debug!( "[LC-Reader] batch_id={}, selected_rows={}, cols={}, hits={}, misses={}, fallback={}", *self.current_batch_id, selected_rows, diff --git a/src/datafusion/src/reader/runtime/liquid_stream.rs b/src/datafusion/src/reader/runtime/liquid_stream.rs index a056a68d2..34ec325fc 100644 --- a/src/datafusion/src/reader/runtime/liquid_stream.rs +++ b/src/datafusion/src/reader/runtime/liquid_stream.rs @@ -116,7 +116,7 @@ impl ReaderFactory { }; if row_group_idx == 0 { - log::info!( + log::debug!( "[LC-Stream] plan_row_group: rg={}, cache_cols={:?}, predicate_cols={:?}, \ has_filter={}, rows={}", row_group_idx, From 9de5160d22b902064fa3ec90b082f4db3a88c152 Mon Sep 17 00:00:00 2001 From: Tanvir Alam Date: Wed, 10 Jun 2026 21:54:29 +0530 Subject: [PATCH 12/12] feat: add with_predicate_no_page_pruning() to LiquidParquetSource Adds a method that sets the predicate for row_filter (per-row filtering during decode) and RG-level pruning predicate, but deliberately skips page_pruning_predicate. This is used by the indexed path where the BoolNode's RowSelection is authoritative and page-level statistics must not override it. --- src/datafusion/src/reader/plantime/source.rs | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 31c0823ee..626deaf83 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -253,6 +253,37 @@ 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();