From 7ee12821e63d2ca25b2675c2bbfc47864c745ed0 Mon Sep 17 00:00:00 2001 From: lihongyu <1162914749@qq.com> Date: Tue, 14 Jul 2026 14:44:24 +0800 Subject: [PATCH] perf(index): use dot distance for cosine in FlatFloatStorage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For IVF indexes with cosine distance, vectors are already normalized by the IVF transform pipeline (NormalizeTransformer), so cosine distance is mathematically equivalent to dot distance over normalized vectors. However, cosine still computes the L2 norm of each vector during distance calculation — redundant work that dot avoids. This converts Cosine to Dot in FlatFloatStorage::try_from_batch, consistent with how PQ and RQ already convert Cosine to L2 for the same reason. Adds tests verifying: - Cosine is converted to Dot - L2 and Dot are preserved unchanged - Partition serde round-trip adapts to the conversion --- rust/lance-index/src/vector/flat/storage.rs | 82 ++++++++++++++++++- .../src/index/vector/ivf/partition_serde.rs | 9 +- 2 files changed, 89 insertions(+), 2 deletions(-) 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); } }