perf(index): avoid HEAD call when opening vector indexes (#7064)

Closes #6944

## Problem

Opening a vector index made an extra network request to find out how big
the index file was. We do not need that request. The file size is
already saved in the manifest, in `IndexMetadata.files`.

`Dataset::open_vector_index` opens `index.idx` just to read the footer
and work out the format version. The plain `open` call does not know the
file size, so the reader sends a HEAD request to get it. This happens
for every vector index, not only HNSW. On the modern reader path the
file is then opened a second time with the size already known, so the
first HEAD was wasted. The same wasted HEAD also happened on the HNSW
auxiliary file open (the case named in the issue), on the legacy remap
path, and on legacy detail inference.

##  Changes

Added one small helper, `open_index_file`. It reads the size from the
manifest (`IndexMetadata::file_size_map()`) and opens the file with
`open_with_size`. If the size is not recorded, which is the case for
older indices, it falls back to the plain `open`.

The helper is now used everywhere a vector index file is opened:

- `index.rs`, the main open that detects the format version (IVF_PQ,
IVF_RQ, FLAT). This also removes a duplicate `file_size_map()` call in
the same path.
- `vector.rs`, the `IVF_HNSW_PQ` and `IVF_HNSW_SQ` auxiliary file opens.
- `ivf.rs`, the legacy v1 `remap_index_file`.
- `details.rs`, the legacy `infer_vector_index_details` fallback.

## Testing

- New test `test_open_index_file_skips_head_when_size_known`. It wraps
the store in a proxy that counts metadata reads against the index file.
The result is 0 HEAD requests when the size is known and 1 HEAD on the
older fallback path.
- `cargo fmt --all`
- `cargo clippy -p lance --tests -- -D warnings`, clean.
- `cargo test -p lance --lib index::vector::`, 208 passed.

Author: gstamatakis95

rust/lance/src/index.rs

@@ -1877,7 +1877,15 @@ impl DatasetIndexInternalExt for Dataset {
         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.clone().join(uuid).join(INDEX_FILE_NAME);
-        let reader: Arc<dyn Reader> = object_store.open(&index_file).await?.into();
+        let file_sizes = index_meta.file_size_map();
+        let reader: Arc<dyn Reader> = vector::open_index_file(
+            object_store.as_ref(),
+            &index_file,
+            INDEX_FILE_NAME,
+            &file_sizes,
+        )
+        .await?
+        .into();
 
         let tailing_bytes = read_last_block(reader.as_ref()).await?;
         let (major_version, minor_version) = read_version(&tailing_bytes)?;
@@ -1944,7 +1952,6 @@ impl DatasetIndexInternalExt for Dataset {
                     self.object_store.clone(),
                     SchedulerConfig::max_bandwidth(&self.object_store),
                 );
-                let file_sizes = index_meta.file_size_map();
                 let cached_size = file_sizes
                     .get(INDEX_FILE_NAME)
                     .map(|&size| CachedFileSize::new(size))

rust/lance/src/index/vector.rs

@@ -52,6 +52,7 @@ use lance_index::vector::{
     sq::{ScalarQuantizer, builder::SQBuildParams},
 };
 use lance_index::{INDEX_AUXILIARY_FILE_NAME, INDEX_METADATA_SCHEMA_KEY, IndexType};
+use lance_io::object_store::ObjectStore;
 use lance_io::traits::Reader;
 use lance_linalg::distance::*;
 use lance_table::format::{IndexMetadata, list_index_files_with_sizes};
@@ -1588,6 +1589,24 @@ pub(crate) async fn open_vector_index(
     Ok(idx)
 }
 
+/// Open an index file without a HEAD request when the size is already known.
+///
+/// `file_sizes` maps a file name to its size in bytes (see
+/// `IndexMetadata::file_size_map`). If `file_name` is missing, which is the case
+/// for older indices that did not record sizes, this falls back to `open`, which
+/// issues a HEAD to learn the size.
+pub(crate) async fn open_index_file(
+    object_store: &ObjectStore,
+    path: &Path,
+    file_name: &str,
+    file_sizes: &HashMap<String, u64>,
+) -> Result<Box<dyn Reader>> {
+    match file_sizes.get(file_name) {
+        Some(&size) => object_store.open_with_size(path, size as usize).await,
+        None => object_store.open(path).await,
+    }
+}
+
 #[instrument(level = "debug", skip(dataset, reader))]
 pub(crate) async fn open_vector_index_v2(
     dataset: Arc<Dataset>,
@@ -1612,11 +1631,18 @@ pub(crate) async fn open_vector_index_v2(
         .ok_or_else(|| Error::index(format!("Index with id {} does not exist", uuid)))?;
     let index_dir = dataset.indice_files_dir(&index_meta)?;
     let object_store = dataset.object_store_for_index(&index_meta).await?;
+    let file_sizes = index_meta.file_size_map();
 
     let index: Arc<dyn VectorIndex> = match index_metadata.index_type.as_str() {
         "IVF_HNSW_PQ" => {
             let aux_path = index_dir.clone().join(uuid).join(INDEX_AUXILIARY_FILE_NAME);
-            let aux_reader = object_store.open(&aux_path).await?;
+            let aux_reader = open_index_file(
+                object_store.as_ref(),
+                &aux_path,
+                INDEX_AUXILIARY_FILE_NAME,
+                &file_sizes,
+            )
+            .await?;
 
             let ivf_data = IvfModel::load(&reader).await?;
             let options = HNSWIndexOptions { use_residual: true };
@@ -1643,7 +1669,13 @@ pub(crate) async fn open_vector_index_v2(
 
         "IVF_HNSW_SQ" => {
             let aux_path = index_dir.clone().join(uuid).join(INDEX_AUXILIARY_FILE_NAME);
-            let aux_reader = object_store.open(&aux_path).await?;
+            let aux_reader = open_index_file(
+                object_store.as_ref(),
+                &aux_path,
+                INDEX_AUXILIARY_FILE_NAME,
+                &file_sizes,
+            )
+            .await?;
 
             let ivf_data = IvfModel::load(&reader).await?;
             let options = HNSWIndexOptions {
@@ -1960,6 +1992,125 @@ mod tests {
     use lance_index::metrics::NoOpMetricsCollector;
     use lance_linalg::distance::MetricType;
 
+    /// `open_index_file` skips the HEAD when the size is known and still falls
+    /// back to a HEAD for older indices that did not record sizes. A HEAD is
+    /// issued as a `get_opts` call with `head = true`, so a proxy store counts
+    /// those against the index file.
+    ///
+    /// Regression test for <https://github.com/lance-format/lance/issues/6944>.
+    #[tokio::test]
+    async fn test_open_index_file_skips_head_when_size_known() {
+        use lance_index::INDEX_FILE_NAME;
+        use lance_io::assert_io_eq;
+        use lance_io::object_store::{ObjectStoreParams, ObjectStoreRegistry};
+
+        let (store, base) = ObjectStore::from_uri_and_params(
+            Arc::new(ObjectStoreRegistry::default()),
+            "memory:///",
+            &ObjectStoreParams::default(),
+        )
+        .await
+        .unwrap();
+
+        let path = base.join(INDEX_FILE_NAME);
+        // Larger than the block size so size discovery needs a separate HEAD.
+        let data = vec![7u8; 2 * store.block_size()];
+        store.put(&path, &data).await.unwrap();
+
+        let file_sizes = HashMap::from([(INDEX_FILE_NAME.to_string(), data.len() as u64)]);
+
+        // Size recorded in the manifest, so reading the size issues no HEAD.
+        let _ = store.io_stats_incremental(); // reset
+        let reader = open_index_file(store.as_ref(), &path, INDEX_FILE_NAME, &file_sizes)
+            .await
+            .unwrap();
+        assert_eq!(reader.size().await.unwrap(), data.len());
+        let stats = store.io_stats_incremental();
+        assert_io_eq!(
+            stats,
+            read_iops,
+            0,
+            "a known file size must not trigger a HEAD request"
+        );
+
+        // Size unknown, as in an older index, so it falls back to a HEAD.
+        let _ = store.io_stats_incremental(); // reset
+        let reader = open_index_file(store.as_ref(), &path, INDEX_FILE_NAME, &HashMap::new())
+            .await
+            .unwrap();
+        assert_eq!(reader.size().await.unwrap(), data.len());
+        let stats = store.io_stats_incremental();
+        assert_io_eq!(
+            stats,
+            read_iops,
+            1,
+            "an unknown file size must fall back to exactly one HEAD request"
+        );
+    }
+
+    /// `open_index_file` looks up sizes in `IndexMetadata::file_size_map()` by
+    /// bare file name. This pins that a freshly created HNSW index records both
+    /// the main and auxiliary files under those exact names with nonzero sizes,
+    /// which is what lets the open path skip the HEAD.
+    #[tokio::test]
+    async fn test_hnsw_index_records_file_sizes() {
+        use lance_index::{INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME};
+
+        let test_dir = TempStrDir::default();
+        let uri = format!("{}/ds", test_dir.as_str());
+
+        let reader = lance_datagen::gen_batch()
+            .col("vector", array::rand_vec::<Float32Type>(32.into()))
+            .into_reader_rows(RowCount::from(400), BatchCount::from(1));
+        let mut dataset = Dataset::write(reader, &uri, None).await.unwrap();
+
+        let params = VectorIndexParams::with_ivf_hnsw_pq_params(
+            MetricType::L2,
+            IvfBuildParams {
+                num_partitions: Some(8),
+                ..Default::default()
+            },
+            HnswBuildParams {
+                max_level: 6,
+                m: 24,
+                ef_construction: 120,
+                prefetch_distance: None,
+            },
+            PQBuildParams {
+                num_sub_vectors: 8,
+                num_bits: 8,
+                ..Default::default()
+            },
+        );
+        dataset
+            .create_index(
+                &["vector"],
+                IndexType::Vector,
+                Some("hnsw".to_string()),
+                &params,
+                false,
+            )
+            .await
+            .unwrap();
+
+        let indices = dataset.load_indices().await.unwrap();
+        let index = indices.iter().find(|idx| idx.name == "hnsw").unwrap();
+        let file_sizes = index.file_size_map();
+
+        assert!(
+            file_sizes.get(INDEX_FILE_NAME).copied().unwrap_or(0) > 0,
+            "manifest should record a nonzero {INDEX_FILE_NAME} size, got {file_sizes:?}"
+        );
+        assert!(
+            file_sizes
+                .get(INDEX_AUXILIARY_FILE_NAME)
+                .copied()
+                .unwrap_or(0)
+                > 0,
+            "manifest should record a nonzero {INDEX_AUXILIARY_FILE_NAME} size, got {file_sizes:?}"
+        );
+    }
+
     #[tokio::test]
     async fn test_initialize_vector_index_ivf_pq() {
         let test_dir = TempStrDir::default();

rust/lance/src/index/vector/details.rs

@@ -503,7 +503,15 @@ pub async fn infer_vector_index_details(
     let index_dir = dataset.indice_files_dir(index)?;
     let file_dir = index_dir.clone().join(uuid.as_str());
     let index_file = file_dir.clone().join(INDEX_FILE_NAME);
-    let reader: Arc<dyn Reader> = dataset.object_store.open(&index_file).await?.into();
+    let file_sizes = index.file_size_map();
+    let reader: Arc<dyn Reader> = super::open_index_file(
+        dataset.object_store.as_ref(),
+        &index_file,
+        INDEX_FILE_NAME,
+        &file_sizes,
+    )
+    .await?
+    .into();
 
     let tailing_bytes = read_last_block(reader.as_ref()).await?;
     let (major_version, minor_version) = read_version(&tailing_bytes)?;

rust/lance/src/index/vector/ivf.rs

@@ -13,7 +13,9 @@ use super::{
     utils::PartitionLoadLock,
 };
 use crate::dataset::index::dataset_format_version;
+use crate::index::DatasetIndexExt;
 use crate::index::DatasetIndexInternalExt;
+use crate::index::vector::open_index_file;
 use crate::index::vector::utils::{get_vector_dim, get_vector_type};
 use crate::{
     dataset::Dataset,
@@ -1842,7 +1844,15 @@ pub(crate) async fn remap_index_file(
     let old_path = dataset.indices_dir().join(old_uuid).join(INDEX_FILE_NAME);
     let new_path = dataset.indices_dir().join(new_uuid).join(INDEX_FILE_NAME);
 
-    let reader: Arc<dyn Reader> = object_store.open(&old_path).await?.into();
+    let file_sizes = dataset
+        .load_index(old_uuid)
+        .await?
+        .map(|index| index.file_size_map())
+        .unwrap_or_default();
+    let reader: Arc<dyn Reader> =
+        open_index_file(object_store, &old_path, INDEX_FILE_NAME, &file_sizes)
+            .await?
+            .into();
     let mut writer = object_store.create(&new_path).await?;
 
     let tasks = generate_remap_tasks(&index.ivf.offsets, &index.ivf.lengths)?;