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 a7525529db4..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); @@ -3572,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,