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()); + } }