From 08be5ae34945f43af60c60347bde6fa6e8a539a1 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:32:45 +0000 Subject: [PATCH 1/2] fix(index): search legacy truncated PQ indexes --- rust/lance-index/src/vector/pq.rs | 30 ++++++++++++++++++++++ rust/lance-index/src/vector/pq/distance.rs | 26 ++++++++++++++++--- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index 5749e56ed31..3025f4144a5 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/rust/lance-index/src/vector/pq.rs @@ -615,6 +615,7 @@ mod tests { use lance_linalg::kernels::argmin; use lance_testing::datagen::generate_random_array; use num_traits::Zero; + use rstest::rstest; use storage::transpose; #[test] @@ -690,6 +691,35 @@ mod tests { }); } + #[rstest] + #[case::l2(DistanceType::L2)] + #[case::dot(DistanceType::Dot)] + fn test_distance_with_legacy_truncated_dimension(#[case] distance_type: DistanceType) { + const DIM: usize = 64; + const NUM_SUB_VECTORS: usize = 14; + const NUM_BITS: u32 = 4; + const NUM_CENTROIDS: usize = 1 << NUM_BITS; + const SUB_VECTOR_DIM: usize = DIM / NUM_SUB_VECTORS; + + // 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 codebook = + Float32Array::from(vec![0.0; NUM_SUB_VECTORS * NUM_CENTROIDS * SUB_VECTOR_DIM]); + let pq = ProductQuantizer::new( + NUM_SUB_VECTORS, + NUM_BITS, + DIM, + FixedSizeListArray::try_new_from_values(codebook, DIM as i32).unwrap(), + distance_type, + ); + let query = Float32Array::from(vec![0.0; DIM]); + let code = UInt8Array::from(vec![0; NUM_SUB_VECTORS / 2]); + + let distances = pq.compute_distances(&query, &code).unwrap(); + assert_eq!(distances.len(), 1); + } + #[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( From fa3d326e6ed3c84b48b5aaae899fc33211f516e2 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:36:57 +0000 Subject: [PATCH 2/2] test(index): cover legacy non-divisible PQ index --- rust/lance-index/src/vector/pq.rs | 90 +++++++++++++++--- rust/lance/src/index/vector/ivf/v2.rs | 28 ++++++ test_data/readme.md | 4 + test_data/v0.10.15/datagen.py | 42 ++++++++ .../index.idx | Bin 0 -> 57725 bytes .../non_divisible_pq/_latest.manifest | Bin 0 -> 298 bytes ...0-573045e8-48b7-4662-941d-d30af281d50f.txn | Bin 0 -> 177 bytes ...1-11123412-ba3d-4e23-8edd-5b0d5a1e84f6.txn | Bin 0 -> 102 bytes .../non_divisible_pq/_versions/1.manifest | Bin 0 -> 233 bytes .../non_divisible_pq/_versions/2.manifest | Bin 0 -> 298 bytes ...2f1a33f2-fdcd-40f7-a361-844659965e1e.lance | Bin 0 -> 631 bytes 11 files changed, 149 insertions(+), 15 deletions(-) create mode 100644 test_data/v0.10.15/datagen.py create mode 100644 test_data/v0.10.15/non_divisible_pq/_indices/be068df8-322d-4309-8347-51afc73e8d3f/index.idx create mode 100644 test_data/v0.10.15/non_divisible_pq/_latest.manifest create mode 100644 test_data/v0.10.15/non_divisible_pq/_transactions/0-573045e8-48b7-4662-941d-d30af281d50f.txn create mode 100644 test_data/v0.10.15/non_divisible_pq/_transactions/1-11123412-ba3d-4e23-8edd-5b0d5a1e84f6.txn create mode 100644 test_data/v0.10.15/non_divisible_pq/_versions/1.manifest create mode 100644 test_data/v0.10.15/non_divisible_pq/_versions/2.manifest create mode 100644 test_data/v0.10.15/non_divisible_pq/data/2f1a33f2-fdcd-40f7-a361-844659965e1e.lance diff --git a/rust/lance-index/src/vector/pq.rs b/rust/lance-index/src/vector/pq.rs index 3025f4144a5..4f6ba450f7e 100644 --- a/rust/lance-index/src/vector/pq.rs +++ b/rust/lance-index/src/vector/pq.rs @@ -615,7 +615,6 @@ mod tests { use lance_linalg::kernels::argmin; use lance_testing::datagen::generate_random_array; use num_traits::Zero; - use rstest::rstest; use storage::transpose; #[test] @@ -691,33 +690,94 @@ mod tests { }); } - #[rstest] - #[case::l2(DistanceType::L2)] - #[case::dot(DistanceType::Dot)] - fn test_distance_with_legacy_truncated_dimension(#[case] distance_type: DistanceType) { + #[test] + fn test_distance_with_legacy_truncated_dimension() { const DIM: usize = 64; const NUM_SUB_VECTORS: usize = 14; - const NUM_BITS: u32 = 4; + 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 codebook = - Float32Array::from(vec![0.0; NUM_SUB_VECTORS * NUM_CENTROIDS * SUB_VECTOR_DIM]); - let pq = ProductQuantizer::new( + 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(codebook, DIM as i32).unwrap(), - distance_type, + FixedSizeListArray::try_new_from_values( + Float32Array::from(codebook.clone()), + DIM as i32, + ) + .unwrap(), + DistanceType::L2, ); - let query = Float32Array::from(vec![0.0; DIM]); - let code = UInt8Array::from(vec![0; NUM_SUB_VECTORS / 2]); + 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 distances = pq.compute_distances(&query, &code).unwrap(); - assert_eq!(distances.len(), 1); + 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] 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 0000000000000000000000000000000000000000..6af76b538b445ae5758c6e45fa03acb2db574e52 GIT binary patch literal 57725 zcmeI*p=(uP7zg0*oFj4=BA0{6z&x}M>_Z;l0UqE19^e5UsLunly!qFY&16i}`r$md59Wb+Xdl>zJir4y zzymzM13XZl2X^!3-%Fk+%S5do&V&149+-#rfqlpWJir4yzymzM1NC{}P2T+b$%o`H zQR|2E;69iK=AnIHAMyYX@Bk0+01xm$eIEFfH~;74OEM*D{cs-K2lK!@v=8h<9^e5U z-~k@s0UoH&1K;!J|B?JmR*70aoCo*8JTMRK1N)E%cz_3ZfCqSh2kP^{@4Wf{B!81~ zqZGA%I1lcFd0-ye2lgQk@Bk0+01xm057g&@Q+f0E$(dx9sP)5na39PA^Uyx94|#wG zcz_3ZfCqS>J`bGBn}42MNY)dzemD>AgLz;c+6VR_5AXmF@Bk0+01wpXflGPwUrw$h zn~7RKoCo*8JTMRK1N)E%cz_3ZfCqSh2kP^{)x7zyCD)TNQR|2E;69iK=AnIHAMyYX z@Bk0+01xm$eIB@#H~;NqC)rKZ`r$md59Wb+Xdl>zJir4yzymzM13XZl2kzy~e?NJU s>?Laba30(T^T0f`59~u8-~k@s0UqE19;nR&=krf3j=!3+vwi30KOiO%M*si- literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..3bb9ca51bc779c13366ff9fa7c8e19f3e50c3f2a GIT binary patch literal 298 zcmYk1y-LJD6oqF38+EVpg&%ajn76RR6dB(*0*iHwn4$}BC4EH%s|t?C+sJpdyZdK8q$hoA4;4?p;NetPuk zm=VSp7m{&OD=}IO7bG`YleDHfRm|klW=TJEHvyVWvICm1@zTZq!5tvDJil800%(j> AHUIzs literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..c7a0feb2315b5c12fd1a2642c622462541ed632b GIT binary patch literal 177 zcmYkzy$ZrG6hL89utUd8#i@f+;YxmP(zGvAVv`Gj(n1?Vd;#Cb7t}!n&-@*JHk}q| zq62i-4$xW!J<0$<3s1_)Ulw&;|RXo-yOI-`RO_2{M{@EL;@l+fQ?f1KW57t5Eg(pg&%ajn76RR6dB(*0*iHwn4$}BC4EH%s|t?C+sJpdyZdK8q$hoA4;4?p;NetPuk zm=VSp7m{&OD=}IO7bG`YleDHfRm|klW=TJEHvyVWvICm1@zTZq!5tvDJil800%(j> AHUIzs literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..8fda283a059f9dc877ed3e9f5f6e81f52454892c GIT binary patch literal 631 zcmajbKS%;m9KiAS&Oh1FAV<_A8XOu#A}+EJ3DRn4kZ@>d2=q8ly0}i!Q$!kWbZl&N zY;0_FY;0_FtgW%FzL5vw(vSDyy&u1Kzx(loh$1}VRAM2-7!yn(%0i@(fr%2TsH2S# zBi!QwV?5#s&zRr^uXw{0@A$+Q#IX=@q_BcCR-v$o47QNPHcaGEz#&RF#ThE7;sTer zLLCiU!$TVZ2Dn3rzmT&%wUE!!f4A(cLDa`w$O{X(Lq~nX(A-AmI|`XgxPEYuH)TQU zdq%SF*g>nK