diff --git a/rust/lance-index/src/vector/flat/storage.rs b/rust/lance-index/src/vector/flat/storage.rs index ff8756cbe41..ca81e66826b 100644 --- a/rust/lance-index/src/vector/flat/storage.rs +++ b/rust/lance-index/src/vector/flat/storage.rs @@ -53,6 +53,14 @@ impl QuantizerStorage for FlatFloatStorage { distance_type: DistanceType, frag_reuse_index: Option>, ) -> Result { + // For cosine distance, vectors have already been normalized by the IVF + // transform pipeline, so cosine distance is equivalent to dot distance + // over normalized vectors. Use dot to avoid redundant norm computation. + let distance_type = match distance_type { + DistanceType::Cosine => DistanceType::Dot, + DistanceType::L2 | DistanceType::Dot | DistanceType::Hamming => distance_type, + }; + let batch = if let Some(frag_reuse_index_ref) = frag_reuse_index.as_ref() { frag_reuse_index_ref.remap_row_ids_record_batch(batch, 0)? } else { @@ -534,9 +542,10 @@ impl DistCalculator for FlatFloatDistanceCalc<'_> { mod tests { use super::*; - use arrow_array::{Float16Array, Float64Array}; + use arrow_array::{Float16Array, Float32Array, Float64Array}; use half::f16; use lance_arrow::FixedSizeListArrayExt; + use rstest::rstest; fn make_f16_storage() -> FlatFloatStorage { let values = Float16Array::from(vec![ @@ -583,4 +592,75 @@ mod tests { assert_eq!(distances[0], 0.0); assert!((distances[1] - 25.0).abs() < 1e-6); } + + fn make_flat_test_batch(vectors: FixedSizeListArray) -> RecordBatch { + let num_rows = vectors.len(); + RecordBatch::try_from_iter(vec![ + ( + ROW_ID, + Arc::new(UInt64Array::from_iter_values(0..num_rows as u64)) as ArrayRef, + ), + (FLAT_COLUMN, Arc::new(vectors) as ArrayRef), + ]) + .unwrap() + } + + #[test] + fn test_try_from_batch_converts_cosine_to_dot() { + // For cosine distance, vectors are normalized by the IVF transform + // pipeline, so cosine is equivalent to dot over normalized vectors. + // try_from_batch should convert Cosine to Dot to avoid redundant + // norm computation during distance calculation. + let values = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]); + let vectors = FixedSizeListArray::try_new_from_values(values, 4).unwrap(); + let batch = make_flat_test_batch(vectors); + let metadata = FlatMetadata { dim: 4 }; + + let storage = + FlatFloatStorage::try_from_batch(batch, &metadata, DistanceType::Cosine, None).unwrap(); + + assert_eq!(storage.distance_type(), DistanceType::Dot); + } + + #[rstest] + #[case::l2(DistanceType::L2)] + #[case::dot(DistanceType::Dot)] + fn test_try_from_batch_keeps_non_cosine_distance_types(#[case] distance_type: DistanceType) { + let values = Float32Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]); + let vectors = FixedSizeListArray::try_new_from_values(values, 4).unwrap(); + let batch = make_flat_test_batch(vectors); + let metadata = FlatMetadata { dim: 4 }; + + let storage = + FlatFloatStorage::try_from_batch(batch, &metadata, distance_type, None).unwrap(); + assert_eq!(storage.distance_type(), distance_type); + } + + #[test] + fn test_cosine_dot_equivalence_on_normalized_vectors() { + // For unit-normalized vectors, cosine distance and dot distance + // produce identical values: cosine = 1 - dot/(|x||y|) = 1 - dot + // (when |x| = |y| = 1), and dot_distance = 1 - dot. This equivalence + // is the correctness premise for converting Cosine to Dot. + let dim: i32 = 4; + let v1: Vec = vec![0.5, 0.5, 0.5, 0.5]; // norm = 1.0 + let v2: Vec = vec![1.0, 0.0, 0.0, 0.0]; // norm = 1.0 + let values = Float32Array::from([v1, v2].concat()); + let vectors = FixedSizeListArray::try_new_from_values(values, dim).unwrap(); + + let query: ArrayRef = Arc::new(Float32Array::from(vec![0.5, 0.5, 0.5, 0.5])); + + let cosine_storage = FlatFloatStorage::new(vectors.clone(), DistanceType::Cosine); + let dot_storage = FlatFloatStorage::new(vectors, DistanceType::Dot); + + let cosine_dist = cosine_storage + .dist_calculator(query.clone(), 0.0) + .distance_all(2); + let dot_dist = dot_storage.dist_calculator(query, 0.0).distance_all(2); + + assert_eq!(cosine_dist.len(), dot_dist.len()); + for (c, d) in cosine_dist.iter().zip(dot_dist.iter()) { + assert!((c - d).abs() < 1e-6, "cosine {} != dot {}", c, d); + } + } } diff --git a/rust/lance/src/index/vector/ivf/partition_serde.rs b/rust/lance/src/index/vector/ivf/partition_serde.rs index ad737620a94..fa5430ace93 100644 --- a/rust/lance/src/index/vector/ivf/partition_serde.rs +++ b/rust/lance/src/index/vector/ivf/partition_serde.rs @@ -742,7 +742,14 @@ mod tests { }; let bytes = ser_body(&entry); let restored = de_body::>(bytes).unwrap(); - assert_eq!(restored.storage.distance_type(), dt); + // Cosine is converted to Dot during deserialization (try_from_batch) + // because vectors are normalized by the IVF transform pipeline, + // making cosine equivalent to dot over normalized vectors. + let expected_dt = match dt { + DistanceType::Cosine => DistanceType::Dot, + other => other, + }; + assert_eq!(restored.storage.distance_type(), expected_dt); } }