perf(index): fetch the IVF and quantizer global buffers together on index open

rust/lance-index/src/vector/storage.rs

@@ -594,26 +594,39 @@ impl<Q: Quantization> IvfQuantizationStorage<Q> {
             .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<String> = 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<String> = 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)?;
         }
 

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<dyn object_store::ObjectStore>,
+        in_flight: Arc<AtomicUsize>,
+        max_in_flight: Arc<AtomicUsize>,
+    }
+
+    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<object_store::PutResult> {
+            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<Box<dyn object_store::MultipartUpload>> {
+            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<object_store::GetResult> {
+            self.enter().await;
+            self.target.get_opts(location, options).await
+        }
+
+        async fn get_ranges(
+            &self,
+            location: &object_store::path::Path,
+            ranges: &[Range<u64>],
+        ) -> object_store::Result<Vec<bytes::Bytes>> {
+            self.enter().await;
+            self.target.get_ranges(location, ranges).await
+        }
+
+        fn delete_stream(
+            &self,
+            locations: futures::stream::BoxStream<
+                'static,
+                object_store::Result<object_store::path::Path>,
+            >,
+        ) -> futures::stream::BoxStream<'static, object_store::Result<object_store::path::Path>>
+        {
+            self.target.delete_stream(locations)
+        }
+
+        fn list(
+            &self,
+            prefix: Option<&object_store::path::Path>,
+        ) -> futures::stream::BoxStream<'static, object_store::Result<object_store::ObjectMeta>>
+        {
+            self.target.list(prefix)
+        }
+
+        async fn list_with_delimiter(
+            &self,
+            prefix: Option<&object_store::path::Path>,
+        ) -> object_store::Result<object_store::ListResult> {
+            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::<Float32Type>("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, &params, 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::<Float32Type>(test_uri, 0.0..1.0).await;
         let params = VectorIndexParams::with_ivf_hnsw_sq_params(