fix: return cosine-scaled ANN distances

rust/lancedb/src/table/query.rs

@@ -17,7 +17,11 @@ use arrow::array::{AsArray, FixedSizeListBuilder, Float32Builder};
 use arrow::datatypes::{Float32Type, UInt8Type};
 use arrow_array::Array;
 use arrow_schema::{DataType, Schema};
+use datafusion_common::ScalarValue;
+use datafusion_expr::Operator;
+use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
 use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::PhysicalExpr;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::union::UnionExec;
@@ -25,7 +29,12 @@ use futures::future::try_join_all;
 use lance::dataset::mem_wal::DatasetMemWalExt;
 use lance::dataset::scanner::DatasetRecordBatchStream;
 use lance::dataset::scanner::Scanner;
+use lance::index::DatasetIndexInternalExt;
+use lance::io::exec::{ANNIvfSubIndexExec, KNNVectorDistanceExec};
 use lance_datafusion::exec::{analyze_plan as lance_analyze_plan, execute_plan};
+use lance_index::metrics::NoOpMetricsCollector;
+use lance_index::vector::{DIST_COL, quantizer::QuantizationType};
+use lance_linalg::distance::DistanceType as LanceDistanceType;
 use lance_namespace::LanceNamespace;
 use lance_namespace::models::{
     QueryTableRequest as NsQueryTableRequest, QueryTableRequestColumns,
@@ -324,11 +333,103 @@ pub async fn create_plan(
         scanner.order_by(Some(order_by.clone()))?;
     }
 
-    Ok(scanner.create_plan().await?)
+    let mut plan = scanner.create_plan().await?;
+    if let Some(scale) = approximate_cosine_distance_scale(plan.as_ref()).await? {
+        // Lance's ANN range filter sees the quantizer's internal normalized squared-L2
+        // scores. Translate the public cosine bounds before rebuilding the plan.
+        if query.lower_bound.is_some() || query.upper_bound.is_some() {
+            scanner.distance_range(
+                query.lower_bound.map(|bound| bound / scale),
+                query.upper_bound.map(|bound| bound / scale),
+            );
+            plan = scanner.create_plan().await?;
+        }
+        plan = scale_distance_column(plan, scale)?;
+    }
+
+    Ok(plan)
 }
 
 //Helper functions below
 
+/// Return the conversion from an ANN index's internal score to its public cosine scale.
+///
+/// Cosine PQ/SQ/RQ indices normalize their vectors and use squared L2 internally.  This
+/// preserves ranking, but squared L2 over unit vectors is twice the cosine distance.  Flat
+/// cosine indices calculate cosine directly and exact/refined plans already recompute the
+/// public distance, so those plans must not be adjusted.
+async fn approximate_cosine_distance_scale(plan: &dyn ExecutionPlan) -> Result<Option<f32>> {
+    if contains_plan::<KNNVectorDistanceExec>(plan) {
+        return Ok(None);
+    }
+
+    let Some(ann) = find_plan::<ANNIvfSubIndexExec>(plan) else {
+        return Ok(None);
+    };
+    if ann.query().metric_type != Some(LanceDistanceType::Cosine) {
+        return Ok(None);
+    }
+
+    let Some(index) = ann.indices().first() else {
+        return Ok(None);
+    };
+    let vector_index = ann
+        .dataset()
+        .open_vector_index(&ann.query().column, &index.uuid, &NoOpMetricsCollector)
+        .await?;
+    let (_, quantization_type) = vector_index.sub_index_type();
+    if matches!(
+        quantization_type,
+        QuantizationType::Product | QuantizationType::Scalar | QuantizationType::Rabit
+    ) {
+        Ok(Some(0.5))
+    } else {
+        Ok(None)
+    }
+}
+
+fn find_plan<T: 'static>(plan: &dyn ExecutionPlan) -> Option<&T> {
+    if let Some(plan) = (plan as &dyn std::any::Any).downcast_ref::<T>() {
+        return Some(plan);
+    }
+    plan.children()
+        .into_iter()
+        .find_map(|child| find_plan::<T>(child.as_ref()))
+}
+
+fn contains_plan<T: 'static>(plan: &dyn ExecutionPlan) -> bool {
+    find_plan::<T>(plan).is_some()
+}
+
+fn scale_distance_column(
+    plan: Arc<dyn ExecutionPlan>,
+    scale: f32,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let schema = plan.schema();
+    if schema.column_with_name(DIST_COL).is_none() {
+        return Ok(plan);
+    }
+
+    let expressions: Vec<(Arc<dyn PhysicalExpr>, String)> = schema
+        .fields()
+        .iter()
+        .enumerate()
+        .map(|(index, field)| {
+            let column: Arc<dyn PhysicalExpr> = Arc::new(Column::new(field.name(), index));
+            let expression = if field.name() == DIST_COL {
+                let scale: Arc<dyn PhysicalExpr> =
+                    Arc::new(Literal::new(ScalarValue::Float32(Some(scale))));
+                Arc::new(BinaryExpr::new(column, Operator::Multiply, scale))
+                    as Arc<dyn PhysicalExpr>
+            } else {
+                column
+            };
+            (expression, field.name().clone())
+        })
+        .collect();
+    Ok(Arc::new(ProjectionExec::try_new(expressions, plan)?))
+}
+
 // Take many execution plans and map them into a single plan that adds
 // a query_index column and unions them.
 pub(crate) fn create_multi_vector_plan(
@@ -693,7 +794,7 @@ mod tests {
     };
 
     use super::*;
-    use crate::query::{QueryExecutionOptions, QueryRequest};
+    use crate::query::{ExecutableQuery, QueryBase, QueryExecutionOptions, QueryRequest};
 
     fn fixed_size_list_array(values: Vec<f32>, dimension: i32) -> FixedSizeListArray {
         FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap()
@@ -1071,6 +1172,127 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_cosine_pq_distance_uses_public_cosine_scale() {
+        use arrow_array::{Int32Array, RecordBatch, types::Float32Type};
+        use arrow_schema::{DataType, Field, Schema};
+
+        use crate::connect;
+        use crate::index::{Index, vector::IvfPqIndexBuilder};
+
+        fn normalized_vector(state: &mut u64, dimension: usize) -> Vec<f32> {
+            let mut vector = (0..dimension)
+                .map(|_| {
+                    *state = state
+                        .wrapping_mul(6_364_136_223_846_793_005)
+                        .wrapping_add(1);
+                    ((*state >> 32) as u32 as f32 / u32::MAX as f32) * 2.0 - 1.0
+                })
+                .collect::<Vec<_>>();
+            let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
+            vector.iter_mut().for_each(|value| *value /= norm);
+            vector
+        }
+
+        fn distances(batches: &[RecordBatch]) -> Vec<f32> {
+            batches
+                .iter()
+                .flat_map(|batch| {
+                    batch[DIST_COL]
+                        .as_primitive::<Float32Type>()
+                        .values()
+                        .to_vec()
+                })
+                .collect()
+        }
+
+        let conn = connect("memory://").execute().await.unwrap();
+        let dimension = 8;
+        let num_rows = 256;
+        let mut state = 42;
+        let values = (0..num_rows)
+            .flat_map(|_| normalized_vector(&mut state, dimension))
+            .collect::<Vec<_>>();
+        let query_vector = normalized_vector(&mut state, dimension);
+        let vectors = Arc::new(fixed_size_list_array(values, dimension as i32));
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("vector", vectors.data_type().clone(), false),
+        ]));
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(Int32Array::from_iter_values(0..num_rows)), vectors],
+        )
+        .unwrap();
+        let table = conn
+            .create_table("test_cosine_pq_distance", batch)
+            .execute()
+            .await
+            .unwrap();
+        table
+            .create_index(
+                &["vector"],
+                Index::IvfPq(
+                    IvfPqIndexBuilder::default()
+                        .distance_type(crate::DistanceType::Cosine)
+                        .num_partitions(1)
+                        .num_sub_vectors(1),
+                ),
+            )
+            .execute()
+            .await
+            .unwrap();
+
+        let approximate = table
+            .vector_search(query_vector.as_slice())
+            .unwrap()
+            .limit(5)
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let refined = table
+            .vector_search(query_vector.as_slice())
+            .unwrap()
+            .limit(5)
+            .refine_factor(1)
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let approximate_distances = distances(&approximate);
+        let refined_distances = distances(&refined);
+        assert_eq!(approximate_distances.len(), refined_distances.len());
+        for (approximate, refined) in approximate_distances.iter().zip(&refined_distances) {
+            assert!(
+                (approximate - refined).abs() < 1e-5,
+                "approximate cosine distance {approximate} did not use the public scale; refined distance was {refined}"
+            );
+        }
+
+        // Distance range bounds are public cosine distances too. Lance applies them to
+        // internal ANN scores, so the planner must translate the bounds before execution.
+        let nearest = approximate_distances[0];
+        let ranged = table
+            .vector_search(query_vector.as_slice())
+            .unwrap()
+            .limit(1)
+            .distance_range(Some(nearest - 1e-5), Some(nearest + 1e-5))
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let ranged_distances = distances(&ranged);
+        assert_eq!(ranged_distances.len(), 1);
+        assert!((ranged_distances[0] - nearest).abs() < 1e-5);
+    }
+
     #[tokio::test]
     async fn test_create_plan_applies_approx_mode_to_ann_query() {
         use arrow_array::RecordBatch;