From efcf735212076bb258f2ee90eb2fc21d5ef92ccd Mon Sep 17 00:00:00 2001 From: Vova Kolmakov Date: Wed, 8 Jul 2026 10:37:01 +0700 Subject: [PATCH] test(index): make flaky IVF_RQ recall test deterministic --- rust/lance/src/index/vector/ivf/v2.rs | 185 +++++++++++++++++++++++++- 1 file changed, 182 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 0faa79dce5f..549fefbc4ab 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -1921,8 +1921,8 @@ mod tests { use arrow::datatypes::{Float64Type, UInt8Type, UInt64Type}; use arrow::{array::AsArray, datatypes::Float32Type}; use arrow_array::{ - Array, ArrayRef, ArrowPrimitiveType, FixedSizeListArray, Float32Array, Int64Array, - ListArray, RecordBatch, RecordBatchIterator, UInt64Array, + Array, ArrayRef, ArrowPrimitiveType, FixedSizeListArray, Float32Array, Float64Array, + Int64Array, ListArray, RecordBatch, RecordBatchIterator, UInt64Array, }; use arrow_buffer::OffsetBuffer; use arrow_schema::{DataType, Field, Schema, SchemaRef}; @@ -1943,7 +1943,7 @@ mod tests { use crate::utils::test::copy_test_data_to_tmp; use crate::{ Dataset, - index::vector::{VectorIndex, VectorIndexParams}, + index::vector::{StageParams, VectorIndex, VectorIndexParams}, }; use crate::{ dataset::optimize::{CompactionOptions, compact_files}, @@ -1988,6 +1988,14 @@ mod tests { const NUM_ROWS: usize = 512; const DIM: usize = 32; + /// Seed for the deterministic IVF_RQ rotation used by `determinize_rq_params`. + /// RaBitQ picks a random rotation at build time (unseeded), which makes the + /// recall of `test_build_ivf_rq` non-deterministic and borderline against its + /// `>= 0.5` bar. Pinning the rotation (and the IVF centroids) makes recall + /// reproducible. This is the single knob to bump if a case ever dips below its + /// bar. + const RQ_SEED: u64 = 42; + lance_testing::define_stage_event_progress!(RecordingProgress, IndexBuildProgress, Result<()>); #[test] @@ -3978,6 +3986,171 @@ mod tests { assert_lt!(sq_meta.bounds.end, FRAGMENT_OFFSETS[1] as f64); } + /// Deterministic IVF centroids for the RQ recall tests: partition `k` is the + /// mean of every row whose index satisfies `row_index % nlist == k`. For + /// `nlist == 1` this is the global mean (what k-means converges to for a single + /// cluster); for `nlist == 4` it yields distinct, near-central centroids. The + /// output value type matches the dataset so the IVF stage accepts it verbatim. + fn fixed_rq_centroids(fsl: &FixedSizeListArray, nlist: usize) -> Arc { + let n = fsl.len(); + let mut sums = vec![0.0f64; nlist * DIM]; + let mut counts = vec![0usize; nlist]; + match fsl.value_type() { + DataType::Float32 => { + let values = fsl.values().as_primitive::(); + for row in 0..n { + let k = row % nlist; + counts[k] += 1; + for d in 0..DIM { + sums[k * DIM + d] += values.value(row * DIM + d) as f64; + } + } + } + DataType::Float64 => { + let values = fsl.values().as_primitive::(); + for row in 0..n { + let k = row % nlist; + counts[k] += 1; + for d in 0..DIM { + sums[k * DIM + d] += values.value(row * DIM + d); + } + } + } + dt => unreachable!("RQ centroids: unexpected value type {dt:?}"), + } + for k in 0..nlist { + let count = counts[k].max(1) as f64; + for d in 0..DIM { + sums[k * DIM + d] /= count; + } + } + let centroids = match fsl.value_type() { + DataType::Float32 => FixedSizeListArray::try_new_from_values( + Float32Array::from(sums.iter().map(|&x| x as f32).collect::>()), + DIM as i32, + ) + .unwrap(), + DataType::Float64 => { + FixedSizeListArray::try_new_from_values(Float64Array::from(sums), DIM as i32) + .unwrap() + } + dt => unreachable!("RQ centroids: unexpected value type {dt:?}"), + }; + Arc::new(centroids) + } + + /// A deterministic, genuinely mixing (Haar-like) orthogonal matrix, built from + /// seeded Gaussian entries (Box-Muller) orthonormalized by modified + /// Gram-Schmidt over its rows. Production builds this with an unseeded + /// `random_orthogonal` (private to `lance-index`); we reproduce an equivalent + /// deterministically in-crate. The value type must match the dataset because + /// the quantize path downcasts `rotate_mat` to the dataset's primitive type. + fn deterministic_orthogonal_fsl(value_type: &DataType, seed: u64) -> FixedSizeListArray { + let mut rng = StdRng::seed_from_u64(seed); + let mut matrix = vec![0.0f64; DIM * DIM]; + for entry in matrix.iter_mut() { + let u1: f64 = rng.random::().max(1e-12); + let u2: f64 = rng.random::(); + *entry = (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos(); + } + // Modified Gram-Schmidt over the DIM rows. + for i in 0..DIM { + for j in 0..i { + let dot: f64 = (0..DIM) + .map(|k| matrix[i * DIM + k] * matrix[j * DIM + k]) + .sum(); + for k in 0..DIM { + matrix[i * DIM + k] -= dot * matrix[j * DIM + k]; + } + } + let norm: f64 = (0..DIM) + .map(|k| matrix[i * DIM + k] * matrix[i * DIM + k]) + .sum::() + .sqrt(); + for k in 0..DIM { + matrix[i * DIM + k] /= norm; + } + } + match value_type { + DataType::Float32 => FixedSizeListArray::try_new_from_values( + Float32Array::from(matrix.iter().map(|&x| x as f32).collect::>()), + DIM as i32, + ) + .unwrap(), + DataType::Float64 => { + FixedSizeListArray::try_new_from_values(Float64Array::from(matrix), DIM as i32) + .unwrap() + } + dt => unreachable!("RQ rotate_mat: unexpected value type {dt:?}"), + } + } + + /// Build a fixed `RabitQuantizationMetadata` so the RQ quantizer reuses this + /// rotation instead of generating a fresh random one. Mirrors the field layout + /// of `RabitQuantizer::new_with_rotation`. + fn fixed_rq_rotation_metadata( + rotation_type: RQRotationType, + num_bits: u8, + value_type: &DataType, + seed: u64, + ) -> RabitQuantizationMetadata { + match rotation_type { + RQRotationType::Fast => { + // 4 rounds, one sign bit per dimension per round; see + // `fast_rotation_signs_len` in lance-index bq/rotation.rs. + let mut rng = StdRng::seed_from_u64(seed); + let mut signs = vec![0u8; 4 * DIM.div_ceil(8)]; + rng.fill(&mut signs[..]); + RabitQuantizationMetadata { + rotate_mat: None, + rotate_mat_position: None, + fast_rotation_signs: Some(signs), + rotation_type: RQRotationType::Fast, + code_dim: DIM as u32, + num_bits, + packed: false, + query_estimator: RabitQueryEstimator::RawQuery, + } + } + RQRotationType::Matrix => RabitQuantizationMetadata { + rotate_mat: Some(deterministic_orthogonal_fsl(value_type, seed)), + rotate_mat_position: None, + fast_rotation_signs: None, + rotation_type: RQRotationType::Matrix, + code_dim: DIM as u32, + num_bits, + packed: false, + query_estimator: RabitQueryEstimator::RawQuery, + }, + } + } + + /// Make an IVF_RQ build deterministic by pinning both the IVF centroids and the + /// RaBitQ rotation via the existing build-param seams. Non-RQ params are + /// returned unchanged, so the SQ/PQ/flat/HNSW callers of the shared test + /// helpers are unaffected. + fn determinize_rq_params( + params: &VectorIndexParams, + nlist: usize, + fsl: &FixedSizeListArray, + ) -> VectorIndexParams { + let Some(StageParams::RQ(rq)) = params.stages.last() else { + return params.clone(); + }; + let mut rq = rq.clone(); + rq.rotation = Some(fixed_rq_rotation_metadata( + rq.rotation_type, + rq.num_bits, + &fsl.value_type(), + RQ_SEED, + )); + let ivf = + IvfBuildParams::try_with_centroids(nlist, fixed_rq_centroids(fsl, nlist)).unwrap(); + let mut params = params.clone(); + params.stages = vec![StageParams::Ivf(ivf), StageParams::RQ(rq)]; + params + } + async fn test_index( params: VectorIndexParams, nlist: usize, @@ -4028,6 +4201,7 @@ mod tests { Some((dataset, vectors)) => (dataset, vectors), None => generate_test_dataset::(test_uri, range).await, }; + let params = determinize_rq_params(¶ms, nlist, &vectors); let vector_column = "vector"; dataset @@ -4111,6 +4285,7 @@ mod tests { let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); let (mut dataset, vectors) = generate_test_dataset::(test_uri, range.clone()).await; + let params = determinize_rq_params(¶ms, nlist, &vectors); let vector_column = "vector"; dataset @@ -4685,6 +4860,10 @@ mod tests { let test_uri = test_dir.as_str(); let (mut dataset, vectors) = generate_multivec_test_dataset::(test_uri, range).await; + // The multivec column is a List>; centroids/rotation are + // computed over the inner per-vector FixedSizeList. + let inner_vectors = vectors.values(); + let params = determinize_rq_params(¶ms, nlist, inner_vectors.as_fixed_size_list()); dataset .create_index(