From 5f3afb3294dfaa73baab924839a45ec929092401 Mon Sep 17 00:00:00 2001 From: kid Date: Mon, 27 Jul 2026 17:21:07 +0800 Subject: [PATCH 1/2] fix(index): rebind vector readers to current object store - cache only portable IVF state and file metadata, rebuilding readers with the object store supplied by each dataset open - resolve the IVF_RQ index directory from the index metadata so shallow clones and external index bases read the source index path - preserve reusable IVF partition entries across credential rotations while routing reader I/O through the current object store - keep in-memory caching of legacy (v0.1/v0.2) live vector indices and document that their store-bound readers assume internally refreshing credentials Closes #7904 --- rust/lance/src/index.rs | 34 +- rust/lance/src/index/vector/ivf/v2.rs | 669 ++++++++++++++++++++++---- 2 files changed, 591 insertions(+), 112 deletions(-) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index a11295d0fb6..041a079ebb4 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -703,6 +703,11 @@ impl CacheKey for LegacyVectorIndexCacheKey<'_> { /// 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). +/// +/// Note: legacy entries hold live readers bound to the object store that +/// opened them, so they are only valid for credential setups that refresh +/// internally (e.g. a credentials provider). Deployments that pass fresh +/// static credentials per dataset open should use v0.3+ index formats. #[derive(Debug, Clone)] pub(crate) struct IvfIndexStateCacheKey<'a> { uuid: &'a Uuid, @@ -2438,7 +2443,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. + // Namespace the index cache by the UUID of the index. v2+ partition + // entries are store-free and remain reusable across object-store + // generations alongside their serializable state. let index_cache = self.index_cache.for_index(uuid, frag_reuse_uuid.as_ref()); // Extract the cacheable state before type-erasing to Arc. @@ -2497,8 +2504,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) @@ -2532,7 +2539,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, @@ -2545,7 +2552,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, @@ -2564,7 +2571,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, @@ -2578,7 +2585,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, @@ -2592,8 +2599,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(), @@ -2607,7 +2614,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, @@ -2620,7 +2627,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, @@ -2635,7 +2642,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, @@ -2649,7 +2656,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, @@ -2684,7 +2691,6 @@ impl DatasetIndexInternalExt for Dataset { 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; diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 06239887a08..344d7f4e3b4 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -489,49 +489,6 @@ impl CacheKey for FileMetadataCacheKey { fn write_key(&self, _builder: &mut KeyBuilder) {} } -/// 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: Uuid, -} - -impl CacheKey for CachedIndexReadersKey { - type ValueType = CachedIndexReaders; - fn type_name() -> &'static str { - "CachedIndexReaders" - } - fn key(&self) -> std::borrow::Cow<'_, str> { - self.uuid.to_string().into() - } - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.index.ivf-cached-readers-key", 1) - } - fn write_key(&self, builder: &mut KeyBuilder) { - builder.write_fixed_bytes(self.uuid.as_bytes()); - } - // No codec() override → in-memory only -} - /// Open a FileReader, reusing cached file metadata if available. async fn open_reader_cached( scheduler: &Arc, @@ -559,14 +516,20 @@ async fn open_reader_cached( .await } else { let file_scheduler = scheduler.open_file(path, &cached_size).await?; - FileReader::try_open( + let reader = FileReader::try_open( file_scheduler, None, Arc::::default(), cache, FileReaderOptions::default(), ) - .await + .await?; + // File metadata is store-free, so it outlives the reader opened here: + // cache it to spare later reconstructions the footer read. + file_cache + .insert_with_key(&FileMetadataCacheKey, reader.metadata().clone()) + .await; + Ok(reader) } } @@ -1117,17 +1080,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 }, - 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); @@ -1987,43 +1939,25 @@ async fn reconstruct_typed( let dir: Path = parts.into_iter().collect(); let aux_path = dir.clone().join(INDEX_AUXILIARY_FILE_NAME); - let parsed_uuid = Uuid::parse_str(&state.uuid) - .map_err(|e| Error::index(format!("Invalid UUID in IvfIndexState: {e}")))?; - let readers_key = CachedIndexReadersKey { uuid: parsed_uuid }; - - 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, @@ -2034,6 +1968,8 @@ async fn reconstruct_typed( ); let rq_search_cache = IVFIndex::::rq_search_cache_from_state(state, &storage)?; + let parsed_uuid = Uuid::parse_str(&state.uuid) + .map_err(|e| Error::index(format!("Invalid UUID in IvfIndexState: {e}")))?; let index = IVFIndex::::from_cached_state( to_local_path(&index_path), index_path.to_string(), @@ -2052,9 +1988,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}; @@ -2078,7 +2020,9 @@ mod tests { use crate::dataset::{InsertBuilder, UpdateBuilder, WriteMode, WriteParams}; use crate::index::DatasetIndexExt; use crate::index::DatasetIndexInternalExt; - use crate::index::vector::ivf::v2::{IvfFlatIndex, IvfPq, IvfStateEntryBox, PartitionEntry}; + use crate::index::vector::ivf::v2::{ + IVFPartitionKey, IvfFlatIndex, IvfPq, IvfStateEntryBox, PartitionEntry, + }; use crate::utils::test::copy_test_data_to_tmp; use crate::{ Dataset, @@ -2088,7 +2032,7 @@ mod tests { dataset::optimize::{CompactionOptions, compact_files}, index::vector::IndexFileVersion, }; - use lance_core::cache::{CacheCodecImpl, LanceCache}; + use lance_core::cache::{CacheBackend, CacheCodecImpl, LanceCache}; use lance_core::utils::tempfile::TempStrDir; use lance_core::{ROW_ID, Result}; use lance_encoding::decoder::DecoderPlugins; @@ -2112,7 +2056,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, }; @@ -6540,6 +6484,181 @@ mod tests { assert_io_eq!(stats, read_iops, 0, "second prewarm should not perform IO"); } + /// Index-cache backend that can drop partition entries on demand. + /// + /// Used to simulate cache invalidation after a credential rotation: + /// partition keys are opaque digests, so the test supplies the exact + /// [`InternalCacheKey`] set to bypass once the index identity is known + /// (see [`ivf_partition_cache_keys`]). + #[derive(Debug)] + struct PartitionBypassCacheBackend { + inner: lance_core::cache::MokaCacheBackend, + partition_keys: std::sync::Mutex>, + bypass_partitions: AtomicBool, + partition_hits: AtomicUsize, + } + + impl PartitionBypassCacheBackend { + fn new() -> Self { + Self { + inner: lance_core::cache::MokaCacheBackend::with_capacity(256 * 1024 * 1024), + partition_keys: std::sync::Mutex::new(HashSet::new()), + bypass_partitions: AtomicBool::new(false), + partition_hits: AtomicUsize::new(0), + } + } + + fn set_partition_keys(&self, partition_keys: HashSet) { + *self.partition_keys.lock().unwrap() = partition_keys; + } + + fn is_partition(&self, key: &lance_core::cache::InternalCacheKey) -> bool { + self.partition_keys.lock().unwrap().contains(key) + } + + 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) + } + + /// Whether the backend currently holds an entry for `key`. + async fn contains(&self, key: &lance_core::cache::InternalCacheKey) -> bool { + self.inner.get(key, None).await.is_some() + } + + fn partition_hits(&self) -> usize { + self.partition_hits.load(Ordering::Relaxed) + } + } + + /// Derive the internal cache keys of the IVF partition entries for an + /// index, replicating the namespace path + /// `dataset URI -> index UUID -> frag-reuse UUID` used when opening the + /// index. V3 partitions use [`IVFPartitionKey`]; legacy (v0.1/v0.2) + /// indices use `LegacyIVFPartitionKey`. + fn ivf_partition_cache_keys( + dataset_uri: &str, + uuid: &uuid::Uuid, + fri_uuid: Option<&uuid::Uuid>, + num_partitions: usize, + index_version: &IndexFileVersion, + ) -> HashSet { + use lance_core::cache::{CacheKey, CacheNamespace, KeyBuilder, UnsizedCacheKey}; + + let mut namespace = CacheNamespace::root().child(dataset_uri); + namespace = namespace.child(uuid.as_hyphenated().to_string().as_str()); + if let Some(fri_uuid) = fri_uuid { + namespace = namespace.child(fri_uuid.as_hyphenated().to_string().as_str()); + } + + (0..num_partitions) + .map(|partition_id| { + if matches!(index_version, IndexFileVersion::V3) { + let cache_key = + IVFPartitionKey::::new(partition_id); + let mut builder = KeyBuilder::new( + namespace, + IVFPartitionKey::::stable_type_id(), + IVFPartitionKey::::schema(), + ); + cache_key.write_key(&mut builder); + builder.finish() + } else { + let cache_key = + crate::index::vector::ivf::LegacyIVFPartitionKey::new(partition_id); + let mut builder = KeyBuilder::new( + namespace, + crate::index::vector::ivf::LegacyIVFPartitionKey::stable_type_id(), + crate::index::vector::ivf::LegacyIVFPartitionKey::schema(), + ); + cache_key.write_key(&mut builder); + builder.finish() + } + }) + .collect() + } + + #[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 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 + } + + fn approx_num_entries(&self) -> usize { + self.inner.approx_num_entries() + } + + fn approx_size_bytes(&self) -> usize { + self.inner.approx_size_bytes() + } + } + /// 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 @@ -6697,4 +6816,358 @@ mod tests { "warmed IVF query should not perform IO after backend restart" ); } + + #[rstest] + #[case::v3(IndexFileVersion::V3)] + #[case::legacy(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 partition_keys = ivf_partition_cache_keys( + dataset.uri(), + &index_meta.uuid, + frag_reuse_uuid.as_ref(), + 4, + &index_version, + ); + cache_backend.set_partition_keys(partition_keys.clone()); + for partition_key in &partition_keys { + assert!( + cache_backend.contains(partition_key).await, + "the first query should populate portable partition entries" + ); + } + let index_entries_after_a = dataset.session().index_cache_stats().await.num_entries; + let metadata_entries_after_a = dataset.session().metadata_cache_stats().await.num_entries; + 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(); + if matches!(index_version, IndexFileVersion::V3) { + 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:#?}" + ); + } else { + // Legacy live indices are shared across dataset opens: their + // readers stay bound to the object store that first populated the + // cache, so the second dataset keeps reading through the old store. + assert!( + old_store_index_reads > 0, + "the cached legacy index should keep reading through the original object store: {old_store_stats:#?}" + ); + assert_eq!( + new_store_index_reads, 0, + "the cached legacy index must not reopen 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: V3 portable state is rebound to + // the store supplied by each reconstruction, while the cached legacy + // index keeps reading through its original store. Either way the second + // store must not be touched. + 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 read the index through its 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:#?}" + ); + + // Cache keys are opaque digests, so they cannot embed credential + // material by construction. What rotation must not do is mint new + // entries: the same portable state, partitions, and file metadata + // serve both object-store generations. + let index_entries_after_rotation = dataset.session().index_cache_stats().await.num_entries; + let metadata_entries_after_rotation = + dataset.session().metadata_cache_stats().await.num_entries; + assert_eq!( + index_entries_after_rotation, index_entries_after_a, + "credential rotation must not create new index cache entries" + ); + assert_eq!( + metadata_entries_after_rotation, metadata_entries_after_a, + "credential rotation must not create new metadata cache entries" + ); + + 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); + } } From f27fc1339c75693ad7485827ccf4be1b431bc641 Mon Sep 17 00:00:00 2001 From: kid Date: Thu, 30 Jul 2026 22:50:04 +0800 Subject: [PATCH 2/2] ci: rerun checks