From 97322ee6eff04d4d83a220f8724240ead6e39c90 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 15 Oct 2026 17:03:14 +0530 Subject: [PATCH] fix(cache): serve cached entries only to the file they came from A cache key keeps 16 bits of the file id, while file ids come from a counter that only climbs. The 65,537th distinct file a process registers therefore computes the same keys as the first, and reads the first file's cached data: silently wrong rows when the column types match, a panic in `RecordBatch::try_new` when they do not. The narrowing is guarded by `debug_assert!` only, so release builds carry on. Record the full file id ("identity") with every cached entry, in the index and as an 8-byte trailer on every object written to disk, and compare it on every lookup. A mismatch reads as a miss, so the caller falls back to parquet and gets correct data. An insert replaces whatever held the key, so the newer file can still cache its own data. The trailer is what makes the disk side safe without changing how objects are keyed: an eviction or flush that lands after the key changed hands leaves bytes that name the old identity, and the new owner's read of them misses. It follows the payload so the payload keeps its alignment, and the disk budget is charged for it. `LiquidCache::insert`, `get`, `eval_predicate`, `prefetch`, `try_read_liquid`, `contains`, `read_entry` and `eval_predicate_on_entry` take the identity as a new argument. Callers whose keys are already unique can pass any constant. The `should_panic` test for an oversized file id pinned the debug-only guard; it is replaced by one pinning the narrowing the identity check relies on. --- README.md | 11 +- examples/core.rs | 10 +- src/core/README.md | 16 +- src/core/src/cache/builders.rs | 27 +- src/core/src/cache/core.rs | 422 +++++++++++++++++++-------- src/core/src/cache/index.rs | 89 ++++-- src/core/src/cache/tests/policies.rs | 24 +- src/core/study/cache_storage.rs | 4 +- src/datafusion/src/cache/column.rs | 30 +- src/datafusion/src/cache/id.rs | 14 +- src/datafusion/src/cache/mod.rs | 73 ++++- 11 files changed, 528 insertions(+), 192 deletions(-) diff --git a/README.md b/README.md index 1cf6210b..fcd42aca 100644 --- a/README.md +++ b/README.md @@ -40,18 +40,21 @@ use std::sync::Arc; tokio_test::block_on(async { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(1); + // The identity names the source the data belongs to; reads under any + // other identity miss. One source here, so any constant will do. + let identity = 0; let values = Arc::new(UInt64Array::from(vec![10, 11, 12, 13, 14, 15])); // 1) insert - cache.insert(entry_id, values.clone()).await; + cache.insert(entry_id, identity, values.clone()).await; // 2) get - let all_rows = cache.get(&entry_id).await.expect("entry should exist"); + let all_rows = cache.get(&entry_id, identity).await.expect("entry should exist"); // 3) get filtered (selection pushdown): keep rows 0, 2, 4 let selection = BooleanBuffer::from(vec![true, false, true, false, true, false]); let selected_rows = cache - .get(&entry_id) + .get(&entry_id, identity) .with_selection(&selection) .await .expect("entry should exist"); @@ -63,7 +66,7 @@ tokio_test::block_on(async { Arc::new(Literal::new(ScalarValue::UInt64(Some(12)))), )); let predicate_mask = cache - .eval_predicate(&entry_id, &predicate) + .eval_predicate(&entry_id, identity, &predicate) .await .expect("entry should exist") .expect("predicate should be evaluated in cache"); diff --git a/examples/core.rs b/examples/core.rs index 1df164c7..f5823291 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -18,13 +18,19 @@ async fn main() -> Result<(), Box> { .await; let entry_id = EntryID::from(7); + // The identity names the source this data belongs to. One source here, + // so any constant will do. + let identity = 0; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); - storage.insert(entry_id, arrow_array.clone()).await.unwrap(); + storage + .insert(entry_id, identity, arrow_array.clone()) + .await + .unwrap(); // Move data to disk so the read demonstrates async I/O storage.flush_all_to_disk().await.unwrap(); - let retrieved = storage.get(&entry_id).await.unwrap(); + let retrieved = storage.get(&entry_id, identity).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); Ok(()) diff --git a/src/core/README.md b/src/core/README.md index e0e75212..26acdce7 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -24,10 +24,12 @@ let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(42); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..1000)); -// Insert once; replacement/placement is handled by the cache policy -storage.insert(entry_id, arrow_array.clone()).await; +// Insert once; replacement/placement is handled by the cache policy. +// `0` is the identity: the source this data belongs to. Reads under any +// other identity miss. One source here, so any constant will do. +storage.insert(entry_id, 0, arrow_array.clone()).await; -assert!(storage.contains(&entry_id)); +assert!(storage.contains(&entry_id, 0)); }); ``` @@ -43,13 +45,13 @@ let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(7); let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); -storage.insert(entry_id, arrow_array.clone()).await; +storage.insert(entry_id, 0, arrow_array.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; // Read asynchronously -let retrieved = storage.get(&entry_id).await.unwrap(); +let retrieved = storage.get(&entry_id, 0).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); }); ``` @@ -74,7 +76,7 @@ let entry_id = EntryID::from(8); let data = Arc::new(StringArray::from(vec![ Some("apple"), Some("banana"), None, Some("apple"), Some("cherry"), ])); -storage.insert(entry_id, data.clone()).await; +storage.insert(entry_id, 0, data.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; @@ -93,7 +95,7 @@ let liquid_expr = liquid_cache::cache::LiquidExpr::try_new( // Read with predicate pushdown let mask = storage - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, 0, &liquid_expr) .with_selection(&selection) .await .unwrap(); diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 9bb4a843..4a08b9be 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -175,16 +175,23 @@ pub fn default_max_memory_bytes() -> usize { pub struct Insert<'a> { pub(super) storage: &'a Arc, pub(super) entry_id: EntryID, + pub(super) identity: u64, pub(super) batch: ArrayRef, pub(super) skip_gc: bool, pub(super) lineage: Option>, } impl<'a> Insert<'a> { - pub(super) fn new(storage: &'a Arc, entry_id: EntryID, batch: ArrayRef) -> Self { + pub(super) fn new( + storage: &'a Arc, + entry_id: EntryID, + identity: u64, + batch: ArrayRef, + ) -> Self { Self { storage, entry_id, + identity, batch, skip_gc: false, lineage: None, @@ -213,7 +220,9 @@ impl<'a> Insert<'a> { self.storage.add_lineage(&self.entry_id, lineage); } let batch = CacheEntry::memory_arrow(batch); - self.storage.insert_inner(self.entry_id, batch).await + self.storage + .insert_inner(self.entry_id, self.identity, batch) + .await } } @@ -231,15 +240,17 @@ impl<'a> IntoFuture for Insert<'a> { pub struct Get<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) selection: Option<&'a BooleanBuffer>, pub(super) expression_hint: Option>, } impl<'a> Get<'a> { - pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID) -> Self { + pub(super) fn new(storage: &'a LiquidCache, entry_id: &'a EntryID, identity: u64) -> Self { Self { storage, entry_id, + identity, selection: None, expression_hint: None, } @@ -271,6 +282,7 @@ impl<'a> Get<'a> { self.storage .read_arrow_array( self.entry_id, + self.identity, self.selection, self.expression_hint.as_deref(), ) @@ -326,6 +338,7 @@ fn maybe_gc_view_arrays(array: &ArrayRef) -> Option { pub struct EvaluatePredicate<'a> { pub(super) storage: &'a LiquidCache, pub(super) entry_id: &'a EntryID, + pub(super) identity: u64, pub(super) predicate: &'a LiquidExpr, pub(super) selection: Option<&'a BooleanBuffer>, } @@ -334,11 +347,13 @@ impl<'a> EvaluatePredicate<'a> { pub(super) fn new( storage: &'a LiquidCache, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> Self { Self { storage, entry_id, + identity, predicate, selection: None, } @@ -353,7 +368,7 @@ impl<'a> EvaluatePredicate<'a> { /// Evaluate the predicate against the cached data. pub async fn read(self) -> Option { self.storage - .eval_predicate_internal(self.entry_id, self.selection, self.predicate) + .eval_predicate_internal(self.entry_id, self.identity, self.selection, self.predicate) .await } } @@ -451,9 +466,9 @@ mod tests { let cache = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(123usize); - cache.insert(entry_id, root.clone()).await.unwrap(); + cache.insert(entry_id, 0, root.clone()).await.unwrap(); - let stored = cache.get(&entry_id).await.expect("array present"); + let stored = cache.get(&entry_id, 0).await.expect("array present"); let post_size = stored.get_array_memory_size(); // GC should have compacted the view arrays, reducing memory footprint. diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index d007d316..a1b77d8e 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -25,6 +25,15 @@ use crate::sync::Arc; /// Cache storage for liquid cache. /// +/// # Identity +/// +/// Every entry is cached for an *identity*: a `u64` naming the source the data +/// belongs to. Callers that pack several fields into a fixed-width +/// [`EntryID`] can have two sources compute the same key; the identity is how +/// the cache tells them apart. A lookup under any identity other than the one +/// an entry was cached for is a miss, and an insert replaces whatever held the +/// key. Callers whose keys are already unique can pass any constant. +/// /// Example (async read): /// ```rust /// use liquid_cache::cache::{LiquidCacheBuilder, EntryID}; @@ -36,10 +45,12 @@ use crate::sync::Arc; /// /// let entry_id = EntryID::from(0); /// let arrow_array = Arc::new(UInt64Array::from_iter_values(0..32)); -/// storage.insert(entry_id, arrow_array.clone()).await; +/// // The identity names the source this data belongs to. One source here, +/// // so any constant will do. +/// storage.insert(entry_id, 0, arrow_array.clone()).await; /// /// // Get the arrow array back asynchronously -/// let retrieved = storage.get(&entry_id).await.unwrap(); +/// let retrieved = storage.get(&entry_id, 0).await.unwrap(); /// assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); /// }); /// ``` @@ -65,6 +76,9 @@ pub enum PrefetchResult { Absent, } +/// Length of the identity trailer appended to every object written to disk. +const IDENTITY_TRAILER_LEN: usize = std::mem::size_of::(); + /// Builder returned by [`LiquidCache::insert`] for configuring cache writes. impl LiquidCache { /// Return current cache statistics: counts and resource usage. @@ -80,7 +94,7 @@ impl LiquidCache { let mut memory_arrow_bytes = 0usize; let mut memory_liquid_bytes = 0usize; - self.index.for_each(|_, batch| match batch { + self.index.for_each(|_, slot| match &slot.entry { CacheEntry::MemoryArrow(array) => { memory_arrow_entries += 1; memory_arrow_bytes += array.get_array_memory_size(); @@ -113,52 +127,66 @@ impl LiquidCache { } } - /// Insert a batch into the cache. + /// Insert a batch into the cache for `identity`. pub fn insert<'a>( self: &'a Arc, entry_id: EntryID, + identity: u64, batch_to_cache: ArrayRef, ) -> Insert<'a> { - Insert::new(self, entry_id, batch_to_cache) + Insert::new(self, entry_id, identity, batch_to_cache) } - /// Create a [`Get`] builder for the provided entry. - pub fn get<'a>(&'a self, entry_id: &'a EntryID) -> Get<'a> { - Get::new(self, entry_id) + /// Create a [`Get`] builder for the provided entry, as cached for `identity`. + pub fn get<'a>(&'a self, entry_id: &'a EntryID, identity: u64) -> Get<'a> { + Get::new(self, entry_id, identity) } /// Create an [`EvaluatePredicate`] builder for evaluating predicates on cached data. pub fn eval_predicate<'a>( &'a self, entry_id: &'a EntryID, + identity: u64, predicate: &'a LiquidExpr, ) -> EvaluatePredicate<'a> { - EvaluatePredicate::new(self, entry_id, predicate) + EvaluatePredicate::new(self, entry_id, identity, predicate) } /// Prefetch an entry into a memory-form snapshot without recording an access. - pub async fn prefetch(&self, entry_id: &EntryID) -> PrefetchResult { - let Some(entry) = self.index.get(entry_id) else { + pub async fn prefetch(&self, entry_id: &EntryID, identity: u64) -> PrefetchResult { + let Some(slot) = self.index.get_for(entry_id, identity) else { return PrefetchResult::Absent; }; - match entry.as_ref() { - CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) => { - PrefetchResult::Snapshot(entry) + match &slot.entry { + entry @ (CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_)) => { + PrefetchResult::Snapshot(Arc::new(entry.clone())) } disk @ CacheEntry::DiskArrow { .. } => { - let Some(array) = self.read_disk_arrow_array(entry_id).await else { + let Some(array) = self.read_disk_arrow_array(entry_id, identity).await else { return PrefetchResult::Absent; }; - self.maybe_hydrate(entry_id, disk, MaterializedEntry::Arrow(&array), None) - .await; + self.maybe_hydrate( + entry_id, + identity, + disk, + MaterializedEntry::Arrow(&array), + None, + ) + .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_arrow(array))) } disk @ CacheEntry::DiskLiquid { .. } => { - let Some(array) = self.read_disk_liquid_array(entry_id).await else { + let Some(array) = self.read_disk_liquid_array(entry_id, identity).await else { return PrefetchResult::Absent; }; - self.maybe_hydrate(entry_id, disk, MaterializedEntry::Liquid(&array), None) - .await; + self.maybe_hydrate( + entry_id, + identity, + disk, + MaterializedEntry::Liquid(&array), + None, + ) + .await; PrefetchResult::Snapshot(Arc::new(CacheEntry::memory_liquid(array))) } } @@ -169,30 +197,38 @@ impl LiquidCache { pub async fn try_read_liquid( &self, entry_id: &EntryID, + identity: u64, ) -> Option { self.observer.on_try_read_liquid(); self.trace(InternalEvent::TryReadLiquid { entry: *entry_id }); - let batch = self.index.get(entry_id)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); - match batch.as_ref() { + match &slot.entry { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; Some(liquid) } CacheEntry::DiskArrow { .. } | CacheEntry::MemoryArrow(_) => None, } } - /// Iterate over all entries in the cache. + /// Iterate over all entries in the cache, whichever identity holds them. /// No guarantees are made about the order of the entries. /// Isolation level: read-committed pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { - self.index.for_each(&mut f); + self.index + .for_each(|entry_id, slot| f(entry_id, &slot.entry)); } /// Reset the cache. @@ -201,9 +237,9 @@ impl LiquidCache { self.budget.reset_usage(); } - /// Check whether the cache contains a batch. - pub fn contains(&self, entry_id: &EntryID) -> bool { - self.index.contains(entry_id) + /// Check whether the cache holds a batch under `entry_id` for `identity`. + pub fn contains(&self, entry_id: &EntryID, identity: u64) -> bool { + self.index.contains(entry_id, identity) } /// Get the config of the cache. @@ -234,18 +270,22 @@ impl LiquidCache { /// Flush all entries to disk. pub async fn flush_all_to_disk(&self) -> Result<(), CacheFull> { let mut entires = Vec::new(); - self.for_each_entry(|entry_id, batch| { - entires.push((*entry_id, batch.clone())); + self.index.for_each(|entry_id, slot| { + entires.push((*entry_id, slot.identity, slot.entry.clone())); }); - for (entry_id, batch) in entires { + for (entry_id, identity, batch) in entires { match &batch { CacheEntry::MemoryArrow(array) => { let bytes = arrow_to_bytes(array).expect("failed to convert arrow to bytes"); let disk_bytes = bytes.len(); - match self.write_batch_to_disk(entry_id, &batch, bytes).await { + match self + .write_batch_to_disk(entry_id, identity, &batch, bytes) + .await + { Ok(()) => { self.try_insert( entry_id, + identity, CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), ) .expect("failed to insert disk arrow entry"); @@ -257,12 +297,13 @@ impl LiquidCache { let liquid_bytes = liquid_array.to_bytes(); let disk_bytes = liquid_bytes.len(); match self - .write_batch_to_disk(entry_id, &batch, Bytes::from(liquid_bytes)) + .write_batch_to_disk(entry_id, identity, &batch, Bytes::from(liquid_bytes)) .await { Ok(()) => { self.try_insert( entry_id, + identity, CacheEntry::disk_liquid( liquid_array.original_arrow_data_type(), disk_bytes, @@ -287,6 +328,7 @@ impl LiquidCache { async fn write_in_memory_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: CacheEntry, ) -> Result { match &batch { @@ -302,7 +344,7 @@ impl LiquidCache { unreachable!("memory Arrow eviction cannot remove entry"); }; if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(entry_id, &new_batch, bytes_to_write) + self.write_batch_to_disk(entry_id, identity, &new_batch, bytes_to_write) .await?; } Ok(new_batch) @@ -310,7 +352,7 @@ impl LiquidCache { CacheEntry::MemoryLiquid(liquid_array) => { let liquid_bytes = Bytes::from(liquid_array.to_bytes()); let disk_bytes = liquid_bytes.len(); - self.write_batch_to_disk(entry_id, &batch, liquid_bytes) + self.write_batch_to_disk(entry_id, identity, &batch, liquid_bytes) .await?; Ok(CacheEntry::disk_liquid( liquid_array.original_arrow_data_type(), @@ -327,10 +369,11 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, + identity: u64, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { loop { - let Err(not_inserted) = self.try_insert(entry_id, batch_to_cache) else { + let Err(not_inserted) = self.try_insert(entry_id, identity, batch_to_cache) else { return Ok(()); }; self.trace(InternalEvent::InsertFailed { @@ -344,7 +387,7 @@ impl LiquidCache { // this can happen if the entry to be inserted is too large, in that case, // we write it to disk let on_disk_batch = self - .write_in_memory_batch_to_disk(entry_id, not_inserted) + .write_in_memory_batch_to_disk(entry_id, identity, not_inserted) .await?; batch_to_cache = on_disk_batch; continue; @@ -389,10 +432,15 @@ impl LiquidCache { } } - fn try_insert(&self, entry_id: EntryID, to_insert: CacheEntry) -> Result<(), CacheEntry> { + fn try_insert( + &self, + entry_id: EntryID, + identity: u64, + to_insert: CacheEntry, + ) -> Result<(), CacheEntry> { let new_memory_size = to_insert.memory_usage_bytes(); - let cached_batch_type = if let Some(entry) = self.index.get(&entry_id) { - let old_memory_size = entry.memory_usage_bytes(); + let cached_batch_type = if let Some(slot) = self.index.get(&entry_id) { + let old_memory_size = slot.entry.memory_usage_bytes(); if self .budget .try_update_memory_usage(old_memory_size, new_memory_size) @@ -401,14 +449,14 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + self.index.insert(&entry_id, identity, to_insert); batch_type } else { if self.budget.try_reserve_memory(new_memory_size).is_err() { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + self.index.insert(&entry_id, identity, to_insert); batch_type }; @@ -428,13 +476,13 @@ impl LiquidCache { }; assert!( matches!( - removed.as_ref(), + &removed.entry, CacheEntry::MemoryArrow(_) | CacheEntry::MemoryLiquid(_) ), "flush should only drop memory entries" ); self.budget - .try_update_memory_usage(removed.memory_usage_bytes(), 0) + .try_update_memory_usage(removed.entry.memory_usage_bytes(), 0) .expect("memory release cannot fail"); self.cache_policy.notify_remove(&entry_id); } @@ -443,7 +491,7 @@ impl LiquidCache { let Some(removed) = self.index.remove(&entry_id) else { return; }; - let disk_bytes = match removed.as_ref() { + let disk_bytes = match &removed.entry { CacheEntry::DiskLiquid { disk_bytes, .. } | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, _ => panic!("remove_disk_entry called for non-disk entry"), @@ -452,7 +500,7 @@ impl LiquidCache { .remove(&entry_id_to_key(&entry_id)) .await .expect("disk remove failed"); - self.budget.release_disk(disk_bytes); + self.budget.release_disk(disk_bytes + IDENTITY_TRAILER_LEN); self.cache_policy.notify_remove(&entry_id); self.trace(InternalEvent::DiskEvict { entry: entry_id, @@ -496,15 +544,19 @@ impl LiquidCache { } async fn evict_victim_inner(&self, victim: EntryID) -> Result<(), CacheFull> { - let Some(mut victim_entry) = self.index.get(&victim) else { + let Some(slot) = self.index.get(&victim) else { return Ok(()); }; + // The rewrite is made for the identity the victim was read under. If + // another identity takes the key over meanwhile, the rewrite still + // carries the old identity, so the new owner reads it as a miss. + let identity = slot.identity; + let mut victim_entry = slot.entry.clone(); self.trace(InternalEvent::EvictionVictim { entry: victim }); loop { - let outcome = self.eviction_policy.evict( - victim_entry.as_ref(), - self.metadata.lineage(&victim).as_deref(), - ); + let outcome = self + .eviction_policy + .evict(&victim_entry, self.metadata.lineage(&victim).as_deref()); match outcome { EvictionOutcome::Replace { @@ -512,15 +564,15 @@ impl LiquidCache { bytes_to_write, } => { if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(victim, &new_batch, bytes_to_write) + self.write_batch_to_disk(victim, identity, &new_batch, bytes_to_write) .await?; } - match self.try_insert(victim, new_batch) { + match self.try_insert(victim, identity, new_batch) { Ok(()) => { break; } Err(batch) => { - victim_entry = Arc::new(batch); + victim_entry = batch; } } } @@ -536,6 +588,7 @@ impl LiquidCache { async fn maybe_hydrate( &self, entry_id: &EntryID, + identity: u64, cached: &CacheEntry, materialized: MaterializedEntry<'_>, expression: Option<&CacheExpression>, @@ -553,40 +606,48 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self.insert_inner(*entry_id, new_entry).await; + let _ = self.insert_inner(*entry_id, identity, new_entry).await; } } pub(crate) async fn read_arrow_array( &self, entry_id: &EntryID, + identity: u64, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - let batch = self.index.get(entry_id)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.read_entry_inner(entry_id, batch.as_ref(), selection, expression) + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); + self.read_entry_inner(entry_id, identity, &slot.entry, selection, expression) .await } - /// Read an already-looked-up cache entry. + /// Read an already-looked-up cache entry, cached for `identity`. + /// + /// `entry` is a value the caller already holds, so it is not checked + /// against `identity`: the caller must have obtained it for that identity, + /// e.g. from [`Self::prefetch`]. `identity` is checked when the entry's + /// data has to be read from disk. pub async fn read_entry( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { self.observer.on_get(selection.is_some()); - self.read_entry_inner(entry_id, entry, selection, expression) + self.read_entry_inner(entry_id, identity, entry, selection, expression) .await } async fn read_entry_inner( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, @@ -613,7 +674,7 @@ impl LiquidCache { None => Some(array.to_arrow_array()), }, CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { - self.read_disk_array(entry, entry_id, expression, selection) + self.read_disk_array(entry, entry_id, identity, expression, selection) .await } } @@ -623,6 +684,7 @@ impl LiquidCache { &self, entry: &CacheEntry, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -633,9 +695,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let full_array = self.read_disk_arrow_array(entry_id).await?; + let full_array = self.read_disk_arrow_array(entry_id, identity).await?; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Arrow(&full_array), expression, @@ -655,9 +718,10 @@ impl LiquidCache { { return Some(arrow::array::new_empty_array(data_type)); } - let liquid = self.read_disk_liquid_array(entry_id).await?; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Liquid(&liquid), expression, @@ -672,15 +736,27 @@ impl LiquidCache { } } + /// Write `bytes` to disk for `identity`, followed by the identity itself. + /// + /// Objects are keyed by entry id alone, so a write landing after the key + /// changed hands can leave one identity's bytes where another's entry + /// points. The trailer is what lets a read notice that and miss. It sits + /// after the payload so the payload keeps the alignment it is read with. + /// The disk budget is charged for the trailer as well as the payload. async fn write_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: &CacheEntry, bytes: Bytes, ) -> Result<(), CacheFull> { let len = bytes.len(); loop { - if self.budget.try_reserve_disk(len).is_ok() { + if self + .budget + .try_reserve_disk(len + IDENTITY_TRAILER_LEN) + .is_ok() + { break; } let victims = self.cache_policy.find_disk_victim(8); @@ -696,19 +772,34 @@ impl LiquidCache { kind: CachedBatchType::from(batch), bytes: len, }); + let mut object = Vec::with_capacity(len + IDENTITY_TRAILER_LEN); + object.extend_from_slice(&bytes); + object.extend_from_slice(&identity.to_le_bytes()); self.store - .put(entry_id_to_key(&entry_id), bytes.to_vec()) + .put(entry_id_to_key(&entry_id), object) .await .expect("write failed"); Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { - Ok(bytes) => bytes, + /// The payload stored under `entry_id`, if it was written for `identity`. + async fn read_disk_bytes(&self, entry_id: &EntryID, identity: u64) -> Option { + let mut object = match self.store.get(&entry_id_to_key(entry_id)).await { + Ok(object) => object, Err(t4::Error::NotFound) => return None, Err(error) => panic!("read failed: {error}"), }; + let payload_len = object.len().checked_sub(IDENTITY_TRAILER_LEN)?; + let trailer: [u8; IDENTITY_TRAILER_LEN] = object[payload_len..].try_into().unwrap(); + if u64::from_le_bytes(trailer) != identity { + return None; + } + object.truncate(payload_len); + Some(Bytes::from(object)) + } + + async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> Option { + let bytes = self.read_disk_bytes(entry_id, identity).await?; let bytes_len = bytes.len(); let cursor = std::io::Cursor::new(bytes); let mut reader = @@ -725,51 +816,60 @@ impl LiquidCache { async fn read_disk_liquid_array( &self, entry_id: &EntryID, + identity: u64, ) -> Option { - let bytes = match self.store.get(&entry_id_to_key(entry_id)).await { - Ok(bytes) => bytes, - Err(t4::Error::NotFound) => return None, - Err(error) => panic!("read failed: {error}"), - }; + let bytes = self.read_disk_bytes(entry_id, identity).await?; self.trace(InternalEvent::IoReadLiquid { entry: *entry_id, bytes: bytes.len(), }); Some(Arc::new(crate::liquid_array::LiquidArray::from_bytes( - Bytes::from(bytes), + bytes, ))) } pub(crate) async fn eval_predicate_internal( &self, entry_id: &EntryID, + identity: u64, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - let batch = self.index.get(entry_id)?; + let slot = self.index.get_for(entry_id, identity)?; self.cache_policy - .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); - self.eval_predicate_on_entry_inner(entry_id, batch.as_ref(), selection_opt, predicate) - .await + .notify_access(entry_id, CachedBatchType::from(&slot.entry)); + self.eval_predicate_on_entry_inner( + entry_id, + identity, + &slot.entry, + selection_opt, + predicate, + ) + .await } - /// Evaluate a predicate on an already-looked-up cache entry. + /// Evaluate a predicate on an already-looked-up cache entry, cached for `identity`. + /// + /// As with [`Self::read_entry`], `entry` itself is not checked against + /// `identity`; only data read from disk is. pub async fn eval_predicate_on_entry( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { self.observer.on_eval_predicate(); - self.eval_predicate_on_entry_inner(entry_id, entry, selection_opt, predicate) + self.eval_predicate_on_entry_inner(entry_id, identity, entry, selection_opt, predicate) .await } async fn eval_predicate_on_entry_inner( &self, entry_id: &EntryID, + identity: u64, entry: &CacheEntry, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, @@ -793,9 +893,15 @@ impl LiquidCache { Some(self.eval_predicate_on_array(filtered, predicate)) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) - .await; + let array = self.read_disk_arrow_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Arrow(&array), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(array.len())); @@ -815,9 +921,15 @@ impl LiquidCache { Some(array.try_eval_predicate(predicate, selection)) } entry @ CacheEntry::DiskLiquid { .. } => { - let liquid = self.read_disk_liquid_array(entry_id).await?; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await?; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; let mut owned = None; let selection = selection_opt.unwrap_or_else(|| { owned = Some(BooleanBuffer::new_set(liquid.len())); @@ -897,12 +1009,12 @@ mod tests { let entry_id1: EntryID = EntryID::from(1); let array1 = create_test_array(100); let size1 = array1.memory_usage_bytes(); - store.insert_inner(entry_id1, array1).await.unwrap(); + store.insert_inner(entry_id1, 0, array1).await.unwrap(); // Verify budget usage and data correctness assert_eq!(store.budget.memory_usage_bytes(), size1); let retrieved1 = store.index().get(&entry_id1).unwrap(); - match retrieved1.as_ref() { + match &retrieved1.entry { CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), _ => panic!("Expected ArrowMemory"), } @@ -910,13 +1022,13 @@ mod tests { let entry_id2: EntryID = EntryID::from(2); let array2 = create_test_array(200); let size2 = array2.memory_usage_bytes(); - store.insert_inner(entry_id2, array2).await.unwrap(); + store.insert_inner(entry_id2, 0, array2).await.unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size1 + size2); let array3 = create_test_array(150); let size3 = array3.memory_usage_bytes(); - store.insert_inner(entry_id1, array3).await.unwrap(); + store.insert_inner(entry_id1, 0, array3).await.unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); assert!(store.index().get(&EntryID::from(999)).is_none()); @@ -936,19 +1048,19 @@ mod tests { let store = create_cache_store(8000, Box::new(advisor)).await; // Small budget to force advice store - .insert_inner(entry_id1, create_test_array(800)) + .insert_inner(entry_id1, 0, create_test_array(800)) .await .unwrap(); - match store.index().get(&entry_id1).unwrap().as_ref() { + match &store.index().get(&entry_id1).unwrap().entry { CacheEntry::MemoryArrow(_) => {} other => panic!("Expected ArrowMemory, got {other:?}"), } store - .insert_inner(entry_id2, create_test_array(800)) + .insert_inner(entry_id2, 0, create_test_array(800)) .await .unwrap(); - match store.index().get(&entry_id1).unwrap().as_ref() { + match &store.index().get(&entry_id1).unwrap().entry { CacheEntry::MemoryLiquid(_) => {} other => panic!("Expected LiquidMemory after eviction, got {other:?}"), } @@ -995,7 +1107,7 @@ mod tests { let unique_id = thread_id * ops_per_thread + i; let entry_id: EntryID = EntryID::from(unique_id); let array = create_test_arrow_array(100); - store.insert(entry_id, array).await.unwrap(); + store.insert(entry_id, 0, array).await.unwrap(); } }); })); @@ -1029,8 +1141,14 @@ mod tests { // Insert two small batches let arr1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); let arr2: ArrayRef = Arc::new(Int32Array::from_iter_values(0..128)); - storage.insert(EntryID::from(1usize), arr1).await.unwrap(); - storage.insert(EntryID::from(2usize), arr2).await.unwrap(); + storage + .insert(EntryID::from(1usize), 0, arr1) + .await + .unwrap(); + storage + .insert(EntryID::from(2usize), 0, arr2) + .await + .unwrap(); // Stats after insert: 2 entries, memory usage > 0, disk usage == 0 let s = storage.stats(); @@ -1054,21 +1172,68 @@ mod tests { let entry_id = EntryID::from(321usize); let array = create_test_arrow_array(8); - store.insert(entry_id, array.clone()).await.unwrap(); + store.insert(entry_id, 0, array.clone()).await.unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::DiskArrow { .. })); + assert!(matches!(&entry.entry, CacheEntry::DiskArrow { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::MemoryArrow(_))); + assert!(matches!(&entry.entry, CacheEntry::MemoryArrow(_))); } } + #[tokio::test] + async fn an_entry_is_served_only_to_its_identity() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(321usize); + let array = create_test_arrow_array(8); + + store.insert(entry_id, 1, array.clone()).await.unwrap(); + assert!(store.contains(&entry_id, 1)); + assert!(!store.contains(&entry_id, 2)); + assert!(store.get(&entry_id, 2).await.is_none()); + assert!(matches!( + store.prefetch(&entry_id, 2).await, + PrefetchResult::Absent + )); + + store.flush_all_to_disk().await.unwrap(); + assert!(store.get(&entry_id, 2).await.is_none()); + let result = store.get(&entry_id, 1).await.expect("present"); + assert_eq!(result.as_ref(), array.as_ref()); + + let replacement = create_test_arrow_array(16); + store + .insert(entry_id, 2, replacement.clone()) + .await + .unwrap(); + assert!(!store.contains(&entry_id, 1)); + let result = store.get(&entry_id, 2).await.expect("present"); + assert_eq!(result.as_ref(), replacement.as_ref()); + } + + #[tokio::test] + async fn disk_bytes_written_for_another_identity_are_a_miss() { + let store = create_cache_store(1 << 20, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(321usize); + let array = create_test_arrow_array(8); + + store.insert(entry_id, 1, array.clone()).await.unwrap(); + store.flush_all_to_disk().await.unwrap(); + let disk_entry = store.index().get(&entry_id).unwrap().entry.clone(); + assert!(matches!(disk_entry, CacheEntry::DiskArrow { .. })); + + // A record naming identity 2 over bytes identity 1 wrote: what a + // rewrite landing after the key changed hands leaves behind. + store.index().insert(&entry_id, 2, disk_entry); + assert!(store.get(&entry_id, 2).await.is_none()); + } + #[tokio::test] async fn missing_disk_blob_is_a_cache_miss() { let directory = tempfile::tempdir().unwrap(); @@ -1079,11 +1244,14 @@ mod tests { .await; let id = EntryID::from(320usize); - cache.insert(id, create_test_arrow_array(8)).await.unwrap(); + cache + .insert(id, 0, create_test_arrow_array(8)) + .await + .unwrap(); cache.flush_all_to_disk().await.unwrap(); store.remove(&entry_id_to_key(&id)).await.unwrap(); - assert!(cache.get(&id).await.is_none()); + assert!(cache.get(&id, 0).await.is_none()); } #[tokio::test] @@ -1095,20 +1263,20 @@ mod tests { Arc::new(crate::liquid_array::LiquidArray::from_arrow_array(&arrow_array).unwrap()); store - .insert_inner(entry_id, CacheEntry::memory_liquid(liquid.clone())) + .insert_inner(entry_id, 0, CacheEntry::memory_liquid(liquid.clone())) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::DiskLiquid { .. })); + assert!(matches!(&entry.entry, CacheEntry::DiskLiquid { .. })); } - let result = store.get(&entry_id).await.expect("present"); + let result = store.get(&entry_id, 0).await.expect("present"); assert_eq!(result.as_ref(), arrow_array.as_ref()); { let entry = store.index().get(&entry_id).unwrap(); - assert!(matches!(entry.as_ref(), CacheEntry::MemoryLiquid(_))); + assert!(matches!(&entry.entry, CacheEntry::MemoryLiquid(_))); } } @@ -1122,18 +1290,18 @@ mod tests { .await; let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let err = cache.insert(EntryID::from(900usize), array).await; + let err = cache.insert(EntryID::from(900usize), 0, array).await; assert_eq!(err, Err(CacheFull)); - assert!(!cache.contains(&EntryID::from(900usize))); + assert!(!cache.contains(&EntryID::from(900usize), 0)); } #[tokio::test] async fn insert_until_disk_full_then_evicts_oldest_disk_entry() { let first_array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let second_array: ArrayRef = Arc::new(Int32Array::from_iter_values(16..32)); - let first_bytes = arrow_to_bytes(&first_array).unwrap().len(); - let second_bytes = arrow_to_bytes(&second_array).unwrap().len(); + let first_bytes = arrow_to_bytes(&first_array).unwrap().len() + IDENTITY_TRAILER_LEN; + let second_bytes = arrow_to_bytes(&second_array).unwrap().len() + IDENTITY_TRAILER_LEN; let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(first_bytes.max(second_bytes)) @@ -1144,16 +1312,16 @@ mod tests { let first = EntryID::from(910usize); let second = EntryID::from(911usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(cache.contains(&first)); + assert!(cache.contains(&first, 0)); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.contains(&first)); + assert!(!cache.contains(&first, 0)); assert!(matches!( - cache.index().get(&second).unwrap().as_ref(), + &cache.index().get(&second).unwrap().entry, CacheEntry::DiskArrow { .. } )); } @@ -1172,19 +1340,19 @@ mod tests { .await; let first = EntryID::from(912usize); let second = EntryID::from(913usize); - cache.insert(first, first_array).await.unwrap(); + cache.insert(first, 0, first_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - cache.insert(second, second_array).await.unwrap(); + cache.insert(second, 0, second_array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.contains(&first) || !cache.contains(&second)); + assert!(!cache.contains(&first, 0) || !cache.contains(&second, 0)); } #[tokio::test] async fn disk_eviction_releases_budget() { let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - let disk_bytes = arrow_to_bytes(&array).unwrap().len(); + let disk_bytes = arrow_to_bytes(&array).unwrap().len() + IDENTITY_TRAILER_LEN; let cache = LiquidCacheBuilder::new() .with_max_memory_bytes(1 << 20) .with_max_disk_bytes(disk_bytes) @@ -1193,14 +1361,14 @@ mod tests { .build() .await; let entry = EntryID::from(914usize); - cache.insert(entry, array).await.unwrap(); + cache.insert(entry, 0, array).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); let before = cache.stats().disk_usage_bytes; cache.remove_disk_entry(entry).await; assert_eq!(cache.stats().disk_usage_bytes, before - disk_bytes); - assert!(!cache.contains(&entry)); + assert!(!cache.contains(&entry, 0)); } #[tokio::test] @@ -1213,11 +1381,11 @@ mod tests { .await; let entry_id = EntryID::from(901usize); let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); - cache.insert(entry_id, array).await.unwrap(); + cache.insert(entry_id, 0, array).await.unwrap(); let result = cache.flush_all_to_disk().await; assert_eq!(result, Ok(())); - assert!(!cache.contains(&entry_id)); + assert!(!cache.contains(&entry_id, 0)); } } diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 62acca36..8607acec 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -9,8 +9,27 @@ use std::{ use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; +/// What the index holds for a key: the entry, and the identity of the source +/// it was cached for. +/// +/// A key does not name its source uniquely: callers pack fields into fixed +/// widths, so two sources can compute the same key. The identity is what tells +/// them apart, so a read under any other identity must treat the slot as +/// absent. +#[derive(Debug)] +pub(crate) struct Slot { + pub(crate) identity: u64, + pub(crate) entry: CacheEntry, +} + +impl Slot { + pub(crate) fn belongs_to(&self, identity: u64) -> bool { + self.identity == identity + } +} + pub(crate) struct ArtIndex { - art: CongeeArc, + art: CongeeArc, entry_count: AtomicUsize, } @@ -22,36 +41,46 @@ impl Debug for ArtIndex { impl ArtIndex { pub(crate) fn new() -> Self { - let art: CongeeArc = CongeeArc::new(); + let art: CongeeArc = CongeeArc::new(); Self { art, entry_count: AtomicUsize::new(0), } } - pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { + /// The slot under `entry_id`, whichever identity holds it. + pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); - let batch = self.art.get(*entry_id, &guard)?; - Some(batch) + self.art.get(*entry_id, &guard) } - pub(crate) fn contains(&self, entry_id: &EntryID) -> bool { - let guard = self.art.pin(); - self.art.get(*entry_id, &guard).is_some() + /// The slot under `entry_id`, only if `identity` holds it. + pub(crate) fn get_for(&self, entry_id: &EntryID, identity: u64) -> Option> { + self.get(entry_id).filter(|slot| slot.belongs_to(identity)) + } + + pub(crate) fn contains(&self, entry_id: &EntryID, identity: u64) -> bool { + self.get_for(entry_id, identity).is_some() } - pub(crate) fn insert(&self, entry_id: &EntryID, batch: CacheEntry) { + /// Store `batch` under `entry_id` for `identity`, replacing whatever held + /// the key, including a slot held by another identity. + pub(crate) fn insert(&self, entry_id: &EntryID, identity: u64, batch: CacheEntry) { let guard = self.art.pin(); + let slot = Slot { + identity, + entry: batch, + }; let existing = self .art - .insert(*entry_id, Arc::new(batch), &guard) + .insert(*entry_id, Arc::new(slot), &guard) .expect("Insertion failed"); if existing.is_none() { self.entry_count.fetch_add(1, Ordering::Relaxed); } } - pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { + pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); let removed = self.art.remove(*entry_id, &guard); if removed.is_some() { @@ -68,7 +97,7 @@ impl ArtIndex { self.entry_count.store(0, Ordering::Relaxed); } - pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &CacheEntry)) { + pub(crate) fn for_each(&self, mut f: impl FnMut(&EntryID, &Slot)) { let guard = self.art.pin(); for id in self.art.keys().into_iter() { f( @@ -106,21 +135,21 @@ mod tests { let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.contains(&entry_id1)); - assert!(!store.contains(&entry_id2)); + assert!(!store.contains(&entry_id1, 0)); + assert!(!store.contains(&entry_id2, 0)); assert!(store.get(&entry_id1).is_none()); // Insert an entry and verify it's cached { - store.insert(&entry_id1, array1.clone()); + store.insert(&entry_id1, 0, array1.clone()); } - assert!(store.contains(&entry_id1)); - assert!(!store.contains(&entry_id2)); + assert!(store.contains(&entry_id1, 0)); + assert!(!store.contains(&entry_id2, 0)); // Get should return the cached value - match store.get(&entry_id1) { - Some(batch) => match batch.as_ref() { + match store.get_for(&entry_id1, 0) { + Some(slot) => match &slot.entry { CacheEntry::MemoryArrow(arr) => assert_eq!(arr.len(), 100), _ => panic!("Expected ArrowMemory batch"), }, @@ -134,13 +163,29 @@ mod tests { let entry_id: EntryID = EntryID::from(1); let array = create_test_array(100); - store.insert(&entry_id, array.clone()); + store.insert(&entry_id, 0, array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.contains(&entry_id)); + assert!(store.contains(&entry_id, 0)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.contains(&entry_id)); + assert!(!store.contains(&entry_id, 0)); + } + + #[test] + fn a_slot_is_visible_only_to_its_identity() { + let store = ArtIndex::new(); + let entry_id = EntryID::from(1); + store.insert(&entry_id, 7, create_test_array(10)); + + assert!(store.contains(&entry_id, 7)); + assert!(!store.contains(&entry_id, 8)); + assert!(store.get_for(&entry_id, 8).is_none()); + + store.insert(&entry_id, 8, create_test_array(20)); + assert!(!store.contains(&entry_id, 7)); + assert!(store.contains(&entry_id, 8)); + assert_eq!(store.entry_count(), 1); } } diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index ee3681a4..e074adbd 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -26,12 +26,12 @@ async fn default_policies() { for i in 0..5 { let entry_id = EntryID::from(i); - cache.insert(entry_id, test_array.clone()).await.unwrap(); + cache.insert(entry_id, 0, test_array.clone()).await.unwrap(); } for i in 0..5 { let entry_id = EntryID::from(i); - let array = cache.get(&entry_id).read().await.unwrap(); + let array = cache.get(&entry_id, 0).read().await.unwrap(); assert_eq!(array.len(), test_array.len()); } @@ -54,17 +54,17 @@ async fn insert_wont_fit_cache() { .build() .await; cache - .insert(EntryID::from(0), test_array.clone()) + .insert(EntryID::from(0), 0, test_array.clone()) .await .unwrap(); let array_3x = arrow::compute::concat(&[&test_array, &test_array, &test_array]).unwrap(); let array_9x = arrow::compute::concat(&[&array_3x, &array_3x, &array_3x]).unwrap(); let array_27x = arrow::compute::concat(&[&array_9x, &array_9x, &array_9x]).unwrap(); cache - .insert(EntryID::from(1), array_27x.clone()) + .insert(EntryID::from(1), 0, array_27x.clone()) .await .unwrap(); - cache.get(&EntryID::from(1)).read().await.unwrap(); + cache.get(&EntryID::from(1), 0).read().await.unwrap(); let trace = cache.consume_event_trace(); let json_trace = serde_json::to_string(&trace).unwrap(); @@ -100,7 +100,7 @@ async fn liquid_eviction_reads_memory_and_disk() { .into_iter() .enumerate() { - cache.insert(EntryID::from(id), array).await.unwrap(); + cache.insert(EntryID::from(id), 0, array).await.unwrap(); } let mut states = Vec::new(); @@ -113,13 +113,13 @@ async fn liquid_eviction_reads_memory_and_disk() { let id = EntryID::from(0); assert_eq!( - cache.get(&id).read().await.unwrap().as_ref(), + cache.get(&id, 0).read().await.unwrap().as_ref(), integers.as_ref() ); let selection = BooleanBuffer::from_iter((0..integers.len()).map(|index| index % 3 == 0)); let selected = cache - .get(&id) + .get(&id, 0) .with_selection(&selection) .read() .await @@ -137,7 +137,11 @@ async fn liquid_eviction_reads_memory_and_disk() { Arc::new(Literal::new(ScalarValue::Int64(Some(1_000)))), )); let predicate = LiquidExpr::try_new(physical, &DataType::Int64).unwrap(); - let actual = cache.eval_predicate(&id, &predicate).read().await.unwrap(); + let actual = cache + .eval_predicate(&id, 0, &predicate) + .read() + .await + .unwrap(); let expected = arrow::array::BooleanArray::from_iter( integers .as_any() @@ -157,7 +161,7 @@ async fn liquid_eviction_reads_memory_and_disk() { ))); assert!(disk_states.contains(&CachedBatchType::DiskLiquid)); assert_eq!( - cache.get(&id).read().await.unwrap().as_ref(), + cache.get(&id, 0).read().await.unwrap().as_ref(), integers.as_ref() ); } diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index 18e01d5c..b2571753 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -108,7 +108,7 @@ fn main() { continue; }; if storage - .eval_predicate(id, &liquid_expr) + .eval_predicate(id, 0, &liquid_expr) .with_selection(&selection) .await .is_some() @@ -161,7 +161,7 @@ fn load_and_insert_referer( let id = EntryID::from(idx); ids.push(id); total_size += array.get_array_memory_size(); - storage.insert(id, array).await.unwrap(); + storage.insert(id, 0, array).await.unwrap(); idx += 1; } diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index c08cf9af..a0cdaba6 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -22,6 +22,11 @@ pub struct CachedColumn { cache_store: Arc, field: Arc, column_path: ColumnAccessPath, + /// The unnarrowed id of the file this column belongs to. `column_path` + /// keeps only its low 16 bits, so another file can compute the same + /// entry ids; the cache serves an entry only to the identity it was + /// cached for. + file_identity: u64, expression: Option>, snapshots: Arc, } @@ -55,6 +60,7 @@ impl CachedColumn { field: Arc, cache_store: Arc, column_access_path: ColumnAccessPath, + file_identity: u64, expression: Option>, is_predicate_column: bool, snapshots: Arc, @@ -79,6 +85,7 @@ impl CachedColumn { field, cache_store, column_path: column_access_path, + file_identity, expression, snapshots, } @@ -89,8 +96,13 @@ impl CachedColumn { self.column_path.entry_id(batch_id) } + pub(crate) fn file_identity(&self) -> u64 { + self.file_identity + } + pub(crate) fn contains(&self, batch_id: BatchID) -> bool { - self.cache_store.contains(&self.entry_id(batch_id).into()) + self.cache_store + .contains(&self.entry_id(batch_id).into(), self.file_identity) } pub(crate) fn snapshot_entry(&self, batch_id: BatchID) -> Option> { @@ -131,6 +143,7 @@ impl CachedColumn { self.cache_store .eval_predicate_on_entry( &entry_id, + self.file_identity, entry.as_ref(), Some(filter), &liquid_expr, @@ -139,7 +152,7 @@ impl CachedColumn { } None => { self.cache_store - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, self.file_identity, &liquid_expr) .with_selection(filter) .await } @@ -191,6 +204,7 @@ impl CachedColumn { .cache_store .read_entry( &entry_id, + self.file_identity, entry.as_ref(), Some(filter), self.expression.as_deref(), @@ -198,7 +212,7 @@ impl CachedColumn { .await; } self.cache_store - .get(&entry_id) + .get(&entry_id, self.file_identity) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() @@ -208,7 +222,7 @@ impl CachedColumn { #[cfg(test)] pub(crate) async fn get_arrow_array_test_only(&self, batch_id: BatchID) -> Option { let entry_id = self.entry_id(batch_id).into(); - self.cache_store.get(&entry_id).await + self.cache_store.get(&entry_id, self.file_identity).await } /// Insert an array into the cache. @@ -222,7 +236,7 @@ impl CachedColumn { } self.cache_store - .insert(self.entry_id(batch_id).into(), array) + .insert(self.entry_id(batch_id).into(), self.file_identity, array) .await?; Ok(()) } @@ -239,7 +253,11 @@ impl CachedColumn { if self.snapshots.get(&entry_id).is_some() { return PrefetchOutcome::AlreadySnapshotted; } - match self.cache_store.prefetch(&entry_id).await { + match self + .cache_store + .prefetch(&entry_id, self.file_identity) + .await + { PrefetchResult::Snapshot(entry) => { self.snapshots.insert(entry_id, entry); PrefetchOutcome::Snapshotted diff --git a/src/datafusion/src/cache/id.rs b/src/datafusion/src/cache/id.rs index b52379de..84674589 100644 --- a/src/datafusion/src/cache/id.rs +++ b/src/datafusion/src/cache/id.rs @@ -51,8 +51,11 @@ const _: () = assert!(std::mem::align_of::() == 8); impl ParquetArrayID { /// Creates a new CacheEntryID. + /// + /// Only the low 16 bits of `file_id` are kept, so distinct files can + /// produce the same id. The cache tells them apart by the full file id it + /// is given alongside the entry id. pub fn new(file_id: u64, row_group_id: u64, column_id: u64, batch_id: BatchID) -> Self { - debug_assert!(file_id <= u16::MAX as u64); debug_assert!(row_group_id <= u16::MAX as u64); debug_assert!(column_id <= u16::MAX as u64); Self { @@ -146,8 +149,9 @@ pub struct ColumnAccessPath { impl ColumnAccessPath { /// Create a new instance of ColumnAccessPath. + /// + /// Only the low 16 bits of `file_id` are kept; see [`ParquetArrayID::new`]. pub fn new(file_id: u64, row_group_id: u64, column_id: u64) -> Self { - debug_assert!(file_id <= u16::MAX as u64); debug_assert!(row_group_id <= u16::MAX as u64); debug_assert!(column_id <= u16::MAX as u64); Self { @@ -226,9 +230,9 @@ mod tests { } #[test] - #[should_panic] - fn test_cache_entry_id_new_panic_file_id() { - ParquetArrayID::new((u16::MAX as u64) + 1, 0, 0, BatchID::from_raw(0)); + fn test_cache_entry_id_new_keeps_low_16_bits_of_file_id() { + let wrapped = ParquetArrayID::new((u16::MAX as u64) + 6, 2, 3, BatchID::from_raw(4)); + assert_eq!(wrapped, ParquetArrayID::new(5, 2, 3, BatchID::from_raw(4))); } #[test] diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index d51bc611..e4f751ff 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -119,6 +119,7 @@ impl CachedRowGroup { Arc::clone(field), Arc::clone(&cache_store), column_access_path, + file_idx, expression.clone(), *is_predicate_column, Arc::clone(&snapshots), @@ -218,7 +219,11 @@ impl CachedRowGroup { let entry_id = column.entry_id(batch_id).into(); let liquid_array = match snapshot_liquid { Some(array) => Some(array), - None => self.cache_store.try_read_liquid(&entry_id).await, + None => { + self.cache_store + .try_read_liquid(&entry_id, column.file_identity()) + .await + } }; let liquid_array = match liquid_array { None => { @@ -928,4 +933,70 @@ mod tests { let expected = BooleanBuffer::collect_bool(batch_size, |i| i == 1 || i == 3).into(); assert_eq!(result, expected); } + + /// A cache key holds 16 bits of file id, so the 65,537th file registered + /// computes the same keys as the first. Its reads must not be served the + /// first file's data, and its own data must still be cacheable. + #[tokio::test] + async fn files_whose_keys_collide_do_not_share_entries() { + let batch_size = 64; + let tmp_dir = tempfile::tempdir().unwrap(); + let store = t4::mount(tmp_dir.path().join("liquid_cache.t4")) + .await + .unwrap(); + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let identity = + |path: String| ParquetFileIdentity::new(ObjectStoreUrl::local_filesystem(), path); + let batch_id = BatchID::from_row_id(0, batch_size); + + let first = cache.register_or_get_file(identity("first".to_string()), schema.clone()); + let first_column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + let first_data: ArrayRef = Arc::new(Int32Array::from_iter_values(0..batch_size as i32)); + first_column + .insert(batch_id, first_data.clone()) + .await + .unwrap(); + + for i in 1..=u16::MAX as u32 { + cache.register_or_get_file(identity(format!("filler-{i}")), schema.clone()); + } + + let colliding = cache.register_or_get_file(identity("colliding".to_string()), schema); + let colliding_column = colliding.create_row_group(0, vec![]).get_column(0).unwrap(); + assert_eq!( + colliding_column.entry_id(batch_id), + first_column.entry_id(batch_id), + "the 65,537th file should compute the first file's key" + ); + + assert!( + colliding_column + .get_arrow_array_test_only(batch_id) + .await + .is_none(), + "a file was served another file's cached data" + ); + + let colliding_data: ArrayRef = + Arc::new(Int32Array::from_iter_values((0..batch_size as i32).rev())); + colliding_column + .insert(batch_id, colliding_data.clone()) + .await + .expect("a file could not cache its own data"); + let read = colliding_column + .get_arrow_array_test_only(batch_id) + .await + .unwrap(); + assert_eq!(read.as_ref(), colliding_data.as_ref()); + } }