diff --git a/rust/lance-index/src/vector/storage.rs b/rust/lance-index/src/vector/storage.rs index bbb478f4862..260b2db6507 100644 --- a/rust/lance-index/src/vector/storage.rs +++ b/rust/lance-index/src/vector/storage.rs @@ -594,26 +594,39 @@ 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( + // 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 .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(); + + // 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(pos) = metadata.buffer_index() { - let bytes = reader.read_global_buffer(pos).await?; + 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(