diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 4b85a6ca8..24c5f55b3 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -2372,4 +2372,325 @@ mod lsm_tests { sstables.len() ); } + + #[tokio::test] + async fn lsm_cosine_distance_scale_and_mixed_tier_ordering() { + use arrow::array::{ + FixedSizeListArray, FixedSizeListBuilder, Float32Array, Float32Builder, + }; + use arrow::datatypes::Float32Type; + + use crate::index::Index; + use crate::index::vector::IvfPqIndexBuilder; + use datafusion_physical_plan::{ExecutionPlan, displayable}; + use lance::io::exec::ANNIvfSubIndexExec; + use std::any::Any; + + const DIM: usize = 8; + const N: usize = 256; + + fn normalized_vector(state: &mut u64) -> Vec { + let mut vector = (0..DIM) + .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 + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new( + "vec", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + DIM as i32, + ), + false, + ), + ])); + let make_batch = |rows: Vec<(i64, Vec)>| { + let ids = rows.iter().map(|(id, _)| *id).collect::>(); + let mut vectors = FixedSizeListBuilder::new(Float32Builder::new(), DIM as i32); + for (_, vector) in &rows { + vectors.values().append_slice(vector); + vectors.append(true); + } + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(ids)), Arc::new(vectors.finish())], + ) + .unwrap() + }; + let first_result = |batches: &[RecordBatch]| { + let batch = &batches[0]; + let id = batch["id"].as_primitive::().value(0); + let distance = batch["_distance"].as_primitive::().value(0); + (id, distance) + }; + + let mut state = 42; + let base_rows = (0..N) + .map(|id| (id as i64, normalized_vector(&mut state))) + .collect::>(); + let query = normalized_vector(&mut state); + + let dir = tempdir().unwrap(); + let conn = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let base = make_batch(base_rows.clone()); + let reader: Box = + Box::new(RecordBatchIterator::new(vec![Ok(base)], schema.clone())); + let table = conn + .create_table("cosine_lsm", reader) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .create_index( + &["vec"], + Index::IvfPq( + IvfPqIndexBuilder::default() + .distance_type(crate::DistanceType::Cosine) + .num_partitions(1) + .num_sub_vectors(1), + ), + ) + .name("vec_cosine".to_string()) + .execute() + .await + .unwrap(); + table + .enable_rte( + RTEWriteSpec::unsharded().with_maintained_indexes(vec!["vec_cosine".to_string()]), + ) + .await + .unwrap(); + + let base_only = table + .query() + .nearest_to(query.as_slice()) + .unwrap() + .limit(1) + .use_rte(false) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let (base_id, public_distance) = first_result(&base_only); + + let lsm = table + .query() + .nearest_to(query.as_slice()) + .unwrap() + .limit(1) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let (lsm_id, lsm_distance) = first_result(&lsm); + assert_eq!(lsm_id, base_id); + assert!( + (lsm_distance - public_distance).abs() < 1e-5, + "LSM cosine distance {lsm_distance} did not use the public scale {public_distance}" + ); + + // Add an exact memtable result whose distance lies between the public ANN + // score and its doubled internal score. Correctly normalized plans still + // rank the ANN row first; mixed units would incorrectly rank this row first. + assert!(public_distance > 0.0 && public_distance < 4.0 / 3.0); + let memtable_distance = public_distance * 1.5; + let cosine_similarity = 1.0 - memtable_distance; + let mut orthogonal = normalized_vector(&mut state); + let projection = orthogonal + .iter() + .zip(&query) + .map(|(left, right)| left * right) + .sum::(); + for (value, query_value) in orthogonal.iter_mut().zip(&query) { + *value -= projection * query_value; + } + let norm = orthogonal + .iter() + .map(|value| value * value) + .sum::() + .sqrt(); + orthogonal.iter_mut().for_each(|value| *value /= norm); + let sine = (1.0 - cosine_similarity * cosine_similarity).sqrt(); + let memtable_vector = query + .iter() + .zip(&orthogonal) + .map(|(query_value, orthogonal_value)| { + cosine_similarity * query_value + sine * orthogonal_value + }) + .collect::>(); + + let mut merge = table.merge_insert(&[]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let memtable = make_batch(vec![(N as i64, memtable_vector)]); + merge + .execute(Box::new(RecordBatchIterator::new( + vec![Ok(memtable)], + schema.clone(), + ))) + .await + .unwrap(); + + let mixed = table + .query() + .nearest_to(query.as_slice()) + .unwrap() + .limit(1) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let (mixed_id, mixed_distance) = first_result(&mixed); + assert_eq!( + mixed_id, base_id, + "mixed LSM tiers must compare ANN and exact distances in public units" + ); + assert!((mixed_distance - public_distance).abs() < 1e-5); + + // Re-ingest every base row so the flushed SSTable generation owns a + // complete copy of the dataset. ANN ranking of the rebuilt index is an + // independent approximation, so exact row parity with the base scan is + // not guaranteed. + let mut merge = table.merge_insert(&[]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let upsert_batch = make_batch(base_rows.clone()); + merge + .execute(Box::new(RecordBatchIterator::new( + vec![Ok(upsert_batch)], + schema.clone(), + ))) + .await + .unwrap(); + assert!( + base_rows.iter().any(|(id, _)| *id == base_id), + "upsert source must still contain the nearest base row" + ); + + table.close_rte_writers().await.unwrap(); + let sstables = published_sstables(&table).await; + assert_eq!(sstables.len(), 1, "expected one published SSTable"); + assert_eq!( + sstables[0].count_rows(None).await.unwrap(), + N + 1, + "published SSTable must contain the complete replacement fixture" + ); + + // The public query sets no explicit refine factor. Because the upsert + // rewrote every base id, a newer generation exists, so the engine must + // auto-refine ANN results over the published SSTable. + let public_query = table.query().nearest_to(query.as_slice()).unwrap().limit(1); + + fn collect_ann_nodes<'a>( + node: &'a dyn ExecutionPlan, + target_uri: &str, + out: &mut Vec<&'a ANNIvfSubIndexExec>, + ) { + // Local datasets may report their path with a different separator on Windows. + if let Some(ann) = (node as &dyn Any).downcast_ref::() + && std::path::Path::new(ann.dataset().uri()) == std::path::Path::new(target_uri) + { + out.push(ann); + } + for child in node.children() { + collect_ann_nodes(child.as_ref(), target_uri, out); + } + } + + let plan = public_query + .create_plan(crate::query::QueryExecutionOptions::default()) + .await + .unwrap(); + let mut ann_nodes = Vec::new(); + collect_ann_nodes(plan.as_ref(), sstables[0].uri(), &mut ann_nodes); + assert!( + !ann_nodes.is_empty(), + "public plan must contain an ANN sub-index over the published SSTable" + ); + assert_eq!( + ann_nodes[0].query().refine_factor, + Some(1), + "engine must auto-refine the SSTable ANN results" + ); + println!( + "RTE LSM public query plan:\n{}", + displayable(plan.as_ref()).indent(true) + ); + + // Execute the same public query. + let published = public_query + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let (published_id, published_distance) = first_result(&published); + assert!( + (0_i64..=N as i64).contains(&published_id), + "SSTable result id {published_id} must belong to the complete 0..={N} replacement fixture" + ); + + // The maintained ANN index may rank near-tied rows differently than + // the base scan, so the durable contract is distance correctness: the + // auto-refined result carries its own stored vector, and the query + // returns the analytical cosine distance of that exact vector. + let expected_distance = { + let published_vec = published[0]["vec"] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_any() + .downcast_ref::() + .unwrap(); + // The query is a unit vector, so cosine distance is 1 - cos(q, v). + let dot = (0..DIM) + .map(|index| query[index] * published_vec.value(index)) + .sum::(); + let norm = (0..DIM) + .map(|index| published_vec.value(index) * published_vec.value(index)) + .sum::() + .sqrt(); + 1.0 - dot / norm + }; + assert!( + expected_distance > 0.0 && expected_distance.is_finite(), + "independent cosine oracle is degenerate: {expected_distance}" + ); + assert!( + (published_distance - expected_distance).abs() < 1e-5, + "SSTable ANN distance {published_distance} must match the analytical cosine distance {expected_distance} of its returned vector" + ); + println!( + "sstable_uri={} ann_nodes={} refine_factor={:?} observed_distance={} expected_cosine={}", + sstables[0].uri(), + ann_nodes.len(), + ann_nodes[0].query().refine_factor, + published_distance, + expected_distance, + ); + } } diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 4ad840350..dbb3dbc68 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::sync::Arc; +use std::{ + collections::{HashSet, VecDeque}, + sync::Arc, +}; mod lsm; @@ -17,15 +20,24 @@ use arrow::array::{AsArray, FixedSizeListBuilder, Float32Builder}; use arrow::datatypes::{Float32Type, UInt8Type}; use arrow_array::Array; use arrow_schema::{DataType, Schema}; -use datafusion_common::{Column, DataFusionError, SchemaError}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_common::{Column, DataFusionError, ScalarValue, SchemaError}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::{BinaryExpr, Column as PhysicalColumn, Literal}; +use datafusion_physical_plan::PhysicalExpr; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::union::UnionExec; +use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary}; use lance::dataset::mem_wal::DatasetMemWalExt; use lance::dataset::scanner::DatasetRecordBatchStream; use lance::dataset::scanner::Scanner; +use lance::index::DatasetIndexInternalExt; +use lance::io::exec::{ANNIvfBatchExec, ANNIvfSubIndexExec}; use lance_datafusion::exec::{analyze_plan as lance_analyze_plan, execute_plan}; +use lance_index::IndexType; +use lance_index::metrics::NoOpMetricsCollector; +use lance_index::vector::DIST_COL; +use lance_linalg::distance::DistanceType as LanceDistanceType; use lance_namespace::LanceNamespace; use lance_namespace::models::{ QueryTableRequest as NsQueryTableRequest, QueryTableRequestColumns, @@ -409,10 +421,30 @@ pub async fn create_plan( scanner.order_by(Some(order_by.clone()))?; } - scanner + let mut plan = scanner .create_plan() .await - .map_err(|error| enrich_lance_field_not_found(error, schema)) + .map_err(|error| enrich_lance_field_not_found(error, schema))?; + let normalized_l2_indices = normalized_l2_ann_indices(plan.as_ref()).await?; + if !normalized_l2_indices.is_empty() { + // Rebuild only the affected ANN nodes with internal normalized squared-L2 + // bounds. Exact branches keep the public cosine bounds from `plan`. + let internal_plan = if query.lower_bound.is_some() || query.upper_bound.is_some() { + scanner.distance_range( + query.lower_bound.map(|bound| bound / COSINE_ANN_SCALE), + query.upper_bound.map(|bound| bound / COSINE_ANN_SCALE), + ); + scanner + .create_plan() + .await + .map_err(|error| enrich_lance_field_not_found(error, schema))? + } else { + plan.clone() + }; + plan = normalize_ann_branches(plan, internal_plan, &normalized_l2_indices)?; + } + + Ok(plan) } /// Replace DataFusion's top-level field candidates with qualified leaf paths. @@ -504,6 +536,242 @@ fn leaf_field_paths(schema: &Schema) -> Vec { //Helper functions below +const COSINE_ANN_SCALE: f32 = 0.5; + +#[derive(Clone, Copy)] +enum AnnExecRef<'a> { + Single(&'a ANNIvfSubIndexExec), + Batch(&'a ANNIvfBatchExec), +} + +impl<'a> AnnExecRef<'a> { + fn from_plan(plan: &'a dyn ExecutionPlan) -> Option { + if let Some(single) = plan.downcast_ref::() { + return Some(Self::Single(single)); + } + plan.downcast_ref::().map(Self::Batch) + } + + fn query(self) -> &'a lance_index::vector::Query { + match self { + Self::Single(ann) => ann.query(), + Self::Batch(ann) => ann.query(), + } + } + + fn dataset(self) -> &'a Arc { + match self { + Self::Single(ann) => ann.dataset(), + Self::Batch(ann) => ann.dataset(), + } + } + + fn indices(self) -> &'a [lance_table::format::IndexMetadata] { + match self { + Self::Single(ann) => ann.indices(), + Self::Batch(ann) => ann.indices(), + } + } + + fn query_count(self) -> Option { + match self { + Self::Single(_) => None, + Self::Batch(ann) => Some(ann.query_count()), + } + } + + fn same_kind(self, other: Self) -> bool { + matches!( + (self, other), + (Self::Single(_), Self::Single(_)) | (Self::Batch(_), Self::Batch(_)) + ) + } +} + +/// Find ANN index segments whose scores use normalized squared L2 for cosine search. +/// +/// 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, so they are not included. +async fn normalized_l2_ann_indices(plan: &dyn ExecutionPlan) -> Result> { + let mut ann_plans = Vec::new(); + find_ann_plans(plan, &mut ann_plans); + + let mut checked = HashSet::new(); + let mut normalized_l2 = HashSet::new(); + for ann in ann_plans { + if ann.query().metric_type != Some(LanceDistanceType::Cosine) { + continue; + } + for index in ann.indices() { + let uuid = index.uuid.to_string(); + if !checked.insert(uuid.clone()) { + continue; + } + let vector_index = ann + .dataset() + .open_vector_index(&ann.query().column, &index.uuid, &NoOpMetricsCollector) + .await?; + if matches!( + vector_index.index_type(), + IndexType::IvfPq + | IndexType::IvfSq + | IndexType::IvfRq + | IndexType::IvfHnswPq + | IndexType::IvfHnswSq + ) { + normalized_l2.insert(uuid); + } + } + } + Ok(normalized_l2) +} + +/// Normalize affected ANN outputs before their parent plan nodes consume them. +/// +/// This is used by planners that do not support distance ranges, such as the MemWAL +/// LSM planner. The standard scanner path rebuilds a second plan when it also needs +/// to translate range bounds, then calls [`normalize_ann_branches`] directly. +pub(super) async fn normalize_cosine_ann_branches( + plan: Arc, +) -> Result> { + let normalized_l2_indices = normalized_l2_ann_indices(plan.as_ref()).await?; + if normalized_l2_indices.is_empty() { + return Ok(plan); + } + normalize_ann_branches(plan.clone(), plan, &normalized_l2_indices) +} + +fn find_ann_plans<'a>(plan: &'a dyn ExecutionPlan, ann_plans: &mut Vec>) { + if let Some(ann) = AnnExecRef::from_plan(plan) { + ann_plans.push(ann); + return; + } + for child in plan.children() { + find_ann_plans(child.as_ref(), ann_plans); + } +} + +fn collect_ann_plans( + plan: &Arc, + ann_plans: &mut VecDeque>, +) { + if AnnExecRef::from_plan(plan.as_ref()).is_some() { + ann_plans.push_back(plan.clone()); + return; + } + for child in plan.children() { + collect_ann_plans(child, ann_plans); + } +} + +/// Replace normalized-L2 ANN nodes with equivalent nodes that use internal bounds, then +/// convert their output to the public cosine scale before any generic plan node consumes it. +fn normalize_ann_branches( + public_plan: Arc, + internal_plan: Arc, + normalized_l2_indices: &HashSet, +) -> Result> { + let mut internal_ann_plans = VecDeque::new(); + collect_ann_plans(&internal_plan, &mut internal_ann_plans); + let normalized = + replace_ann_branches(public_plan, &mut internal_ann_plans, normalized_l2_indices)?; + if !internal_ann_plans.is_empty() { + return Err(Error::Runtime { + message: "internal and public vector plans contained different ANN branches" + .to_string(), + }); + } + Ok(normalized) +} + +fn replace_ann_branches( + public_plan: Arc, + internal_ann_plans: &mut VecDeque>, + normalized_l2_indices: &HashSet, +) -> Result> { + if let Some(public_ann) = AnnExecRef::from_plan(public_plan.as_ref()) { + let internal_plan = internal_ann_plans + .pop_front() + .ok_or_else(|| Error::Runtime { + message: "internal vector plan was missing an ANN branch".to_string(), + })?; + let internal_ann = + AnnExecRef::from_plan(internal_plan.as_ref()).ok_or_else(|| Error::Runtime { + message: "internal vector plan contained a non-ANN branch where an ANN branch was expected" + .to_string(), + })?; + let same_indices = public_ann + .indices() + .iter() + .map(|index| &index.uuid) + .eq(internal_ann.indices().iter().map(|index| &index.uuid)); + if !public_ann.same_kind(internal_ann) + || public_ann.query_count() != internal_ann.query_count() + || public_ann.query().column != internal_ann.query().column + || public_ann.query().metric_type != internal_ann.query().metric_type + || !same_indices + { + return Err(Error::Runtime { + message: "internal and public vector plans had mismatched ANN branches".to_string(), + }); + } + + let normalized_count = public_ann + .indices() + .iter() + .filter(|index| normalized_l2_indices.contains(&index.uuid.to_string())) + .count(); + if normalized_count == 0 { + return Ok(public_plan); + } + if normalized_count != public_ann.indices().len() { + return Err(Error::Runtime { + message: "one ANN branch mixed public and normalized-L2 distance scales" + .to_string(), + }); + } + return scale_distance_column(internal_plan, COSINE_ANN_SCALE); + } + + let children = public_plan + .children() + .into_iter() + .cloned() + .map(|child| replace_ann_branches(child, internal_ann_plans, normalized_l2_indices)) + .collect::>>()?; + Ok(with_new_children_if_necessary(public_plan, children)?) +} + +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(PhysicalColumn::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( @@ -1622,6 +1890,354 @@ mod tests { "plan should add query_index column, got:\n{display}" ); } + #[rstest::rstest] + #[case::modern(false)] + #[case::legacy(true)] + #[tokio::test] + async fn test_cosine_pq_distance_uses_public_cosine_scale(#[case] legacy: bool) { + use arrow_array::{ + Float32Array, Int32Array, RecordBatch, + types::{Float32Type, Int32Type}, + }; + use arrow_schema::{DataType, Field, Schema}; + + use crate::connect; + use crate::index::{Index, vector::IvfPqIndexBuilder}; + use lance::index::{ + DatasetIndexExt, + vector::{IndexFileVersion, VectorIndexParams}, + }; + + 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.clone(), + 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(); + if legacy { + let mut dataset = (*table.dataset().unwrap().get().await.unwrap()).clone(); + dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + VectorIndexParams::ivf_pq(1, 8, 1, LanceDistanceType::Cosine, 10) + .version(IndexFileVersion::Legacy), + true, + ) + .await + .unwrap(); + table.dataset().unwrap().update(dataset); + } else { + table + .create_index( + &["vector"], + Index::IvfPq( + IvfPqIndexBuilder::default() + .distance_type(crate::DistanceType::Cosine) + .num_partitions(1) + .num_sub_vectors(1), + ), + ) + .execute() + .await + .unwrap(); + } + + let expected_index_version = if legacy { "Legacy" } else { "V3" }; + let index_file_version = { + let dataset = table.dataset().unwrap().get().await.unwrap(); + let metadata = dataset + .load_indices() + .await + .unwrap() + .first() + .unwrap() + .clone(); + let index = dataset + .open_vector_index("vector", &metadata.uuid, &NoOpMetricsCollector) + .await + .unwrap(); + index.statistics().unwrap()["index_file_version"] + .as_str() + .unwrap() + .to_string() + }; + assert_eq!(index_file_version, expected_index_version); + + let plan = table + .dataset() + .unwrap() + .get() + .await + .unwrap() + .scan() + .nearest("vector", &Float32Array::from(query_vector.clone()), 5) + .unwrap() + .create_plan() + .await + .unwrap(); + let mut ann_plans = Vec::new(); + find_ann_plans(plan.as_ref(), &mut ann_plans); + assert!(!ann_plans.is_empty()); + + 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(), 5); + 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); + + if !legacy { + fn batch_distances_by_query( + batches: &[RecordBatch], + query_count: usize, + ) -> Vec> { + let mut distances = vec![Vec::new(); query_count]; + for batch in batches { + let query_index_column = batch["query_index"].as_primitive::(); + let distance_column = batch[DIST_COL].as_primitive::(); + for row in 0..batch.num_rows() { + let query_index = query_index_column.value(row) as usize; + distances[query_index].push(distance_column.value(row)); + } + } + distances + } + + fn sorted(mut values: Vec) -> Vec { + values.sort_by(|left, right| left.partial_cmp(right).unwrap()); + values + } + + let batch_query = table + .vector_search(query_vector.as_slice()) + .unwrap() + .add_query_vector(query_vector.as_slice()) + .unwrap() + .nprobes(1) + .unwrap() + .limit(5); + let batch_plan = batch_query.explain_plan(false).await.unwrap(); + assert!(batch_plan.contains("ANNIvfBatch"), "{batch_plan}"); + let batched = batch_query + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batched_distances = batch_distances_by_query(&batched, 2); + let expected_distances = sorted(approximate_distances.clone()); + for (query_index, actual) in batched_distances.into_iter().enumerate() { + assert_eq!( + actual.len(), + expected_distances.len(), + "query {query_index} returned the wrong number of rows" + ); + for (actual, expected) in sorted(actual).into_iter().zip(&expected_distances) { + assert!( + (actual - expected).abs() < 1e-5, + "batch distance {actual} for query {query_index} must use single-query public scale {expected}" + ); + } + } + + let batched_ranged = table + .vector_search(query_vector.as_slice()) + .unwrap() + .add_query_vector(query_vector.as_slice()) + .unwrap() + .nprobes(1) + .unwrap() + .limit(1) + .distance_range(None, Some(nearest + 1e-5)) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batched_ranged_distances = batch_distances_by_query(&batched_ranged, 2); + for (query_index, query_distances) in batched_ranged_distances.into_iter().enumerate() { + assert_eq!( + query_distances.len(), + 1, + "public cosine range must return one hit for query {query_index}" + ); + assert!( + (query_distances[0] - nearest).abs() < 1e-5, + "ranged batch distance {} for query {query_index} must match {nearest}", + query_distances[0] + ); + } + } + + let refined_ranged = table + .vector_search(query_vector.as_slice()) + .unwrap() + .limit(1) + .refine_factor(1) + .distance_range(None, Some(nearest + 1e-5)) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + distances(&refined_ranged).len(), + 1, + "refinement must not apply public cosine bounds to internal ANN scores" + ); + + let aliased = table + .vector_search(query_vector.as_slice()) + .unwrap() + .limit(1) + .select(Select::dynamic(&[("aliased_distance", "_distance")])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = &aliased[0]; + let aliased_distance = batch["aliased_distance"] + .as_primitive::() + .value(0); + let public_distance = batch[DIST_COL].as_primitive::().value(0); + assert!( + (aliased_distance - public_distance).abs() < 1e-5, + "distance aliases and auto-projected distances must use the same public scale" + ); + + // Appended rows take an exact fallback branch. Its public range filter must stay + // independent of the translated ANN bounds before both branches are merged. + let mut orthogonal = normalized_vector(&mut state, dimension); + let projection = orthogonal + .iter() + .zip(&query_vector) + .map(|(left, right)| left * right) + .sum::(); + for (value, query_value) in orthogonal.iter_mut().zip(&query_vector) { + *value -= projection * query_value; + } + let norm = orthogonal + .iter() + .map(|value| value * value) + .sum::() + .sqrt(); + orthogonal.iter_mut().for_each(|value| *value /= norm); + let appended_vectors = Arc::new(fixed_size_list_array(orthogonal, dimension as i32)); + let appended = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![num_rows])), appended_vectors], + ) + .unwrap(); + table.add(appended).execute().await.unwrap(); + + let mixed = table + .vector_search(query_vector.as_slice()) + .unwrap() + .limit(5) + .distance_range(None, Some(nearest + 1e-5)) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let mixed_distances = distances(&mixed); + assert_eq!(mixed_distances.len(), 1); + assert!((mixed_distances[0] - nearest).abs() < 1e-5); + } #[tokio::test] async fn test_create_plan_applies_approx_mode_to_ann_query() { diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index 0cc240291..7af52e572 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -130,6 +130,10 @@ pub(super) async fn create_lsm_plan( .await? }; + // Normalize cosine ANN arms before LSM merge and sort nodes compare their + // distances with exact SSTable and memtable arms. + let plan = super::normalize_cosine_ann_branches(plan).await?; + // Lance appends the primary-key columns internally for dedup and keeps them in // the output; drop the ones the user did not request so the projection matches. restore_projection(plan, &query, &pk_columns)