Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 20 additions & 70 deletions rust/lance/src/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -652,40 +652,12 @@ fn segment_has_rtree_details(segment: &IndexMetadata) -> bool {
.is_some_and(|details| details.type_url.ends_with("RTreeIndexDetails"))
}

// Cache keys for different index types
#[derive(Debug, Clone)]
pub(crate) struct LegacyVectorIndexCacheKey<'a> {
uuid: &'a Uuid,
fri_uuid: Option<&'a Uuid>,
}

impl<'a> LegacyVectorIndexCacheKey<'a> {
fn new(uuid: &'a Uuid, fri_uuid: Option<&'a Uuid>) -> Self {
Self { uuid, fri_uuid }
}
}

impl CacheKey for LegacyVectorIndexCacheKey<'_> {
type ValueType = CachedLegacyVectorIndex;

fn key(&self) -> std::borrow::Cow<'_, str> {
if let Some(fri_uuid) = self.fri_uuid {
format!("{}-{}", self.uuid, fri_uuid).into()
} else {
self.uuid.to_string().into()
}
}

fn type_name() -> &'static str {
"LegacyVectorIndex"
}
}

/// Sized cache key for `IvfIndexState`.
///
/// Used for v0.3+ indices that support serialization. This key has a codec,
/// so custom cache backends can serialize the state to disk/Redis/etc.
/// Legacy indices use `LegacyVectorIndexCacheKey` instead (in-memory only).
/// Legacy live indices are not cached because their readers are bound to the
/// object store used to open them.
#[derive(Debug, Clone)]
pub(crate) struct IvfIndexStateCacheKey<'a> {
uuid: &'a Uuid,
Expand Down Expand Up @@ -718,11 +690,6 @@ impl CacheKey for IvfIndexStateCacheKey<'_> {
}
}

/// Wrapper that stores a live VectorIndex in the cache.
/// Used for v0.1/v0.2 indices that don't support serializable caching.
#[derive(Debug, lance_core::deepsize::DeepSizeOf)]
pub(crate) struct CachedLegacyVectorIndex(Arc<dyn VectorIndex>);

#[derive(Debug, Clone)]
pub struct FragReuseIndexCacheKey<'a> {
pub uuid: &'a Uuid,
Expand Down Expand Up @@ -2246,12 +2213,6 @@ impl DatasetIndexInternalExt for Dataset {
return Ok(index.as_index());
}

// Fallback: in-memory cache for legacy indices.
let vector_cache_key = LegacyVectorIndexCacheKey::new(uuid, frag_reuse_uuid.as_ref());
if let Some(cached) = self.index_cache.get_with_key(&vector_cache_key).await {
return Ok(cached.0.clone().as_index());
}

let frag_reuse_cache_key = FragReuseIndexCacheKey::new(uuid, frag_reuse_uuid.as_ref());
if let Some(index) = self.index_cache.get_with_key(&frag_reuse_cache_key).await {
return Ok(Arc::new(FragReuseIndexHandle(index)).as_index());
Expand Down Expand Up @@ -2339,12 +2300,6 @@ impl DatasetIndexInternalExt for Dataset {
.await;
}

// Fallback: in-memory cache for legacy indices.
let cache_key = LegacyVectorIndexCacheKey::new(uuid, frag_reuse_uuid.as_ref());
if let Some(cached) = self.index_cache.get_with_key(&cache_key).await {
return Ok(cached.0.clone());
}

let frag_reuse_index = self.open_frag_reuse_index(metrics).await?;
let index_dir = self.indice_files_dir(&index_meta)?;
let index_file = index_dir
Expand All @@ -2364,8 +2319,9 @@ impl DatasetIndexInternalExt for Dataset {
let tailing_bytes = read_last_block(reader.as_ref()).await?;
let (major_version, minor_version) = read_version(&tailing_bytes)?;

// Namespace the index cache by the UUID of the index.
let index_cache = self.index_cache.with_key_prefix(&cache_key.key());
// v2+ partition entries are store-free and remain reusable across
// object-store generations alongside their serializable state.
let index_cache = self.index_cache.with_key_prefix(&state_key.key());

// Extract the cacheable state before type-erasing to Arc<dyn VectorIndex>.
fn wrap_ivf<S: IvfSubIndex + 'static, Q: Quantization + 'static>(
Expand Down Expand Up @@ -2423,8 +2379,8 @@ impl DatasetIndexInternalExt for Dataset {

(0, 3) | (2, _) => {
let scheduler = ScanScheduler::new(
self.object_store.clone(),
SchedulerConfig::max_bandwidth(&self.object_store),
object_store.clone(),
SchedulerConfig::max_bandwidth(&object_store),
);
let cached_size = file_sizes
.get(INDEX_FILE_NAME)
Expand Down Expand Up @@ -2458,7 +2414,7 @@ impl DatasetIndexInternalExt for Dataset {
"IVF_FLAT" => match element_type {
DataType::Float16 | DataType::Float32 | DataType::Float64 => {
let ivf = IVFIndex::<FlatIndex, FlatQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2471,7 +2427,7 @@ impl DatasetIndexInternalExt for Dataset {
}
DataType::UInt8 => {
let ivf = IVFIndex::<FlatIndex, FlatBinQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2490,7 +2446,7 @@ impl DatasetIndexInternalExt for Dataset {

"IVF_PQ" => {
let ivf = IVFIndex::<FlatIndex, ProductQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2504,7 +2460,7 @@ impl DatasetIndexInternalExt for Dataset {

"IVF_SQ" => {
let ivf = IVFIndex::<FlatIndex, ScalarQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2518,8 +2474,8 @@ impl DatasetIndexInternalExt for Dataset {

"IVF_RQ" => {
let ivf = IVFIndex::<FlatIndex, RabitQuantizer>::try_new(
self.object_store.clone(),
self.indices_dir(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
self.metadata_cache.as_ref(),
Expand All @@ -2533,7 +2489,7 @@ impl DatasetIndexInternalExt for Dataset {
"IVF_HNSW_FLAT" => match element_type {
DataType::UInt8 => {
let ivf = IVFIndex::<HNSW, FlatBinQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2546,7 +2502,7 @@ impl DatasetIndexInternalExt for Dataset {
}
_ => {
let ivf = IVFIndex::<HNSW, FlatQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2561,7 +2517,7 @@ impl DatasetIndexInternalExt for Dataset {

"IVF_HNSW_SQ" => {
let ivf = IVFIndex::<HNSW, ScalarQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2575,7 +2531,7 @@ impl DatasetIndexInternalExt for Dataset {

"IVF_HNSW_PQ" => {
let ivf = IVFIndex::<HNSW, ProductQuantizer>::try_new(
self.object_store.clone(),
object_store.clone(),
index_dir,
uuid.to_owned(),
frag_reuse_index,
Expand All @@ -2600,24 +2556,18 @@ impl DatasetIndexInternalExt for Dataset {
};
let (index, ivf_entry) = result?;
metrics.record_index_load();
// Attribute the one-time index-open I/O (file footers, IVF centroids,
// quantization metadata) to this query's metrics. This runs only on a
// real open; cache hits return earlier, so a warm query reports zero
// index-open I/O.
// Attribute index-open I/O (file footers, IVF centroids, quantization
// metadata) to this query's metrics. Cacheable V2+ indices return
// earlier after their portable state has been populated.
if let Some(io_stats) = metrics.io_stats()
&& let Some(open_stats) = index.open_io_stats()
{
io_stats.add_scan_stats(&open_stats);
}
if let Some(ivf_entry) = ivf_entry {
let state_key = IvfIndexStateCacheKey::new(uuid, frag_reuse_uuid.as_ref());
self.index_cache
.insert_with_key(&state_key, Arc::new(ivf_entry))
.await;
} else {
self.index_cache
.insert_with_key(&cache_key, Arc::new(CachedLegacyVectorIndex(index.clone())))
.await;
}
Ok(index)
}
Expand Down
62 changes: 42 additions & 20 deletions rust/lance/src/index/vector/ivf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6602,8 +6602,6 @@ mod tests {

#[tokio::test]
async fn test_prewarm_ivf_legacy() {
use lance_io::assert_io_eq;

let test_dir = TempStrDir::default();
let test_uri = test_dir.as_str();

Expand Down Expand Up @@ -6648,8 +6646,10 @@ mod tests {
stats.read_iops > 0,
"prewarm should have read from disk, but read_iops was 0"
);
let cache_stats_after_prewarm = dataset.session().index_cache_stats().await;

// Can query index without IO
// Legacy live indices are store-bound and must be reopened for each
// query. The prewarmed, store-free partitions remain reusable.
let q = Float32Array::from_iter_values(repeat_n(0.0, DIM));
dataset
.scan()
Expand All @@ -6661,23 +6661,32 @@ mod tests {
.await
.unwrap();
let stats = dataset.object_store.as_ref().io_stats_incremental();
assert_io_eq!(
stats,
read_iops,
0,
"query should not perform IO after prewarm"
assert!(
stats.read_iops > 0,
"query should reopen the legacy live index through the current store"
);
let cache_stats_after_query = dataset.session().index_cache_stats().await;
assert!(
cache_stats_after_query.hits > cache_stats_after_prewarm.hits,
"query should reuse partitions populated by prewarm"
);

// Second prewarm should not need IO (already cached)
// A second prewarm reopens the live index but reuses all partitions.
dataset.prewarm_index("my_idx").await.unwrap();
let stats = dataset.object_store.as_ref().io_stats_incremental();
assert_io_eq!(stats, read_iops, 0, "second prewarm should not perform IO");
assert!(
stats.read_iops > 0,
"second prewarm should reopen the legacy live index"
);
let cache_stats_after_second_prewarm = dataset.session().index_cache_stats().await;
assert!(
cache_stats_after_second_prewarm.hits > cache_stats_after_query.hits,
"second prewarm should reuse cached partitions"
);
}

#[tokio::test]
async fn test_prewarm_ivf_legacy_multiple_deltas() {
use lance_io::assert_io_eq;

let test_dir = copy_test_data_to_tmp("v0.21.0/bad_index_fragment_bitmap").unwrap();
let test_uri = test_dir.path_str();
let test_uri = &test_uri;
Expand Down Expand Up @@ -6723,8 +6732,10 @@ mod tests {
stats.read_iops > 0,
"prewarm should have read from disk, but read_iops was 0"
);
let cache_stats_after_prewarm = dataset.session().index_cache_stats().await;

// Query should not perform index IO after prewarm of all deltas.
// Each legacy delta is reopened through the current store, while its
// prewarmed, store-free partitions remain reusable.
dataset
.scan()
.nearest("vector", &q, 10)
Expand All @@ -6735,17 +6746,28 @@ mod tests {
.await
.unwrap();
let stats = dataset.object_store.as_ref().io_stats_incremental();
assert_io_eq!(
stats,
read_iops,
0,
"query should not perform IO after prewarm"
assert!(
stats.read_iops > 0,
"query should reopen each legacy live index through the current store"
);
let cache_stats_after_query = dataset.session().index_cache_stats().await;
assert!(
cache_stats_after_query.hits > cache_stats_after_prewarm.hits,
"query should reuse prewarmed partitions from every index delta"
);

// Second prewarm should not need IO (already cached).
// A second prewarm reopens the live indices but reuses all partitions.
dataset.prewarm_index("vector_idx").await.unwrap();
let stats = dataset.object_store.as_ref().io_stats_incremental();
assert_io_eq!(stats, read_iops, 0, "second prewarm should not perform IO");
assert!(
stats.read_iops > 0,
"second prewarm should reopen each legacy live index"
);
let cache_stats_after_second_prewarm = dataset.session().index_cache_stats().await;
assert!(
cache_stats_after_second_prewarm.hits > cache_stats_after_query.hits,
"second prewarm should reuse cached partitions from every index delta"
);
}

#[tokio::test]
Expand Down
Loading
Loading