From d4dfdafd2b3122bc2497f2724540dbbdc1320a98 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 29 Aug 2026 23:52:09 +0800 Subject: [PATCH] fix(index): reject RabitQ indices whose dimension is not a multiple of 8 --- rust/lance-index/src/vector/bq/storage.rs | 49 +++++++++++++++++++++++ rust/lance-linalg/src/simd/dist_table.rs | 1 + 2 files changed, 50 insertions(+) diff --git a/rust/lance-index/src/vector/bq/storage.rs b/rust/lance-index/src/vector/bq/storage.rs index b96517d03b3..b54ae3b5113 100644 --- a/rust/lance-index/src/vector/bq/storage.rs +++ b/rust/lance-index/src/vector/bq/storage.rs @@ -2429,6 +2429,24 @@ impl QuantizerStorage for RabitQuantizationStorage { _ => distance_type, }; validate_rq_num_bits(metadata.num_bits)?; + // The FastScan LUT is `4 * rotated_dim` bytes while the kernels index it + // as `BATCH_SIZE * rotated_dim.div_ceil(8)`, so the two agree only when + // the dimension is a multiple of 8. `RabitQuantizer::build` has rejected + // a non-multiple since #6024, but an index written before that still + // loads here, and the AVX-512, AVX2 and NEON kernels read the LUT + // through unchecked raw pointers. Only the scalar fallback panics. + // + // This has to go through `rotated_dim()`, not `metadata.code_dim`: + // `code_dim` was added by #6024 itself, so it deserializes to 0 for the + // very indices this rejects, and `rotated_dim()` recovers the real + // dimension from the rotation matrix that `parse_buffer` backfills. + let rotated_dim = metadata.rotated_dim(); + if rotated_dim % 8 != 0 { + return Err(Error::invalid_input(format!( + "RabitQ vector dimension must be divisible by 8, got {rotated_dim}. \ + Rebuild the index." + ))); + } let row_ids = batch[ROW_ID].as_primitive::().clone(); let codes = batch[RABIT_CODE_COLUMN].as_fixed_size_list().clone(); let expected_code_bytes = metadata.binary_code_bytes(); @@ -2982,6 +3000,37 @@ mod tests { .metadata(None) } + /// Indices written before #6024 could carry a `code_dim` that is not a + /// multiple of 8. The FastScan LUT is sized `4 * code_dim` while the kernels + /// index it as `BATCH_SIZE * code_dim.div_ceil(8)`, so loading one made the + /// kernels read past the LUT through raw pointers. + #[test] + fn test_try_from_batch_rejects_dim_not_multiple_of_eight() { + let code_dim = 12usize; + let metadata = make_test_metadata(code_dim); + assert_eq!(metadata.code_dim as usize, code_dim); + let code_bytes = metadata.binary_code_bytes(); + let codes = FixedSizeListArray::try_new_from_values( + UInt8Array::from(vec![0u8; 2 * code_bytes]), + code_bytes as i32, + ) + .unwrap(); + let batch = make_test_batch(codes); + + let err = + RabitQuantizationStorage::try_from_batch(batch, &metadata, DistanceType::L2, None) + .expect_err("a dimension that is not a multiple of 8 must be rejected"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "unexpected variant: {err:?}" + ); + let message = err.to_string(); + assert!( + message.contains("divisible by 8") && message.contains("12"), + "unexpected error: {message}" + ); + } + #[test] fn test_rabit_metadata_defaults_old_indexes_to_residual_query() { let metadata: RabitQuantizationMetadata = serde_json::from_str( diff --git a/rust/lance-linalg/src/simd/dist_table.rs b/rust/lance-linalg/src/simd/dist_table.rs index 00bc9143cf0..1d46da1d878 100644 --- a/rust/lance-linalg/src/simd/dist_table.rs +++ b/rust/lance-linalg/src/simd/dist_table.rs @@ -67,6 +67,7 @@ pub unsafe fn sum_4bit_dist_table_uninit( debug_assert!(n.is_multiple_of(BATCH_SIZE)); debug_assert!(dists.len() >= n); debug_assert!(codes.len() >= n * code_len); + debug_assert!(dist_table.len() >= BATCH_SIZE * code_len); match *SIMD_SUPPORT { #[cfg(all(kernel_support = "avx512_dist_table", target_arch = "x86_64"))]