diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index 5749e56ed31..4f6ba450f7e 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/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::>(); + 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::>(); + 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::::from( + codebook + .iter() + .map(|value| *value as f64) + .collect::>(), + ), + DIM as i32, + ) + .unwrap(), + DistanceType::L2, + ); + assert!(generic_l2.l2_targets.is_none()); + let distances = generic_l2 + .compute_distances( + &PrimitiveArray::::from( + query.iter().map(|value| *value as f64).collect::>(), + ), + &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::(); + 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; diff --git a/rust/lance-index/src/vector/pq/distance.rs b/rust/lance-index/src/vector/pq/distance.rs index b341ba98af7..905594ff534 100644 --- a/rust/lance-index/src/vector/pq/distance.rs +++ b/rust/lance-index/src/vector/pq/distance.rs @@ -42,7 +42,13 @@ pub fn build_distance_table_l2_impl( 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::(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( 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::(codebook, dimension, num_sub_vectors, i); result.extend(dot_distance_batch( diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 10df6ed52d7..9c0abbabd25 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/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::>(), + ); + + 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::().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(); diff --git a/test_data/readme.md b/test_data/readme.md index 69f153c22d1..a299c48850d 100644 --- a/test_data/readme.md +++ b/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 diff --git a/test_data/v0.10.15/datagen.py b/test_data/v0.10.15/datagen.py new file mode 100644 index 00000000000..f864cfa0ee8 --- /dev/null +++ b/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, +) diff --git a/test_data/v0.10.15/non_divisible_pq/_indices/be068df8-322d-4309-8347-51afc73e8d3f/index.idx b/test_data/v0.10.15/non_divisible_pq/_indices/be068df8-322d-4309-8347-51afc73e8d3f/index.idx new file mode 100644 index 00000000000..6af76b538b4 Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/_indices/be068df8-322d-4309-8347-51afc73e8d3f/index.idx differ diff --git a/test_data/v0.10.15/non_divisible_pq/_latest.manifest b/test_data/v0.10.15/non_divisible_pq/_latest.manifest new file mode 100644 index 00000000000..3bb9ca51bc7 Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/_latest.manifest differ diff --git a/test_data/v0.10.15/non_divisible_pq/_transactions/0-573045e8-48b7-4662-941d-d30af281d50f.txn b/test_data/v0.10.15/non_divisible_pq/_transactions/0-573045e8-48b7-4662-941d-d30af281d50f.txn new file mode 100644 index 00000000000..c7a0feb2315 Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/_transactions/0-573045e8-48b7-4662-941d-d30af281d50f.txn differ diff --git a/test_data/v0.10.15/non_divisible_pq/_transactions/1-11123412-ba3d-4e23-8edd-5b0d5a1e84f6.txn b/test_data/v0.10.15/non_divisible_pq/_transactions/1-11123412-ba3d-4e23-8edd-5b0d5a1e84f6.txn new file mode 100644 index 00000000000..4a5161722fc Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/_transactions/1-11123412-ba3d-4e23-8edd-5b0d5a1e84f6.txn differ diff --git a/test_data/v0.10.15/non_divisible_pq/_versions/1.manifest b/test_data/v0.10.15/non_divisible_pq/_versions/1.manifest new file mode 100644 index 00000000000..acb27411508 Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/_versions/1.manifest differ diff --git a/test_data/v0.10.15/non_divisible_pq/_versions/2.manifest b/test_data/v0.10.15/non_divisible_pq/_versions/2.manifest new file mode 100644 index 00000000000..3bb9ca51bc7 Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/_versions/2.manifest differ diff --git a/test_data/v0.10.15/non_divisible_pq/data/2f1a33f2-fdcd-40f7-a361-844659965e1e.lance b/test_data/v0.10.15/non_divisible_pq/data/2f1a33f2-fdcd-40f7-a361-844659965e1e.lance new file mode 100644 index 00000000000..8fda283a059 Binary files /dev/null and b/test_data/v0.10.15/non_divisible_pq/data/2f1a33f2-fdcd-40f7-a361-844659965e1e.lance differ