fix: remove dup cache consts and respect index size param (#2948)

- remove dup cache size consts - replace the index size param --------- Signed-off-by: BubbleCal <[email protected]>
lancedb · Sep 29, 2024 · 396b4d1 · 396b4d1
1 parent 681db8c
commit 396b4d1
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/rust/lance-core/src/cache.rs b/rust/lance-core/src/cache.rs
@@ -14,9 +14,6 @@ use object_store::path::Path;
 use crate::utils::path::LancePathExt;
 use crate::Result;
 
-pub const DEFAULT_INDEX_CACHE_SIZE: usize = 128;
-pub const DEFAULT_METADATA_CACHE_SIZE: usize = 128;
-
 type ArcAny = Arc<dyn Any + Send + Sync>;
 
 #[derive(Clone)]

diff --git a/rust/lance/src/index/cache.rs b/rust/lance/src/index/cache.rs
@@ -10,6 +10,8 @@ use moka::sync::{Cache, ConcurrentCacheExt};
 
 use std::sync::atomic::{AtomicU64, Ordering};
 
+use crate::dataset::DEFAULT_INDEX_CACHE_SIZE;
+
 #[derive(Debug, Default, DeepSizeOf)]
 struct CacheStats {
     hits: AtomicU64,
@@ -71,6 +73,13 @@ impl IndexCache {
         }
     }
 
+    pub(crate) fn capacity(&self) -> u64 {
+        self.vector_cache
+            .policy()
+            .max_capacity()
+            .unwrap_or(DEFAULT_INDEX_CACHE_SIZE as u64)
+    }
+
     #[allow(dead_code)]
     pub(crate) fn len_vector(&self) -> usize {
         self.vector_cache.sync();

diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs
@@ -23,7 +23,7 @@ use futures::prelude::stream::{self, StreamExt, TryStreamExt};
 use lance_arrow::RecordBatchExt;
 use lance_core::cache::FileMetadataCache;
 use lance_core::utils::tokio::get_num_compute_intensive_cpus;
-use lance_core::{cache::DEFAULT_INDEX_CACHE_SIZE, Error, Result};
+use lance_core::{Error, Result};
 use lance_encoding::decoder::{DecoderMiddlewareChain, FilterExpression};
 use lance_file::v2::reader::FileReader;
 use lance_index::vector::flat::index::{FlatIndex, FlatQuantizer};
@@ -122,6 +122,7 @@ impl<S: IvfSubIndex + 'static, Q: Quantization> IVFIndex<S, Q> {
             .upgrade()
             .map(|sess| sess.file_metadata_cache.clone())
             .unwrap_or_else(FileMetadataCache::no_cache);
+        let index_cache_capacity = session.upgrade().unwrap().index_cache.capacity();
         let index_reader = FileReader::try_open(
             scheduler
                 .open_file(&index_dir.child(uuid.as_str()).child(INDEX_FILE_NAME))
@@ -191,7 +192,7 @@ impl<S: IvfSubIndex + 'static, Q: Quantization> IVFIndex<S, Q> {
             ivf,
             reader: index_reader,
             storage,
-            partition_cache: Cache::new(DEFAULT_INDEX_CACHE_SIZE as u64),
+            partition_cache: Cache::new(index_cache_capacity),
             partition_locks: PartitionLoadLock::new(num_partitions),
             sub_index_metadata,
             distance_type,