fix(index): search legacy truncated PQ indexes

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

@@ -690,6 +690,96 @@ mod tests {
             });
     }
 
+    #[test]
+    fn test_distance_with_legacy_truncated_dimension() {
+        const DIM: usize = 64;
+        const NUM_SUB_VECTORS: usize = 14;
+        const NUM_BITS: u32 = 8;
+        const NUM_CENTROIDS: usize = 1 << NUM_BITS;
+        const SUB_VECTOR_DIM: usize = DIM / NUM_SUB_VECTORS;
+        const PERSISTED_DIM: usize = NUM_SUB_VECTORS * SUB_VECTOR_DIM;
+
+        // Older writers silently omitted the tail when the dimension was not
+        // divisible by the number of sub-vectors. Preserve searches over those
+        // indexes even though current writers reject this configuration.
+        let indexed_vector = (1..=DIM).map(|value| value as f32).collect::<Vec<_>>();
+        let mut codebook = Vec::with_capacity(NUM_SUB_VECTORS * NUM_CENTROIDS * SUB_VECTOR_DIM);
+        for sub_vector in indexed_vector[..PERSISTED_DIM].chunks_exact(SUB_VECTOR_DIM) {
+            for _ in 0..NUM_CENTROIDS {
+                codebook.extend_from_slice(sub_vector);
+            }
+        }
+        let query = indexed_vector
+            .iter()
+            .enumerate()
+            .map(|(idx, value)| value + if idx < PERSISTED_DIM { 1.0 } else { 1_000.0 })
+            .collect::<Vec<_>>();
+        let code = UInt8Array::from(vec![0; NUM_SUB_VECTORS]);
+
+        let prepared_l2 = ProductQuantizer::new(
+            NUM_SUB_VECTORS,
+            NUM_BITS,
+            DIM,
+            FixedSizeListArray::try_new_from_values(
+                Float32Array::from(codebook.clone()),
+                DIM as i32,
+            )
+            .unwrap(),
+            DistanceType::L2,
+        );
+        assert!(prepared_l2.l2_targets.is_some());
+        let distances = prepared_l2
+            .compute_distances(&Float32Array::from(query.clone()), &code)
+            .unwrap();
+        assert_relative_eq!(distances.value(0), PERSISTED_DIM as f32, epsilon = 1e-4);
+
+        let generic_l2 = ProductQuantizer::new(
+            NUM_SUB_VECTORS,
+            NUM_BITS,
+            DIM,
+            FixedSizeListArray::try_new_from_values(
+                PrimitiveArray::<datatypes::Float64Type>::from(
+                    codebook
+                        .iter()
+                        .map(|value| *value as f64)
+                        .collect::<Vec<_>>(),
+                ),
+                DIM as i32,
+            )
+            .unwrap(),
+            DistanceType::L2,
+        );
+        assert!(generic_l2.l2_targets.is_none());
+        let distances = generic_l2
+            .compute_distances(
+                &PrimitiveArray::<datatypes::Float64Type>::from(
+                    query.iter().map(|value| *value as f64).collect::<Vec<_>>(),
+                ),
+                &code,
+            )
+            .unwrap();
+        assert_relative_eq!(distances.value(0), PERSISTED_DIM as f32, epsilon = 1e-4);
+
+        let dot = ProductQuantizer::new(
+            NUM_SUB_VECTORS,
+            NUM_BITS,
+            DIM,
+            FixedSizeListArray::try_new_from_values(Float32Array::from(codebook), DIM as i32)
+                .unwrap(),
+            DistanceType::Dot,
+        );
+        let expected_dot_distance = 1.0
+            - indexed_vector[..PERSISTED_DIM]
+                .iter()
+                .zip(&query[..PERSISTED_DIM])
+                .map(|(left, right)| left * right)
+                .sum::<f32>();
+        let distances = dot
+            .compute_distances(&Float32Array::from(query), &code)
+            .unwrap();
+        assert_relative_eq!(distances.value(0), expected_dot_distance, epsilon = 1e-4);
+    }
+
     #[test]
     fn test_pq_transform() {
         const DIM: usize = 16;

rust/lance-index/src/vector/pq/distance.rs

@@ -42,7 +42,13 @@ pub fn build_distance_table_l2_impl<const NUM_BITS: u32, T: L2>(
     let sub_vector_length = dimension / num_sub_vectors;
     let num_centroids = 2_usize.pow(NUM_BITS);
     let mut result = Vec::with_capacity(num_sub_vectors * num_centroids);
-    for (i, sub_vec) in query.chunks_exact(sub_vector_length).enumerate() {
+    // Legacy writers allowed non-divisible dimensions and truncated the tail.
+    // Limit iteration to the sub-vectors that were persisted by those writers.
+    for (i, sub_vec) in query
+        .chunks_exact(sub_vector_length)
+        .take(num_sub_vectors)
+        .enumerate()
+    {
         let subvec_centroids =
             get_sub_vector_centroids::<NUM_BITS, _>(codebook, dimension, num_sub_vectors, i);
         result.extend(l2_distance_batch(
@@ -63,8 +69,14 @@ pub fn build_distance_table_l2_prepared(l2_targets: &[L2Prepared], query: &[f32]
     let num_targets = l2_targets[0].num_targets();
 
     let mut result = vec![0.0f32; l2_targets.len() * num_targets];
-    for (i, sub_vec) in query.chunks_exact(sub_dim).enumerate() {
-        l2_targets[i].distances_into(sub_vec, &mut result[i * num_targets..][..num_targets]);
+    // The target count also bounds legacy codebooks whose writers truncated
+    // a non-divisible vector tail.
+    for (i, (target, sub_vec)) in l2_targets
+        .iter()
+        .zip(query.chunks_exact(sub_dim))
+        .enumerate()
+    {
+        target.distances_into(sub_vec, &mut result[i * num_targets..][..num_targets]);
     }
     result
 }
@@ -94,7 +106,13 @@ pub fn build_distance_table_dot_impl<const NUM_BITS: u32, T: Dot>(
     let sub_vector_length = dimension / num_sub_vectors;
     let num_centroids = 2_usize.pow(NUM_BITS);
     let mut result = Vec::with_capacity(num_sub_vectors * num_centroids);
-    for (i, sub_vec) in query.chunks_exact(sub_vector_length).enumerate() {
+    // Legacy writers allowed non-divisible dimensions and truncated the tail.
+    // Limit iteration to the sub-vectors that were persisted by those writers.
+    for (i, sub_vec) in query
+        .chunks_exact(sub_vector_length)
+        .take(num_sub_vectors)
+        .enumerate()
+    {
         let subvec_centroids =
             get_sub_vector_centroids::<NUM_BITS, _>(codebook, dimension, num_sub_vectors, i);
         result.extend(dot_distance_batch(

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

@@ -5587,6 +5587,34 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_legacy_non_divisible_pq_search() {
+        const DIM: usize = 64;
+        const PERSISTED_DIM: usize = 56;
+
+        let test_dir = copy_test_data_to_tmp("v0.10.15/non_divisible_pq").unwrap();
+        let dataset = Dataset::open(&test_dir.path_str()).await.unwrap();
+        let query = Float32Array::from(
+            (1..=DIM)
+                .map(|value| value as f32 + if value <= PERSISTED_DIM { 1.0 } else { 1_000.0 })
+                .collect::<Vec<_>>(),
+        );
+
+        let result = dataset
+            .scan()
+            .nearest("vector", &query, 1)
+            .unwrap()
+            .try_into_batch()
+            .await
+            .unwrap();
+
+        assert_eq!(result.num_rows(), 1);
+        assert_eq!(
+            result[DIST_COL].as_primitive::<Float32Type>().values(),
+            &[PERSISTED_DIM as f32]
+        );
+    }
+
     #[tokio::test]
     async fn test_pq_storage_backwards_compat() {
         let test_dir = copy_test_data_to_tmp("v0.27.1/pq_in_schema").unwrap();

test_data/readme.md

@@ -27,6 +27,10 @@ folder contains a `datagen.py` script that generates one or more lance datasets.
   correctly, so there are duplicate field ids in the schema. There aren't great
   workarounds for readers. Writers should make sure to check the field ids in
   the schema and re-compute them if necessary.
+* `v0.10.15/non_divisible_pq`: This dataset has an 8-bit IVF-PQ index whose
+  64-dimensional vectors were divided into 14 sub-vectors. Writers at this
+  version silently omitted the final eight dimensions from the PQ codebook.
+  Readers should preserve that prefix-only search behavior.
 * `v0.27.1/pq_in_schema`: This dataset uses the old method of storing the PQ
   metadata in the schema metadata in the index file. We switched to storing them
   in a global buffer in https://github.com/lancedb/lance/pull/3829, but still

test_data/v0.10.15/datagen.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The Lance Authors
+
+import shutil
+
+import lance
+import numpy as np
+import pyarrow as pa
+
+# To generate the test file, we should be running this version of lance.
+assert lance.__version__ == "0.10.15"
+
+name = "non_divisible_pq"
+dimension = 64
+num_sub_vectors = 14
+sub_vector_dimension = dimension // num_sub_vectors
+
+shutil.rmtree(name, ignore_errors=True)
+
+vector = np.arange(1, dimension + 1, dtype=np.float32)
+data = pa.table(
+    {
+        "id": pa.array([0]),
+        "vector": pa.FixedSizeListArray.from_arrays(pa.array(vector), dimension),
+    }
+)
+dataset = lance.write_dataset(data, name)
+
+ivf_centroids = np.zeros((1, dimension), dtype=np.float32)
+persisted_prefix = vector[: num_sub_vectors * sub_vector_dimension].reshape(
+    num_sub_vectors, sub_vector_dimension
+)
+pq_codebook = np.repeat(persisted_prefix[:, np.newaxis, :], 256, axis=1)
+dataset.create_index(
+    "vector",
+    "IVF_PQ",
+    metric="l2",
+    num_partitions=1,
+    ivf_centroids=ivf_centroids,
+    num_sub_vectors=num_sub_vectors,
+    pq_codebook=pq_codebook,
+)