diff --git a/Cargo.lock b/Cargo.lock index 791030cff2d..37153836907 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4586,6 +4586,7 @@ dependencies = [ "pin-project", "proptest", "prost", + "quick_cache", "rand 0.9.5", "roaring", "rstest", @@ -7016,6 +7017,18 @@ dependencies = [ "serde", ] +[[package]] +name = "quick_cache" +version = "0.6.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9c6658afe513a3b484e3abfdaa0d03ef3c0bbf017542c178dd55f94eb3051f9" +dependencies = [ + "ahash", + "equivalent", + "hashbrown 0.16.1", + "parking_lot", +] + [[package]] name = "quinn" version = "0.11.9" diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 68f0528b0eb..5b253cc9b5a 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -3807,6 +3807,7 @@ dependencies = [ "object_store", "pin-project", "prost", + "quick_cache", "rand 0.9.5", "roaring", "serde_json", @@ -5515,6 +5516,18 @@ dependencies = [ "serde", ] +[[package]] +name = "quick_cache" +version = "0.6.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9c6658afe513a3b484e3abfdaa0d03ef3c0bbf017542c178dd55f94eb3051f9" +dependencies = [ + "ahash", + "equivalent", + "hashbrown 0.16.1", + "parking_lot", +] + [[package]] name = "quinn" version = "0.11.9" diff --git a/python/Cargo.lock b/python/Cargo.lock index dbaa6aeebad..f4e533ee0eb 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4134,6 +4134,7 @@ dependencies = [ "object_store", "pin-project", "prost", + "quick_cache", "rand 0.9.5", "roaring", "serde_json", @@ -6193,6 +6194,18 @@ dependencies = [ "serde", ] +[[package]] +name = "quick_cache" +version = "0.6.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9c6658afe513a3b484e3abfdaa0d03ef3c0bbf017542c178dd55f94eb3051f9" +dependencies = [ + "ahash", + "equivalent", + "hashbrown 0.16.1", + "parking_lot", +] + [[package]] name = "quinn" version = "0.11.9" diff --git a/rust/lance-core/Cargo.toml b/rust/lance-core/Cargo.toml index 2f6183be8a1..601e83dde23 100644 --- a/rust/lance-core/Cargo.toml +++ b/rust/lance-core/Cargo.toml @@ -28,6 +28,7 @@ itertools.workspace = true libc.workspace = true libm.workspace = true moka.workspace = true +quick_cache = "0.6" num_cpus = "1.0" object_store = { workspace = true } pin-project.workspace = true diff --git a/rust/lance-core/src/cache/mod.rs b/rust/lance-core/src/cache/mod.rs index cd7476011bf..dc64b79f047 100644 --- a/rust/lance-core/src/cache/mod.rs +++ b/rust/lance-core/src/cache/mod.rs @@ -49,6 +49,7 @@ pub mod backend; pub mod codec; mod entry_io; mod moka; +mod quick; pub use backend::{CacheBackend, CacheEntry, CacheKeyIterator, InternalCacheKey}; pub use codec::{ @@ -56,6 +57,7 @@ pub use codec::{ }; pub use entry_io::{CacheEntryReader, CacheEntryWriter}; pub use moka::MokaCacheBackend; +pub use quick::{QuickCacheBackend, recommended_cache_shards}; use std::borrow::Cow; use std::sync::{ diff --git a/rust/lance-core/src/cache/moka.rs b/rust/lance-core/src/cache/moka.rs index 024d11cdce8..e3d413b349d 100644 --- a/rust/lance-core/src/cache/moka.rs +++ b/rust/lance-core/src/cache/moka.rs @@ -23,7 +23,7 @@ struct MokaCacheEntry { /// Per-entry key cost for eviction: the struct plus the unique `key` bytes. /// Excludes the shared `prefix` `Arc`, which isn't freed per eviction. -fn key_footprint(key: &InternalCacheKey) -> usize { +pub(super) fn key_footprint(key: &InternalCacheKey) -> usize { std::mem::size_of::() + key.key().len() } diff --git a/rust/lance-core/src/cache/quick.rs b/rust/lance-core/src/cache/quick.rs new file mode 100644 index 00000000000..83b0754122e --- /dev/null +++ b/rust/lance-core/src/cache/quick.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! [`CacheBackend`] backed by [quick_cache](https://crates.io/crates/quick_cache), +//! whose hit path is one atomic bit — no read-op channel or inline +//! housekeeping. Used for the session index cache, which sees thousands of +//! cache reads per query. + +use std::pin::Pin; + +use async_trait::async_trait; +use futures::Future; + +use super::CacheCodec; +use super::backend::{CacheBackend, CacheEntry, CacheKeyIterator, InternalCacheKey}; +use super::moka::key_footprint; +use crate::Result; + +#[derive(Clone)] +struct QuickEntry { + entry: CacheEntry, + size_bytes: usize, +} + +#[derive(Clone)] +struct EntryWeighter; + +impl quick_cache::Weighter for EntryWeighter { + fn weight(&self, key: &InternalCacheKey, value: &QuickEntry) -> u64 { + // Same accounting as the moka backend. + (key_footprint(key) + value.size_bytes).max(1) as u64 + } +} + +pub struct QuickCacheBackend { + cache: quick_cache::sync::Cache, +} + +impl std::fmt::Debug for QuickCacheBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("QuickCacheBackend") + .field("entry_count", &self.cache.len()) + .finish() + } +} + +/// Minimum weight budget (4 GiB) per shard: shards don't borrow capacity, and +/// an entry heavier than ~its shard's budget is silently refused admission. +const MIN_SHARD_SHARE: usize = 4 << 30; + +/// Recommended shard count: `min(cpus / 2, capacity / 4 GiB)`, power of two +/// in `[1, 1024]`. The cpu term bounds lock contention; the capacity term +/// keeps each shard's budget >= 4 GiB so large entries stay admissible. +/// Rounded down because quick_cache rounds requests up. +pub fn recommended_cache_shards(capacity: usize) -> usize { + let by_cpu = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(2) + / 2; + let shards = (capacity / MIN_SHARD_SHARE).min(by_cpu).max(1); + let shards = if shards.is_power_of_two() { + shards + } else { + shards.next_power_of_two() / 2 + }; + shards.clamp(1, 1024) +} + +/// Assumed average entry size for pre-allocation sizing. +const ESTIMATED_AVG_ENTRY_BYTES: usize = 64 << 10; + +impl QuickCacheBackend { + /// Create a backend holding up to `capacity` bytes of weighted entries + /// (weight = key footprint + declared size), sharded per + /// [`recommended_cache_shards`]. + pub fn with_capacity(capacity: usize) -> Self { + let shards = recommended_cache_shards(capacity); + // Floor protects the shard count from quick_cache's items-per-shard + // heuristic; ceiling bounds pre-allocation. + let estimated_items = (capacity / ESTIMATED_AVG_ENTRY_BYTES).clamp(shards * 32, 1_000_000); + let options = quick_cache::OptionsBuilder::new() + .estimated_items_capacity(estimated_items) + .weight_capacity(capacity as u64) + .shards(shards) + .build() + // Only errors when weight/item capacity is missing; both are set. + .expect("quick_cache options"); + let cache = quick_cache::sync::Cache::with_options( + options, + EntryWeighter, + Default::default(), + Default::default(), + ); + Self { cache } + } +} + +#[async_trait] +impl CacheBackend for QuickCacheBackend { + async fn get(&self, key: &InternalCacheKey, _codec: Option) -> Option { + self.cache.get(key).map(|v| v.entry) + } + + async fn insert( + &self, + key: &InternalCacheKey, + entry: CacheEntry, + size_bytes: usize, + _codec: Option, + ) { + self.cache + .insert(key.clone(), QuickEntry { entry, size_bytes }); + } + + async fn get_or_insert<'a>( + &self, + key: &InternalCacheKey, + loader: Pin> + Send + 'a>>, + _codec: Option, + ) -> Result<(CacheEntry, bool)> { + match self.cache.get_value_or_guard_async(key).await { + Ok(value) => Ok((value.entry, true)), + Err(guard) => { + let (entry, size_bytes) = loader.await?; + let _ = guard.insert(QuickEntry { + entry: entry.clone(), + size_bytes, + }); + Ok((entry, false)) + } + } + } + + async fn invalidate_prefix(&self, prefix: &str) { + let matching: Vec = self + .cache + .iter() + .filter(|(k, _)| k.starts_with(prefix)) + .map(|(k, _)| k) + .collect(); + for key in matching { + self.cache.remove(&key); + } + } + + async fn clear(&self) { + self.cache.clear(); + } + + async fn keys(&self) -> Option> { + Some(Box::new(self.cache.iter().map(|(key, _)| key))) + } + + async fn num_entries(&self) -> usize { + self.cache.len() + } + + async fn size_bytes(&self) -> usize { + self.cache.weight() as usize + } + + fn approx_num_entries(&self) -> usize { + self.cache.len() + } + + fn approx_size_bytes(&self) -> usize { + self.cache.weight() as usize + } +} + +#[cfg(test)] +mod tests { + use std::marker::PhantomData; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + use crate::cache::{CacheKey, LanceCache}; + + struct TestKey { + key: String, + _phantom: PhantomData, + } + + impl TestKey { + fn new(key: &str) -> Self { + Self { + key: key.to_string(), + _phantom: PhantomData, + } + } + } + + impl CacheKey for TestKey { + type ValueType = T; + fn key(&self) -> std::borrow::Cow<'_, str> { + std::borrow::Cow::Borrowed(&self.key) + } + fn type_name() -> &'static str { + std::any::type_name::() + } + } + + #[tokio::test] + async fn test_quick_backend_roundtrip_singleflight_and_eviction() { + // Capacity must be large relative to one entry: quick_cache shards + // its weight budget, and an entry heavier than its shard's share is + // not admitted at all. + const CAPACITY: usize = 1 << 20; + let item = Arc::new(vec![1u8, 2, 3]); + let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(CAPACITY))); + + // insert + get roundtrip and weighted accounting + cache + .insert_with_key(&TestKey::>::new("a"), item.clone()) + .await; + assert_eq!( + cache + .get_with_key(&TestKey::>::new("a")) + .await + .as_deref(), + Some(&vec![1u8, 2, 3]) + ); + assert_eq!(cache.approx_size(), 1); + assert!(cache.size_bytes().await > 0); + + // get_or_insert runs the loader only on a miss + let loads = Arc::new(AtomicUsize::new(0)); + for _ in 0..2 { + let loads = loads.clone(); + let value = cache + .get_or_insert_with_key(TestKey::>::new("b"), || async move { + loads.fetch_add(1, Ordering::SeqCst); + Ok(vec![7u8]) + }) + .await + .unwrap(); + assert_eq!(value.as_ref(), &vec![7u8]); + } + assert_eq!(loads.load(Ordering::SeqCst), 1); + + // capacity is enforced: overfill with 4x capacity of 16KiB entries + // and confirm eviction kept the weighted size within budget + for i in 0..256 { + cache + .insert_with_key( + &TestKey::>::new(&format!("fill-{i}")), + Arc::new(vec![0u8; 16 << 10]), + ) + .await; + } + assert!(cache.size_bytes().await <= CAPACITY); + assert!(cache.size().await < 258); + + cache.clear().await; + assert_eq!(cache.size().await, 0); + } + + #[tokio::test] + async fn test_quick_backend_tiny_capacity() { + // A tiny cache must not over-provision item metadata and must still + // admit and evict correctly within its weight budget. + const CAPACITY: usize = 64 << 10; + let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(CAPACITY))); + for i in 0..64 { + cache + .insert_with_key( + &TestKey::>::new(&format!("k-{i}")), + Arc::new(vec![0u8; 4 << 10]), + ) + .await; + } + assert!(cache.size_bytes().await <= CAPACITY); + assert!(cache.size().await >= 1); + let hit = cache + .get_with_key(&TestKey::>::new("k-63")) + .await + .is_some() + || cache + .get_with_key(&TestKey::>::new("k-62")) + .await + .is_some(); + assert!(hit, "recently inserted entries should be resident"); + } +} diff --git a/rust/lance/src/session.rs b/rust/lance/src/session.rs index c6214298341..c707586929e 100644 --- a/rust/lance/src/session.rs +++ b/rust/lance/src/session.rs @@ -4,7 +4,7 @@ use std::collections::HashMap; use std::sync::Arc; -use lance_core::cache::{CacheBackend, CacheKeyIterator, LanceCache}; +use lance_core::cache::{CacheBackend, CacheKeyIterator, LanceCache, QuickCacheBackend}; use lance_core::deepsize::DeepSizeOf; use lance_core::{Error, Result}; use lance_index::IndexType; @@ -95,7 +95,8 @@ impl Session { /// /// Parameters: /// - /// - ***index_cache_size***: the size of the index cache. + /// - ***index_cache_size***: the size of the index cache, backed by + /// [`QuickCacheBackend`]. /// - ***metadata_cache_size***: the size of the metadata cache. /// - ***store_registry***: the object store registry to use when opening /// datasets. This determines which schemes are available, and also allows @@ -106,7 +107,9 @@ impl Session { store_registry: Arc, ) -> Self { Self { - index_cache: GlobalIndexCache(LanceCache::with_capacity(index_cache_size)), + index_cache: GlobalIndexCache(LanceCache::with_backend(Arc::new( + QuickCacheBackend::with_capacity(index_cache_size), + ))), metadata_cache: GlobalMetadataCache(LanceCache::with_capacity(metadata_cache_size)), index_extensions: HashMap::new(), store_registry,