diff --git a/rust/lance-index/src/vector/bq/transform.rs b/rust/lance-index/src/vector/bq/transform.rs index c87695e14cd..bd522ee7c10 100644 --- a/rust/lance-index/src/vector/bq/transform.rs +++ b/rust/lance-index/src/vector/bq/transform.rs @@ -356,8 +356,15 @@ impl Transformer for RQTransformer { DistanceType::L2 => res_norm_square.clone(), DistanceType::Dot => { // for dot, the add factor is `1 - v*c + |c|^2 = dist_v_c + |c|^2` - let part_ids = &batch[PART_ID_COLUMN]; - let part_ids = part_ids.as_primitive::(); + let part_ids = batch + .column_by_name(PART_ID_COLUMN) + .ok_or_else(|| { + Error::index(format!( + "RQ Transform: column {} not found in batch", + PART_ID_COLUMN + )) + })? + .as_primitive::(); let centroids_norm_square = self.centroids_norm_square.as_ref().ok_or( Error::index("RQ Transform: centroids norm square not found".to_string()), )?; @@ -427,7 +434,15 @@ impl Transformer for RQTransformer { )); } - let part_ids = batch[PART_ID_COLUMN].as_primitive::(); + let part_ids = batch + .column_by_name(PART_ID_COLUMN) + .ok_or_else(|| { + Error::index(format!( + "RQ Transform: column {} not found in batch", + PART_ID_COLUMN + )) + })? + .as_primitive::(); let rotated_centroids = self.rotated_centroids.as_ref().ok_or_else(|| { Error::internal("RabitQ raw-query transformer is missing rotated centroids") })?; @@ -500,12 +515,50 @@ mod tests { use crate::vector::bq::storage::RABIT_BLOCKED_EX_CODE_COLUMN; use crate::vector::transform::Transformer; use crate::vector::{CENTROID_DIST_COLUMN, PART_ID_COLUMN}; + use lance_core::Error; use super::{ ADD_FACTORS_COLUMN, ERROR_FACTORS_COLUMN, EX_ADD_FACTORS_COLUMN, EX_SCALE_FACTORS_COLUMN, RQTransformer, compute_raw_query_factors, error_factor_value, }; + /// `batch[PART_ID_COLUMN]` panics when the column is absent, and the + /// transformer is public, so a caller that forgets the partition column + /// gets a panic out of a library call instead of the same + /// column-not-found error the vector and centroid-distance columns give. + #[test] + fn test_rq_transformer_reports_missing_partition_column() { + let rq = RabitQuantizer::new_with_rotation::(4, 8, RQRotationType::Fast); + let centroids = + FixedSizeListArray::try_new_from_values(Float32Array::from(vec![0.0f32; 8]), 8) + .unwrap(); + let transformer = RQTransformer::new(rq, DistanceType::L2, centroids, "vector").unwrap(); + + let residual_vectors = FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![1.0f32, -2.0, 3.0, -4.0, 1.5, -2.5, 3.5, -4.5]), + 8, + ) + .unwrap(); + let batch = RecordBatch::try_from_iter(vec![ + ("vector", Arc::new(residual_vectors) as ArrayRef), + ( + CENTROID_DIST_COLUMN, + Arc::new(Float32Array::from(vec![73.0f32])) as ArrayRef, + ), + ]) + .unwrap(); + + let err = transformer.transform(&batch).unwrap_err(); + assert!( + matches!(err, Error::Index { .. }), + "expected an Index error, got: {err:?}" + ); + assert!( + err.to_string().contains(PART_ID_COLUMN), + "the error should name the missing column, got: {err}" + ); + } + #[test] fn test_rq_transformer_writes_multi_bit_ex_factors() { let rq = RabitQuantizer::new_with_rotation::(4, 8, RQRotationType::Fast);