diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 8b871001acf..d7badda05df 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -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, @@ -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); - #[derive(Debug, Clone)] pub struct FragReuseIndexCacheKey<'a> { pub uuid: &'a Uuid, @@ -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()); @@ -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 @@ -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. fn wrap_ivf( @@ -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) @@ -2458,7 +2414,7 @@ impl DatasetIndexInternalExt for Dataset { "IVF_FLAT" => match element_type { DataType::Float16 | DataType::Float32 | DataType::Float64 => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2471,7 +2427,7 @@ impl DatasetIndexInternalExt for Dataset { } DataType::UInt8 => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2490,7 +2446,7 @@ impl DatasetIndexInternalExt for Dataset { "IVF_PQ" => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2504,7 +2460,7 @@ impl DatasetIndexInternalExt for Dataset { "IVF_SQ" => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2518,8 +2474,8 @@ impl DatasetIndexInternalExt for Dataset { "IVF_RQ" => { let ivf = IVFIndex::::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(), @@ -2533,7 +2489,7 @@ impl DatasetIndexInternalExt for Dataset { "IVF_HNSW_FLAT" => match element_type { DataType::UInt8 => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2546,7 +2502,7 @@ impl DatasetIndexInternalExt for Dataset { } _ => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2561,7 +2517,7 @@ impl DatasetIndexInternalExt for Dataset { "IVF_HNSW_SQ" => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -2575,7 +2531,7 @@ impl DatasetIndexInternalExt for Dataset { "IVF_HNSW_PQ" => { let ivf = IVFIndex::::try_new( - self.object_store.clone(), + object_store.clone(), index_dir, uuid.to_owned(), frag_reuse_index, @@ -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) } diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 1d56300fc4b..8e307950d04 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -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(); @@ -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() @@ -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; @@ -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) @@ -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] diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index c8d8a15de55..7629af70ac0 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -488,43 +488,6 @@ impl CacheKey for FileMetadataCacheKey { } } -/// Cached open file readers for the index and aux files. -/// -/// Stored in `file_metadata_cache` to avoid re-opening files on every reconstruction. -/// Not serializable (no codec); a cache miss just triggers a re-open. -struct CachedIndexReaders { - index_reader: Arc, - aux_reader: Arc, -} - -impl lance_core::deepsize::DeepSizeOf for CachedIndexReaders { - fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize { - // FileReader doesn't impl DeepSizeOf. We approximate by counting the - // fixed struct size for each reader plus the Arc - // heap contents. The metadata Arcs are also held by FileMetadataCacheKey - // entries, so this may over-count across cache entries, but - // over-counting is safer than under-counting for eviction purposes. - std::mem::size_of::() * 2 - + self.index_reader.metadata().deep_size_of_children(context) - + self.aux_reader.metadata().deep_size_of_children(context) - } -} - -struct CachedIndexReadersKey { - uuid: String, -} - -impl CacheKey for CachedIndexReadersKey { - type ValueType = CachedIndexReaders; - fn type_name() -> &'static str { - "CachedIndexReaders" - } - fn key(&self) -> std::borrow::Cow<'_, str> { - self.uuid.as_str().into() - } - // No codec() override → in-memory only -} - /// Open a FileReader, reusing cached file metadata if available. async fn open_reader_cached( scheduler: &Arc, @@ -1096,19 +1059,6 @@ impl IVFIndex { .insert_with_key(&FileMetadataCacheKey, storage.reader().metadata().clone()) .await; - // Cache open readers so the first reconstruction also skips file opens. - file_metadata_cache - .insert_with_key( - &CachedIndexReadersKey { - uuid: uuid_str.clone(), - }, - Arc::new(CachedIndexReaders { - index_reader: Arc::new(index_reader.clone()), - aux_reader: Arc::new(storage.reader().clone()), - }), - ) - .await; - let scratch_pool = Arc::new(Self::query_scratch_pool(&ivf, &storage)); let use_query_residual = Self::use_query_residual(&storage, distance_type); let use_residual_scratch = Self::use_residual_scratch(&ivf, use_query_residual); @@ -1962,43 +1912,25 @@ async fn reconstruct_typed( let dir: Path = parts.into_iter().collect(); let aux_path = dir.clone().join(INDEX_AUXILIARY_FILE_NAME); - let readers_key = CachedIndexReadersKey { - uuid: state.uuid.clone(), - }; - - let (index_reader, aux_reader) = - if let Some(cached) = file_metadata_cache.get_with_key(&readers_key).await { - // Warm path: reuse the cached readers directly, no file opens needed. - ((*cached.index_reader).clone(), (*cached.aux_reader).clone()) - } else { - // Cold path: open files, then cache the readers for future reconstructions. - let scheduler_config = SchedulerConfig::max_bandwidth(&object_store); - let scheduler = ScanScheduler::new(object_store, scheduler_config); - let index_reader = open_reader_cached( - &scheduler, - &index_path, - file_metadata_cache, - state.index_file_size, - ) - .await?; - let aux_reader = open_reader_cached( - &scheduler, - &aux_path, - file_metadata_cache, - state.aux_file_size, - ) - .await?; - file_metadata_cache - .insert_with_key( - &readers_key, - Arc::new(CachedIndexReaders { - index_reader: Arc::new(index_reader.clone()), - aux_reader: Arc::new(aux_reader.clone()), - }), - ) - .await; - (index_reader, aux_reader) - }; + // Readers carry a scheduler bound to an object store, so they cannot be + // shared across dataset opens. Reuse only portable file metadata and bind + // fresh readers to the object store supplied for this reconstruction. + let scheduler_config = SchedulerConfig::max_bandwidth(&object_store); + let scheduler = ScanScheduler::new(object_store, scheduler_config); + let index_reader = open_reader_cached( + &scheduler, + &index_path, + file_metadata_cache, + state.index_file_size, + ) + .await?; + let aux_reader = open_reader_cached( + &scheduler, + &aux_path, + file_metadata_cache, + state.aux_file_size, + ) + .await?; let storage = IvfQuantizationStorage::from_cached( aux_reader, @@ -2029,9 +1961,15 @@ async fn reconstruct_typed( #[cfg(test)] mod tests { - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; use std::iter::repeat_n; - use std::{ops::Range, sync::Arc}; + use std::{ + ops::Range, + sync::{ + Arc, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }, + }; use all_asserts::{assert_ge, assert_le, assert_lt}; use arrow::datatypes::{Float64Type, UInt8Type, UInt64Type}; @@ -2088,7 +2026,7 @@ mod tests { }; use lance_index::{INDEX_AUXILIARY_FILE_NAME, metrics::NoOpMetricsCollector}; use lance_io::{ - object_store::ObjectStore, + object_store::{ObjectStore, ObjectStoreParams, StorageOptionsAccessor}, scheduler::{ScanScheduler, SchedulerConfig}, utils::CachedFileSize, }; @@ -6635,6 +6573,117 @@ mod tests { } } + #[derive(Debug)] + struct PartitionBypassCacheBackend { + inner: lance_core::cache::MokaCacheBackend, + bypass_partitions: AtomicBool, + partition_hits: AtomicUsize, + } + + impl PartitionBypassCacheBackend { + fn new() -> Self { + Self { + inner: lance_core::cache::MokaCacheBackend::with_capacity(256 * 1024 * 1024), + bypass_partitions: AtomicBool::new(false), + partition_hits: AtomicUsize::new(0), + } + } + + fn is_partition(key: &lance_core::cache::InternalCacheKey) -> bool { + key.type_name().contains("PartitionEntry") || key.type_name() == "LegacyIVFPartition" + } + + fn set_bypass_partitions(&self, bypass_partitions: bool) { + self.bypass_partitions + .store(bypass_partitions, Ordering::Relaxed); + } + + fn should_bypass(&self, key: &lance_core::cache::InternalCacheKey) -> bool { + self.bypass_partitions.load(Ordering::Relaxed) && Self::is_partition(key) + } + + fn partition_hits(&self) -> usize { + self.partition_hits.load(Ordering::Relaxed) + } + } + + #[async_trait::async_trait] + impl lance_core::cache::CacheBackend for PartitionBypassCacheBackend { + async fn get( + &self, + key: &lance_core::cache::InternalCacheKey, + codec: Option, + ) -> Option { + if self.should_bypass(key) { + None + } else { + let entry = self.inner.get(key, codec).await; + if entry.is_some() && Self::is_partition(key) { + self.partition_hits.fetch_add(1, Ordering::Relaxed); + } + entry + } + } + + async fn insert( + &self, + key: &lance_core::cache::InternalCacheKey, + entry: lance_core::cache::CacheEntry, + size_bytes: usize, + codec: Option, + ) { + if !self.should_bypass(key) { + self.inner.insert(key, entry, size_bytes, codec).await; + } + } + + async fn get_or_insert<'a>( + &self, + key: &lance_core::cache::InternalCacheKey, + loader: std::pin::Pin< + Box< + dyn futures::Future> + + Send + + 'a, + >, + >, + codec: Option, + ) -> Result<(lance_core::cache::CacheEntry, bool)> { + if self.should_bypass(key) { + let (entry, _) = loader.await?; + Ok((entry, false)) + } else { + let result = self.inner.get_or_insert(key, loader, codec).await; + if result.as_ref().is_ok_and(|(_, is_cache_hit)| *is_cache_hit) + && Self::is_partition(key) + { + self.partition_hits.fetch_add(1, Ordering::Relaxed); + } + result + } + } + + async fn invalidate_prefix(&self, prefix: &str) { + self.inner.invalidate_prefix(prefix).await; + } + + async fn clear(&self) { + self.inner.clear().await; + } + + async fn num_entries(&self) -> usize { + self.inner.num_entries().await + } + + async fn size_bytes(&self) -> usize { + self.inner.size_bytes().await + } + + async fn keys(&self) -> Option> { + self.inner.keys().await + } + } + /// Integration test: create a vector index, prewarm it through a /// serializing cache backend, then query. Verifies that entries are /// serialized to bytes and that queries produce correct results after @@ -6714,4 +6763,364 @@ mod tests { assert!(w[1] >= w[0], "distances should be sorted ascending"); } } + + #[rstest] + #[case(IndexFileVersion::V3)] + #[case(IndexFileVersion::Legacy)] + #[tokio::test] + async fn test_vector_cache_uses_current_object_store(#[case] index_version: IndexFileVersion) { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let (mut dataset, vectors) = generate_test_dataset::(test_uri, 0.0..1.0).await; + append_dataset::(&mut dataset, NUM_ROWS, 0.0..1.0).await; + assert_eq!(dataset.get_fragments().len(), 2); + + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams::new(4), + PQBuildParams::default(), + ) + .version(index_version.clone()) + .clone(); + dataset + .create_index( + &["vector"], + IndexType::Vector, + Some("credential_rotation_idx".to_owned()), + ¶ms, + true, + ) + .await + .unwrap(); + let index_meta = dataset + .load_indices_by_name("credential_rotation_idx") + .await + .unwrap() + .pop() + .unwrap(); + let query = vectors.value(0); + let ground_truth = ground_truth(&dataset, "vector", &query, 20, DistanceType::L2).await; + + let cache_backend = Arc::new(PartitionBypassCacheBackend::new()); + let session = Arc::new(crate::session::Session::with_index_cache_backend( + cache_backend.clone(), + 128 * 1024 * 1024, + Arc::new(lance_io::object_store::ObjectStoreRegistry::default()), + )); + let dataset = crate::DatasetBuilder::from_uri(test_uri) + .with_session(session) + .load() + .await + .unwrap(); + + let store_params_a = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([( + "credential_generation".to_owned(), + "secret-generation-a".to_owned(), + )]), + ))), + ..Default::default() + }; + let (store_a, _) = ObjectStore::from_uri_and_params( + dataset.session().store_registry(), + dataset.uri(), + &store_params_a, + ) + .await + .unwrap(); + let dataset_a = dataset.with_object_store(store_a.clone(), Some(store_params_a)); + + let store_params_b = ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([( + "credential_generation".to_owned(), + "secret-generation-b".to_owned(), + )]), + ))), + ..Default::default() + }; + let (store_b, _) = ObjectStore::from_uri_and_params( + dataset.session().store_registry(), + dataset.uri(), + &store_params_b, + ) + .await + .unwrap(); + assert!(!Arc::ptr_eq(&store_a, &store_b)); + let dataset_b = dataset.with_object_store(store_b.clone(), Some(store_params_b)); + + let _ = store_a.io_stats_incremental(); + let _ = store_b.io_stats_incremental(); + + dataset_a + .scan() + .nearest("vector", &query, 20) + .unwrap() + .minimum_nprobes(4) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + + let frag_reuse_uuid = dataset_a.frag_reuse_index_uuid().await; + let state_cache_key = + crate::index::IvfIndexStateCacheKey::new(&index_meta.uuid, frag_reuse_uuid.as_ref()); + let cached_state = if matches!(index_version, IndexFileVersion::V3) { + Some( + dataset_a + .index_cache + .get_with_key(&state_cache_key) + .await + .expect("V3 IVF state should be cached"), + ) + } else { + None + }; + let index_path_fragment = format!("_indices/{}", index_meta.uuid); + let first_store_stats = store_a.io_stats_incremental(); + assert!( + first_store_stats + .requests + .iter() + .any(|request| request.path.as_ref().contains(&index_path_fragment)), + "the first query should read the index through the first object store: {first_store_stats:#?}" + ); + let mut index_keys_after_a = dataset + .session() + .index_cache_keys() + .await + .unwrap() + .map(|key| format!("{}{}{}", key.prefix(), key.key(), key.type_name())) + .collect::>(); + index_keys_after_a.sort_unstable(); + assert!( + index_keys_after_a.iter().any(|key| { + key.contains("PartitionEntry") || key.contains("LegacyIVFPartition") + }), + "the first query should populate portable partition entries" + ); + let mut metadata_keys_after_a = dataset + .session() + .metadata_cache_keys() + .await + .unwrap() + .map(|key| format!("{}{}{}", key.prefix(), key.key(), key.type_name())) + .collect::>(); + metadata_keys_after_a.sort_unstable(); + let _ = store_b.io_stats_incremental(); + + cache_backend.set_bypass_partitions(true); + let results = dataset_b + .scan() + .nearest("vector", &query, 20) + .unwrap() + .minimum_nprobes(4) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + let row_ids = results[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + let recall = row_ids.intersection(&ground_truth).count() as f32 / 20.0; + assert_ge!(recall, 0.5); + + let old_store_stats = store_a.io_stats_incremental(); + let old_store_index_reads = old_store_stats + .requests + .iter() + .filter(|request| request.path.as_ref().contains(&index_path_fragment)) + .count(); + let new_store_stats = store_b.io_stats_incremental(); + let new_store_index_reads = new_store_stats + .requests + .iter() + .filter(|request| request.path.as_ref().contains(&index_path_fragment)) + .count(); + assert_eq!( + old_store_index_reads, 0, + "the new dataset query must not use readers bound to the old object store: {old_store_stats:#?}" + ); + assert!( + new_store_index_reads > 0, + "the new dataset query should read the index through the new object store: {new_store_stats:#?}" + ); + if let Some(cached_state) = cached_state { + let state_after_rotation = dataset_b + .index_cache + .get_with_key(&state_cache_key) + .await + .expect("V3 IVF state should remain cached after rotation"); + assert!( + Arc::ptr_eq(&cached_state, &state_after_rotation), + "store-free IVF state should be reused across object-store generations" + ); + } + + // Re-query through the first dataset to prove portable cache state is + // rebound to the object store supplied by each reconstruction. + let _ = store_a.io_stats_incremental(); + let _ = store_b.io_stats_incremental(); + dataset_a + .scan() + .nearest("vector", &query, 20) + .unwrap() + .minimum_nprobes(4) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + let store_a_stats = store_a.io_stats_incremental(); + let store_a_index_reads = store_a_stats + .requests + .iter() + .filter(|request| request.path.as_ref().contains(&index_path_fragment)) + .count(); + let store_b_stats = store_b.io_stats_incremental(); + let store_b_index_reads = store_b_stats + .requests + .iter() + .filter(|request| request.path.as_ref().contains(&index_path_fragment)) + .count(); + assert!( + store_a_index_reads > 0, + "re-querying the first dataset should use its current object store: {store_a_stats:#?}" + ); + assert_eq!( + store_b_index_reads, 0, + "re-querying the first dataset must not use the second object store: {store_b_stats:#?}" + ); + + let mut index_keys_after_rotation = dataset + .session() + .index_cache_keys() + .await + .unwrap() + .map(|key| format!("{}{}{}", key.prefix(), key.key(), key.type_name())) + .collect::>(); + index_keys_after_rotation.sort_unstable(); + let mut metadata_keys_after_rotation = dataset + .session() + .metadata_cache_keys() + .await + .unwrap() + .map(|key| format!("{}{}{}", key.prefix(), key.key(), key.type_name())) + .collect::>(); + metadata_keys_after_rotation.sort_unstable(); + assert_eq!( + index_keys_after_rotation, index_keys_after_a, + "credential rotation must not create a new index cache namespace" + ); + assert_eq!( + metadata_keys_after_rotation, metadata_keys_after_a, + "credential rotation must not create a new metadata cache namespace" + ); + for key in index_keys_after_rotation + .iter() + .chain(metadata_keys_after_rotation.iter()) + { + assert!(!key.contains("secret-generation-a")); + assert!(!key.contains("secret-generation-b")); + } + + cache_backend.set_bypass_partitions(false); + let partition_hits_before = cache_backend.partition_hits(); + dataset_b + .scan() + .nearest("vector", &query, 20) + .unwrap() + .minimum_nprobes(4) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + assert!( + cache_backend.partition_hits() > partition_hits_before, + "the second store should reuse portable partitions populated by the first" + ); + } + + #[tokio::test] + async fn test_shallow_clone_ivf_rq_uses_resolved_index_directory() { + let test_dir = TempStrDir::default(); + let source_uri = format!("{}/source", test_dir.as_str()); + let clone_uri = format!("{}/clone", test_dir.as_str()); + let (mut source, vectors) = + generate_test_dataset::(&source_uri, 0.0..1.0).await; + append_dataset::(&mut source, NUM_ROWS, 0.0..1.0).await; + assert_eq!(source.get_fragments().len(), 2); + + let params = VectorIndexParams::ivf_rq(4, 5, DistanceType::L2); + source + .create_index( + &["vector"], + IndexType::Vector, + Some("ivf_rq_idx".to_owned()), + ¶ms, + true, + ) + .await + .unwrap(); + + let query = vectors.value(0); + let ground_truth = ground_truth(&source, "vector", &query, 20, DistanceType::L2).await; + source + .tags() + .create("with_ivf_rq", source.version().version) + .await + .unwrap(); + let cloned = source + .shallow_clone(&clone_uri, "with_ivf_rq", None) + .await + .unwrap(); + + let index_meta = cloned + .load_indices_by_name("ivf_rq_idx") + .await + .unwrap() + .pop() + .unwrap(); + assert!( + index_meta.base_id.is_some(), + "a shallow-cloned index should reference its source base" + ); + assert_eq!( + cloned.indice_files_dir(&index_meta).unwrap(), + source.indices_dir(), + "the cloned index should resolve its path through the source base" + ); + assert_ne!( + cloned.indice_files_dir(&index_meta).unwrap(), + cloned.indices_dir(), + "the cloned index should not use the clone's primary index directory" + ); + + let cloned = crate::DatasetBuilder::from_uri(&clone_uri) + .with_session(Arc::new(crate::session::Session::default())) + .load() + .await + .unwrap(); + + let results = cloned + .scan() + .nearest("vector", &query, 20) + .unwrap() + .minimum_nprobes(4) + .with_row_id() + .try_into_batch() + .await + .unwrap(); + let row_ids = results[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>(); + let recall = row_ids.intersection(&ground_truth).count() as f32 / 20.0; + assert_ge!(recall, 0.5); + } }