From 064c0bb8b3b04616c4f0d7b7b011b6501ce5a14e Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 11 Sep 2026 23:47:42 +0800 Subject: [PATCH 1/2] fix(index): enforce precomputed buffers contracts in IvfBuildParams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The field docs promise that precomputed_shuffle_buffers is mutually exclusive with precomputed_partitions_file and requires centroids, but neither was validated. Combining both silently used two different assignment sources, and building without centroids trained fresh ones while the buffers were produced against different centroids — silently wrong search results. Enforce both contracts at the start of build(). Assisted-by: GLM-5.3 --- rust/lance/src/index/vector/builder.rs | 42 ++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index a7525529db4..37a7b279d9d 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -487,6 +487,28 @@ impl IvfIndexBuilder // build the index and return the files created by the writer. pub async fn build(&mut self) -> Result { + // Enforce the documented contracts on IvfBuildParams that the Python + // layer cannot check. + if let Some(ivf_params) = &self.ivf_params + && ivf_params.precomputed_shuffle_buffers.is_some() + { + if ivf_params.precomputed_partitions_file.is_some() { + return Err(Error::invalid_input( + "precomputed_shuffle_buffers and precomputed_partitions_file are \ + mutually exclusive; both were set" + .to_string(), + )); + } + if ivf_params.centroids.is_none() { + return Err(Error::invalid_input( + "precomputed_shuffle_buffers requires centroids to be set; the buffers \ + were produced against specific centroids and cannot be combined with \ + freshly trained ones" + .to_string(), + )); + } + } + let progress = self.progress.clone(); // step 1. train IVF & quantizer @@ -3129,6 +3151,26 @@ pub(crate) fn index_type_string(sub_index: SubIndexType, quantizer: Quantization #[cfg(test)] mod tests { + + #[test] + fn test_precomputed_buffers_contract_validation() { + // The documented contracts on IvfBuildParams are enforced at the + // start of build(); verify the rejection logic directly. + let mut params = IvfBuildParams::new(4); + params.precomputed_shuffle_buffers = Some((object_store::path::Path::from("/tmp"), vec![])); + params.precomputed_partitions_file = Some("/tmp/parts".to_string()); + // Both set -> mutually exclusive error (validated in build(), which + // requires a dataset; the check itself is pure field inspection). + assert!( + params.precomputed_shuffle_buffers.is_some() + && params.precomputed_partitions_file.is_some() + ); + + let mut params = IvfBuildParams::new(4); + params.precomputed_shuffle_buffers = Some((object_store::path::Path::from("/tmp"), vec![])); + // No centroids -> requires-centroids error + assert!(params.precomputed_shuffle_buffers.is_some() && params.centroids.is_none()); + } use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; From 8e0c8b64d42d309edbd7bb19dd594436d8481108 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 18 Sep 2026 16:45:09 +0800 Subject: [PATCH 2/2] fix(index): enforce the precomputed-input contracts in the current IVF build path --- rust/lance-index/src/vector/ivf/builder.rs | 110 +++++++++++++++++++++ rust/lance/src/index/vector/builder.rs | 91 +++++++++-------- 2 files changed, 159 insertions(+), 42 deletions(-) diff --git a/rust/lance-index/src/vector/ivf/builder.rs b/rust/lance-index/src/vector/ivf/builder.rs index b8b6e0ec4cc..aa59a18c8f6 100644 --- a/rust/lance-index/src/vector/ivf/builder.rs +++ b/rust/lance-index/src/vector/ivf/builder.rs @@ -140,6 +140,37 @@ impl IvfBuildParams { ..Default::default() }) } + + /// Check the field combinations the precomputed inputs above declare. + /// + /// Both precomputed inputs carry partition ids that were assigned against + /// one particular set of centroids, so neither says anything about an index + /// whose centroids are trained from the data instead, and supplying both at + /// once means one of the two assignments is silently unused. + pub fn validate(&self) -> Result<()> { + if self.precomputed_shuffle_buffers.is_some() && self.precomputed_partitions_file.is_some() + { + return Err(Error::invalid_input( + "precomputed_shuffle_buffers and precomputed_partitions_file are mutually \ + exclusive, but both were set", + )); + } + if self.centroids.is_none() { + if self.precomputed_shuffle_buffers.is_some() { + return Err(Error::invalid_input( + "precomputed_shuffle_buffers requires centroids to be set: the buffers hold \ + partition ids assigned against the centroids they were built with", + )); + } + if self.precomputed_partitions_file.is_some() { + return Err(Error::invalid_input( + "precomputed_partitions_file requires centroids to be set: the file holds \ + partition ids assigned against the centroids it was built with", + )); + } + } + Ok(()) + } } pub fn recommended_num_partitions(num_rows: usize, target_partition_size: usize) -> usize { @@ -180,3 +211,82 @@ pub async fn load_precomputed_partitions( Ok(partition_lookup) } + +#[cfg(test)] +mod tests { + use arrow_array::Float32Array; + use lance_arrow::FixedSizeListArrayExt; + use rstest::rstest; + + use super::*; + + fn centroids(num_partitions: usize) -> Arc { + let values = Float32Array::from(vec![0.0_f32; num_partitions * 2]); + Arc::new(FixedSizeListArray::try_new_from_values(values, 2).unwrap()) + } + + fn buffers() -> (Path, Vec) { + (Path::from("buffers/data"), vec!["buffer1.lance".to_owned()]) + } + + #[rstest] + #[case::buffers_and_partitions_file(true, true, true, "mutually exclusive")] + #[case::buffers_without_centroids( + true, + false, + false, + "precomputed_shuffle_buffers requires centroids" + )] + #[case::partitions_file_without_centroids( + false, + true, + false, + "precomputed_partitions_file requires centroids" + )] + fn test_validate_rejects_precomputed_inputs( + #[case] with_buffers: bool, + #[case] with_partitions_file: bool, + #[case] with_centroids: bool, + #[case] expected: &str, + ) { + let mut params = IvfBuildParams::new(2); + if with_buffers { + params.precomputed_shuffle_buffers = Some(buffers()); + } + if with_partitions_file { + params.precomputed_partitions_file = Some("partitions.lance".to_owned()); + } + if with_centroids { + params.centroids = Some(centroids(2)); + } + + let err = params.validate().unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput, got: {err:?}" + ); + assert!( + err.to_string().contains(expected), + "expected the message to mention {expected:?}, got: {err}" + ); + } + + #[rstest] + #[case::nothing_precomputed(false, false)] + #[case::buffers_with_centroids(true, false)] + #[case::partitions_file_with_centroids(false, true)] + fn test_validate_accepts_supported_combinations( + #[case] with_buffers: bool, + #[case] with_partitions_file: bool, + ) { + let mut params = IvfBuildParams::try_with_centroids(2, centroids(2)).unwrap(); + if with_buffers { + params.precomputed_shuffle_buffers = Some(buffers()); + } + if with_partitions_file { + params.precomputed_partitions_file = Some("partitions.lance".to_owned()); + } + + params.validate().unwrap(); + } +} diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index 37a7b279d9d..86f3ce1e70d 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -385,6 +385,13 @@ impl IvfIndexBuilder sub_index_params: S::BuildParams, frag_reuse_index: Option>, ) -> Result { + if let Some(ivf_params) = ivf_params.as_ref() { + // The legacy IVF_PQ writer checks these combinations in + // `sanity_check_ivf_params`; this path had no equivalent, so a + // precomputed input paired with trained centroids was accepted and + // produced an index whose partition ids belong to other centroids. + ivf_params.validate()?; + } let temp_dir = TempStdDir::default(); let temp_dir_path = Path::from_filesystem_path(&temp_dir)?; let format_version = dataset_format_version(&dataset); @@ -487,28 +494,6 @@ impl IvfIndexBuilder // build the index and return the files created by the writer. pub async fn build(&mut self) -> Result { - // Enforce the documented contracts on IvfBuildParams that the Python - // layer cannot check. - if let Some(ivf_params) = &self.ivf_params - && ivf_params.precomputed_shuffle_buffers.is_some() - { - if ivf_params.precomputed_partitions_file.is_some() { - return Err(Error::invalid_input( - "precomputed_shuffle_buffers and precomputed_partitions_file are \ - mutually exclusive; both were set" - .to_string(), - )); - } - if ivf_params.centroids.is_none() { - return Err(Error::invalid_input( - "precomputed_shuffle_buffers requires centroids to be set; the buffers \ - were produced against specific centroids and cannot be combined with \ - freshly trained ones" - .to_string(), - )); - } - } - let progress = self.progress.clone(); // step 1. train IVF & quantizer @@ -3151,26 +3136,6 @@ pub(crate) fn index_type_string(sub_index: SubIndexType, quantizer: Quantization #[cfg(test)] mod tests { - - #[test] - fn test_precomputed_buffers_contract_validation() { - // The documented contracts on IvfBuildParams are enforced at the - // start of build(); verify the rejection logic directly. - let mut params = IvfBuildParams::new(4); - params.precomputed_shuffle_buffers = Some((object_store::path::Path::from("/tmp"), vec![])); - params.precomputed_partitions_file = Some("/tmp/parts".to_string()); - // Both set -> mutually exclusive error (validated in build(), which - // requires a dataset; the check itself is pure field inspection). - assert!( - params.precomputed_shuffle_buffers.is_some() - && params.precomputed_partitions_file.is_some() - ); - - let mut params = IvfBuildParams::new(4); - params.precomputed_shuffle_buffers = Some((object_store::path::Path::from("/tmp"), vec![])); - // No centroids -> requires-centroids error - assert!(params.precomputed_shuffle_buffers.is_some() && params.centroids.is_none()); - } use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; @@ -3614,6 +3579,48 @@ mod tests { .collect() } + /// The V3 builder is the current write path, so the contract the legacy + /// IVF_PQ writer enforces has to hold here too: precomputed partition ids + /// are only meaningful next to the centroids they were assigned against. + #[tokio::test] + async fn test_new_rejects_precomputed_buffers_without_centroids() { + use lance_index::vector::v3::shuffler::IvfShuffler; + + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_str().unwrap(); + let dataset = write_clusters(uri, &[(8, 0.0)]).await; + let index_dir = dataset.indices_dir().join("idx"); + + let mut ivf_params = IvfBuildParams::new(1); + ivf_params.precomputed_shuffle_buffers = + Some((Path::from("buffers/data"), vec!["buffer1.lance".to_owned()])); + + let builder = IvfIndexBuilder::::new( + dataset, + "vec".to_owned(), + index_dir.clone(), + DistanceType::L2, + Box::new(IvfShuffler::new(index_dir, 1)), + Some(ivf_params), + Some(()), + (), + None, + ); + + let Err(err) = builder else { + panic!("expected the constructor to reject the params"); + }; + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected InvalidInput, got: {err:?}" + ); + assert!( + err.to_string() + .contains("precomputed_shuffle_buffers requires centroids"), + "unexpected message: {err}" + ); + } + fn cluster_batch( schema: &Arc, num_rows: usize,