From 4e3504558e1c775ac0e63949050eaff25fdbb16a Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 10:13:40 +0530 Subject: [PATCH 01/12] fix(cache): serve entries only to their own file A cache key packs the file id into 16 bits, so the 65,537th distinct file a process registers aliases the first. Entries recorded nothing about where they came from and the read API took no expected identity, so an aliased lookup returned the other file's data: a panic when the column types differed, silently wrong rows when they matched. Record the unnarrowed file id alongside each entry and compare it on read. A mismatch reads as a miss, so the caller re-reads from its source and gets correct data, and it is counted. Maintenance that rewrites a key in place keeps the identity already recorded, and can no longer resurrect a key that has since been removed. Lease file ids instead of assigning them permanently. A lease is held by the file handle and by every row group and column derived from it, and returns to a FIFO pool when the last holder drops, so the id space tracks the files being read rather than every file ever read. Entries deliberately hold no lease: an id reused while its old entries are resident leaves them unreachable rather than readable, which is what lets the release stay out of index removal. Stop panicking when a cached entry cannot answer a predicate. try_eval_predicate now returns Option, which every caller already treated as "materialize from the source", so an entry built for a different column degrades to a re-read instead of unwinding the stream mid-flight. Drop the debug_assert on the narrowing conversions. It made debug builds panic where release builds truncate, which left the shipped behaviour untestable; a test pins the wrap down instead. --- src/core/src/cache/builders.rs | 27 +- src/core/src/cache/core.rs | 214 +++++++++------ src/core/src/cache/index.rs | 174 +++++++++++-- src/core/src/cache/observer/stats.rs | 7 + src/core/src/cache/tests/policies.rs | 10 +- src/core/src/cache/tests/squeezed.rs | 23 +- .../src/liquid_array/byte_view_array/mod.rs | 16 +- src/core/src/liquid_array/decimal_array.rs | 21 +- src/core/src/liquid_array/float_array.rs | 29 ++- .../liquid_array/hybrid_primitive_array.rs | 60 +++-- .../src/liquid_array/linear_integer_array.rs | 6 +- src/core/src/liquid_array/mod.rs | 33 ++- src/core/src/liquid_array/primitive_array.rs | 12 +- .../src/liquid_array/squeezed_date32_array.rs | 2 +- src/core/src/liquid_array/tests.rs | 5 +- src/core/study/cache_storage.rs | 4 +- src/core/study/squeeze_integer.rs | 6 +- src/core/tests/memory_footprint.rs | 4 +- src/datafusion/src/cache/column.rs | 38 ++- src/datafusion/src/cache/file_id.rs | 243 ++++++++++++++++++ src/datafusion/src/cache/id.rs | 52 ++-- src/datafusion/src/cache/mod.rs | 193 ++++++++++++-- src/datafusion/src/cache/stats.rs | 12 +- .../src/reader/runtime/liquid_cache_reader.rs | 7 +- 24 files changed, 930 insertions(+), 268 deletions(-) create mode 100644 src/datafusion/src/cache/file_id.rs diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index af0572dac..824c0fc0e 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) squeeze_hint: 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, squeeze_hint: None, @@ -214,7 +221,9 @@ impl<'a> Insert<'a> { } let batch = CacheEntry::memory_arrow(batch); self.storage.supersede_disk_copy(self.entry_id).await; - self.storage.insert_inner(self.entry_id, batch).await + self.storage + .insert_inner(self.entry_id, Some(self.identity), batch) + .await } } @@ -232,15 +241,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, } @@ -273,6 +284,7 @@ impl<'a> Get<'a> { self.storage .read_arrow_array( self.entry_id, + self.identity, self.selection, self.expression_hint.as_deref(), ) @@ -328,6 +340,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>, } @@ -336,11 +349,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, } @@ -355,7 +370,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 } } @@ -453,9 +468,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 2f85a1a7d..851495cac 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -91,10 +91,10 @@ impl DiskCopy { /// /// 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; +/// 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()); /// }); /// ``` @@ -161,6 +161,7 @@ impl LiquidCache { memory_arrow_bytes, memory_liquid_bytes, memory_squeezed_liquid_bytes, + identity_mismatches: self.index.identity_mismatches(), memory_usage_bytes, disk_usage_bytes, max_memory_bytes: self.config.max_memory_bytes(), @@ -173,23 +174,25 @@ impl LiquidCache { 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) + 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) } /// Try to read a liquid array from the cache. @@ -197,10 +200,11 @@ 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 batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); @@ -238,8 +242,8 @@ impl LiquidCache { } /// Check if a batch is cached. - pub fn is_cached(&self, entry_id: &EntryID) -> bool { - self.index.is_cached(entry_id) + pub fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { + self.index.is_cached(entry_id, identity) } /// Get the config of the cache. @@ -287,6 +291,7 @@ impl LiquidCache { Ok(()) => { self.try_insert( entry_id, + None, CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), ) .expect("failed to insert disk arrow entry"); @@ -304,7 +309,7 @@ impl LiquidCache { // Hydrated from disk and never modified since: the // bytes are already there, flip the index rather // than re-serialising and rewriting them. - self.try_insert(entry_id, CacheEntry::disk_liquid(data_type, bytes)) + self.try_insert(entry_id, None, CacheEntry::disk_liquid(data_type, bytes)) .expect("failed to insert disk liquid entry"); continue; } @@ -317,6 +322,7 @@ impl LiquidCache { Ok(()) => { self.try_insert( entry_id, + None, CacheEntry::disk_liquid(data_type, disk_bytes), ) .expect("failed to insert disk liquid entry"); @@ -327,7 +333,7 @@ impl LiquidCache { CacheEntry::MemorySqueezedLiquid(array) => { // We don't have to do anything, because it's already on disk let disk_entry = Self::disk_entry_from_squeezed(array); - self.try_insert(entry_id, disk_entry) + self.try_insert(entry_id, None, disk_entry) .expect("failed to insert disk entry"); } CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { @@ -406,10 +412,11 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, + identity: Option, 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 { @@ -568,7 +575,18 @@ impl LiquidCache { } } - fn try_insert(&self, entry_id: EntryID, to_insert: CacheEntry) -> Result<(), CacheEntry> { + /// `identity` names whose data this is for a caller-originated insert, and + /// is `None` for maintenance rewriting a key in place — see + /// [`ArtIndex::insert`]. A declined write is not an error: the key belongs + /// to someone else, or the entry being rewritten has since been removed. + /// Neither is worth retrying, so the reservation is handed back and the + /// call reports success with nothing stored. + fn try_insert( + &self, + entry_id: EntryID, + identity: Option, + 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(); @@ -580,14 +598,24 @@ impl LiquidCache { return Err(to_insert); } let batch_type = CachedBatchType::from(&to_insert); - self.index.insert(&entry_id, to_insert); + if !self.index.insert(&entry_id, identity, to_insert) { + self.budget + .try_update_memory_usage(new_memory_size, old_memory_size) + .expect("memory release cannot fail"); + return Ok(()); + } 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); + if !self.index.insert(&entry_id, identity, to_insert) { + self.budget + .try_update_memory_usage(new_memory_size, 0) + .expect("memory release cannot fail"); + return Ok(()); + } batch_type }; @@ -718,7 +746,7 @@ impl LiquidCache { self.write_batch_to_disk(to_squeeze, &new_batch, bytes_to_write) .await?; } - match self.try_insert(to_squeeze, new_batch) { + match self.try_insert(to_squeeze, None, new_batch) { Ok(()) => { break; } @@ -766,19 +794,20 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self.insert_inner(*entry_id, new_entry).await; + let _ = self.insert_inner(*entry_id, None, new_entry).await; } } pub(crate) async fn read_arrow_array( &self, entry_id: &EntryID, + identity: u64, selection: Option<&BooleanBuffer>, expression: Option<&CacheExpression>, ) -> Option { use arrow::array::BooleanArray; - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); self.trace(InternalEvent::Read { @@ -1060,13 +1089,14 @@ impl LiquidCache { pub(crate) async fn eval_predicate_internal( &self, entry_id: &EntryID, + identity: u64, selection_opt: Option<&BooleanBuffer>, predicate: &LiquidExpr, ) -> Option { use arrow::array::BooleanArray; self.observer.on_eval_predicate(); - let batch = self.index.get(entry_id)?; + let batch = self.index.get_checked(entry_id, identity)?; self.cache_policy .notify_access(entry_id, CachedBatchType::from(batch.as_ref())); self.trace(InternalEvent::EvalPredicate { @@ -1083,9 +1113,8 @@ impl LiquidCache { owned.as_ref().unwrap() }); let selection_array = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::filter(array, &selection_array) - .expect("selection must match array length"); - Some(self.eval_predicate_on_array(filtered, predicate)) + let filtered = arrow::compute::filter(array, &selection_array).ok()?; + self.eval_predicate_on_array(filtered, predicate) } entry @ CacheEntry::DiskArrow { .. } => { let array = self.read_disk_arrow_array(entry_id).await; @@ -1097,9 +1126,8 @@ impl LiquidCache { owned.as_ref().unwrap() }); let selection_array = BooleanArray::new(selection.clone(), None); - let filtered = arrow::compute::filter(&array, &selection_array) - .expect("selection must match array length"); - Some(self.eval_predicate_on_array(filtered, predicate)) + let filtered = arrow::compute::filter(&array, &selection_array).ok()?; + self.eval_predicate_on_array(filtered, predicate) } CacheEntry::MemoryLiquid(array) => { let mut owned = None; @@ -1107,7 +1135,7 @@ impl LiquidCache { owned = Some(BooleanBuffer::new_set(array.len())); owned.as_ref().unwrap() }); - Some(array.try_eval_predicate(predicate, selection)) + array.try_eval_predicate(predicate, selection) } entry @ CacheEntry::DiskLiquid { .. } => { let liquid = self.read_disk_liquid_array(entry_id).await; @@ -1118,7 +1146,7 @@ impl LiquidCache { owned = Some(BooleanBuffer::new_set(liquid.len())); owned.as_ref().unwrap() }); - Some(liquid.try_eval_predicate(predicate, selection)) + liquid.try_eval_predicate(predicate, selection) } CacheEntry::MemorySqueezedLiquid(array) => { self.eval_predicate_on_squeezed(array, selection_opt, predicate) @@ -1138,25 +1166,25 @@ impl LiquidCache { owned = Some(BooleanBuffer::new_set(array.len())); owned.as_ref().unwrap() }); - Some(array.try_eval_predicate(predicate, selection).await) + array.try_eval_predicate(predicate, selection).await } - fn eval_predicate_on_array(&self, array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { + /// `None` when the cached array cannot answer the predicate. See the + /// free function of the same name in `liquid_array`. + fn eval_predicate_on_array( + &self, + array: ArrayRef, + predicate: &LiquidExpr, + ) -> Option { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", array.data_type().clone(), true, )])); - let record_batch = - RecordBatch::try_new(schema, vec![array]).expect("single-column predicate batch"); - let result = predicate - .physical_expr() - .evaluate(&record_batch) - .expect("validated LiquidExpr must evaluate"); - let boolean_array = result - .into_array(record_batch.num_rows()) - .expect("predicate output must be an array"); - boolean_array.as_boolean().clone() + let record_batch = RecordBatch::try_new(schema, vec![array]).ok()?; + let result = predicate.physical_expr().evaluate(&record_batch).ok()?; + let boolean_array = result.into_array(record_batch.num_rows()).ok()?; + Some(boolean_array.as_boolean().clone()) } } @@ -1218,7 +1246,10 @@ 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, Some(0), array1) + .await + .unwrap(); // Verify budget usage and data correctness assert_eq!(store.budget.memory_usage_bytes(), size1); @@ -1231,13 +1262,19 @@ 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, Some(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, Some(0), array3) + .await + .unwrap(); assert_eq!(store.budget.memory_usage_bytes(), size3 + size2); assert!(store.index().get(&EntryID::from(999)).is_none()); @@ -1256,6 +1293,7 @@ mod tests { store .insert_inner( entry_id, + Some(0), CacheEntry::memory_squeezed_liquid(squeezed.clone()), ) .await @@ -1263,7 +1301,7 @@ mod tests { let expr = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); let result = store - .get(&entry_id) + .get(&entry_id, 0) .with_expression_hint(expr) .read() .await @@ -1294,7 +1332,7 @@ 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, Some(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -1303,7 +1341,7 @@ mod tests { } store - .insert_inner(entry_id2, create_test_array(800)) + .insert_inner(entry_id2, Some(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -1353,7 +1391,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(); } }); })); @@ -1387,8 +1425,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(); @@ -1412,14 +1456,14 @@ 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 { .. })); } - 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(); @@ -1436,7 +1480,7 @@ mod tests { let liquid = transcode_liquid_inner(&arrow_array, &compressor).unwrap(); store - .insert_inner(entry_id, CacheEntry::memory_liquid(liquid.clone())) + .insert_inner(entry_id, Some(0), CacheEntry::memory_liquid(liquid.clone())) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); @@ -1445,7 +1489,7 @@ mod tests { assert!(matches!(entry.as_ref(), 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(); @@ -1463,10 +1507,10 @@ 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.is_cached(&EntryID::from(900usize))); + assert!(!cache.is_cached(&EntryID::from(900usize), 0)); } #[tokio::test] @@ -1485,14 +1529,14 @@ 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.is_cached(&first)); + assert!(cache.is_cached(&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.is_cached(&first)); + assert!(!cache.is_cached(&first, 0)); assert!(matches!( cache.index().get(&second).unwrap().as_ref(), CacheEntry::DiskArrow { .. } @@ -1513,13 +1557,13 @@ 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.is_cached(&first) || !cache.is_cached(&second)); + assert!(!cache.is_cached(&first, 0) || !cache.is_cached(&second, 0)); } #[tokio::test] @@ -1534,14 +1578,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.is_cached(&entry)); + assert!(!cache.is_cached(&entry, 0)); } #[tokio::test] @@ -1554,12 +1598,12 @@ 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.is_cached(&entry_id)); + assert!(!cache.is_cached(&entry_id, 0)); } async fn hydrating_cache() -> Arc { @@ -1594,15 +1638,15 @@ mod tests { let v1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let v2: ArrayRef = Arc::new(Int32Array::from_iter_values(100..164)); - cache.insert(id, v1).await.unwrap(); + cache.insert(id, 0, v1).await.unwrap(); let v1_disk_bytes = demote_to_disk(&cache, id).await; assert_eq!(cache.budget().disk_usage_bytes(), v1_disk_bytes); - cache.insert(id, v2.clone()).await.unwrap(); + cache.insert(id, 0, v2.clone()).await.unwrap(); let disk_after_overwrite = cache.budget().disk_usage_bytes(); let v2_disk_bytes = demote_to_disk(&cache, id).await; - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), v2.as_ref(), "read back the superseded value"); assert_eq!( disk_after_overwrite, 0, @@ -1621,9 +1665,9 @@ mod tests { let v1: ArrayRef = Arc::new(Int32Array::from_iter_values(0..16)); let v2: ArrayRef = Arc::new(Int32Array::from_iter_values(100..164)); - cache.insert(id, v1.clone()).await.unwrap(); + cache.insert(id, 0, v1.clone()).await.unwrap(); let v1_disk_bytes = demote_to_disk(&cache, id).await; - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), v1.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1631,11 +1675,11 @@ mod tests { )); assert_eq!(cache.budget().disk_usage_bytes(), v1_disk_bytes); - cache.insert(id, v2.clone()).await.unwrap(); + cache.insert(id, 0, v2.clone()).await.unwrap(); let disk_after_overwrite = cache.budget().disk_usage_bytes(); let v2_disk_bytes = demote_to_disk(&cache, id).await; - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), v2.as_ref(), "read back the superseded value"); assert_eq!(disk_after_overwrite, 0); assert_eq!(cache.budget().disk_usage_bytes(), v2_disk_bytes); @@ -1651,13 +1695,13 @@ mod tests { let id = EntryID::from(922usize); let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); - cache.insert(id, array.clone()).await.unwrap(); + cache.insert(id, 0, array.clone()).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), CacheEntry::DiskArrow { .. } )); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), array.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1666,7 +1710,7 @@ mod tests { let disk_bytes = demote_to_disk(&cache, id).await; assert_eq!(cache.budget().disk_usage_bytes(), disk_bytes); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), array.as_ref()); } @@ -1686,7 +1730,7 @@ mod tests { let expr = Arc::new(CacheExpression::extract_date32(Date32Field::Year)); cache - .insert(id, dates.clone()) + .insert(id, 0, dates.clone()) .with_squeeze_hint(expr.clone()) .await .unwrap(); @@ -1702,7 +1746,7 @@ mod tests { // Drain the IO counters so the count below covers only the re-eviction. let _ = cache.observer().runtime_snapshot(); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), dates.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1725,7 +1769,7 @@ mod tests { assert_eq!(cache.budget().disk_usage_bytes(), disk_bytes); let years = cache - .get(&id) + .get(&id, 0) .with_expression_hint(expr) .read() .await @@ -1766,7 +1810,7 @@ mod tests { let expr = expr.clone(); async move { cache - .insert(id, dates) + .insert(id, 0, dates) .with_squeeze_hint(expr) .await .unwrap(); @@ -1797,11 +1841,11 @@ mod tests { // Too big for memory, and the disk tier is full with no victims. let too_big: ArrayRef = Arc::new(Int32Array::from_iter_values(0..(1 << 16))); - let result = cache.insert(id, too_big).await; + let result = cache.insert(id, 0, too_big).await; assert_eq!(result, Err(CacheFull)); - assert!(!cache.is_cached(&id)); - assert!(cache.get(&id).await.is_none()); + assert!(!cache.is_cached(&id, 0)); + assert!(cache.get(&id, 0).await.is_none()); assert_eq!(cache.budget().disk_usage_bytes(), 0); assert_eq!(cache.budget().memory_usage_bytes(), 0); } @@ -1823,9 +1867,9 @@ mod tests { .await; let id = EntryID::from(925usize); - cache.insert(id, array.clone()).await.unwrap(); + cache.insert(id, 0, array.clone()).await.unwrap(); cache.flush_all_to_disk().await.unwrap(); - let read = cache.get(&id).await.expect("present"); + let read = cache.get(&id, 0).await.expect("present"); assert_eq!(read.as_ref(), array.as_ref()); assert!(matches!( cache.index().get(&id).unwrap().as_ref(), @@ -1837,7 +1881,7 @@ mod tests { // that is full with the entry's own copy, so the entry is dropped. cache.flush_all_to_disk().await.unwrap(); - assert!(!cache.is_cached(&id)); + assert!(!cache.is_cached(&id, 0)); assert_eq!( cache.budget().disk_usage_bytes(), 0, diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index cf8881ab5..1cb3a8b53 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -1,7 +1,7 @@ use congee::CongeeArc; use std::{ fmt::{Debug, Formatter}, - sync::atomic::{AtomicUsize, Ordering}, + sync::atomic::{AtomicU64, AtomicUsize, Ordering}, }; use crate::cache::{cached_batch::CacheEntry, utils::EntryID}; @@ -20,25 +20,40 @@ use crate::sync::{Arc, RwLock}; /// So the tree stores a small slot and the payload is taken out of it the /// moment the index gives the entry up. The deferred drop then reclaims only /// an empty shell, and the array dies with the last caller-held reference. -struct Slot(RwLock>>); +/// +/// The slot also records the identity of what it holds. `EntryID` is a packed +/// integer whose fields are narrower than the values they encode, so two +/// distinct sources can compute the same key; the key alone therefore cannot +/// answer "is this the entry I asked for?". `identity` is the caller's +/// unnarrowed name for the data, compared on every read through +/// [`ArtIndex::get_checked`], which turns such aliasing into a miss rather +/// than a wrong answer. +struct Slot { + identity: u64, + entry: RwLock>>, +} impl Slot { - fn new(entry: CacheEntry) -> Arc { - Arc::new(Self(RwLock::new(Some(Arc::new(entry))))) + fn new(identity: u64, entry: CacheEntry) -> Arc { + Arc::new(Self { + identity, + entry: RwLock::new(Some(Arc::new(entry))), + }) } fn load(&self) -> Option> { - self.0.read().unwrap().clone() + self.entry.read().unwrap().clone() } fn take(&self) -> Option> { - self.0.write().unwrap().take() + self.entry.write().unwrap().take() } } pub(crate) struct ArtIndex { art: CongeeArc, entry_count: AtomicUsize, + identity_mismatches: AtomicU64, } impl Debug for ArtIndex { @@ -52,9 +67,16 @@ impl ArtIndex { Self { art: CongeeArc::new(), entry_count: AtomicUsize::new(0), + identity_mismatches: AtomicU64::new(0), } } + /// Look up an entry without checking whose it is. + /// + /// This is for maintenance that acts on whatever currently occupies a key — + /// eviction, squeezing, disk supersession, iteration for stats. A read + /// serving a caller must use [`Self::get_checked`] instead, so that a key + /// collision cannot return one caller another's data. pub(crate) fn get(&self, entry_id: &EntryID) -> Option> { let guard = self.art.pin(); // An empty slot means the entry was removed or replaced between the @@ -69,15 +91,68 @@ impl ArtIndex { self.art.get(*entry_id, &guard)?.load() } - pub(crate) fn is_cached(&self, entry_id: &EntryID) -> bool { - self.get(entry_id).is_some() + /// Look up an entry, returning it only if it is the one `identity` names. + /// + /// A mismatch reads as a miss, so the caller re-reads from its source and + /// gets correct data. It is also counted: the tally is expected to stay at + /// zero, and a non-zero value means two sources are computing the same + /// `EntryID`. + pub(crate) fn get_checked(&self, entry_id: &EntryID, identity: u64) -> Option> { + let guard = self.art.pin(); + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + if let Some(entry) = slot.load() { + return Some(entry); + } + // Re-read as in `get`: a replace leaves the key present under a new + // slot, which carries its own identity and must be checked again. + let slot = self.art.get(*entry_id, &guard)?; + if slot.identity != identity { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return None; + } + slot.load() + } + + pub(crate) fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { + self.get_checked(entry_id, identity).is_some() } - pub(crate) fn insert(&self, entry_id: &EntryID, batch: CacheEntry) { + /// Store `batch` under `entry_id`, returning whether it was stored. + /// + /// `identity` is `Some` for a caller-originated insert, naming whose data + /// this is. A key already held by a *different* identity is refused rather + /// than overwritten, so two sources colliding on one key do not evict each + /// other's data back and forth; the loser reads from its own source. + /// + /// `identity` is `None` for maintenance that rewrites a key in place — + /// transcode, squeeze, hydrate, spill to disk — which keeps the identity + /// already recorded. Maintenance cannot change whose an entry is, and + /// cannot resurrect a key that has since been removed: with nothing to + /// preserve there is no identity to record, so the write is dropped. + pub(crate) fn insert( + &self, + entry_id: &EntryID, + identity: Option, + batch: CacheEntry, + ) -> bool { let guard = self.art.pin(); + let existing_identity = self.art.get(*entry_id, &guard).map(|slot| slot.identity); + let identity = match (identity, existing_identity) { + (Some(new), Some(old)) if new != old => { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + return false; + } + (Some(new), _) => new, + (None, Some(old)) => old, + (None, None) => return false, + }; let existing = self .art - .insert(*entry_id, Slot::new(batch), &guard) + .insert(*entry_id, Slot::new(identity, batch), &guard) .expect("Insertion failed"); match existing { Some(replaced) => drop(replaced.take()), @@ -85,6 +160,7 @@ impl ArtIndex { self.entry_count.fetch_add(1, Ordering::Relaxed); } } + true } pub(crate) fn remove(&self, entry_id: &EntryID) -> Option> { @@ -117,6 +193,15 @@ impl ArtIndex { pub(crate) fn entry_count(&self) -> usize { self.entry_count.load(Ordering::Relaxed) } + + /// How many lookups or inserts found a key held by a different identity. + /// + /// Expected to stay at zero. A non-zero value means two sources compute the + /// same `EntryID`, and every one of them was served correctly only because + /// the check turned it into a miss. + pub(crate) fn identity_mismatches(&self) -> u64 { + self.identity_mismatches.load(Ordering::Relaxed) + } } #[cfg(test)] @@ -134,17 +219,17 @@ mod tests { let array1 = create_test_array(100); // Initially, entries should not be cached - assert!(!store.is_cached(&entry_id1)); - assert!(!store.is_cached(&entry_id2)); + assert!(!store.is_cached(&entry_id1, 0)); + assert!(!store.is_cached(&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, Some(0), array1.clone()); } - assert!(store.is_cached(&entry_id1)); - assert!(!store.is_cached(&entry_id2)); + assert!(store.is_cached(&entry_id1, 0)); + assert!(!store.is_cached(&entry_id2, 0)); // Get should return the cached value match store.get(&entry_id1) { @@ -162,14 +247,14 @@ 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, Some(0), array.clone()); let entry_id: EntryID = EntryID::from(1); - assert!(store.is_cached(&entry_id)); + assert!(store.is_cached(&entry_id, 0)); store.reset(); let entry_id: EntryID = EntryID::from(1); - assert!(!store.is_cached(&entry_id)); + assert!(!store.is_cached(&entry_id, 0)); } /// The array behind a removed or replaced entry must die with the last @@ -184,8 +269,8 @@ mod tests { unreachable!() }; let weak_first = Arc::downgrade(first_array); - store.insert(&id, first); - store.insert(&id, create_test_array(200)); + store.insert(&id, Some(0), first); + store.insert(&id, Some(0), create_test_array(200)); assert!( weak_first.upgrade().is_none(), "replaced entry still alive: held by the index's deferred drop" @@ -204,4 +289,53 @@ mod tests { ); assert_eq!(store.entry_count(), 0); } + + /// Two files whose ids narrow to the same `EntryID` must not read each + /// other's data. Before the identity check this returned the incumbent's + /// array, which is a wrong answer whenever the two happen to share a type. + #[test] + fn an_entry_is_never_served_to_a_different_identity() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(7); + + assert!(store.insert(&key, Some(1), create_test_array(100))); + + // The colliding file asks for the same key and is told nothing is there. + assert!(store.get_checked(&key, 2).is_none()); + assert!(!store.is_cached(&key, 2)); + assert_eq!(store.identity_mismatches(), 2); + + // The owner still reads its own entry. + assert!(store.get_checked(&key, 1).is_some()); + + // And the colliding file cannot displace it by writing over the key, + // so the two do not take turns evicting each other. + assert!(!store.insert(&key, Some(2), create_test_array(200))); + match store.get_checked(&key, 1).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!(array.len(), 100), + other => panic!("expected the owner's array, found {other}"), + } + } + + /// Maintenance rewrites a key in place and must neither change whose the + /// entry is nor bring back one that has been removed — a stale reader that + /// misses goes on to insert what it read, and that write must not land + /// under a key nobody owns any more. + #[test] + fn maintenance_preserves_identity_and_cannot_resurrect_a_removed_key() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(9); + + assert!(store.insert(&key, Some(5), create_test_array(10))); + assert!(store.insert(&key, None, create_test_array(20))); + assert!( + store.get_checked(&key, 5).is_some(), + "rewriting in place kept the identity" + ); + + store.remove(&key); + assert!(!store.insert(&key, None, create_test_array(30))); + assert!(store.get(&key).is_none()); + assert_eq!(store.entry_count(), 0); + } } diff --git a/src/core/src/cache/observer/stats.rs b/src/core/src/cache/observer/stats.rs index fa0c3d9ae..bd827b752 100644 --- a/src/core/src/cache/observer/stats.rs +++ b/src/core/src/cache/observer/stats.rs @@ -142,6 +142,13 @@ pub struct CacheStats { pub memory_liquid_bytes: usize, /// Total size of in-memory Squeezed-Liquid entries in bytes. pub memory_squeezed_liquid_bytes: usize, + /// Lookups and inserts that found a key held by a different identity. + /// + /// Expected to stay at zero. A non-zero value means two sources compute + /// the same `EntryID` — each was served correctly, because the check + /// turns the collision into a miss, but the cache is not holding what + /// either of them could use. + pub identity_mismatches: u64, /// Total memory usage of the cache. pub memory_usage_bytes: usize, /// Total disk usage of the cache. diff --git a/src/core/src/cache/tests/policies.rs b/src/core/src/cache/tests/policies.rs index 29c849902..b0e159606 100644 --- a/src/core/src/cache/tests/policies.rs +++ b/src/core/src/cache/tests/policies.rs @@ -18,12 +18,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()); } @@ -44,17 +44,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(); diff --git a/src/core/src/cache/tests/squeezed.rs b/src/core/src/cache/tests/squeezed.rs index 3d7c152e7..0f3e75e2f 100644 --- a/src/core/src/cache/tests/squeezed.rs +++ b/src/core/src/cache/tests/squeezed.rs @@ -41,7 +41,7 @@ async fn read_squeezed_date_time() { for i in 0..4 { let entry_id = EntryID::from(i); cache - .insert(entry_id, array.clone()) + .insert(entry_id, 0, array.clone()) .with_squeeze_hint(expression.clone()) .await .unwrap(); @@ -50,14 +50,14 @@ async fn read_squeezed_date_time() { for i in 0..4 { let entry_id = EntryID::from(i); let array = cache - .get(&entry_id) + .get(&entry_id, 0) .with_expression_hint(expression.clone()) .await .unwrap(); assert_eq!(array.len(), array.len()); } cache - .get(&EntryID::from(1)) + .get(&EntryID::from(1), 0) .with_expression_hint(Arc::new(CacheExpression::extract_date32( Date32Field::Month, ))) @@ -111,14 +111,14 @@ async fn read_squeezed_variant_path() { for i in 0..3 { let entry_id = EntryID::from(i); cache - .insert(entry_id, variant_array.clone()) + .insert(entry_id, 0, variant_array.clone()) .with_squeeze_hint(name_expr.clone()) .await .unwrap(); } let squeezed = cache - .get(&EntryID::from(0)) + .get(&EntryID::from(0), 0) .with_expression_hint(name_expr.clone()) .read() .await @@ -126,13 +126,13 @@ async fn read_squeezed_variant_path() { assert_eq!(squeezed.len(), variant_array.len()); cache - .get(&EntryID::from(0)) + .get(&EntryID::from(0), 0) .with_expression_hint(age_expr.clone()) .read() .await .unwrap(); cache - .get(&EntryID::from(1)) + .get(&EntryID::from(1), 0) .with_expression_hint(zipcode_expr.clone()) .read() .await @@ -171,19 +171,22 @@ async fn read_squeezed_int64_array() { let entry_id = EntryID::from(i); if i % 2 == 0 { cache - .insert(entry_id, int64_array.clone()) + .insert(entry_id, 0, int64_array.clone()) .with_squeeze_hint(expression.clone()) .await .unwrap(); } else { - cache.insert(entry_id, int64_array.clone()).await.unwrap(); + cache + .insert(entry_id, 0, int64_array.clone()) + .await + .unwrap(); } } for i in 0..4 { let entry_id = EntryID::from(i); let array = cache - .get(&entry_id) + .get(&entry_id, 0) .with_expression_hint(expression.clone()) .read() .await diff --git a/src/core/src/liquid_array/byte_view_array/mod.rs b/src/core/src/liquid_array/byte_view_array/mod.rs index 1d484148b..0be3fcf6b 100644 --- a/src/core/src/liquid_array/byte_view_array/mod.rs +++ b/src/core/src/liquid_array/byte_view_array/mod.rs @@ -354,11 +354,15 @@ impl LiquidArray for LiquidByteViewArray { Arc::new(dict) } - fn try_eval_predicate(&self, expr: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + expr: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = helpers::filter_inner(self, filter); helpers::try_eval_predicate_in_memory(expr.physical_expr(), &filtered) - .unwrap_or_else(|| eval_predicate_on_array(filtered.to_arrow_array(), expr)) + .or_else(|| eval_predicate_on_array(filtered.to_arrow_array(), expr)) } fn to_bytes(&self) -> Vec { @@ -487,13 +491,17 @@ impl LiquidSqueezedArray for LiquidByteViewArray { /// /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. /// The returned boolean mask is nullable if the the original array is nullable. - async fn try_eval_predicate(&self, expr: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + async fn try_eval_predicate( + &self, + expr: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { // Reuse generic filter path first to reduce input rows if any let filtered = helpers::filter_inner(self, filter); if let Some(mask) = helpers::try_eval_predicate_on_disk(expr.physical_expr(), &filtered).await { - mask + Some(mask) } else { eval_predicate_on_array(filtered.to_arrow_array().await, expr) } diff --git a/src/core/src/liquid_array/decimal_array.rs b/src/core/src/liquid_array/decimal_array.rs index 374ff9cc7..4e7b7e43b 100644 --- a/src/core/src/liquid_array/decimal_array.rs +++ b/src/core/src/liquid_array/decimal_array.rs @@ -545,7 +545,7 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { let filtered = self.filter_inner(filter); let expr = if let Some(expr) = unwrap_dynamic_filter(liquid_expr.physical_expr()) { @@ -575,7 +575,7 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { match filtered.try_eval_predicate_inner(&op, literal) { Ok(Some(mask)) => { self.io.trace_io_saved(); - return mask; + return Some(mask); } Ok(None) => { let fallback = self.filter(filter).await; @@ -588,8 +588,7 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs = ColumnarValue::Array(filtered_arr); @@ -606,12 +605,11 @@ impl LiquidSqueezedArray for LiquidDecimalQuantizedArray { return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() + // A comparison that cannot evaluate means this array is not what the + // predicate was built for. Report it as unanswerable rather than + // asserting: the caller falls back to the source. + let result = result.ok()?; + Some(result.into_array(filtered_len).ok()?.as_boolean().clone()) } } @@ -685,7 +683,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = BooleanArray::from(vec![Some(true), Some(true), None, Some(true)]); assert_eq!(got, expected); assert_eq!(io.reads(), 0); diff --git a/src/core/src/liquid_array/float_array.rs b/src/core/src/liquid_array/float_array.rs index 683dece88..6c4a2b1a2 100644 --- a/src/core/src/liquid_array/float_array.rs +++ b/src/core/src/liquid_array/float_array.rs @@ -993,7 +993,7 @@ where &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { // Apply selection first to reduce input rows let filtered = self.filter_inner(filter); let expr = liquid_expr.physical_expr(); @@ -1005,7 +1005,7 @@ where let supported_op = Operator::from_datafusion(op); if let Some(supported_op) = supported_op { match filtered.try_eval_predicate_inner(&supported_op, literal) { - Ok(Some(mask)) => return mask, + Ok(Some(mask)) => return Some(mask), Ok(None) => { let fallback = self.filter(filter).await; return eval_predicate_on_array(fallback, liquid_expr); @@ -1018,8 +1018,7 @@ where let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs = ColumnarValue::Array(filtered_arr); @@ -1036,12 +1035,10 @@ where return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - return result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone(); + // Unanswerable rather than fatal: this array is not what the + // predicate was built for, so the caller falls back. + let result = result.ok()?; + return Some(result.into_array(filtered_len).ok()?.as_boolean().clone()); } } let fallback = self.filter(filter).await; @@ -1291,7 +1288,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1323,7 +1321,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1386,7 +1385,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1418,7 +1418,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { diff --git a/src/core/src/liquid_array/hybrid_primitive_array.rs b/src/core/src/liquid_array/hybrid_primitive_array.rs index f2d0845bb..65563cdc1 100644 --- a/src/core/src/liquid_array/hybrid_primitive_array.rs +++ b/src/core/src/liquid_array/hybrid_primitive_array.rs @@ -338,7 +338,7 @@ where &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { // Apply selection first to reduce input rows let filtered = self.filter_inner(filter); @@ -374,7 +374,7 @@ where match filtered.try_eval_predicate_inner(&supported_op, literal) { Ok(Some(mask)) => { self.io.trace_io_saved(); - return mask; + return Some(mask); } Ok(None) => { let fallback = self.filter(filter).await; @@ -389,8 +389,7 @@ where let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs_array = match lhs_kind { PredicateLhs::Plain => filtered_arr, @@ -415,12 +414,11 @@ where return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() + // A comparison that cannot evaluate means this array is not what the + // predicate was built for. Report it as unanswerable rather than + // asserting: the caller falls back to the source. + let result = result.ok()?; + Some(result.into_array(filtered_len).ok()?.as_boolean().clone()) } } @@ -701,7 +699,7 @@ where &self, liquid_expr: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { // Apply selection first to reduce input rows let filtered = self.filter_inner(filter); @@ -737,7 +735,7 @@ where match filtered.try_eval_predicate_inner(&supported_op, literal) { Ok(Some(mask)) => { self.io.trace_io_saved(); - return mask; + return Some(mask); } Ok(None) => { let fallback = self.filter(filter).await; @@ -752,8 +750,7 @@ where let full = self.hydrate_full_arrow().await; let selection_array = BooleanArray::new(filter.clone(), None); - let filtered_arr = arrow::compute::filter(&full, &selection_array) - .expect("selection must match array length"); + let filtered_arr = arrow::compute::filter(&full, &selection_array).ok()?; let filtered_len = filtered_arr.len(); let lhs_array = match lhs_kind { PredicateLhs::Plain => filtered_arr, @@ -778,12 +775,11 @@ where return eval_predicate_on_array(fallback, liquid_expr); } }; - let result = result.expect("validated LiquidExpr comparison must evaluate"); - result - .into_array(filtered_len) - .expect("comparison output must be an array") - .as_boolean() - .clone() + // A comparison that cannot evaluate means this array is not what the + // predicate was built for. Report it as unanswerable rather than + // asserting: the caller falls back to the source. + let result = result.ok()?; + Some(result.into_array(filtered_len).ok()?.as_boolean().clone()) } } @@ -994,7 +990,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert_eq!(io.reads(), 0); assert_eq!(got, expected); @@ -1015,7 +1012,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert!(io.reads() > 0); assert_eq!(got, expected); @@ -1081,7 +1079,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert_eq!(io.reads(), 0); assert_eq!(got, expected); @@ -1101,7 +1100,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = expected_for(op, k); assert!(io.reads() > 0); assert_eq!(got, expected); @@ -1143,7 +1143,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1175,7 +1176,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1225,7 +1227,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { @@ -1257,7 +1260,8 @@ mod tests { let got = block_on(hybrid.try_eval_predicate( &crate::cache::LiquidExpr::new_unchecked(expr_eq_present.clone()), &mask, - )); + )) + .expect("predicate must evaluate in this test"); let expected = { let vals: Vec> = (0..arr.len()) .map(|i| { diff --git a/src/core/src/liquid_array/linear_integer_array.rs b/src/core/src/liquid_array/linear_integer_array.rs index 4a8bc14ed..bbbcee7ef 100644 --- a/src/core/src/liquid_array/linear_integer_array.rs +++ b/src/core/src/liquid_array/linear_integer_array.rs @@ -351,7 +351,11 @@ where filter::filter(&arr, &selection).unwrap() } - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let arr = self.filter(filter); eval_predicate_on_array(arr, predicate) } diff --git a/src/core/src/liquid_array/mod.rs b/src/core/src/liquid_array/mod.rs index 7776c7640..cb5b3d00c 100644 --- a/src/core/src/liquid_array/mod.rs +++ b/src/core/src/liquid_array/mod.rs @@ -124,7 +124,11 @@ pub trait LiquidArray: std::fmt::Debug + Send + Sync { /// /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable. /// The returned boolean mask is nullable if the the original array is nullable. - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } @@ -252,7 +256,7 @@ pub trait LiquidSqueezedArray: std::fmt::Debug + Send + Sync { &self, predicate: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { let filtered = self.filter(filter).await; eval_predicate_on_array(filtered, predicate) } @@ -262,21 +266,26 @@ pub trait LiquidSqueezedArray: std::fmt::Debug + Send + Sync { fn disk_backing(&self) -> SqueezedBacking; } -pub(crate) fn eval_predicate_on_array(array: ArrayRef, predicate: &LiquidExpr) -> BooleanArray { +/// Evaluate `predicate` against a one-column batch built from `array`. +/// +/// Returns `None` when the array cannot answer the predicate — a data type the +/// expression was not built for, or a batch it cannot be evaluated against. +/// That is reachable whenever a cached entry is not the one the predicate was +/// built for, so it must not be an assertion: the caller treats `None` as "the +/// cache cannot answer", materializes from the source and evaluates there. +pub(crate) fn eval_predicate_on_array( + array: ArrayRef, + predicate: &LiquidExpr, +) -> Option { let schema = Arc::new(Schema::new(vec![Field::new( "liquid_predicate_col", array.data_type().clone(), true, )])); - let record_batch = RecordBatch::try_new(schema, vec![array]).expect("predicate input batch"); - let result = predicate - .physical_expr() - .evaluate(&record_batch) - .expect("validated LiquidExpr must evaluate"); - let boolean_array = result - .into_array(record_batch.num_rows()) - .expect("predicate output must be an array"); - boolean_array.as_boolean().clone() + let record_batch = RecordBatch::try_new(schema, vec![array]).ok()?; + let result = predicate.physical_expr().evaluate(&record_batch).ok()?; + let boolean_array = result.into_array(record_batch.num_rows()).ok()?; + Some(boolean_array.as_boolean().clone()) } /// A trait to read the backing bytes of a squeezed array from disk. diff --git a/src/core/src/liquid_array/primitive_array.rs b/src/core/src/liquid_array/primitive_array.rs index c651112b5..eb0b2eb78 100644 --- a/src/core/src/liquid_array/primitive_array.rs +++ b/src/core/src/liquid_array/primitive_array.rs @@ -373,7 +373,11 @@ where arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() } - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } @@ -573,7 +577,11 @@ where arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap() } - fn try_eval_predicate(&self, predicate: &LiquidExpr, filter: &BooleanBuffer) -> BooleanArray { + fn try_eval_predicate( + &self, + predicate: &LiquidExpr, + filter: &BooleanBuffer, + ) -> Option { let filtered = self.filter(filter); eval_predicate_on_array(filtered, predicate) } diff --git a/src/core/src/liquid_array/squeezed_date32_array.rs b/src/core/src/liquid_array/squeezed_date32_array.rs index e2497e556..75b817340 100644 --- a/src/core/src/liquid_array/squeezed_date32_array.rs +++ b/src/core/src/liquid_array/squeezed_date32_array.rs @@ -480,7 +480,7 @@ impl LiquidSqueezedArray for SqueezedDate32Array { &self, predicate: &LiquidExpr, filter: &BooleanBuffer, - ) -> BooleanArray { + ) -> Option { let filtered = self.filter(filter).await; eval_predicate_on_array(filtered, predicate) } diff --git a/src/core/src/liquid_array/tests.rs b/src/core/src/liquid_array/tests.rs index de3475737..ac04d8c9c 100644 --- a/src/core/src/liquid_array/tests.rs +++ b/src/core/src/liquid_array/tests.rs @@ -185,8 +185,9 @@ mod byte_view_tests { let expr: Arc = Arc::new(BinaryExpr::new(col, Operator::Eq, lit)); let liquid = make_byte_view(&input); - let result = - liquid.try_eval_predicate(&crate::cache::LiquidExpr::new_unchecked(expr), &mask); + let result = liquid + .try_eval_predicate(&crate::cache::LiquidExpr::new_unchecked(expr), &mask) + .expect("predicate must evaluate in this test"); let expected = BooleanArray::from(vec![ Some(true), None, diff --git a/src/core/study/cache_storage.rs b/src/core/study/cache_storage.rs index e4d6533a3..b4c4c5a91 100644 --- a/src/core/study/cache_storage.rs +++ b/src/core/study/cache_storage.rs @@ -92,7 +92,7 @@ fn main() { continue; }; if storage - .eval_predicate(id, &liquid_expr) + .eval_predicate(id, 0, &liquid_expr) .with_selection(&selection) .await .is_some() @@ -145,7 +145,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/core/study/squeeze_integer.rs b/src/core/study/squeeze_integer.rs index cfdb89cc1..a19a99d3f 100644 --- a/src/core/study/squeeze_integer.rs +++ b/src/core/study/squeeze_integer.rs @@ -400,8 +400,10 @@ fn try_eval_or_fetch( ) -> (BooleanArray, usize) { io.reset_bytes_read(); let maybe_expr = LiquidExpr::try_new(expr.clone(), &hybrid.original_arrow_data_type(), None); - if let Some(liquid_expr) = maybe_expr { - let mask = futures::executor::block_on(hybrid.try_eval_predicate(&liquid_expr, filter)); + if let Some(liquid_expr) = maybe_expr + && let Some(mask) = + futures::executor::block_on(hybrid.try_eval_predicate(&liquid_expr, filter)) + { return (mask, io.bytes_read()); } // Not supported in hybrid form: materialize from full bytes and compute via Arrow. diff --git a/src/core/tests/memory_footprint.rs b/src/core/tests/memory_footprint.rs index 836255880..bc10719ea 100644 --- a/src/core/tests/memory_footprint.rs +++ b/src/core/tests/memory_footprint.rs @@ -168,7 +168,7 @@ async fn heap_footprint_tracks_budget_for_oversized_working_set() { // column batch as arrow, dropping the caller's copy right after. for i in 0..ENTRIES { let arr = make_entry(i as u64, ROWS); - cache.insert(EntryID::from(i), arr).await.unwrap(); + cache.insert(EntryID::from(i), 0, arr).await.unwrap(); } report(&cache, "after fill", baseline); let idle_after_fill = live() - baseline; @@ -189,7 +189,7 @@ async fn heap_footprint_tracks_budget_for_oversized_working_set() { reset_peak(); for _pass in 0..2 { for i in 0..ENTRIES { - let arr = cache.get(&EntryID::from(i)).await.unwrap(); + let arr = cache.get(&EntryID::from(i), 0).await.unwrap(); assert_eq!(arr.len(), ROWS); drop(arr); } diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index 1f582fd86..1ea8d3cfa 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -10,7 +10,7 @@ use parquet::arrow::arrow_reader::ArrowPredicate; use crate::{ LiquidPredicate, - cache::{BatchID, ColumnAccessPath, ParquetArrayID}, + cache::{BatchID, ColumnAccessPath, ParquetArrayID, file_id::FileId}, }; use std::sync::Arc; @@ -20,6 +20,15 @@ pub struct CachedColumn { cache_store: Arc, field: Arc, column_path: ColumnAccessPath, + /// The file id before it is narrowed into `column_path`. Two files whose + /// ids differ only in the bits `ColumnAccessPath` drops share every + /// `EntryID` this column computes; the cache compares this value to tell + /// them apart and treat the other file's data as a miss. + /// + /// Held as a lease rather than copied: a row group outlives the + /// `CachedFile` it came from, and the id must stay allocated for as long + /// as anything can still compute a key from it. + file_id: Arc, expression: Option>, } @@ -46,6 +55,7 @@ impl CachedColumn { field: Arc, cache_store: Arc, column_access_path: ColumnAccessPath, + file_id: Arc, expression: Option>, is_predicate_column: bool, ) -> Self { @@ -69,6 +79,7 @@ impl CachedColumn { field, cache_store, column_path: column_access_path, + file_id, expression, } } @@ -78,8 +89,14 @@ impl CachedColumn { self.column_path.entry_id(batch_id) } + /// The unnarrowed file id this column belongs to. + pub(crate) fn identity(&self) -> u64 { + self.file_id.get() + } + pub(crate) fn is_cached(&self, batch_id: BatchID) -> bool { - self.cache_store.is_cached(&self.entry_id(batch_id).into()) + self.cache_store + .is_cached(&self.entry_id(batch_id).into(), self.identity()) } /// Returns the Arrow field metadata for this cached column. @@ -92,9 +109,12 @@ impl CachedColumn { self.expression.clone() } - fn array_to_record_batch(&self, array: ArrayRef) -> RecordBatch { + /// `None` when the array does not match this column's field — the cache + /// returned something built for a different column. The caller treats that + /// as "cannot answer from cache" and reads the source instead. + fn array_to_record_batch(&self, array: ArrayRef) -> Option { let schema = Arc::new(Schema::new(vec![self.field.clone()])); - RecordBatch::try_new(schema, vec![array]).unwrap() + RecordBatch::try_new(schema, vec![array]).ok() } /// Evaluates a predicate on a cached column. @@ -114,7 +134,7 @@ impl CachedColumn { if let Some(liquid_expr) = liquid_expr && let Some(boolean_array) = self .cache_store - .eval_predicate(&entry_id, &liquid_expr) + .eval_predicate(&entry_id, self.identity(), &liquid_expr) .with_selection(filter) .await { @@ -126,7 +146,7 @@ impl CachedColumn { } let array = self.get_arrow_array_with_filter(batch_id, filter).await?; - let record_batch = self.array_to_record_batch(array); + let record_batch = self.array_to_record_batch(array)?; let boolean_array = match predicate.evaluate(record_batch) { Ok(arr) => arr, Err(err) => return Some(Err(err)), @@ -160,7 +180,7 @@ impl CachedColumn { ) -> Option { let entry_id = self.entry_id(batch_id).into(); self.cache_store - .get(&entry_id) + .get(&entry_id, self.identity()) .with_selection(filter) .with_optional_expression_hint(self.expression()) .read() @@ -170,7 +190,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.identity()).await } /// Insert an array into the cache. @@ -184,7 +204,7 @@ impl CachedColumn { } self.cache_store - .insert(self.entry_id(batch_id).into(), array) + .insert(self.entry_id(batch_id).into(), self.identity(), array) .await?; Ok(()) } diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs new file mode 100644 index 000000000..277aef035 --- /dev/null +++ b/src/datafusion/src/cache/file_id.rs @@ -0,0 +1,243 @@ +//! Allocation of the file ids that name cached data. +//! +//! An id is the part of a cache key that says which file an entry came from. +//! It is narrowed into 16 bits by [`crate::cache::ColumnAccessPath`], so the +//! supply of *distinct* keys is finite while the number of files a process +//! opens is not. Handing ids out from a counter that only ever climbs means a +//! long-lived process eventually reuses a key while its previous owner's data +//! is still cached. +//! +//! So an id is a lease rather than a permanent assignment. It is held by +//! everything that can still compute a key from it — the file handle, the row +//! groups and columns derived from it — and returns to the pool when the last +//! of them is dropped. The live id count is then bounded by what is actually +//! being read, not by everything that has ever been read. +//! +//! Cache *entries* deliberately do not hold a lease. An id can be reused while +//! entries keyed from it are still resident, and those entries are simply +//! unreachable: each records the identity of the file it came from, so the new +//! owner's reads miss and its writes are refused (see +//! `liquid_cache::cache::ArtIndex::get_checked`). The cost is cache space held +//! by data nobody will read until it is evicted; the alternative — releasing +//! ids from inside index removal — would take a process-wide lock underneath a +//! crossbeam-epoch pin, and would deadlock against `reset`. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicU64, Ordering}; + +use ahash::AHashMap; +use std::sync::{Arc, Mutex, Weak}; + +/// A leased file id. The id returns to its pool when this is dropped. +#[derive(Debug)] +pub(crate) struct FileId { + id: u64, + path: String, + pool: Arc, +} + +impl FileId { + /// The id itself, as the cache key and the entry identity use it. + pub(crate) fn get(&self) -> u64 { + self.id + } +} + +impl Drop for FileId { + fn drop(&mut self) { + self.pool.release(&self.path, self.id); + } +} + +/// Hands out file ids and takes them back. +#[derive(Debug, Default)] +pub(crate) struct FileIdPool { + inner: Mutex, + /// Ids ever allocated that did not fit the key's 16-bit file field. A + /// non-zero count means keys are aliasing and the cache is refusing to + /// serve entries across the alias, which is correct but costs hit rate. + over_key_width: AtomicU64, +} + +#[derive(Debug, Default)] +struct PoolInner { + /// Live leases by path, so concurrent readers of one file share an id. + /// Entries are weak: the map never keeps a file alive on its own, and a + /// path is removed when its lease is released. + leases: AHashMap>, + /// Released ids, reused oldest-first. FIFO rather than LIFO on purpose: a + /// just-released id is the one whose entries are most likely still + /// resident, and reusing it last gives them the longest window to be + /// evicted before anything keys over them. + free: VecDeque, + next: u64, +} + +impl FileIdPool { + pub(crate) fn new() -> Arc { + Arc::new(Self::default()) + } + + /// The lease for `path`, shared with any reader already holding one. + pub(crate) fn acquire(self: &Arc, path: &str) -> Arc { + let mut inner = self.inner.lock().unwrap(); + if let Some(existing) = inner.leases.get(path).and_then(Weak::upgrade) { + return existing; + } + let id = match inner.free.pop_front() { + Some(id) => id, + None => { + let id = inner.next; + inner.next += 1; + id + } + }; + if id > u16::MAX as u64 { + self.over_key_width.fetch_add(1, Ordering::Relaxed); + } + let lease = Arc::new(FileId { + id, + path: path.to_string(), + pool: Arc::clone(self), + }); + inner + .leases + .insert(path.to_string(), Arc::downgrade(&lease)); + lease + } + + fn release(&self, path: &str, id: u64) { + let Ok(mut inner) = self.inner.lock() else { + // A poisoned pool means some other thread panicked holding it. + // Losing one id is better than panicking again inside a drop. + return; + }; + // Only drop the path if it still points at the lease being released. + // A new lease for the same path may already have replaced it, and + // removing that one would hand the same file two live ids. + if inner + .leases + .get(path) + .is_some_and(|weak| weak.strong_count() == 0) + { + inner.leases.remove(path); + } + inner.free.push_back(id); + } + + /// Ids currently leased. Bounded by what is being read, which is what + /// keeps the 16-bit key field from running out. + pub(crate) fn live_count(&self) -> usize { + self.inner.lock().map(|i| i.leases.len()).unwrap_or(0) + } + + /// Ids handed out that do not fit the key's file field. Expected to stay + /// at zero. + pub(crate) fn over_key_width(&self) -> u64 { + self.over_key_width.load(Ordering::Relaxed) + } + + /// Forget every lease and start ids from zero again. + /// + /// Only valid when nothing holds a lease; callers that still do would keep + /// computing keys from ids this pool is free to hand out again. + pub(crate) fn reset(&self) { + if let Ok(mut inner) = self.inner.lock() { + inner.leases.clear(); + inner.free.clear(); + inner.next = 0; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn concurrent_readers_of_one_path_share_a_lease() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let second = pool.acquire("a.parquet"); + assert_eq!(first.get(), second.get()); + assert_eq!(pool.live_count(), 1); + } + + #[test] + fn an_id_returns_to_the_pool_when_its_last_holder_drops() { + let pool = FileIdPool::new(); + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + assert_eq!((a.get(), b.get()), (0, 1)); + assert_eq!(pool.live_count(), 2); + + drop(a); + assert_eq!(pool.live_count(), 1, "the released path is forgotten"); + + // Reused rather than climbing to 2: the supply tracks what is being + // read, which is the whole point. + let c = pool.acquire("c.parquet"); + assert_eq!(c.get(), 0); + } + + #[test] + fn a_second_holder_keeps_the_id_alive() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let second = pool.acquire("a.parquet"); + drop(first); + assert_eq!(pool.live_count(), 1); + // Bound, not a temporary: an unheld lease is released the moment the + // expression ends, which would put its id back before the next call. + let other = pool.acquire("b.parquet"); + assert_eq!(other.get(), 1, "id 0 is still leased"); + drop(second); + let recycled = pool.acquire("c.parquet"); + assert_eq!(recycled.get(), 0); + } + + #[test] + fn released_ids_are_reused_oldest_first() { + let pool = FileIdPool::new(); + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + drop(a); + drop(b); + // 0 was released first, so it is handed out first — giving b's entries + // the longer eviction window. + assert_eq!(pool.acquire("x.parquet").get(), 0); + assert_eq!(pool.acquire("y.parquet").get(), 1); + } + + #[test] + fn ids_beyond_the_key_width_are_counted() { + let pool = FileIdPool::new(); + { + let mut inner = pool.inner.lock().unwrap(); + inner.next = u16::MAX as u64; + } + let _fits = pool.acquire("fits.parquet"); + assert_eq!(pool.over_key_width(), 0); + let _over = pool.acquire("over.parquet"); + assert_eq!(pool.over_key_width(), 1); + } + + /// A path re-registered while its old lease is being dropped must not lose + /// the new lease's entry in the map — that would give one file two live + /// ids and split its cache. + #[test] + fn releasing_a_stale_lease_leaves_a_newer_one_alone() { + let pool = FileIdPool::new(); + let first = pool.acquire("a.parquet"); + let first_id = first.get(); + drop(first); + let second = pool.acquire("a.parquet"); + assert_eq!(second.get(), first_id, "the id came back round"); + assert_eq!(pool.live_count(), 1); + assert_eq!( + pool.acquire("a.parquet").get(), + second.get(), + "the live lease is still the one the map points at" + ); + } +} diff --git a/src/datafusion/src/cache/id.rs b/src/datafusion/src/cache/id.rs index b52379dea..97464a30e 100644 --- a/src/datafusion/src/cache/id.rs +++ b/src/datafusion/src/cache/id.rs @@ -51,10 +51,15 @@ const _: () = assert!(std::mem::align_of::() == 8); impl ParquetArrayID { /// Creates a new CacheEntryID. + /// + /// Each field is narrowed to the width the packed key gives it. That is + /// lossy above `u16::MAX`, and deliberately not an assertion: a process + /// that outlives 65,536 distinct files goes on working, because the cache + /// compares each entry's recorded identity and treats an aliased key as a + /// miss. Asserting here would panic in debug builds on a case release + /// builds handle, which would leave the shipped behaviour untestable. + /// `LiquidCacheParquet` counts ids it hands out that will not fit. 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 { file_id: file_id as u16, rg_id: row_group_id as u16, @@ -146,10 +151,10 @@ pub struct ColumnAccessPath { impl ColumnAccessPath { /// Create a new instance of ColumnAccessPath. + /// + /// Narrowing above `u16::MAX` is lossy and handled rather than asserted — + /// 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 { file_id: file_id as u16, rg_id: row_group_id as u16, @@ -225,22 +230,27 @@ mod tests { assert_eq!(entry_id.batch_id_inner(), *batch_id as u64); } + /// Each field wraps rather than panicking above `u16::MAX`, in every build. + /// The consequence — two sources computing one key — is caught by the + /// identity the cache records alongside each entry, not here. This test + /// pins the wrap down so the aliasing it produces stays a known, + /// reproducible condition rather than a debug-only assertion that the + /// shipped binary never evaluates. #[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)); - } - - #[test] - #[should_panic] - fn test_cache_entry_id_new_panic_row_group_id() { - ParquetArrayID::new(0, (u16::MAX as u64) + 1, 0, BatchID::from_raw(0)); - } - - #[test] - #[should_panic] - fn test_cache_entry_id_new_panic_column_id() { - ParquetArrayID::new(0, 0, (u16::MAX as u64) + 1, BatchID::from_raw(0)); + fn fields_wrap_rather_than_panicking_above_their_width() { + let wrapped = ParquetArrayID::new( + (u16::MAX as u64) + 1, + (u16::MAX as u64) + 2, + (u16::MAX as u64) + 3, + BatchID::from_raw(0), + ); + assert_eq!(wrapped.file_id_inner(), 0); + assert_eq!(wrapped.row_group_id_inner(), 1); + assert_eq!(wrapped.column_id_inner(), 2); + + // Which is exactly the aliasing the identity check exists to absorb. + let first = ParquetArrayID::new(0, 1, 2, BatchID::from_raw(0)); + assert_eq!(usize::from(wrapped), usize::from(first)); } #[test] diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 814d0d324..49b5ae193 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -3,7 +3,6 @@ use crate::io::ParquetCacheMetadata; use crate::reader::{LiquidPredicate, extract_multi_column_or}; -use crate::sync::Mutex; use ahash::AHashMap; use arrow::array::{BooleanArray, RecordBatch, RecordBatchOptions}; use arrow::buffer::BooleanBuffer; @@ -16,12 +15,14 @@ use parquet::arrow::arrow_reader::ArrowPredicate; use std::collections::HashMap; use std::path::Path; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; mod column; +mod file_id; mod id; mod stats; +use file_id::{FileId, FileIdPool}; + pub(crate) use column::InsertArrowArrayError; pub use column::{CachedColumn, CachedColumnRef}; pub(crate) use id::ColumnAccessPath; @@ -58,16 +59,18 @@ impl CachedRowGroup { fn new( cache_store: Arc, row_group_idx: u64, - file_idx: u64, + file_id: Arc, columns: &[CachedColumnSpec], ) -> Self { let mut column_maps = ColumnMaps::default(); for (column_id, field, expression, is_predicate_column) in columns { - let column_access_path = ColumnAccessPath::new(file_idx, row_group_idx, *column_id); + let column_access_path = + ColumnAccessPath::new(file_id.get(), row_group_idx, *column_id); let column = Arc::new(CachedColumn::new( Arc::clone(field), Arc::clone(&cache_store), column_access_path, + Arc::clone(&file_id), expression.clone(), *is_predicate_column, )); @@ -138,7 +141,10 @@ impl CachedRowGroup { } }; let entry_id = column.entry_id(batch_id).into(); - let liquid_array = self.cache_store.try_read_liquid(&entry_id).await; + let liquid_array = self + .cache_store + .try_read_liquid(&entry_id, column.identity()) + .await; let liquid_array = match liquid_array { None => { combined_buffer = None; @@ -146,7 +152,7 @@ impl CachedRowGroup { } Some(array) => array, }; - let buffer = liquid_array.try_eval_predicate(&liquid_expr, selection); + let buffer = liquid_array.try_eval_predicate(&liquid_expr, selection)?; combined_buffer = Some(match combined_buffer { None => buffer, @@ -190,7 +196,9 @@ pub(crate) type CachedRowGroupRef = Arc; #[derive(Debug)] pub struct CachedFile { cache_store: Arc, - file_id: u64, + /// Held, not copied: the id stays allocated for as long as anything can + /// still compute a cache key from it. + file_id: Arc, file_schema: SchemaRef, squeeze_hints: Arc, } @@ -198,7 +206,7 @@ pub struct CachedFile { impl CachedFile { fn new( cache_store: Arc, - file_id: u64, + file_id: Arc, file_schema: SchemaRef, squeeze_hints: Arc, ) -> Self { @@ -236,11 +244,17 @@ impl CachedFile { Arc::new(CachedRowGroup::new( self.cache_store.clone(), row_group_id, - self.file_id, + Arc::clone(&self.file_id), &columns, )) } + /// The leased id this file's cache keys are built from. + #[cfg(test)] + pub(crate) fn file_id(&self) -> u64 { + self.file_id.get() + } + /// Return the configured cache batch size. pub fn batch_size(&self) -> usize { self.cache_store.config().batch_size() @@ -258,12 +272,12 @@ pub(crate) type CachedFileRef = Arc; /// The main cache structure. #[derive(Debug)] pub struct LiquidCacheParquet { - /// Map file path to file id. - files: Mutex>, + /// Leases the file ids that name cached data. Ids come back when nothing + /// is reading the file any more, so the number in use tracks what is being + /// read rather than everything ever read — see [`file_id`]. + file_ids: Arc, cache_store: Arc, - - current_file_id: AtomicU64, } /// A reference to the main cache structure. @@ -322,9 +336,8 @@ impl LiquidCacheParquet { .await; LiquidCacheParquet { - files: Mutex::new(AHashMap::new()), + file_ids: FileIdPool::new(), cache_store: cache_storage, - current_file_id: AtomicU64::new(0), } } @@ -345,15 +358,9 @@ impl LiquidCacheParquet { full_file_schema: SchemaRef, squeeze_hints: Arc, ) -> CachedFileRef { - let mut files = self.files.lock().unwrap(); - let file_id = *files - .entry(file_path.clone()) - .or_insert_with(|| self.current_file_id.fetch_add(1, Ordering::Relaxed)); - drop(files); - Arc::new(CachedFile::new( self.cache_store.clone(), - file_id, + self.file_ids.acquire(&file_path), full_file_schema, squeeze_hints, )) @@ -384,6 +391,34 @@ impl LiquidCacheParquet { self.cache_store.budget().disk_usage_bytes() } + /// How many file ids are currently leased. + /// + /// This tracks the files being read, not the files ever read. It is the + /// number that has to stay under the cache key's 16-bit file field, so it + /// is worth watching: rising without bound means leases are being held by + /// something that should have let go. + pub fn leased_file_ids(&self) -> usize { + self.file_ids.live_count() + } + + /// How many ids have been handed out that do not fit the cache key's file + /// field. + /// + /// Expected to stay at zero. Above zero, distinct files are computing the + /// same keys — served correctly, because each entry records which file it + /// came from, but unable to share the cache. + pub fn file_ids_over_key_width(&self) -> u64 { + self.file_ids.over_key_width() + } + + /// How many cache lookups or writes found a key held by another file. + /// + /// The consequence of the counter above, and the one that proves the + /// aliasing is being caught rather than served. + pub fn identity_mismatches(&self) -> u64 { + self.cache_store.stats().identity_mismatches + } + /// Flush the cache trace to a file. pub fn flush_trace(&self, to_file: impl AsRef) { self.cache_store.observer().flush_cache_trace(to_file); @@ -405,8 +440,7 @@ impl LiquidCacheParquet { /// This is unsafe because resetting the cache while other threads are using the cache may cause undefined behavior. /// You should only call this when no one else is using the cache. pub unsafe fn reset(&self) { - let mut files = self.files.lock().unwrap(); - files.clear(); + self.file_ids.reset(); self.cache_store.reset(); } @@ -437,7 +471,7 @@ mod tests { use super::*; use crate::cache::{CachedRowGroupRef, LiquidCacheParquet}; use crate::reader::FilterCandidateBuilder; - use arrow::array::{Array, Int32Array}; + use arrow::array::{Array, ArrayRef, Int32Array}; use arrow::buffer::BooleanBuffer; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; @@ -470,6 +504,115 @@ mod tests { file.create_row_group(0, vec![]) } + /// What part of the fix is actually for: a process that reads far more + /// files than it holds open at once must not exhaust the key's 16-bit file + /// field. Before ids were leased this counter only ever climbed, so a + /// long-lived instance wrapped it purely by having *seen* enough files. + #[tokio::test] + async fn reading_files_one_after_another_does_not_consume_the_id_space() { + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let tmp_dir = tempfile::tempdir().unwrap(); + let store = crate::test_utils::mount_test_store(tmp_dir.path()).await; + let cache = LiquidCacheParquet::new( + 8, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + + // Well past the ceiling in total, but only ever one open at a time. + for i in 0..(u16::MAX as usize + 1_000) { + let file = cache.register_or_get_file(format!("scan-{i}.parquet"), Arc::clone(&schema)); + assert_eq!( + file.file_id(), + 0, + "each file should reuse the id the previous one gave back" + ); + } + + // And a file opened now still fits the key field. + let after = cache.register_or_get_file("after.parquet".to_string(), schema); + assert!(after.file_id() <= u16::MAX as u64); + } + + /// The bug in its real shape, walked through the actual registration path. + /// + /// `ColumnAccessPath` narrows the file id to 16 bits, so the 65,537th + /// distinct file a process registers is keyed identically to the first. + /// Before entries recorded their identity, the newcomer read the + /// incumbent's data — a panic when the column types differed, silently + /// wrong rows when they matched. + #[tokio::test] + async fn a_file_past_the_key_ceiling_does_not_read_the_first_file_s_data() { + let batch_size = 8; + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let tmp_dir = tempfile::tempdir().unwrap(); + let store = crate::test_utils::mount_test_store(tmp_dir.path()).await; + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + + let batch_id = BatchID::from_row_id(0, batch_size); + let filter = BooleanBuffer::new_set(batch_size); + + // File id 0, with data in the cache. + let first = cache.register_or_get_file("first.parquet".to_string(), Arc::clone(&schema)); + let first_column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + let first_data: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); + first_column + .insert(batch_id, Arc::clone(&first_data)) + .await + .unwrap(); + + // Burn the rest of the 16-bit id space. The handles are held: ids are + // leased, so files that are opened and closed hand their id straight + // back and the ceiling is only reachable with this many files open at + // once. + let _fillers: Vec<_> = (1..=u16::MAX as usize) + .map(|i| cache.register_or_get_file(format!("filler-{i}.parquet"), Arc::clone(&schema))) + .collect(); + + // Id 65536, which narrows to 0. + let wrapped = cache.register_or_get_file("wrapped.parquet".to_string(), schema); + let wrapped_column = wrapped.create_row_group(0, vec![]).get_column(0).unwrap(); + + assert_eq!( + usize::from(wrapped_column.entry_id(batch_id)), + usize::from(first_column.entry_id(batch_id)), + "the packed keys must actually collide, or this test proves nothing" + ); + + // The newcomer must not be handed the incumbent's rows. + assert!(!wrapped_column.is_cached(batch_id)); + assert!( + wrapped_column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .is_none(), + "a colliding key must read as a miss, not as the other file's data" + ); + + // And the incumbent still reads its own. + let got = first_column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .expect("the owner's entry is still there"); + assert_eq!(got.as_ref(), first_data.as_ref()); + } + /// Issue #19: `NOT (s = s)` simplifies to `s IS NULL AND NULL`, so a conjunct /// that reads no column reaches the row filter. It has to survive candidate /// building and then evaluate against the selection's row count — an diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index 474abcb2d..8f17e6e51 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -205,9 +205,15 @@ mod tests { let mut row_start_id_sum = 0; let mut row_count_sum = 0; let mut memory_size_sum = 0; - for file_no in 0..8 { - let file_name = format!("test_{file_no}.parquet"); - let file = cache.register_or_get_file(file_name, schema.clone()); + // Held for the whole loop, not per iteration: a file id is leased and + // comes back when its handle drops, so releasing each file before + // opening the next would hand them all the same id. + let files: Vec<_> = (0..8) + .map(|file_no| { + cache.register_or_get_file(format!("test_{file_no}.parquet"), schema.clone()) + }) + .collect(); + for file in &files { for rg in 0..8 { let row_group = file.create_row_group(rg, vec![]); for col in 0..8 { diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index fadc8e70d..778717b30 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -385,9 +385,10 @@ impl LiquidCacheReaderInner { arrays.push(array); } - Ok(Some( - RecordBatch::try_new(self.schema.clone(), arrays).unwrap(), - )) + // A batch that does not match the declared schema is a cache that + // handed back something other than what was asked for. Report it; + // unwinding here aborts the stream mid-flight with no error to show. + Ok(Some(RecordBatch::try_new(self.schema.clone(), arrays)?)) } async fn read_parquet_batch_and_fill_cache( From ddeae5854ee5dc02953b1d448704400e708b8acd Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 10:44:34 +0530 Subject: [PATCH 02/12] fix(cache): keep a file's identity when its id is recycled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Leasing file ids reintroduced the aliasing the identity check exists to prevent: the identity was the id itself, so a file that inherited a released id was indistinguishable from the one that gave it back and read the entries it left behind. TPC-H caught it as a wrong column. Split the two. The id stays narrow and recycled, because the key's file field is 16 bits and has to be reusable. The identity is a separate u64 that is never handed to a different file, so an inherited id carries no claim on what the previous holder cached. A released id remembers the path and identity it had, so re-opening the same file keeps both and its cached entries stay readable — a file read twice is a cache hit, not a collision. Any other path gets a fresh identity. Also fix the callers outside the three crates I had been building: examples/core.rs and the core README doctests. --- examples/core.rs | 11 +++- src/core/README.md | 18 ++++-- src/datafusion/src/cache/column.rs | 4 +- src/datafusion/src/cache/file_id.rs | 97 ++++++++++++++++++++++++++--- src/datafusion/src/cache/mod.rs | 59 ++++++++++++++++++ 5 files changed, 171 insertions(+), 18 deletions(-) diff --git a/examples/core.rs b/examples/core.rs index d172f601a..79644b639 100644 --- a/examples/core.rs +++ b/examples/core.rs @@ -18,13 +18,20 @@ async fn main() -> Result<(), Box> { .await; let entry_id = EntryID::from(7); + // Names whose data this is. Entry ids are packed and can alias between + // sources; the cache compares this and treats a mismatch as a miss, so a + // caller only ever reads back what it put in. + let identity = 1; 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 d263aab51..e776a9ef7 100644 --- a/src/core/README.md +++ b/src/core/README.md @@ -22,12 +22,16 @@ tokio_test::block_on(async { let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(42); +// Names whose data this is. Entry ids are packed and can alias between +// sources, so the cache compares this on every read and treats a mismatch as +// a miss — a caller only ever reads back what it put in. +let identity = 1; 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; +storage.insert(entry_id, identity, arrow_array.clone()).await; -assert!(storage.is_cached(&entry_id)); +assert!(storage.is_cached(&entry_id, identity)); }); ``` @@ -42,14 +46,15 @@ tokio_test::block_on(async { let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(7); +let identity = 1; let arrow_array = Arc::new(UInt64Array::from_iter_values(0..16)); -storage.insert(entry_id, arrow_array.clone()).await; +storage.insert(entry_id, identity, 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, identity).await.unwrap(); assert_eq!(retrieved.as_ref(), arrow_array.as_ref()); }); ``` @@ -71,10 +76,11 @@ tokio_test::block_on(async { let storage = LiquidCacheBuilder::new().build().await; let entry_id = EntryID::from(8); +let identity = 1; 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, identity, data.clone()).await; // Move data to disk so the read will demonstrate async I/O storage.flush_all_to_disk().await; @@ -94,7 +100,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, identity, &liquid_expr) .with_selection(&selection) .await .unwrap(); diff --git a/src/datafusion/src/cache/column.rs b/src/datafusion/src/cache/column.rs index 1ea8d3cfa..f0c2592fd 100644 --- a/src/datafusion/src/cache/column.rs +++ b/src/datafusion/src/cache/column.rs @@ -89,9 +89,9 @@ impl CachedColumn { self.column_path.entry_id(batch_id) } - /// The unnarrowed file id this column belongs to. + /// The never-reused name of the file this column belongs to. pub(crate) fn identity(&self) -> u64 { - self.file_id.get() + self.file_id.identity() } pub(crate) fn is_cached(&self, batch_id: BatchID) -> bool { diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs index 277aef035..f75d4b3d3 100644 --- a/src/datafusion/src/cache/file_id.rs +++ b/src/datafusion/src/cache/file_id.rs @@ -29,23 +29,41 @@ use ahash::AHashMap; use std::sync::{Arc, Mutex, Weak}; /// A leased file id. The id returns to its pool when this is dropped. +/// +/// Two numbers, because they answer different questions and only one of them +/// can be recycled: +/// +/// * `id` goes into the cache key, whose file field is 16 bits wide. It has to +/// be recycled or a long-lived process runs out. +/// * `identity` names *which file* an entry came from, and is never reused. +/// It cannot be the recycled id: a file that inherits id 0 from a file that +/// has finished would otherwise be indistinguishable from it, and would read +/// the entries it left behind — the exact aliasing the identity exists to +/// catch. #[derive(Debug)] pub(crate) struct FileId { id: u64, + identity: u64, path: String, pool: Arc, } impl FileId { - /// The id itself, as the cache key and the entry identity use it. + /// The narrow, recycled id the cache key is built from. pub(crate) fn get(&self) -> u64 { self.id } + + /// The wide, never-reused name for this file, recorded alongside every + /// entry so a recycled key cannot serve one file another's data. + pub(crate) fn identity(&self) -> u64 { + self.identity + } } impl Drop for FileId { fn drop(&mut self) { - self.pool.release(&self.path, self.id); + self.pool.release(&self.path, self.id, self.identity); } } @@ -69,8 +87,24 @@ struct PoolInner { /// just-released id is the one whose entries are most likely still /// resident, and reusing it last gives them the longest window to be /// evicted before anything keys over them. - free: VecDeque, + /// + /// Each carries the path that released it and the identity it had. If the + /// same path comes back it keeps that identity, so its cached entries are + /// still its own and still readable — a file read twice is a cache hit, + /// not a collision. Any other path gets a fresh identity, so it cannot + /// read what the previous holder left behind. + free: VecDeque, next: u64, + /// Only ever climbs. A `u64` of these is not a resource worth reclaiming: + /// at one a microsecond it outlasts the hardware. + next_identity: u64, +} + +#[derive(Debug)] +struct Released { + id: u64, + path: String, + identity: u64, } impl FileIdPool { @@ -84,19 +118,29 @@ impl FileIdPool { if let Some(existing) = inner.leases.get(path).and_then(Weak::upgrade) { return existing; } - let id = match inner.free.pop_front() { - Some(id) => id, + let (id, reusable_identity) = match inner.free.pop_front() { + Some(released) if released.path == path => (released.id, Some(released.identity)), + Some(released) => (released.id, None), None => { let id = inner.next; inner.next += 1; - id + (id, None) } }; if id > u16::MAX as u64 { self.over_key_width.fetch_add(1, Ordering::Relaxed); } + let identity = match reusable_identity { + Some(identity) => identity, + None => { + let identity = inner.next_identity; + inner.next_identity += 1; + identity + } + }; let lease = Arc::new(FileId { id, + identity, path: path.to_string(), pool: Arc::clone(self), }); @@ -106,7 +150,7 @@ impl FileIdPool { lease } - fn release(&self, path: &str, id: u64) { + fn release(&self, path: &str, id: u64, self_identity: u64) { let Ok(mut inner) = self.inner.lock() else { // A poisoned pool means some other thread panicked holding it. // Losing one id is better than panicking again inside a drop. @@ -122,7 +166,11 @@ impl FileIdPool { { inner.leases.remove(path); } - inner.free.push_back(id); + inner.free.push_back(Released { + id, + path: path.to_string(), + identity: self_identity, + }); } /// Ids currently leased. Bounded by what is being read, which is what @@ -196,6 +244,39 @@ mod tests { assert_eq!(recycled.get(), 0); } + /// The two numbers have to move independently. Reusing an id is how the + /// key space stays bounded; reusing an *identity* for a different file is + /// how one file reads another's entries. Re-opening the same path must + /// keep its identity, or every lease boundary silently empties the cache. + #[test] + fn identity_follows_the_path_while_the_id_is_recycled() { + let pool = FileIdPool::new(); + + let first = pool.acquire("a.parquet"); + let (a_id, a_identity) = (first.get(), first.identity()); + drop(first); + + // Same file again: same id and the same name, so its cached entries + // are still its own. + let reopened = pool.acquire("a.parquet"); + assert_eq!(reopened.get(), a_id); + assert_eq!( + reopened.identity(), + a_identity, + "re-opening a file must keep its identity, or its cache is dead" + ); + drop(reopened); + + // A different file inherits the id but must not inherit the name. + let other = pool.acquire("b.parquet"); + assert_eq!(other.get(), a_id, "the id is recycled"); + assert_ne!( + other.identity(), + a_identity, + "a different file must not be able to read what the last one left" + ); + } + #[test] fn released_ids_are_reused_oldest_first() { let pool = FileIdPool::new(); diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 49b5ae193..f76a9a7ed 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -504,6 +504,65 @@ mod tests { file.create_row_group(0, vec![]) } + /// Recycling a file id must not recycle a file's *name*. + /// + /// The id is narrow and reused so the key space cannot run out. If the + /// identity recorded against each entry were that same id, the next file + /// to inherit it would be indistinguishable from the one that gave it + /// back, and would read the entries it left behind — reintroducing the + /// aliasing the identity exists to catch, at every lease boundary rather + /// than only past 65,536 files. + #[tokio::test] + async fn a_file_inheriting_a_recycled_id_does_not_read_its_predecessors_data() { + let batch_size = 8; + let schema: SchemaRef = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let tmp_dir = tempfile::tempdir().unwrap(); + let store = crate::test_utils::mount_test_store(tmp_dir.path()).await; + let cache = LiquidCacheParquet::new( + batch_size, + usize::MAX, + usize::MAX, + store, + Box::new(LiquidPolicy::new()), + Box::new(TranscodeSqueezeEvict), + Box::new(AlwaysHydrate::new()), + ) + .await; + + let batch_id = BatchID::from_row_id(0, batch_size); + let filter = BooleanBuffer::new_set(batch_size); + let first_data: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8])); + + let first_id = { + let first = + cache.register_or_get_file("first.parquet".to_string(), Arc::clone(&schema)); + let column = first.create_row_group(0, vec![]).get_column(0).unwrap(); + column + .insert(batch_id, Arc::clone(&first_data)) + .await + .unwrap(); + first.file_id() + }; // lease dropped here, so the id goes back to the pool + + let second = cache.register_or_get_file("second.parquet".to_string(), schema); + assert_eq!( + second.file_id(), + first_id, + "the id must actually be recycled, or this test proves nothing" + ); + + let column = second.create_row_group(0, vec![]).get_column(0).unwrap(); + assert!(!column.is_cached(batch_id)); + assert!( + column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .is_none(), + "inheriting an id must not inherit the entries keyed from it" + ); + } + /// What part of the fix is actually for: a process that reads far more /// files than it holds open at once must not exhaust the key's 16-bit file /// field. Before ids were leased this counter only ever climbed, so a From da69d9a38d769645ec1f321047b8ce7e8b7a80fd Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 10:48:33 +0530 Subject: [PATCH 03/12] fix(cache): let a recycled key be taken over by its new owner Refusing an insert whose key is held by another identity left the key occupied by data nobody can read: the holder is a file that has since let its id go, and the file that now owns the id cannot read it either. On a cache below its budget nothing evicts it, so that file would never cache the key again. Overwrite instead, and keep counting. The incumbent loses an entry it could not have served anyway. Also stop asserting that restoring a memory reservation succeeds. When the entry being replaced was larger, restoring it grows the reservation again and can fail on a full cache, which would panic. Nothing is stored on that path either way, so leave the budget under-counted rather than bring the process down. --- src/core/src/cache/core.rs | 10 ++++++--- src/core/src/cache/index.rs | 36 +++++++++++++++++++++------------ src/datafusion/src/cache/mod.rs | 15 ++++++++++++++ 3 files changed, 45 insertions(+), 16 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 851495cac..cb07b263b 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -599,9 +599,13 @@ impl LiquidCache { } let batch_type = CachedBatchType::from(&to_insert); if !self.index.insert(&entry_id, identity, to_insert) { - self.budget - .try_update_memory_usage(new_memory_size, old_memory_size) - .expect("memory release cannot fail"); + // Restoring the reservation *grows* it again when the entry we + // were replacing was larger, so this can legitimately fail on + // a full cache. Nothing is stored either way; the budget is + // left under-counted rather than the process brought down. + let _ = self + .budget + .try_update_memory_usage(new_memory_size, old_memory_size); return Ok(()); } batch_type diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 1cb3a8b53..64ce193ce 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -124,9 +124,13 @@ impl ArtIndex { /// Store `batch` under `entry_id`, returning whether it was stored. /// /// `identity` is `Some` for a caller-originated insert, naming whose data - /// this is. A key already held by a *different* identity is refused rather - /// than overwritten, so two sources colliding on one key do not evict each - /// other's data back and forth; the loser reads from its own source. + /// this is. A key held by a *different* identity is overwritten and the + /// collision counted. Refusing instead would be worse: whoever owns the + /// key now cannot read the incumbent's entry either, so refusing leaves + /// the key occupied by data nobody can use, and on a cache below its + /// budget nothing evicts it — the new owner would never cache that key + /// again. Overwriting costs the incumbent its entry, which it could not + /// have read anyway. /// /// `identity` is `None` for maintenance that rewrites a key in place — /// transcode, squeeze, hydrate, spill to disk — which keeps the identity @@ -142,11 +146,13 @@ impl ArtIndex { let guard = self.art.pin(); let existing_identity = self.art.get(*entry_id, &guard).map(|slot| slot.identity); let identity = match (identity, existing_identity) { - (Some(new), Some(old)) if new != old => { - self.identity_mismatches.fetch_add(1, Ordering::Relaxed); - return false; + (Some(new), Some(old)) => { + if new != old { + self.identity_mismatches.fetch_add(1, Ordering::Relaxed); + } + new } - (Some(new), _) => new, + (Some(new), None) => new, (None, Some(old)) => old, (None, None) => return false, }; @@ -308,12 +314,16 @@ mod tests { // The owner still reads its own entry. assert!(store.get_checked(&key, 1).is_some()); - // And the colliding file cannot displace it by writing over the key, - // so the two do not take turns evicting each other. - assert!(!store.insert(&key, Some(2), create_test_array(200))); - match store.get_checked(&key, 1).unwrap().as_ref() { - CacheEntry::MemoryArrow(array) => assert_eq!(array.len(), 100), - other => panic!("expected the owner's array, found {other}"), + // The colliding file takes the key over. It has to: it cannot read + // what is there, so leaving it would cost the key to both of them. + assert!(store.insert(&key, Some(2), create_test_array(200))); + assert!( + store.get_checked(&key, 1).is_none(), + "the displaced file reads a miss, never the other file's rows" + ); + match store.get_checked(&key, 2).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!(array.len(), 200), + other => panic!("expected the new owner's array, found {other}"), } } diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index f76a9a7ed..faecb72c5 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -561,6 +561,21 @@ mod tests { .is_none(), "inheriting an id must not inherit the entries keyed from it" ); + + // It must also be able to cache. The predecessor's entries are keyed + // where this file's belong and nobody can read them any more, so they + // give way — otherwise a cache under its budget, where nothing is ever + // evicted, would leave this file permanently uncacheable. + let second_data: ArrayRef = Arc::new(Int32Array::from(vec![9, 9, 9, 9, 9, 9, 9, 9])); + column + .insert(batch_id, Arc::clone(&second_data)) + .await + .expect("the inheriting file must be able to cache"); + let got = column + .get_arrow_array_with_filter(batch_id, &filter) + .await + .expect("the new owner reads back its own rows"); + assert_eq!(got.as_ref(), second_data.as_ref()); } /// What part of the fix is actually for: a process that reads far more From f4dd8a53dc6ff3983899d90a8a48515bb0e77edc Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 11:18:42 +0530 Subject: [PATCH 04/12] fix(cache): drop a rewrite whose key was taken over since it was read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Maintenance builds its payload from an entry it read earlier and stores it after an await — a squeeze reads, writes bytes to disk, then inserts. Adopting whatever identity held the key by then relabels the old file's data as the new owner's, and the new owner reads those rows as a hit. Carry the identity the entry was read under. A rewrite lands only while the key still holds it, and is dropped otherwise: the payload belongs to a source that no longer owns the key. Replaces the Option on the insert path with an explicit WriteIdentity, so a caller storing its own data and maintenance rewriting someone else's are no longer the same call shape. The index read that feeds maintenance now reports the identity alongside the entry. --- src/core/src/cache/builders.rs | 3 +- src/core/src/cache/core.rs | 119 ++++++++++++++++++++--------- src/core/src/cache/index.rs | 119 +++++++++++++++++++++-------- src/core/tests/memory_footprint.rs | 4 +- src/datafusion/src/cache/stats.rs | 63 +++++++-------- 5 files changed, 205 insertions(+), 103 deletions(-) diff --git a/src/core/src/cache/builders.rs b/src/core/src/cache/builders.rs index 824c0fc0e..e1db0e236 100644 --- a/src/core/src/cache/builders.rs +++ b/src/core/src/cache/builders.rs @@ -11,6 +11,7 @@ use super::core::LiquidCache; use super::io_context::{DefaultCacheMetadata, EntryMetadata}; use super::policies::{CachePolicy, HydrationPolicy, SqueezePolicy, TranscodeSqueezeEvict}; use super::{CacheExpression, CacheFull, EntryID, LiquidExpr, LiquidPolicy}; +use crate::cache::index::WriteIdentity; use crate::sync::Arc; /// Builder for [LiquidCache]. @@ -222,7 +223,7 @@ impl<'a> Insert<'a> { let batch = CacheEntry::memory_arrow(batch); self.storage.supersede_disk_copy(self.entry_id).await; self.storage - .insert_inner(self.entry_id, Some(self.identity), batch) + .insert_inner(self.entry_id, WriteIdentity::Owned(self.identity), batch) .await } } diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index cb07b263b..45cbd99e9 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -18,7 +18,11 @@ use super::{ use crate::cache::DefaultSqueezeIo; use crate::cache::policies::{SqueezeOutcome, SqueezePolicy}; use crate::cache::utils::{LiquidCompressorStates, arrow_to_bytes}; -use crate::cache::{CacheExpression, LiquidExpr, index::ArtIndex, utils::EntryID}; +use crate::cache::{ + CacheExpression, LiquidExpr, + index::{ArtIndex, WriteIdentity}, + utils::EntryID, +}; use crate::cache::{CacheFull, CacheStats, EventTrace}; use crate::liquid_array::{ LiquidSqueezedArrayRef, SqueezeIoHandler, SqueezedBacking, SqueezedDate32Array, @@ -130,7 +134,7 @@ impl LiquidCache { let mut memory_liquid_bytes = 0usize; let mut memory_squeezed_liquid_bytes = 0usize; - self.index.for_each(|_, batch| match batch { + self.index.for_each(|_, _, batch| match batch { CacheEntry::MemoryArrow(array) => { memory_arrow_entries += 1; memory_arrow_bytes += array.get_array_memory_size(); @@ -212,8 +216,14 @@ impl LiquidCache { 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; + self.maybe_hydrate( + entry_id, + identity, + entry, + MaterializedEntry::Liquid(&liquid), + None, + ) + .await; Some(liquid) } CacheEntry::MemorySqueezedLiquid(array) => match array.disk_backing() { @@ -230,7 +240,7 @@ impl LiquidCache { /// Iterate over all entries in the cache. /// 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)) { + pub fn for_each_entry(&self, mut f: impl FnMut(&EntryID, u64, &CacheEntry)) { self.index.for_each(&mut f); } @@ -279,10 +289,10 @@ 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.for_each_entry(|entry_id, identity, batch| { + entires.push((*entry_id, identity, batch.clone())); }); - for (entry_id, batch) in entires { + for (entry_id, flush_identity, batch) in entires { match &batch { CacheEntry::MemoryArrow(array) => { let bytes = arrow_to_bytes(array).expect("failed to convert arrow to bytes"); @@ -291,7 +301,7 @@ impl LiquidCache { Ok(()) => { self.try_insert( entry_id, - None, + WriteIdentity::Rewrite(flush_identity), CacheEntry::disk_arrow(array.data_type().clone(), disk_bytes), ) .expect("failed to insert disk arrow entry"); @@ -309,8 +319,12 @@ impl LiquidCache { // Hydrated from disk and never modified since: the // bytes are already there, flip the index rather // than re-serialising and rewriting them. - self.try_insert(entry_id, None, CacheEntry::disk_liquid(data_type, bytes)) - .expect("failed to insert disk liquid entry"); + self.try_insert( + entry_id, + WriteIdentity::Rewrite(flush_identity), + CacheEntry::disk_liquid(data_type, bytes), + ) + .expect("failed to insert disk liquid entry"); continue; } let liquid_bytes = liquid_array.to_bytes(); @@ -322,7 +336,7 @@ impl LiquidCache { Ok(()) => { self.try_insert( entry_id, - None, + WriteIdentity::Rewrite(flush_identity), CacheEntry::disk_liquid(data_type, disk_bytes), ) .expect("failed to insert disk liquid entry"); @@ -333,7 +347,7 @@ impl LiquidCache { CacheEntry::MemorySqueezedLiquid(array) => { // We don't have to do anything, because it's already on disk let disk_entry = Self::disk_entry_from_squeezed(array); - self.try_insert(entry_id, None, disk_entry) + self.try_insert(entry_id, WriteIdentity::Rewrite(flush_identity), disk_entry) .expect("failed to insert disk entry"); } CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { @@ -412,7 +426,7 @@ impl LiquidCache { pub(crate) async fn insert_inner( &self, entry_id: EntryID, - identity: Option, + identity: WriteIdentity, mut batch_to_cache: CacheEntry, ) -> Result<(), CacheFull> { loop { @@ -575,16 +589,14 @@ impl LiquidCache { } } - /// `identity` names whose data this is for a caller-originated insert, and - /// is `None` for maintenance rewriting a key in place — see - /// [`ArtIndex::insert`]. A declined write is not an error: the key belongs - /// to someone else, or the entry being rewritten has since been removed. - /// Neither is worth retrying, so the reservation is handed back and the - /// call reports success with nothing stored. + /// A declined write is not an error: the entry being rewritten has since + /// been taken over or removed. Neither is worth retrying, so the + /// reservation is handed back and the call reports success with nothing + /// stored. See [`WriteIdentity`]. fn try_insert( &self, entry_id: EntryID, - identity: Option, + identity: WriteIdentity, to_insert: CacheEntry, ) -> Result<(), CacheEntry> { let new_memory_size = to_insert.memory_usage_bytes(); @@ -714,7 +726,9 @@ impl LiquidCache { } async fn squeeze_victim_inner(&self, to_squeeze: EntryID) -> Result<(), CacheFull> { - let Some(mut to_squeeze_batch) = self.index.get(&to_squeeze) else { + let Some((squeezed_identity, mut to_squeeze_batch)) = + self.index.get_with_identity(&to_squeeze) + else { return Ok(()); }; self.trace(InternalEvent::SqueezeVictim { entry: to_squeeze }); @@ -750,7 +764,11 @@ impl LiquidCache { self.write_batch_to_disk(to_squeeze, &new_batch, bytes_to_write) .await?; } - match self.try_insert(to_squeeze, None, new_batch) { + match self.try_insert( + to_squeeze, + WriteIdentity::Rewrite(squeezed_identity), + new_batch, + ) { Ok(()) => { break; } @@ -779,6 +797,7 @@ impl LiquidCache { async fn maybe_hydrate( &self, entry_id: &EntryID, + identity: u64, cached: &CacheEntry, materialized: MaterializedEntry<'_>, expression: Option<&CacheExpression>, @@ -798,7 +817,9 @@ impl LiquidCache { cached: cached_type, new: new_type, }); - let _ = self.insert_inner(*entry_id, None, new_entry).await; + let _ = self + .insert_inner(*entry_id, WriteIdentity::Rewrite(identity), new_entry) + .await; } } @@ -834,11 +855,11 @@ impl LiquidCache { None => Some(array.to_arrow_array()), }, CacheEntry::DiskArrow { .. } | CacheEntry::DiskLiquid { .. } => { - self.read_disk_array(batch.as_ref(), entry_id, expression, selection) + self.read_disk_array(batch.as_ref(), entry_id, identity, expression, selection) .await } CacheEntry::MemorySqueezedLiquid(array) => { - self.read_squeezed_array(array, entry_id, expression, selection) + self.read_squeezed_array(array, entry_id, identity, expression, selection) .await } } @@ -848,6 +869,7 @@ impl LiquidCache { &self, entry: &CacheEntry, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -861,6 +883,7 @@ impl LiquidCache { let full_array = self.read_disk_arrow_array(entry_id).await; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Arrow(&full_array), expression, @@ -883,6 +906,7 @@ impl LiquidCache { let liquid = self.read_disk_liquid_array(entry_id).await; self.maybe_hydrate( entry_id, + identity, entry, MaterializedEntry::Liquid(&liquid), expression, @@ -901,6 +925,7 @@ impl LiquidCache { &self, array: &LiquidSqueezedArrayRef, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -914,7 +939,7 @@ impl LiquidCache { } if let Some(array) = self - .try_read_squeezed_variant_array(array, entry_id, expression, selection) + .try_read_squeezed_variant_array(array, entry_id, identity, expression, selection) .await { self.observer.on_get_squeezed_success(); @@ -959,6 +984,7 @@ impl LiquidCache { &self, array: &LiquidSqueezedArrayRef, entry_id: &EntryID, + identity: u64, expression: Option<&CacheExpression>, selection: Option<&BooleanBuffer>, ) -> Option { @@ -976,6 +1002,7 @@ impl LiquidCache { let full_array = self.read_disk_arrow_array(entry_id).await; self.maybe_hydrate( entry_id, + identity, &batch, MaterializedEntry::Arrow(&full_array), expression, @@ -1122,8 +1149,14 @@ impl LiquidCache { } entry @ CacheEntry::DiskArrow { .. } => { let array = self.read_disk_arrow_array(entry_id).await; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Arrow(&array), None) - .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())); @@ -1143,8 +1176,14 @@ impl LiquidCache { } entry @ CacheEntry::DiskLiquid { .. } => { let liquid = self.read_disk_liquid_array(entry_id).await; - self.maybe_hydrate(entry_id, entry, MaterializedEntry::Liquid(&liquid), None) - .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())); @@ -1251,7 +1290,7 @@ mod tests { let array1 = create_test_array(100); let size1 = array1.memory_usage_bytes(); store - .insert_inner(entry_id1, Some(0), array1) + .insert_inner(entry_id1, WriteIdentity::Owned(0), array1) .await .unwrap(); @@ -1267,7 +1306,7 @@ mod tests { let array2 = create_test_array(200); let size2 = array2.memory_usage_bytes(); store - .insert_inner(entry_id2, Some(0), array2) + .insert_inner(entry_id2, WriteIdentity::Owned(0), array2) .await .unwrap(); @@ -1276,7 +1315,7 @@ mod tests { let array3 = create_test_array(150); let size3 = array3.memory_usage_bytes(); store - .insert_inner(entry_id1, Some(0), array3) + .insert_inner(entry_id1, WriteIdentity::Owned(0), array3) .await .unwrap(); @@ -1297,7 +1336,7 @@ mod tests { store .insert_inner( entry_id, - Some(0), + WriteIdentity::Owned(0), CacheEntry::memory_squeezed_liquid(squeezed.clone()), ) .await @@ -1336,7 +1375,7 @@ mod tests { let store = create_cache_store(8000, Box::new(advisor)).await; // Small budget to force advice store - .insert_inner(entry_id1, Some(0), create_test_array(800)) + .insert_inner(entry_id1, WriteIdentity::Owned(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -1345,7 +1384,7 @@ mod tests { } store - .insert_inner(entry_id2, Some(0), create_test_array(800)) + .insert_inner(entry_id2, WriteIdentity::Owned(0), create_test_array(800)) .await .unwrap(); match store.index().get(&entry_id1).unwrap().as_ref() { @@ -1484,7 +1523,11 @@ mod tests { let liquid = transcode_liquid_inner(&arrow_array, &compressor).unwrap(); store - .insert_inner(entry_id, Some(0), CacheEntry::memory_liquid(liquid.clone())) + .insert_inner( + entry_id, + WriteIdentity::Owned(0), + CacheEntry::memory_liquid(liquid.clone()), + ) .await .unwrap(); store.flush_all_to_disk().await.unwrap(); diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 64ce193ce..6c9794265 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -50,6 +50,22 @@ impl Slot { } } +/// Whose data a write carries, and on what terms. +#[derive(Debug, Clone, Copy)] +pub(crate) enum WriteIdentity { + /// A caller storing its own data. Takes the key over if another identity + /// holds it: that identity belongs to a source that cannot read this key + /// any more, so leaving its entry there would cost the key to both. + Owned(u64), + /// Maintenance rewriting an entry it read earlier — transcode, squeeze, + /// hydrate, spill. It carries the identity the entry was read under and + /// lands only if the key still holds it. Adopting whatever is there + /// instead would relabel one source's data with another's whenever a + /// takeover lands between the read and the write, and the new owner would + /// then read those rows as its own. + Rewrite(u64), +} + pub(crate) struct ArtIndex { art: CongeeArc, entry_count: AtomicUsize, @@ -117,44 +133,49 @@ impl ArtIndex { slot.load() } + /// Look up an entry together with the identity recorded against it, for + /// maintenance that must rewrite it under the identity it observed. + pub(crate) fn get_with_identity(&self, entry_id: &EntryID) -> Option<(u64, Arc)> { + let guard = self.art.pin(); + let slot = self.art.get(*entry_id, &guard)?; + let identity = slot.identity; + if let Some(entry) = slot.load() { + return Some((identity, entry)); + } + let slot = self.art.get(*entry_id, &guard)?; + let identity = slot.identity; + slot.load().map(|entry| (identity, entry)) + } + pub(crate) fn is_cached(&self, entry_id: &EntryID, identity: u64) -> bool { self.get_checked(entry_id, identity).is_some() } /// Store `batch` under `entry_id`, returning whether it was stored. /// - /// `identity` is `Some` for a caller-originated insert, naming whose data - /// this is. A key held by a *different* identity is overwritten and the - /// collision counted. Refusing instead would be worse: whoever owns the - /// key now cannot read the incumbent's entry either, so refusing leaves - /// the key occupied by data nobody can use, and on a cache below its - /// budget nothing evicts it — the new owner would never cache that key - /// again. Overwriting costs the incumbent its entry, which it could not - /// have read anyway. - /// - /// `identity` is `None` for maintenance that rewrites a key in place — - /// transcode, squeeze, hydrate, spill to disk — which keeps the identity - /// already recorded. Maintenance cannot change whose an entry is, and - /// cannot resurrect a key that has since been removed: with nothing to - /// preserve there is no identity to record, so the write is dropped. + /// See [`WriteIdentity`] for the two kinds of write and why they differ. pub(crate) fn insert( &self, entry_id: &EntryID, - identity: Option, + identity: WriteIdentity, batch: CacheEntry, ) -> bool { let guard = self.art.pin(); let existing_identity = self.art.get(*entry_id, &guard).map(|slot| slot.identity); let identity = match (identity, existing_identity) { - (Some(new), Some(old)) => { + (WriteIdentity::Owned(new), Some(old)) => { if new != old { self.identity_mismatches.fetch_add(1, Ordering::Relaxed); } new } - (Some(new), None) => new, - (None, Some(old)) => old, - (None, None) => return false, + (WriteIdentity::Owned(new), None) => new, + // The key still holds what this rewrite was built from. + (WriteIdentity::Rewrite(expected), Some(old)) if expected == old => expected, + // It does not: the entry was taken over or removed while this + // rewrite was in flight, so the payload belongs to a source that + // no longer owns the key. Drop it. + (WriteIdentity::Rewrite(_), _) => return false, }; let existing = self .art @@ -183,10 +204,10 @@ 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, u64, &CacheEntry)) { for id in self.art.keys() { - if let Some(entry) = self.get(&id) { - f(&id, &entry); + if let Some((identity, entry)) = self.get_with_identity(&id) { + f(&id, identity, &entry); } } } @@ -231,7 +252,7 @@ mod tests { // Insert an entry and verify it's cached { - store.insert(&entry_id1, Some(0), array1.clone()); + store.insert(&entry_id1, WriteIdentity::Owned(0), array1.clone()); } assert!(store.is_cached(&entry_id1, 0)); @@ -253,7 +274,7 @@ mod tests { let entry_id: EntryID = EntryID::from(1); let array = create_test_array(100); - store.insert(&entry_id, Some(0), array.clone()); + store.insert(&entry_id, WriteIdentity::Owned(0), array.clone()); let entry_id: EntryID = EntryID::from(1); assert!(store.is_cached(&entry_id, 0)); @@ -275,8 +296,8 @@ mod tests { unreachable!() }; let weak_first = Arc::downgrade(first_array); - store.insert(&id, Some(0), first); - store.insert(&id, Some(0), create_test_array(200)); + store.insert(&id, WriteIdentity::Owned(0), first); + store.insert(&id, WriteIdentity::Owned(0), create_test_array(200)); assert!( weak_first.upgrade().is_none(), "replaced entry still alive: held by the index's deferred drop" @@ -304,7 +325,7 @@ mod tests { let store = ArtIndex::new(); let key: EntryID = EntryID::from(7); - assert!(store.insert(&key, Some(1), create_test_array(100))); + assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); // The colliding file asks for the same key and is told nothing is there. assert!(store.get_checked(&key, 2).is_none()); @@ -316,7 +337,7 @@ mod tests { // The colliding file takes the key over. It has to: it cannot read // what is there, so leaving it would cost the key to both of them. - assert!(store.insert(&key, Some(2), create_test_array(200))); + assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); assert!( store.get_checked(&key, 1).is_none(), "the displaced file reads a miss, never the other file's rows" @@ -327,6 +348,42 @@ mod tests { } } + /// A rewrite is built from an entry read earlier, and the key can be taken + /// over in between — a squeeze reads, awaits a disk write, then stores. If + /// the rewrite adopted whatever identity held the key by then, it would + /// relabel the old file's data as the new owner's, and the new owner would + /// read those rows as a hit. Carrying the identity it read under makes the + /// stale write drop instead. + #[test] + fn a_rewrite_does_not_land_on_a_key_taken_over_since_it_was_read() { + let store = ArtIndex::new(); + let key: EntryID = EntryID::from(11); + + // File A caches, and something begins rewriting that entry. + assert!(store.insert(&key, WriteIdentity::Owned(1), create_test_array(100))); + let (observed, _read) = store.get_with_identity(&key).unwrap(); + assert_eq!(observed, 1); + + // File B takes the key over while that rewrite is in flight. + assert!(store.insert(&key, WriteIdentity::Owned(2), create_test_array(200))); + + // The rewrite lands too late and must be dropped, not relabelled. + assert!(!store.insert( + &key, + WriteIdentity::Rewrite(observed), + create_test_array(100) + )); + + match store.get_checked(&key, 2).unwrap().as_ref() { + CacheEntry::MemoryArrow(array) => assert_eq!( + array.len(), + 200, + "the new owner must still read its own rows, not the rewrite's" + ), + other => panic!("expected the new owner's array, found {other}"), + } + } + /// Maintenance rewrites a key in place and must neither change whose the /// entry is nor bring back one that has been removed — a stale reader that /// misses goes on to insert what it read, and that write must not land @@ -336,15 +393,15 @@ mod tests { let store = ArtIndex::new(); let key: EntryID = EntryID::from(9); - assert!(store.insert(&key, Some(5), create_test_array(10))); - assert!(store.insert(&key, None, create_test_array(20))); + assert!(store.insert(&key, WriteIdentity::Owned(5), create_test_array(10))); + assert!(store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(20))); assert!( store.get_checked(&key, 5).is_some(), "rewriting in place kept the identity" ); store.remove(&key); - assert!(!store.insert(&key, None, create_test_array(30))); + assert!(!store.insert(&key, WriteIdentity::Rewrite(5), create_test_array(30))); assert!(store.get(&key).is_none()); assert_eq!(store.entry_count(), 0); } diff --git a/src/core/tests/memory_footprint.rs b/src/core/tests/memory_footprint.rs index bc10719ea..6d64b711c 100644 --- a/src/core/tests/memory_footprint.rs +++ b/src/core/tests/memory_footprint.rs @@ -103,7 +103,7 @@ fn make_entry(seed: u64, rows: usize) -> ArrayRef { fn indexed_bytes(cache: &LiquidCache) -> usize { let mut sum = 0; - cache.for_each_entry(|_, e| sum += e.memory_usage_bytes()); + cache.for_each_entry(|_, _, e| sum += e.memory_usage_bytes()); sum } @@ -111,7 +111,7 @@ fn indexed_bytes(cache: &LiquidCache) -> usize { /// index rather than the budget, so the budget can be checked against it. fn indexed_disk_bytes(cache: &LiquidCache) -> usize { let mut sum = 0; - cache.for_each_entry(|_, e| { + cache.for_each_entry(|_, _, e| { sum += match e { CacheEntry::DiskLiquid { disk_bytes, .. } | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index 8f17e6e51..1eb8863ce 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -123,37 +123,38 @@ impl LiquidCacheParquet { /// Write the stats of the cache to a parquet file. pub fn write_stats(&self, parquet_file_path: impl AsRef) -> Result<(), ParquetError> { let mut writer = StatsWriter::new(parquet_file_path)?; - self.cache_store.for_each_entry(|entry_id, cached_batch| { - let memory_size = cached_batch.memory_usage_bytes(); - let row_count = match cached_batch { - CacheEntry::MemoryArrow(array) => Some(array.len() as u64), - CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), - CacheEntry::MemorySqueezedLiquid(array) => Some(array.len() as u64), - CacheEntry::DiskLiquid { .. } => None, - CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count - }; - let cache_type = match cached_batch { - CacheEntry::MemoryArrow(_) => "InMemory", - CacheEntry::MemoryLiquid(_) => "LiquidMemory", - CacheEntry::MemorySqueezedLiquid(_) => "LiquidSqueezed", - CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", - CacheEntry::DiskArrow { .. } => "OnDiskArrow", - }; - let reference_count = cached_batch.reference_count(); - let entry_id = ParquetArrayID::from(*entry_id); - writer - .append_entry( - &entry_id.display_path(), - entry_id.row_group_id_inner(), - entry_id.column_id_inner(), - entry_id.batch_id_inner() * self.batch_size() as u64, - row_count, - memory_size as u64, - cache_type, - reference_count as u64, - ) - .unwrap(); - }); + self.cache_store + .for_each_entry(|entry_id, _identity, cached_batch| { + let memory_size = cached_batch.memory_usage_bytes(); + let row_count = match cached_batch { + CacheEntry::MemoryArrow(array) => Some(array.len() as u64), + CacheEntry::MemoryLiquid(array) => Some(array.len() as u64), + CacheEntry::MemorySqueezedLiquid(array) => Some(array.len() as u64), + CacheEntry::DiskLiquid { .. } => None, + CacheEntry::DiskArrow { .. } => None, // We'd need to read it to get the count + }; + let cache_type = match cached_batch { + CacheEntry::MemoryArrow(_) => "InMemory", + CacheEntry::MemoryLiquid(_) => "LiquidMemory", + CacheEntry::MemorySqueezedLiquid(_) => "LiquidSqueezed", + CacheEntry::DiskLiquid { .. } => "OnDiskLiquid", + CacheEntry::DiskArrow { .. } => "OnDiskArrow", + }; + let reference_count = cached_batch.reference_count(); + let entry_id = ParquetArrayID::from(*entry_id); + writer + .append_entry( + &entry_id.display_path(), + entry_id.row_group_id_inner(), + entry_id.column_id_inner(), + entry_id.batch_id_inner() * self.batch_size() as u64, + row_count, + memory_size as u64, + cache_type, + reference_count as u64, + ) + .unwrap(); + }); writer.finish()?; Ok(()) From 5da657f0a804e84c1695bb096316d1b70dca4400 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 11:47:18 +0530 Subject: [PATCH 05/12] fix(cache): scope a disk copy to the entry that wrote it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A declined rewrite has already written its bytes, so the store object and its `disk_copies` record outlive the entry they were built for. The record named only a kind and a length, so the next owner of the key adopted the object on those alone and read the previous owner's rows as a hit — through `reuse_disk_copy` or the spill path. Record the identity that wrote each copy, and treat a copy belonging to anyone else as absent at both adoption sites. The paths that act on whatever object is there regardless of owner — superseding it, discarding it, releasing its reservation — read it unscoped. --- src/core/src/cache/core.rs | 130 ++++++++++++++++++++++++++++++------ src/core/src/cache/index.rs | 9 +++ 2 files changed, 117 insertions(+), 22 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 45cbd99e9..445321b7e 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -44,6 +44,11 @@ use std::collections::HashMap; /// (liquid-cache#43). #[derive(Debug, Clone, Copy)] struct DiskCopy { + /// Whose bytes these are. A key can change hands while a write to it is + /// in flight, and a declined rewrite leaves the object behind; without + /// this the next owner adopts it on kind and length alone and reads the + /// previous owner's rows as its own. + identity: u64, kind: DiskKind, bytes: usize, } @@ -57,22 +62,26 @@ enum DiskKind { impl DiskCopy { /// The store object an entry refers to: a disk stub's bytes, or the /// full serialisation a squeezed entry reads back through. - fn referenced_by(entry: &CacheEntry) -> Option { + fn referenced_by(identity: u64, entry: &CacheEntry) -> Option { match entry { CacheEntry::DiskLiquid { disk_bytes, .. } => Some(Self { + identity, kind: DiskKind::Liquid, bytes: *disk_bytes, }), CacheEntry::DiskArrow { disk_bytes, .. } => Some(Self { + identity, kind: DiskKind::Arrow, bytes: *disk_bytes, }), CacheEntry::MemorySqueezedLiquid(squeezed) => Some(match squeezed.disk_backing() { SqueezedBacking::Liquid(bytes) => Self { + identity, kind: DiskKind::Liquid, bytes, }, SqueezedBacking::Arrow(bytes) => Self { + identity, kind: DiskKind::Arrow, bytes, }, @@ -297,7 +306,10 @@ impl LiquidCache { 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, flush_identity, &batch, bytes) + .await + { Ok(()) => { self.try_insert( entry_id, @@ -314,7 +326,8 @@ impl LiquidCache { if let Some(DiskCopy { kind: DiskKind::Liquid, bytes, - }) = self.disk_copy(&entry_id) + .. + }) = self.disk_copy(&entry_id, flush_identity) { // Hydrated from disk and never modified since: the // bytes are already there, flip the index rather @@ -330,7 +343,12 @@ 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, + flush_identity, + &batch, + Bytes::from(liquid_bytes), + ) .await { Ok(()) => { @@ -364,6 +382,7 @@ impl LiquidCache { async fn write_in_memory_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: CacheEntry, ) -> Result { match &batch { @@ -387,7 +406,7 @@ impl LiquidCache { unreachable!("memory arrow squeeze 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) @@ -397,13 +416,14 @@ impl LiquidCache { if let Some(DiskCopy { kind: DiskKind::Liquid, bytes, - }) = self.disk_copy(&entry_id) + .. + }) = self.disk_copy(&entry_id, identity) { return Ok(CacheEntry::disk_liquid(data_type, bytes)); } 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(data_type, disk_bytes)) } @@ -444,7 +464,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.value(), not_inserted) .await?; batch_to_cache = on_disk_batch; continue; @@ -490,7 +510,21 @@ impl LiquidCache { } } - fn disk_copy(&self, entry_id: &EntryID) -> Option { + /// The store object recorded for `entry_id`, but only if it belongs to + /// `identity`. A copy left by a previous owner reads as absent, so it is + /// never adopted by whoever holds the key now. + fn disk_copy(&self, entry_id: &EntryID, identity: u64) -> Option { + self.disk_copies + .lock() + .unwrap() + .get(entry_id) + .copied() + .filter(|copy| copy.identity == identity) + } + + /// The record regardless of owner, for paths that act on whatever object + /// is there — superseding it, discarding it, releasing its reservation. + fn any_disk_copy(&self, entry_id: &EntryID) -> Option { self.disk_copies.lock().unwrap().get(entry_id).copied() } @@ -505,7 +539,7 @@ impl LiquidCache { /// still land its result after the new one), so this covers the /// sequential case only. pub(crate) async fn supersede_disk_copy(&self, entry_id: EntryID) { - if self.disk_copy(&entry_id).is_none() { + if self.any_disk_copy(&entry_id).is_none() { return; } match self.index.get(&entry_id).as_deref() { @@ -542,7 +576,12 @@ impl LiquidCache { /// If `outcome` demotes an entry to a form backed by a store object whose /// bytes are already there, drop the write and point the entry at the /// existing copy. - fn reuse_disk_copy(&self, entry_id: &EntryID, outcome: SqueezeOutcome) -> SqueezeOutcome { + fn reuse_disk_copy( + &self, + entry_id: &EntryID, + identity: u64, + outcome: SqueezeOutcome, + ) -> SqueezeOutcome { let (entry, bytes) = match outcome { SqueezeOutcome::Replace { entry, @@ -554,9 +593,10 @@ impl LiquidCache { entry, bytes_to_write: Some(bytes), }; - let (Some(copy), Some(wanted)) = - (self.disk_copy(entry_id), DiskCopy::referenced_by(&entry)) - else { + let (Some(copy), Some(wanted)) = ( + self.disk_copy(entry_id, identity), + DiskCopy::referenced_by(identity, &entry), + ) else { return keep_write(entry); }; if copy.kind != wanted.kind { @@ -753,7 +793,7 @@ impl LiquidCache { squeeze_hint, &squeeze_io, ); - let outcome = self.reuse_disk_copy(&to_squeeze, outcome); + let outcome = self.reuse_disk_copy(&to_squeeze, squeezed_identity, outcome); match outcome { SqueezeOutcome::Replace { @@ -761,8 +801,13 @@ impl LiquidCache { bytes_to_write, } => { if let Some(bytes_to_write) = bytes_to_write { - self.write_batch_to_disk(to_squeeze, &new_batch, bytes_to_write) - .await?; + self.write_batch_to_disk( + to_squeeze, + squeezed_identity, + &new_batch, + bytes_to_write, + ) + .await?; } match self.try_insert( to_squeeze, @@ -1028,6 +1073,7 @@ impl LiquidCache { async fn write_batch_to_disk( &self, entry_id: EntryID, + identity: u64, batch: &CacheEntry, bytes: Bytes, ) -> Result<(), CacheFull> { @@ -1063,11 +1109,14 @@ impl LiquidCache { }, CacheEntry::DiskLiquid { .. } | CacheEntry::MemoryLiquid(_) => DiskKind::Liquid, }; - let previous = self - .disk_copies - .lock() - .unwrap() - .insert(entry_id, DiskCopy { kind, bytes: len }); + let previous = self.disk_copies.lock().unwrap().insert( + entry_id, + DiskCopy { + identity, + kind, + bytes: len, + }, + ); if let Some(previous) = previous { // The put replaced the object under this key, so the previous // copy's reservation goes with it. @@ -1544,6 +1593,43 @@ mod tests { } } + /// A rewrite that is declined has already written its bytes, so the store + /// object and its record outlive the entry they were built for. If the + /// next owner of the key could see that record it would adopt the object + /// on kind and length alone and read the previous owner's rows as its own. + #[tokio::test] + async fn a_disk_copy_is_invisible_to_whoever_holds_the_key_next() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(5usize); + + // Identity 1 caches and spills, recording a store object for this key. + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + assert!( + store.disk_copy(&entry_id, 1).is_some(), + "the owner sees the object it wrote" + ); + + // Identity 2 takes the key over. The object is still on disk, and the + // record still names identity 1. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(200)) + .await + .unwrap(); + + assert!( + store.disk_copy(&entry_id, 2).is_none(), + "the new owner must not adopt the object the previous one left" + ); + assert!( + store.any_disk_copy(&entry_id).is_some(), + "the record is still there for the paths that reclaim it" + ); + } + #[tokio::test] async fn insert_returns_cache_full_when_memory_and_disk_are_saturated() { let cache = LiquidCacheBuilder::new() diff --git a/src/core/src/cache/index.rs b/src/core/src/cache/index.rs index 6c9794265..4ec252484 100644 --- a/src/core/src/cache/index.rs +++ b/src/core/src/cache/index.rs @@ -66,6 +66,15 @@ pub(crate) enum WriteIdentity { Rewrite(u64), } +impl WriteIdentity { + /// The identity this write carries, whichever kind it is. + pub(crate) fn value(&self) -> u64 { + match self { + Self::Owned(id) | Self::Rewrite(id) => *id, + } + } +} + pub(crate) struct ArtIndex { art: CongeeArc, entry_count: AtomicUsize, From 9aaf76d9ee52898a21166d7b7a266f531ab8aa84 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 11:58:37 +0530 Subject: [PATCH 06/12] test(cache): put the file id pool under the model checker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pool took its Mutex from std::sync rather than crate::sync, which under the shuttle test feature resolves to shuttle's primitives. A std::sync::Mutex is opaque to the model checker, so the Shuttle job could not explore any interleaving across acquire and release — the pool was excluded from the one job meant to cover it. Import through crate::sync and add a shuttle test for the invariant the scheme rests on: two leases alive at the same time never share an id, and never share an identity. A lease is released from Drop, on whatever thread held it last, so these race by construction. Also correct the module doc, which still said a new owner's writes are refused. They take the key over. --- src/datafusion/src/cache/file_id.rs | 86 ++++++++++++++++++++++++++--- 1 file changed, 78 insertions(+), 8 deletions(-) diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs index f75d4b3d3..a6c362399 100644 --- a/src/datafusion/src/cache/file_id.rs +++ b/src/datafusion/src/cache/file_id.rs @@ -14,19 +14,32 @@ //! being read, not by everything that has ever been read. //! //! Cache *entries* deliberately do not hold a lease. An id can be reused while -//! entries keyed from it are still resident, and those entries are simply -//! unreachable: each records the identity of the file it came from, so the new -//! owner's reads miss and its writes are refused (see -//! `liquid_cache::cache::ArtIndex::get_checked`). The cost is cache space held -//! by data nobody will read until it is evicted; the alternative — releasing -//! ids from inside index removal — would take a process-wide lock underneath a -//! crossbeam-epoch pin, and would deadlock against `reset`. +//! entries keyed from it are still resident, and each entry records the +//! identity of the file it came from, so the new owner's reads miss rather +//! than returning the previous owner's rows. +//! +//! Its writes are not refused, though. A key held by another identity belongs +//! to a file that has already let its id go, so nothing can read that entry +//! any more and the new owner takes the key over +//! (`liquid_cache::cache::ArtIndex::insert`). Refusing instead would leave the +//! key occupied by data nobody can use, and on a cache below its budget +//! nothing evicts it — the new owner would never cache that key again. +//! +//! The alternative, releasing ids from inside index removal, would take a +//! process-wide lock underneath a crossbeam-epoch pin and deadlock against +//! `reset`. Keeping id lifetime and entry lifetime separate is what avoids +//! that, and the identity check is what makes the overlap safe. use std::collections::VecDeque; use std::sync::atomic::{AtomicU64, Ordering}; use ahash::AHashMap; -use std::sync::{Arc, Mutex, Weak}; +// Through `crate::sync`, not `std::sync`: under the shuttle test feature this +// resolves to shuttle's primitives, which is what lets the model checker +// explore interleavings across `acquire` and `release`. A `std::sync::Mutex` +// is opaque to it, so the pool would be excluded from the very job that is +// meant to cover it. +use crate::sync::{Arc, Mutex, Weak}; /// A leased file id. The id returns to its pool when this is dropped. /// @@ -202,6 +215,63 @@ impl FileIdPool { mod tests { use super::*; + /// Two leases alive at the same time must never share an id, and never + /// share an identity. That is the property the whole scheme rests on: + /// a shared id means two files computing one key, and a shared identity + /// means the check that catches it cannot tell them apart. + /// + /// Run under the model checker because `acquire` and `release` race by + /// construction — a lease is released from `Drop`, on whatever thread + /// happened to hold it last. + fn concurrent_leases_stay_distinct() { + let pool = FileIdPool::new(); + let mut threads = Vec::new(); + + for t in 0..3 { + let pool = Arc::clone(&pool); + threads.push(crate::sync::thread::spawn(move || { + for i in 0..3 { + let mine = pool.acquire(&format!("f{t}-{i}.parquet")); + + // Held at the same time, so they cannot be the same file. + let probe = pool.acquire("probe.parquet"); + assert_ne!(mine.get(), probe.get(), "two live leases shared an id"); + assert_ne!( + mine.identity(), + probe.identity(), + "two live leases shared an identity" + ); + drop(probe); + + // The same path always resolves to the same lease. + let again = pool.acquire(&format!("f{t}-{i}.parquet")); + assert_eq!(mine.get(), again.get()); + assert_eq!(mine.identity(), again.identity()); + } + })); + } + + for thread in threads { + thread.join().unwrap(); + } + } + + #[test] + fn concurrent_leases_stay_distinct_single_threaded() { + concurrent_leases_stay_distinct(); + } + + #[cfg(feature = "shuttle")] + #[test] + fn shuttle_concurrent_leases_stay_distinct() { + let mut runner = shuttle::PortfolioRunner::new(true, Default::default()); + let cores = std::thread::available_parallelism().unwrap().get().min(4); + for _ in 0..cores { + runner.add(shuttle::scheduler::PctScheduler::new(10, 1_000)); + } + runner.run(concurrent_leases_stay_distinct); + } + #[test] fn concurrent_readers_of_one_path_share_a_lease() { let pool = FileIdPool::new(); From bef24d091a9e4350f6195adce49a2d0aaa35500c Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 12:29:22 +0530 Subject: [PATCH 07/12] fix(cache): give a re-opened file back its own id, not the queue front A path's identity was restored only when its released record happened to sit at the front of the free queue, and a mismatch discarded that record anyway. Release order is stream completion order and acquire order is partition open order, so for any scan over more than one file the two diverge: a re-read got a fresh identity, orphaned everything it had cached, and the working set was re-read on every query. Match the path against the whole queue and take its own record, falling back to the front only when the path has none. Also run the shuttle job over this package. It does `cd src/core`, which selects the liquid-cache package only, so the pool's model-checked test was never built or run. --- .github/workflows/ci.yml | 7 ++++ src/datafusion/src/cache/file_id.rs | 58 +++++++++++++++++++++++++---- 2 files changed, 58 insertions(+), 7 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cdabc9a08..7b03ebe91 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,6 +139,13 @@ jobs: run: | cd src/core cargo test --features "shuttle" --release -- --test-threads=1 shuttle + # The parquet cache keeps its own concurrent state — the file id + # pool — in this package, and it has its own `shuttle` feature. + # Without this step the model checker never sees it. + - name: Run shuttle test (datafusion) + run: | + cd src/datafusion + cargo test --features "shuttle" --release -- --test-threads=1 shuttle address_san: name: Address Sanitizer diff --git a/src/datafusion/src/cache/file_id.rs b/src/datafusion/src/cache/file_id.rs index a6c362399..e0f96a376 100644 --- a/src/datafusion/src/cache/file_id.rs +++ b/src/datafusion/src/cache/file_id.rs @@ -131,14 +131,26 @@ impl FileIdPool { if let Some(existing) = inner.leases.get(path).and_then(Weak::upgrade) { return existing; } - let (id, reusable_identity) = match inner.free.pop_front() { - Some(released) if released.path == path => (released.id, Some(released.identity)), - Some(released) => (released.id, None), - None => { - let id = inner.next; - inner.next += 1; - (id, None) + // Prefer this path's own released record, wherever it sits in the + // queue. Matching only the front would restore an identity just when + // release order happens to match acquire order — release order is + // stream completion order and acquire order is partition open order, + // so for any scan over more than one file they diverge and every + // re-read would orphan the entries it cached last time. + let mine = inner.free.iter().position(|r| r.path == path); + let (id, reusable_identity) = match mine { + Some(at) => { + let released = inner.free.remove(at).expect("index came from the queue"); + (released.id, Some(released.identity)) } + None => match inner.free.pop_front() { + Some(released) => (released.id, None), + None => { + let id = inner.next; + inner.next += 1; + (id, None) + } + }, }; if id > u16::MAX as u64 { self.over_key_width.fetch_add(1, Ordering::Relaxed); @@ -215,6 +227,38 @@ impl FileIdPool { mod tests { use super::*; + /// A file re-opened after other files have come and gone must still find + /// its own record. Release order is stream completion order and acquire + /// order is partition open order, so the two rarely line up; matching only + /// the front of the queue would hand a re-read a fresh identity and orphan + /// everything it cached before. + #[test] + fn a_reopened_path_finds_its_record_anywhere_in_the_queue() { + let pool = FileIdPool::new(); + + let a = pool.acquire("a.parquet"); + let b = pool.acquire("b.parquet"); + let (a_id, a_identity) = (a.get(), a.identity()); + let (b_id, b_identity) = (b.get(), b.identity()); + + // Released a-then-b, so b's record sits behind a's. + drop(a); + drop(b); + + // Re-open b first: the queue front is a's record, not b's. + let b_again = pool.acquire("b.parquet"); + assert_eq!(b_again.get(), b_id, "b should get its own id back"); + assert_eq!( + b_again.identity(), + b_identity, + "b should keep its identity, or its cached entries are orphaned" + ); + + let a_again = pool.acquire("a.parquet"); + assert_eq!(a_again.get(), a_id); + assert_eq!(a_again.identity(), a_identity); + } + /// Two leases alive at the same time must never share an id, and never /// share an identity. That is the property the whole scheme rests on: /// a shared id means two files computing one key, and a shared identity From 015d6d853c52d34205dc6f3fdc7f5d3fde722066 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 14:38:20 +0530 Subject: [PATCH 08/12] fix(cache): address a store object by its owner, not just its key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scoping the DiskCopy record was not enough. The store object itself was addressed by the packed entry id alone, so two identities that share a key share the object. A write is issued before the index rewrite that would decline it, so it can land after another owner has taken the key over and become disk-backed — overwriting bytes that owner's index entry and record both agree are its own. Silent wrong data, past every check. Put the identity in the store key. The two then address different objects and a late write cannot reach the other's bytes at all, with no ordering to get right. Removes take the identity from the record they are dropping, which is what that record is for. --- src/core/src/cache/core.rs | 83 +++++++++++++++++++++++++------- src/core/src/cache/io_context.rs | 30 ++++++++++-- 2 files changed, 91 insertions(+), 22 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 445321b7e..6f303d9f4 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -224,7 +224,7 @@ impl LiquidCache { match batch.as_ref() { CacheEntry::MemoryLiquid(array) => Some(array.clone()), entry @ CacheEntry::DiskLiquid { .. } => { - 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, @@ -237,7 +237,7 @@ impl LiquidCache { } CacheEntry::MemorySqueezedLiquid(array) => match array.disk_backing() { SqueezedBacking::Liquid(_) => { - let liquid = self.read_disk_liquid_array(entry_id).await; + let liquid = self.read_disk_liquid_array(entry_id, identity).await; Some(liquid) } SqueezedBacking::Arrow(_) => None, @@ -390,6 +390,7 @@ impl LiquidCache { let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( self.store.clone(), entry_id, + identity, self.observer.clone(), )); let outcome = self.squeeze_policy.squeeze( @@ -567,7 +568,7 @@ impl LiquidCache { return; }; self.store - .remove(&entry_id_to_key(&entry_id)) + .remove(&entry_id_to_key(&entry_id, copy.identity)) .await .expect("disk remove failed"); self.budget.release_disk(copy.bytes); @@ -717,11 +718,15 @@ impl LiquidCache { | CacheEntry::DiskArrow { disk_bytes, .. } => *disk_bytes, _ => panic!("remove_disk_entry called for non-disk entry"), }; - self.store - .remove(&entry_id_to_key(&entry_id)) - .await - .expect("disk remove failed"); - self.disk_copies.lock().unwrap().remove(&entry_id); + // Take the record first: it names the owner whose object this is, and + // the key needs it. + let removed_copy = self.disk_copies.lock().unwrap().remove(&entry_id); + if let Some(copy) = removed_copy { + self.store + .remove(&entry_id_to_key(&entry_id, copy.identity)) + .await + .expect("disk remove failed"); + } self.budget.release_disk(disk_bytes); self.cache_policy.notify_remove(&entry_id); self.trace(InternalEvent::DiskEvict { @@ -778,6 +783,7 @@ impl LiquidCache { let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( self.store.clone(), to_squeeze, + squeezed_identity, self.observer.clone(), )); @@ -925,7 +931,7 @@ 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, @@ -948,7 +954,7 @@ 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, @@ -1044,7 +1050,7 @@ impl LiquidCache { let full_array = if !all_paths_present { let batch = CacheEntry::MemorySqueezedLiquid(array.clone()); self.observer.on_get_squeezed_needs_io(); - 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, @@ -1096,7 +1102,7 @@ impl LiquidCache { bytes: len, }); self.store - .put(entry_id_to_key(&entry_id), bytes.to_vec()) + .put(entry_id_to_key(&entry_id, identity), bytes.to_vec()) .await .expect("write failed"); // `bytes` is whatever `batch` serialises to: Arrow IPC for an arrow @@ -1125,10 +1131,10 @@ impl LiquidCache { Ok(()) } - async fn read_disk_arrow_array(&self, entry_id: &EntryID) -> ArrayRef { + async fn read_disk_arrow_array(&self, entry_id: &EntryID, identity: u64) -> ArrayRef { let bytes = self .store - .get(&entry_id_to_key(entry_id)) + .get(&entry_id_to_key(entry_id, identity)) .await .expect("read failed"); let bytes_len = bytes.len(); @@ -1147,10 +1153,11 @@ impl LiquidCache { async fn read_disk_liquid_array( &self, entry_id: &EntryID, + identity: u64, ) -> crate::liquid_array::LiquidArrayRef { let bytes = self .store - .get(&entry_id_to_key(entry_id)) + .get(&entry_id_to_key(entry_id, identity)) .await .expect("read failed"); self.trace(InternalEvent::IoReadLiquid { @@ -1197,7 +1204,7 @@ impl LiquidCache { self.eval_predicate_on_array(filtered, predicate) } entry @ CacheEntry::DiskArrow { .. } => { - let array = self.read_disk_arrow_array(entry_id).await; + let array = self.read_disk_arrow_array(entry_id, identity).await; self.maybe_hydrate( entry_id, identity, @@ -1224,7 +1231,7 @@ impl LiquidCache { array.try_eval_predicate(predicate, selection) } entry @ CacheEntry::DiskLiquid { .. } => { - 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, @@ -1593,6 +1600,48 @@ mod tests { } } + /// The store object has to be owned too, not just the record naming it. + /// + /// A write is issued before the index rewrite that would have declined it, + /// so it can land *after* another owner has taken the key over and become + /// disk-backed. Addressed by the packed id alone, that write overwrites + /// bytes the new owner's index entry and `DiskCopy` record both agree are + /// its own — silent wrong data, past every check. The identity belongs in + /// the store key so the two never address one object. + #[tokio::test] + async fn a_late_write_for_a_previous_owner_cannot_reach_the_current_one() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(21usize); + + // Identity 2 owns the key and is disk-backed. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(200)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + let before = store.read_disk_arrow_array(&entry_id, 2).await; + assert_eq!(before.len(), 200); + + // Identity 1's write lands late, carrying a different array. + let stale_entry = create_test_array(37); + let CacheEntry::MemoryArrow(stale_array) = &stale_entry else { + unreachable!("create_test_array builds an arrow entry") + }; + let stale_bytes = arrow_to_bytes(stale_array).unwrap(); + store + .write_batch_to_disk(entry_id, 1, &stale_entry, stale_bytes) + .await + .unwrap(); + + // The current owner still reads its own rows. + let after = store.read_disk_arrow_array(&entry_id, 2).await; + assert_eq!( + after.len(), + 200, + "a write for a previous owner reached the current owner's object" + ); + } + /// A rewrite that is declined has already written its bytes, so the store /// object and its record outlive the entry they were built for. If the /// next owner of the key could see that record it would adopt the object diff --git a/src/core/src/cache/io_context.rs b/src/core/src/cache/io_context.rs index a9c8974fe..1709fec66 100644 --- a/src/core/src/cache/io_context.rs +++ b/src/core/src/cache/io_context.rs @@ -38,9 +38,21 @@ pub trait EntryMetadata: Debug + Send + Sync { fn get_compressor(&self, entry_id: &EntryID) -> Arc; } -/// Convert an [`EntryID`] to a t4 key (8-byte little-endian representation). -pub(crate) fn entry_id_to_key(entry_id: &EntryID) -> Vec { - usize::from(*entry_id).to_le_bytes().to_vec() +/// Convert an [`EntryID`] and the identity that owns it to a t4 key. +/// +/// Both halves, not just the entry id. `EntryID` is a packed integer whose +/// fields are narrower than the values they encode, so two sources can compute +/// one id — and a store object addressed by that id alone is *shared*. Scoping +/// only the `DiskCopy` record is not enough: a write in flight for one owner +/// can land after another has taken the key over and installed its own disk +/// entry, overwriting bytes the new owner's index entry and record both agree +/// are its own. With the identity in the key the two address different +/// objects, so a late write cannot reach the other's bytes at all. +pub(crate) fn entry_id_to_key(entry_id: &EntryID, identity: u64) -> Vec { + let mut key = Vec::with_capacity(16); + key.extend_from_slice(&usize::from(*entry_id).to_le_bytes()); + key.extend_from_slice(&identity.to_le_bytes()); + key } /// A default implementation of [`EntryMetadata`]. @@ -84,15 +96,23 @@ impl EntryMetadata for DefaultCacheMetadata { pub struct DefaultSqueezeIo { store: t4::Store, entry_id: EntryID, + /// The owner whose object this reads. See [`entry_id_to_key`]. + identity: u64, observer: Arc, } impl DefaultSqueezeIo { /// Create a new instance of [DefaultSqueezeIo]. - pub fn new(store: t4::Store, entry_id: EntryID, observer: Arc) -> Self { + pub fn new( + store: t4::Store, + entry_id: EntryID, + identity: u64, + observer: Arc, + ) -> Self { Self { store, entry_id, + identity, observer, } } @@ -101,7 +121,7 @@ impl DefaultSqueezeIo { #[async_trait::async_trait] impl SqueezeIoHandler for DefaultSqueezeIo { async fn read(&self, range: Option>) -> std::io::Result { - let key = entry_id_to_key(&self.entry_id); + let key = entry_id_to_key(&self.entry_id, self.identity); let bytes = match range { Some(range) => { let len = range.end - range.start; From b5cf5a57e52fceee83b19b0bc60796b067c202b7 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 15:16:19 +0530 Subject: [PATCH 09/12] fix(cache): drop a stale disk write and reclaim the object it supersedes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Putting the identity in the store key stopped a write landing on another owner's object, but left the reservation handling behind: the previous copy's bytes were released while its object stayed on disk, unreachable and uncounted, with nothing that would ever remove it. Removing it whenever the identities differ is wrong in the other direction — a write built before a takeover would delete the live owner's object. The two cases are not distinguishable by identity alone, so ask the index: whoever it says owns the key is the only writer whose bytes can still be read. A write from anyone else is stale and is dropped before it touches the store, which is also why the supersede below can only ever be the current owner reclaiming what a previous one left. The stale case was caught by the test from the previous commit, which started failing when the supersede was added. --- src/core/src/cache/core.rs | 69 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 67 insertions(+), 2 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 6f303d9f4..e7b3abd9b 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -1101,6 +1101,15 @@ impl LiquidCache { kind: CachedBatchType::from(batch), bytes: len, }); + // Whoever the index says owns this key is the only writer whose bytes + // can still be read. A write built before a takeover is stale: writing + // it would leave an object nothing reaches, and superseding on its + // behalf below would delete the live owner's. Drop it instead. + if let Some((current, _)) = self.index.get_with_identity(&entry_id) + && current != identity + { + return Ok(()); + } self.store .put(entry_id_to_key(&entry_id, identity), bytes.to_vec()) .await @@ -1124,8 +1133,22 @@ impl LiquidCache { }, ); if let Some(previous) = previous { - // The put replaced the object under this key, so the previous - // copy's reservation goes with it. + // Same owner: the put replaced that object, so its reservation + // goes with it and there is nothing left to delete. + // + // Different owner: this is the current owner superseding one that + // has let the key go (a stale writer never reaches here). The key + // carries the identity, so the put landed somewhere else and the + // previous object is still there — with no record naming it and + // nothing that would ever reach it. Releasing its reservation + // without removing it would leave disk held by a blob the budget + // has stopped counting. + if previous.identity != identity { + self.store + .remove(&entry_id_to_key(&entry_id, previous.identity)) + .await + .expect("disk remove failed"); + } self.budget.release_disk(previous.bytes); } Ok(()) @@ -1600,6 +1623,48 @@ mod tests { } } + /// Taking a key over must not strand the previous owner's object. + /// + /// Once the identity is part of the store key, a new owner's write lands + /// somewhere else rather than on top — so the old object survives its own + /// record. Releasing its reservation without deleting it leaves disk held + /// by a blob nothing can reach and the budget has stopped counting. + #[tokio::test] + async fn taking_a_key_over_removes_the_previous_owner_s_object() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(31usize); + + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + let disk_after_first = store.budget.disk_usage_bytes(); + assert!(disk_after_first > 0, "the first owner spilled to disk"); + + // A different owner takes the key and spills too. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + + // The first owner's object is gone, not merely unaccounted. + assert!( + store + .store + .get(&crate::cache::io_context::entry_id_to_key(&entry_id, 1)) + .await + .is_err(), + "the previous owner's object outlived its record" + ); + assert_eq!( + store.budget.disk_usage_bytes(), + disk_after_first, + "disk accounting should track one object, not two" + ); + } + /// The store object has to be owned too, not just the record naming it. /// /// A write is issued before the index rewrite that would have declined it, From 8f3da28262f37bd9511c8d4c240ad66878e254fb Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 15:53:27 +0530 Subject: [PATCH 10/12] fix(cache): guard only rewrites, and give back a dropped write's disk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two faults in the stale-write guard, both from taking a bare identity where the kind of write was what mattered. An owned write is the caller taking the key, and the index has not caught up yet by construction — the bytes land before the entry does. Treating that as stale dropped the write and left the caller's own entry pointing at an object that was never written, which panics on the next read. Take a WriteIdentity and drop only a Rewrite. A dropped write also kept the disk it had reserved. No DiskCopy names those bytes, so nothing would ever release them, and repeated takeovers during squeezes would walk the tally to the limit while holding nothing. Release before returning. Both are covered now: an owned write whose index entry lags must still find its bytes, and a dropped rewrite must leave disk usage unchanged. --- src/core/src/cache/core.rs | 125 ++++++++++++++++++++++++++++++++----- 1 file changed, 110 insertions(+), 15 deletions(-) diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index e7b3abd9b..725a8d5ae 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -307,7 +307,12 @@ impl LiquidCache { 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, flush_identity, &batch, bytes) + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(flush_identity), + &batch, + bytes, + ) .await { Ok(()) => { @@ -345,7 +350,7 @@ impl LiquidCache { match self .write_batch_to_disk( entry_id, - flush_identity, + WriteIdentity::Rewrite(flush_identity), &batch, Bytes::from(liquid_bytes), ) @@ -382,7 +387,7 @@ impl LiquidCache { async fn write_in_memory_batch_to_disk( &self, entry_id: EntryID, - identity: u64, + identity: WriteIdentity, batch: CacheEntry, ) -> Result { match &batch { @@ -390,7 +395,7 @@ impl LiquidCache { let squeeze_io: Arc = Arc::new(DefaultSqueezeIo::new( self.store.clone(), entry_id, - identity, + identity.value(), self.observer.clone(), )); let outcome = self.squeeze_policy.squeeze( @@ -418,7 +423,7 @@ impl LiquidCache { kind: DiskKind::Liquid, bytes, .. - }) = self.disk_copy(&entry_id, identity) + }) = self.disk_copy(&entry_id, identity.value()) { return Ok(CacheEntry::disk_liquid(data_type, bytes)); } @@ -465,7 +470,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, identity.value(), not_inserted) + .write_in_memory_batch_to_disk(entry_id, identity, not_inserted) .await?; batch_to_cache = on_disk_batch; continue; @@ -809,7 +814,7 @@ impl LiquidCache { if let Some(bytes_to_write) = bytes_to_write { self.write_batch_to_disk( to_squeeze, - squeezed_identity, + WriteIdentity::Rewrite(squeezed_identity), &new_batch, bytes_to_write, ) @@ -1079,7 +1084,7 @@ impl LiquidCache { async fn write_batch_to_disk( &self, entry_id: EntryID, - identity: u64, + identity: WriteIdentity, batch: &CacheEntry, bytes: Bytes, ) -> Result<(), CacheFull> { @@ -1101,15 +1106,24 @@ impl LiquidCache { kind: CachedBatchType::from(batch), bytes: len, }); - // Whoever the index says owns this key is the only writer whose bytes - // can still be read. A write built before a takeover is stale: writing - // it would leave an object nothing reaches, and superseding on its - // behalf below would delete the live owner's. Drop it instead. - if let Some((current, _)) = self.index.get_with_identity(&entry_id) - && current != identity + // A *rewrite* replays an entry read earlier, so a takeover in the + // meantime makes it stale: its object would be unreachable, and + // superseding on its behalf below would delete the live owner's. Drop + // it, and hand back the reservation taken above — nothing records + // those bytes, so nothing would ever release them. + // + // An *owned* write is the caller taking the key, and the index has not + // caught up yet by construction. Dropping it would leave the caller's + // own index entry pointing at bytes that were never written, and the + // next read of it panics. + if let WriteIdentity::Rewrite(rewriting) = identity + && let Some((current, _)) = self.index.get_with_identity(&entry_id) + && current != rewriting { + self.budget.release_disk(len); return Ok(()); } + let identity = identity.value(); self.store .put(entry_id_to_key(&entry_id, identity), bytes.to_vec()) .await @@ -1623,6 +1637,82 @@ mod tests { } } + /// An owned write is the caller taking the key, and the index has not + /// caught up yet by construction — the write happens first, the index + /// entry second. Dropping it as "stale" would leave the caller's own entry + /// pointing at bytes that were never written, and the next read panics. + #[tokio::test] + async fn an_owned_write_is_not_dropped_because_the_index_lags() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(51usize); + + // Identity 1 holds the key. + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + + // Identity 2 takes it over. Its bytes land before its index entry + // does, so the index still names 1 at write time. + let taking_over = create_test_array(200); + let CacheEntry::MemoryArrow(array) = &taking_over else { + unreachable!("create_test_array builds an arrow entry") + }; + let bytes = arrow_to_bytes(array).unwrap(); + store + .write_batch_to_disk(entry_id, WriteIdentity::Owned(2), &taking_over, bytes) + .await + .unwrap(); + + // The bytes must actually be there, or the entry installed next reads + // a missing object. + let read_back = store.read_disk_arrow_array(&entry_id, 2).await; + assert_eq!( + read_back.len(), + 200, + "an owned write was dropped, so its entry would point at nothing" + ); + } + + /// A dropped write must hand back the disk it reserved. Nothing records + /// those bytes — no `DiskCopy` names them — so no later path would ever + /// release them, and repeated takeovers during squeezes would walk the + /// disk tally up to its limit while holding nothing. + #[tokio::test] + async fn a_dropped_stale_write_releases_its_reservation() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(41usize); + + // Identity 2 owns the key. + store + .insert_inner(entry_id, WriteIdentity::Owned(2), create_test_array(100)) + .await + .unwrap(); + let disk_before = store.budget.disk_usage_bytes(); + + // A rewrite for an identity that has since lost the key is dropped. + let stale_entry = create_test_array(50); + let CacheEntry::MemoryArrow(stale_array) = &stale_entry else { + unreachable!("create_test_array builds an arrow entry") + }; + let stale_bytes = arrow_to_bytes(stale_array).unwrap(); + store + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(1), + &stale_entry, + stale_bytes, + ) + .await + .unwrap(); + + assert_eq!( + store.budget.disk_usage_bytes(), + disk_before, + "a dropped write must not keep the disk it reserved" + ); + } + /// Taking a key over must not strand the previous owner's object. /// /// Once the identity is part of the store key, a new owner's write lands @@ -1694,7 +1784,12 @@ mod tests { }; let stale_bytes = arrow_to_bytes(stale_array).unwrap(); store - .write_batch_to_disk(entry_id, 1, &stale_entry, stale_bytes) + .write_batch_to_disk( + entry_id, + WriteIdentity::Rewrite(1), + &stale_entry, + stale_bytes, + ) .await .unwrap(); From 0637e74b47946fd4b1c25fb42d56f0d22cf46f33 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 18:16:02 +0530 Subject: [PATCH 11/12] fix(cache): keep the assertion on row group and column ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The identity recorded against an entry is the file id alone, so it cannot tell row group 0 from row group 65,536, or column 0 from column 65,536, inside one file. Dropping every assertion left those two narrowing silently with nothing to catch the collision — the doc claimed the identity check absorbed it, which holds for the file id only. Restore the assertion for both. Only the file id wraps unguarded, and it is the one the identity covers. Also leave the multi-column predicate loop rather than the function when an array cannot answer, matching the two arms above it: being unable to answer does not mean the column is unreadable, and the arrow fallback below may still serve it from the cache. --- src/datafusion/src/cache/id.rs | 65 +++++++++++++++++++++------------ src/datafusion/src/cache/mod.rs | 10 ++++- 2 files changed, 50 insertions(+), 25 deletions(-) diff --git a/src/datafusion/src/cache/id.rs b/src/datafusion/src/cache/id.rs index 97464a30e..11f90a07f 100644 --- a/src/datafusion/src/cache/id.rs +++ b/src/datafusion/src/cache/id.rs @@ -52,14 +52,20 @@ const _: () = assert!(std::mem::align_of::() == 8); impl ParquetArrayID { /// Creates a new CacheEntryID. /// - /// Each field is narrowed to the width the packed key gives it. That is - /// lossy above `u16::MAX`, and deliberately not an assertion: a process - /// that outlives 65,536 distinct files goes on working, because the cache - /// compares each entry's recorded identity and treats an aliased key as a - /// miss. Asserting here would panic in debug builds on a case release - /// builds handle, which would leave the shipped behaviour untestable. - /// `LiquidCacheParquet` counts ids it hands out that will not fit. + /// `file_id` is narrowed to the width the packed key gives it, and that is + /// deliberately not an assertion: a process that outlives 65,536 distinct + /// files goes on working, because the cache compares each entry's recorded + /// identity and treats an aliased key as a miss. Asserting would panic in + /// debug builds on a case release builds handle, leaving the shipped + /// behaviour untestable. `LiquidCacheParquet` counts ids that will not fit. + /// + /// The other two keep their assertion, because that identity is the file + /// id *alone*: it cannot tell row group 0 from row group 65,536, or column + /// 0 from column 65,536, within one file. Narrowing those is unguarded and + /// silently wrong, so it stays an assertion rather than a handled case. pub fn new(file_id: u64, row_group_id: u64, column_id: u64, batch_id: BatchID) -> Self { + debug_assert!(row_group_id <= u16::MAX as u64); + debug_assert!(column_id <= u16::MAX as u64); Self { file_id: file_id as u16, rg_id: row_group_id as u16, @@ -152,9 +158,12 @@ pub struct ColumnAccessPath { impl ColumnAccessPath { /// Create a new instance of ColumnAccessPath. /// - /// Narrowing above `u16::MAX` is lossy and handled rather than asserted — - /// see [`ParquetArrayID::new`]. + /// `file_id` narrows unguarded and is handled by the identity check; the + /// other two assert, because the identity cannot distinguish them — see + /// [`ParquetArrayID::new`]. pub fn new(file_id: u64, row_group_id: u64, column_id: u64) -> Self { + debug_assert!(row_group_id <= u16::MAX as u64); + debug_assert!(column_id <= u16::MAX as u64); Self { file_id: file_id as u16, rg_id: row_group_id as u16, @@ -230,29 +239,37 @@ mod tests { assert_eq!(entry_id.batch_id_inner(), *batch_id as u64); } - /// Each field wraps rather than panicking above `u16::MAX`, in every build. - /// The consequence — two sources computing one key — is caught by the - /// identity the cache records alongside each entry, not here. This test - /// pins the wrap down so the aliasing it produces stays a known, - /// reproducible condition rather than a debug-only assertion that the - /// shipped binary never evaluates. + /// `file_id` wraps rather than panicking, in every build. The consequence + /// — two sources computing one key — is caught by the identity the cache + /// records alongside each entry, not here. This pins the wrap down so the + /// aliasing it produces stays a known, reproducible condition rather than + /// a debug-only assertion the shipped binary never evaluates. #[test] - fn fields_wrap_rather_than_panicking_above_their_width() { - let wrapped = ParquetArrayID::new( - (u16::MAX as u64) + 1, - (u16::MAX as u64) + 2, - (u16::MAX as u64) + 3, - BatchID::from_raw(0), - ); + fn a_file_id_wraps_rather_than_panicking_above_its_width() { + let wrapped = ParquetArrayID::new((u16::MAX as u64) + 1, 1, 2, BatchID::from_raw(0)); assert_eq!(wrapped.file_id_inner(), 0); - assert_eq!(wrapped.row_group_id_inner(), 1); - assert_eq!(wrapped.column_id_inner(), 2); // Which is exactly the aliasing the identity check exists to absorb. let first = ParquetArrayID::new(0, 1, 2, BatchID::from_raw(0)); assert_eq!(usize::from(wrapped), usize::from(first)); } + /// The other two fields keep their assertion. The identity recorded + /// against an entry is the file id alone, so it cannot tell row group 0 + /// from row group 65,536 within one file — narrowing those is unguarded + /// and silently wrong, not absorbed. + #[test] + #[should_panic(expected = "row_group_id")] + fn an_over_width_row_group_still_asserts() { + ParquetArrayID::new(0, (u16::MAX as u64) + 1, 0, BatchID::from_raw(0)); + } + + #[test] + #[should_panic(expected = "column_id")] + fn an_over_width_column_still_asserts() { + ParquetArrayID::new(0, 0, (u16::MAX as u64) + 1, BatchID::from_raw(0)); + } + #[test] fn test_cache_entry_id_display_path() { let entry_id = ParquetArrayID::new(1, 2, 3, BatchID::from_raw(4)); diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index faecb72c5..697a37c4f 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -152,7 +152,15 @@ impl CachedRowGroup { } Some(array) => array, }; - let buffer = liquid_array.try_eval_predicate(&liquid_expr, selection)?; + // Leave the loop rather than the function, as the two + // arms above do: an array that cannot answer the predicate + // does not mean the column is unreadable, and the arrow + // fallback below may still serve it from the cache. + let Some(buffer) = liquid_array.try_eval_predicate(&liquid_expr, selection) + else { + combined_buffer = None; + break; + }; combined_buffer = Some(match combined_buffer { None => buffer, From 123cb514cb22230ed8f9f85ae130a05c4aab3b55 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Thu, 17 Sep 2026 18:48:47 +0530 Subject: [PATCH 12/12] fix(cache): delete the store objects a reset forgets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The store key carries the identity that wrote it, so everything written after a reset lands under new keys. Nothing overwrites the old objects, and with disk_copies cleared nothing can find them either — every reset stranded up to max_disk_bytes, unreachable and uncounted. Delete the recorded objects before forgetting the records that name them. reset becomes async, which both callers already are. Before the identity went into the key this was invisible: a new process reused the same keys and overwrote what the old one left. --- benchmark/src/inprocess_runner.rs | 2 +- src/core/src/cache/core.rs | 56 ++++++++++++++++++- src/core/tests/memory_footprint.rs | 2 +- .../src/admin_server/handlers.rs | 2 +- src/datafusion/src/cache/mod.rs | 4 +- 5 files changed, 59 insertions(+), 7 deletions(-) diff --git a/benchmark/src/inprocess_runner.rs b/benchmark/src/inprocess_runner.rs index b45707cf6..03b6c8490 100644 --- a/benchmark/src/inprocess_runner.rs +++ b/benchmark/src/inprocess_runner.rs @@ -657,7 +657,7 @@ impl InProcessBenchmarkRunner { && let Some(cache) = &cache { unsafe { - cache.reset(); + cache.reset().await; } } diff --git a/src/core/src/cache/core.rs b/src/core/src/cache/core.rs index 725a8d5ae..b6bc1c08a 100644 --- a/src/core/src/cache/core.rs +++ b/src/core/src/cache/core.rs @@ -254,10 +254,25 @@ impl LiquidCache { } /// Reset the cache. - pub fn reset(&self) { + /// + /// Deletes the store objects before forgetting the records that name them. + /// The store key carries the identity that wrote it, so everything after a + /// reset writes under new keys and nothing would ever overwrite these + /// again — dropping the records first would strand up to `max_disk_bytes` + /// per reset, unreachable and uncounted. + pub async fn reset(&self) { + let recorded: Vec<(EntryID, DiskCopy)> = { + let mut copies = self.disk_copies.lock().unwrap(); + copies.drain().collect() + }; + for (entry_id, copy) in recorded { + self.store + .remove(&entry_id_to_key(&entry_id, copy.identity)) + .await + .expect("disk remove failed"); + } self.index.reset(); self.budget.reset_usage(); - self.disk_copies.lock().unwrap().clear(); } /// Check if a batch is cached. @@ -1674,6 +1689,43 @@ mod tests { ); } + /// A reset must delete the objects it forgets. The store key carries the + /// identity that wrote it, so everything written after a reset lands under + /// new keys — nothing would overwrite the old objects, and with their + /// records gone nothing would ever find them either. + #[tokio::test] + async fn reset_deletes_the_store_objects_it_forgets() { + let store = create_cache_store(1024 * 1024, Box::new(LiquidPolicy::new())).await; + let entry_id = EntryID::from(61usize); + + store + .insert_inner(entry_id, WriteIdentity::Owned(1), create_test_array(100)) + .await + .unwrap(); + store.flush_all_to_disk().await.unwrap(); + assert!(store.budget.disk_usage_bytes() > 0, "something spilled"); + assert!( + store + .store + .get(&crate::cache::io_context::entry_id_to_key(&entry_id, 1)) + .await + .is_ok(), + "the object is there before the reset" + ); + + store.reset().await; + + assert!( + store + .store + .get(&crate::cache::io_context::entry_id_to_key(&entry_id, 1)) + .await + .is_err(), + "reset left an object nothing can reach and nothing counts" + ); + assert_eq!(store.budget.disk_usage_bytes(), 0); + } + /// A dropped write must hand back the disk it reserved. Nothing records /// those bytes — no `DiskCopy` names them — so no later path would ever /// release them, and repeated takeovers during squeezes would walk the diff --git a/src/core/tests/memory_footprint.rs b/src/core/tests/memory_footprint.rs index 6d64b711c..63da0eb00 100644 --- a/src/core/tests/memory_footprint.rs +++ b/src/core/tests/memory_footprint.rs @@ -249,7 +249,7 @@ async fn heap_footprint_tracks_budget_for_oversized_working_set() { // What survives once the index is emptied is held by the store, the // policy, or the compressor state — not by indexed entries. - cache.reset(); + cache.reset().await; report(&cache, "after reset", baseline); let idle_after_reset = live() - baseline; assert!( diff --git a/src/datafusion-server/src/admin_server/handlers.rs b/src/datafusion-server/src/admin_server/handlers.rs index 1be5b0be7..f3164ba02 100644 --- a/src/datafusion-server/src/admin_server/handlers.rs +++ b/src/datafusion-server/src/admin_server/handlers.rs @@ -51,7 +51,7 @@ pub(crate) async fn reset_cache_handler(State(state): State>) -> J info!("Resetting cache..."); let cache = state.liquid_cache.cache(); unsafe { - cache.reset(); + cache.reset().await; } Json(ApiResponse { diff --git a/src/datafusion/src/cache/mod.rs b/src/datafusion/src/cache/mod.rs index 697a37c4f..404093ed9 100644 --- a/src/datafusion/src/cache/mod.rs +++ b/src/datafusion/src/cache/mod.rs @@ -447,9 +447,9 @@ impl LiquidCacheParquet { /// # Safety /// This is unsafe because resetting the cache while other threads are using the cache may cause undefined behavior. /// You should only call this when no one else is using the cache. - pub unsafe fn reset(&self) { + pub async unsafe fn reset(&self) { self.file_ids.reset(); - self.cache_store.reset(); + self.cache_store.reset().await; } /// Flush all memory-based entries to disk while preserving their format.