From 1d818e384823ae1d3cfe3129304e55a1eb3c46e1 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sun, 6 Sep 2026 22:25:20 +0800 Subject: [PATCH 1/2] fix(index): validate the width and dtype of pre-computed IVF centroids build_ivf_model checked only the flattened length of supplied centroids, so wrongly-shaped arrays passed validation: a half-width array with twice the rows broke the build far from this boundary, and same-width wrong-dtype centroids (e.g. f16 over an f32 column) panicked in a blind downcast inside kmeans. Reject both at the boundary with the offending values. Assisted-by: GLM-5.3 --- rust/lance/src/index/vector/ivf.rs | 79 ++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index e8d0f9e02da..774a1e5eaa1 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -1699,6 +1699,26 @@ pub async fn build_ivf_model( num_partitions * dim, ))); } + // The flattened-length check alone accepts wrongly-shaped centroids + // (e.g. twice the rows at half the dimension), which then break the + // build far away from this boundary. + if centroids.value_length() as usize != dim { + return Err(Error::invalid_input(format!( + "IVF centroids dimension {} does not match the vector column dimension {}", + centroids.value_length(), + dim + ))); + } + // Same-width centroids of a different dtype (e.g. f16 over an f32 + // column) would later hit a blind downcast inside kmeans. + let (_, element_type) = get_vector_type(dataset.schema(), column)?; + if centroids.value_type() != element_type { + return Err(Error::invalid_input(format!( + "IVF centroids type {} does not match the vector column type {}", + centroids.value_type(), + element_type + ))); + } return Ok(IvfModel::new(centroids.clone(), None)); } let sample_size_hint = num_partitions * params.sample_rate; @@ -6319,6 +6339,65 @@ mod tests { ); } + /// Pre-computed centroids whose flattened length matches but whose + /// per-row width disagrees with the column used to pass validation and + /// break the build far away; reject them at this boundary instead. + #[tokio::test] + async fn test_build_ivf_model_rejects_wrong_width_centroids() { + use lance_index::progress::NoopIndexBuildProgress; + + let test_dir = TempStrDir::default(); + let uri = format!("{}/ds", test_dir.as_str()); + let values = generate_random_array_with_seed::(64 * 16, [22; 32]); + let fsl = FixedSizeListArray::try_new_from_values(values, 16).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "vector", + fsl.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(fsl)]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = Dataset::write(reader, &uri, None).await.unwrap(); + + // 64 centroids of width 8 flatten to the same length as 32 rows of + // width 16, so the flattened check alone lets them through. + let values = generate_random_array_with_seed::(64 * 8, [23; 32]); + let centroids = FixedSizeListArray::try_new_from_values(values, 8).unwrap(); + let mut params = IvfBuildParams::new(32); + params.centroids = Some(Arc::new(centroids)); + + let err = build_ivf_model( + &dataset, + "vector", + 16, + MetricType::L2, + ¶ms, + None, + Arc::new(NoopIndexBuildProgress), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("does not match the vector column")); + + // Same width, different dtype: also rejected at the boundary. + let values = generate_random_array_with_seed::(32 * 16, [24; 32]); + let centroids = FixedSizeListArray::try_new_from_values(values, 16).unwrap(); + let mut params = IvfBuildParams::new(32); + params.centroids = Some(Arc::new(centroids)); + let err = build_ivf_model( + &dataset, + "vector", + 16, + MetricType::L2, + ¶ms, + None, + Arc::new(NoopIndexBuildProgress), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("centroids type")); + } + /// Regression test for a hang in the streaming *coreset* trainer /// (`train_streaming_coreset_ivf_model`, taken when `num_partitions > 256`). /// From fc31af1316fe15caa366e829b1180c6e61a07659 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 18 Sep 2026 13:00:07 +0800 Subject: [PATCH 2/2] fix(index): validate the width and dtype of pre-computed IVF centroids --- rust/lance/src/index/vector/ivf.rs | 65 +++++++++++++++++++++++++++--- 1 file changed, 59 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs index 1551f44e2c0..77f5d66a90f 100644 --- a/rust/lance/src/index/vector/ivf.rs +++ b/rust/lance/src/index/vector/ivf.rs @@ -1760,13 +1760,23 @@ pub async fn build_ivf_model( dim ))); } - // Same-width centroids of a different dtype (e.g. f16 over an f32 - // column) would later hit a blind downcast inside kmeans. + // Same-width centroids of a different dtype would be handed to a + // kmeans assignment pair that does not exist. Int8 columns are the + // exception: the trainer converts them to f32 (see the `Int8` arm of + // `train_ivf_kmeans_step`) and the only assignment pair implemented + // for them is (Float32 centroids, Int8 vectors), so f32 is the + // expected centroid type there. let (_, element_type) = get_vector_type(dataset.schema(), column)?; - if centroids.value_type() != element_type { + let expected_centroid_type = if element_type == DataType::Int8 { + DataType::Float32 + } else { + element_type.clone() + }; + if centroids.value_type() != expected_centroid_type { return Err(Error::invalid_input(format!( - "IVF centroids type {} does not match the vector column type {}", + "IVF centroids type {} does not match the expected type {} for a vector column of type {}", centroids.value_type(), + expected_centroid_type, element_type ))); } @@ -6471,7 +6481,10 @@ mod tests { ) .await .unwrap_err(); - assert!(err.to_string().contains("does not match the vector column")); + assert!( + err.to_string().contains("centroids dimension 8"), + "got: {err}" + ); // Same width, different dtype: also rejected at the boundary. let values = generate_random_array_with_seed::(32 * 16, [24; 32]); @@ -6489,7 +6502,47 @@ mod tests { ) .await .unwrap_err(); - assert!(err.to_string().contains("centroids type")); + assert!(err.to_string().contains("centroids type"), "got: {err}"); + } + + /// An Int8 column is trained as f32, and the only assignment pair the + /// kernels implement for it is (Float32 centroids, Int8 vectors), so f32 + /// centroids must be accepted rather than rejected as a dtype mismatch. + #[tokio::test] + async fn test_build_ivf_model_accepts_f32_centroids_for_int8_column() { + use lance_index::progress::NoopIndexBuildProgress; + + let test_dir = TempStrDir::default(); + let uri = format!("{}/ds", test_dir.as_str()); + let values = + arrow_array::Int8Array::from_iter_values((0..64 * 16).map(|i| (i % 127) as i8)); + let fsl = FixedSizeListArray::try_new_from_values(values, 16).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "vector", + fsl.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(fsl)]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = Dataset::write(reader, &uri, None).await.unwrap(); + + let values = generate_random_array_with_seed::(4 * 16, [25; 32]); + let centroids = FixedSizeListArray::try_new_from_values(values, 16).unwrap(); + let mut params = IvfBuildParams::new(4); + params.centroids = Some(Arc::new(centroids)); + + let model = build_ivf_model( + &dataset, + "vector", + 16, + MetricType::L2, + ¶ms, + None, + Arc::new(NoopIndexBuildProgress), + ) + .await + .expect("f32 centroids are the expected type for an Int8 column"); + assert_eq!(model.num_partitions(), 4); } /// Regression test for a hang in the streaming *coreset* trainer