diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 9feb9d5ab..62d65816b 100644 --- a/rust/lancedb/src/table/query.rs +++ b/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> { + if contains_plan::(plan) { + return Ok(None); + } + + let Some(ann) = find_plan::(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(plan: &dyn ExecutionPlan) -> Option<&T> { + if let Some(plan) = (plan as &dyn std::any::Any).downcast_ref::() { + return Some(plan); + } + plan.children() + .into_iter() + .find_map(|child| find_plan::(child.as_ref())) +} + +fn contains_plan(plan: &dyn ExecutionPlan) -> bool { + find_plan::(plan).is_some() +} + +fn scale_distance_column( + plan: Arc, + scale: f32, +) -> Result> { + let schema = plan.schema(); + if schema.column_with_name(DIST_COL).is_none() { + return Ok(plan); + } + + let expressions: Vec<(Arc, String)> = schema + .fields() + .iter() + .enumerate() + .map(|(index, field)| { + let column: Arc = Arc::new(Column::new(field.name(), index)); + let expression = if field.name() == DIST_COL { + let scale: Arc = + Arc::new(Literal::new(ScalarValue::Float32(Some(scale)))); + Arc::new(BinaryExpr::new(column, Operator::Multiply, scale)) + as Arc + } 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, 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 { + 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::>(); + let norm = vector.iter().map(|value| value * value).sum::().sqrt(); + vector.iter_mut().for_each(|value| *value /= norm); + vector + } + + fn distances(batches: &[RecordBatch]) -> Vec { + batches + .iter() + .flat_map(|batch| { + batch[DIST_COL] + .as_primitive::() + .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::>(); + 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::>() + .await + .unwrap(); + let refined = table + .vector_search(query_vector.as_slice()) + .unwrap() + .limit(5) + .refine_factor(1) + .execute() + .await + .unwrap() + .try_collect::>() + .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::>() + .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;