From 04ec64c3a0b0d67467280b50ebc5feb2ad8aa819 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 11 Sep 2026 23:30:18 +0800 Subject: [PATCH 1/2] perf(index): fetch IVF and quantizer buffers concurrently on open The v2 index load path awaited the IVF protobuf read before issuing the quantizer buffer read, two sequential round trips on every cold index open. Parse the quantizer metadata JSON first (schema metadata, no IO), then issue both global-buffer reads together with tokio::join. Assisted-by: GLM-5.3 --- rust/lance-index/src/vector/storage.rs | 29 ++++++++++++++++++-------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index bbb478f4862..143ddb8df50 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -594,26 +594,37 @@ impl IvfQuantizationStorage { .ok_or(Error::index(format!("{} not found", IVF_METADATA_KEY)))? .parse() .map_err(|e| Error::index(format!("Failed to decode IVF metadata: {}", e)))?; - let ivf_bytes = reader.read_global_buffer(ivf_pos).await?; - let ivf = IvfModel::try_from(pb::Ivf::decode(ivf_bytes)?)?; - - let mut metadata: Vec = serde_json::from_str( + // Parse the quantizer metadata JSON before issuing reads so both + // global-buffer reads (IVF protobuf and quantizer buffer) can be + // fetched concurrently — one round trip saved per cold index open. + let mut metadata_strs: Vec = serde_json::from_str( schema .metadata .get(STORAGE_METADATA_KEY) .ok_or(Error::index(format!("{} not found", STORAGE_METADATA_KEY)))? .as_str(), )?; - debug_assert_eq!(metadata.len(), 1); + debug_assert_eq!(metadata_strs.len(), 1); // for now the metadata is the same for all partitions, so we just store one - let metadata = metadata + let metadata_str = metadata_strs .pop() .ok_or(Error::index("metadata is empty".to_string()))?; - let mut metadata: Q::Metadata = serde_json::from_str(&metadata)?; + let mut metadata: Q::Metadata = serde_json::from_str(&metadata_str)?; + let quantizer_buffer_pos = metadata.buffer_index(); + + let ivf_fut = reader.read_global_buffer(ivf_pos); + let quantizer_buffer_fut = async { + match quantizer_buffer_pos { + Some(pos) => reader.read_global_buffer(pos).await.map(Some), + None => Ok(None), + } + }; + let (ivf_bytes, quantizer_bytes) = tokio::join!(ivf_fut, quantizer_buffer_fut); + let ivf = IvfModel::try_from(pb::Ivf::decode(ivf_bytes?)?)?; + // we store large metadata (e.g. PQ codebook) in global buffer, // and the schema metadata just contains a pointer to the buffer - if let Some(pos) = metadata.buffer_index() { - let bytes = reader.read_global_buffer(pos).await?; + if let Some(bytes) = quantizer_bytes? { metadata.parse_buffer(bytes)?; } From 439556c8ffb293226878c83ffa9220f4e48e0fe8 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 18 Sep 2026 17:38:41 +0800 Subject: [PATCH 2/2] perf(index): fetch the IVF and quantizer global buffers together on index open --- rust/lance-index/src/vector/storage.rs | 28 ++--- rust/lance/src/index/vector/ivf/v2.rs | 141 +++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 13 deletions(-) diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index 143ddb8df50..260b2db6507 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -594,9 +594,9 @@ impl IvfQuantizationStorage { .ok_or(Error::index(format!("{} not found", IVF_METADATA_KEY)))? .parse() .map_err(|e| Error::index(format!("Failed to decode IVF metadata: {}", e)))?; - // Parse the quantizer metadata JSON before issuing reads so both - // global-buffer reads (IVF protobuf and quantizer buffer) can be - // fetched concurrently — one round trip saved per cold index open. + // Parsed before the reads below because it holds the position of the + // quantizer buffer, which is what lets that read start alongside the + // IVF one. let mut metadata_strs: Vec = serde_json::from_str( schema .metadata @@ -612,19 +612,21 @@ impl IvfQuantizationStorage { let mut metadata: Q::Metadata = serde_json::from_str(&metadata_str)?; let quantizer_buffer_pos = metadata.buffer_index(); - let ivf_fut = reader.read_global_buffer(ivf_pos); - let quantizer_buffer_fut = async { - match quantizer_buffer_pos { - Some(pos) => reader.read_global_buffer(pos).await.map(Some), - None => Ok(None), - } - }; - let (ivf_bytes, quantizer_bytes) = tokio::join!(ivf_fut, quantizer_buffer_fut); - let ivf = IvfModel::try_from(pb::Ivf::decode(ivf_bytes?)?)?; + // Both positions come from the schema metadata, so the reads do not + // depend on each other: issue them together instead of waiting for the + // IVF protobuf before asking for the quantizer buffer. + let (ivf_bytes, quantizer_bytes) = + futures::try_join!(reader.read_global_buffer(ivf_pos), async { + match quantizer_buffer_pos { + Some(pos) => reader.read_global_buffer(pos).await.map(Some), + None => Ok(None), + } + })?; + let ivf = IvfModel::try_from(pb::Ivf::decode(ivf_bytes)?)?; // we store large metadata (e.g. PQ codebook) in global buffer, // and the schema metadata just contains a pointer to the buffer - if let Some(bytes) = quantizer_bytes? { + if let Some(bytes) = quantizer_bytes { metadata.parse_buffer(bytes)?; } diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index d4f3134c2d4..d9cb1c11649 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -6407,6 +6407,147 @@ mod tests { const HNSW_VECTOR_ID_COL: &str = "__vector_id"; const HNSW_NEIGHBORS_COL: &str = "__neighbors"; + /// Store wrapper that holds every read open for a measurable window and + /// records how many were in flight at once. Instantaneous reads never + /// overlap, so a delay is what makes concurrency observable at all. + #[derive(Debug)] + struct ConcurrencyProbeStore { + target: Arc, + in_flight: Arc, + max_in_flight: Arc, + } + + impl std::fmt::Display for ConcurrencyProbeStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "ConcurrencyProbeStore({})", self.target) + } + } + + impl ConcurrencyProbeStore { + async fn enter(&self) { + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(now, Ordering::SeqCst); + tokio::time::sleep(std::time::Duration::from_millis(25)).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + } + } + + #[async_trait::async_trait] + impl object_store::ObjectStore for ConcurrencyProbeStore { + async fn put_opts( + &self, + location: &object_store::path::Path, + payload: object_store::PutPayload, + opts: object_store::PutOptions, + ) -> object_store::Result { + self.target.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &object_store::path::Path, + opts: object_store::PutMultipartOptions, + ) -> object_store::Result> { + self.target.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &object_store::path::Path, + options: object_store::GetOptions, + ) -> object_store::Result { + self.enter().await; + self.target.get_opts(location, options).await + } + + async fn get_ranges( + &self, + location: &object_store::path::Path, + ranges: &[Range], + ) -> object_store::Result> { + self.enter().await; + self.target.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: futures::stream::BoxStream< + 'static, + object_store::Result, + >, + ) -> futures::stream::BoxStream<'static, object_store::Result> + { + self.target.delete_stream(locations) + } + + fn list( + &self, + prefix: Option<&object_store::path::Path>, + ) -> futures::stream::BoxStream<'static, object_store::Result> + { + self.target.list(prefix) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&object_store::path::Path>, + ) -> object_store::Result { + self.target.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &object_store::path::Path, + to: &object_store::path::Path, + opts: object_store::CopyOptions, + ) -> object_store::Result<()> { + self.target.copy_opts(from, to, opts).await + } + } + + /// Opening the storage reads the IVF protobuf and the quantizer buffer, + /// two independent global buffers whose positions both come from the schema + /// metadata. Reading them one after the other costs an extra round trip on + /// every cold open, so pin that they go out together: with the reads + /// serialized this sees one in flight at a time. + #[tokio::test(flavor = "multi_thread")] + async fn test_storage_open_fetches_ivf_and_quantizer_buffers_together() { + let (mut dataset, _) = generate_test_dataset::("memory://", 0.0..1.0).await; + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + IvfBuildParams::new(16), + PQBuildParams::new(4, 8), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + let indices = dataset.load_indices().await.unwrap(); + + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let mut probed = dataset.object_store.as_ref().clone(); + probed.inner = Arc::new(ConcurrencyProbeStore { + target: probed.inner.clone(), + in_flight: in_flight.clone(), + max_in_flight: max_in_flight.clone(), + }); + let probed = Arc::new(probed); + let scheduler = ScanScheduler::new(probed, SchedulerConfig::default_for_testing()); + let reader = open_rq_aux_reader(&dataset, scheduler, &indices[0].uuid.to_string()).await; + max_in_flight.store(0, Ordering::SeqCst); + let _storage = lance_index::vector::storage::IvfQuantizationStorage::< + lance_index::vector::pq::ProductQuantizer, + >::try_new(reader, None) + .await + .unwrap(); + assert_eq!( + max_in_flight.load(Ordering::SeqCst), + 2, + "both global buffer reads should be in flight at once" + ); + } + async fn build_ivf_hnsw_sq(test_uri: &str, nlist: usize) -> Dataset { let (mut dataset, _) = generate_test_dataset::(test_uri, 0.0..1.0).await; let params = VectorIndexParams::with_ivf_hnsw_sq_params(