perf(index): use dot distance for cosine in FlatFloatStorage

rust/lance-index/src/vector/flat/storage.rs

@@ -53,6 +53,14 @@ impl QuantizerStorage for FlatFloatStorage {
         distance_type: DistanceType,
         frag_reuse_index: Option<Arc<FragReuseIndex>>,
     ) -> Result<Self> {
+        // 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<f32> = vec![0.5, 0.5, 0.5, 0.5]; // norm = 1.0
+        let v2: Vec<f32> = 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);
+        }
+    }
 }

rust/lance/src/index/vector/ivf/partition_serde.rs

@@ -742,7 +742,14 @@ mod tests {
             };
             let bytes = ser_body(&entry);
             let restored = de_body::<PartitionEntry<FlatIndex, FlatQuantizer>>(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);
         }
     }