fix(index): enforce the precomputed-input contracts in the current IVF build path

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<FixedSizeListArray> {
+        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<String>) {
+        (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();
+    }
+}

rust/lance/src/index/vector/builder.rs

@@ -385,6 +385,13 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
         sub_index_params: S::BuildParams,
         frag_reuse_index: Option<Arc<CompactFragReuseIndex>>,
     ) -> Result<Self> {
+        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::<FlatIndex, FlatQuantizer>::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<arrow_schema::Schema>,
         num_rows: usize,