diff --git a/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java b/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java
index b91e5ae4a7d..571bdedc5a5 100644
--- a/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java
+++ b/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java
@@ -117,9 +117,8 @@ private native void nativeCreate(
* @param nprobes number of IVF partitions to probe
* @param columns columns to project; pass {@code null} to return all columns plus {@code
* _distance}
- * @param refineBaseTable when true, the base-table arm re-ranks candidates with exact distances
- * (refine factor 1). Useful when the base table uses an approximate index (IVF-PQ).
- * Auto-enabled whenever stale filtering is on (see {@code overfetchFactor}).
+ * @param refineBaseTable when true, the base-table and SSTable arms re-rank candidates with exact
+ * distances. Auto-enabled when a newer generation blocks the base table.
* @param overfetchFactor over-fetch multiple for sources with rows superseded by newer
* generations. Must be at least {@code 1.0}:
*
@@ -177,7 +176,8 @@ public ExecutionPlan planSearch(
* @param nprobes number of IVF partitions to probe
* @param columns columns to project; pass {@code null} to return all columns plus {@code
* _distance}
- * @param refineBaseTable when true, the base-table arm re-ranks candidates with exact distances.
+ * @param refineBaseTable when true, the base-table and SSTable arms re-rank candidates with exact
+ * distances.
* @return an executable plan
*/
public ExecutionPlan planSearch(
diff --git a/python/python/lance/mem_wal.py b/python/python/lance/mem_wal.py
index b9a8890853a..f0c735eb2e4 100644
--- a/python/python/lance/mem_wal.py
+++ b/python/python/lance/mem_wal.py
@@ -582,10 +582,9 @@ def plan_search(
Columns to project. Returns all columns + ``_distance`` if
omitted.
refine_base_table : bool, optional
- When true, the base-table arm re-ranks its candidates with exact
- distances. Useful when the base index is approximate (e.g.
- IVF-PQ). Memtable arms use exact HNSW and are unaffected.
- Auto-enabled whenever stale filtering is on (see ``overfetch_factor``).
+ When true, the base-table and SSTable arms re-rank their candidates
+ with exact distances. Auto-enabled when a newer generation blocks
+ the base table.
overfetch_factor : float, optional
Over-fetch multiple for sources with rows superseded by newer
generations. Must be at least ``1.0``. At ``1.0`` each affected
diff --git a/rust/lance/src/dataset/mem_wal/index/hnsw.rs b/rust/lance/src/dataset/mem_wal/index/hnsw.rs
index 581814a2309..b98f65ca1eb 100644
--- a/rust/lance/src/dataset/mem_wal/index/hnsw.rs
+++ b/rust/lance/src/dataset/mem_wal/index/hnsw.rs
@@ -569,6 +569,16 @@ async fn write_sq_storage(
};
let row_ids = column(ROW_ID)?;
let vectors = column(FLAT_COLUMN)?;
+ // Cosine search ranks L2 against a normalized query, so stored vectors must
+ // be normalized. The graph is unaffected: raw cosine and normalized L2 order
+ // neighbors identically.
+ let vectors: Arc = if distance_type == DistanceType::Cosine {
+ Arc::new(lance_linalg::kernels::normalize_fsl(
+ vectors.as_fixed_size_list(),
+ )?)
+ } else {
+ vectors
+ };
let vectors = vectors.as_fixed_size_list();
let mut sq = ScalarQuantizer::new(8, dim);
let bounds: Range = sq.update_bounds::(vectors)?;
diff --git a/rust/lance/src/dataset/mem_wal/memtable/flush.rs b/rust/lance/src/dataset/mem_wal/memtable/flush.rs
index 0b52568c733..bbc1e7cf3d0 100644
--- a/rust/lance/src/dataset/mem_wal/memtable/flush.rs
+++ b/rust/lance/src/dataset/mem_wal/memtable/flush.rs
@@ -1878,6 +1878,132 @@ mod tests {
);
}
+ /// A flushed cosine index ranks L2 against a normalized query, so it must
+ /// store normalized vectors.
+ #[tokio::test]
+ async fn test_flusher_cosine_index_ranks_unnormalized_vectors() {
+ use super::super::super::index::IndexStore;
+ use arrow_array::{FixedSizeListArray, Float32Array};
+ use lance_linalg::distance::DistanceType;
+
+ let (store, base_path, base_uri, _temp_dir) = create_local_store().await;
+ let shard_id = Uuid::new_v4();
+ let manifest_store = Arc::new(ShardManifestStore::new(
+ store.clone(),
+ &base_path,
+ shard_id,
+ 2,
+ ));
+ let (epoch, _manifest) = manifest_store.claim_epoch(0).await.unwrap();
+
+ // Norms spread over 0.2..5.0, so ranking by L2 against raw vectors
+ // would disagree with cosine.
+ let (dim, num_vectors) = (32usize, 2000usize);
+ let mut seed: u64 = 42;
+ let mut next = || {
+ seed = seed
+ .wrapping_mul(6364136223846793005)
+ .wrapping_add(1442695040888963407);
+ ((seed >> 33) as f32 / (1u64 << 31) as f32) * 2.0 - 1.0
+ };
+ let mut vectors = Vec::with_capacity(num_vectors * dim);
+ for _ in 0..num_vectors {
+ let row: Vec = (0..dim).map(|_| next()).collect();
+ let norm = row.iter().map(|x| x * x).sum::().sqrt();
+ let target = 0.2 + 4.8 * (next() * 0.5 + 0.5);
+ vectors.extend(row.iter().map(|x| x * target / norm));
+ }
+ let query: Vec = (0..dim).map(|_| next()).collect();
+ let cosine_distance = |id: usize| {
+ lance_linalg::distance::cosine_distance(&query, &vectors[id * dim..][..dim])
+ };
+ let mut by_distance: Vec = (0..num_vectors).collect();
+ by_distance.sort_by(|a, b| cosine_distance(*a).total_cmp(&cosine_distance(*b)));
+ let nearest: std::collections::HashSet =
+ by_distance[..10].iter().map(|id| *id as i32).collect();
+
+ let item = Arc::new(Field::new("item", DataType::Float32, false));
+ let schema = Arc::new(ArrowSchema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new(
+ "vector",
+ DataType::FixedSizeList(item.clone(), dim as i32),
+ false,
+ ),
+ ]));
+ let index_specs = vec![MemIndexSpec::hnsw(
+ "vector_hnsw",
+ 1,
+ "vector",
+ DistanceType::Cosine,
+ )];
+ let mut memtable = MemTable::new(schema.clone(), 1, vec![]).unwrap();
+ let lance_schema = lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap();
+ memtable.set_indexes(
+ IndexStore::from_specs(&index_specs, &lance_schema, num_vectors, 100).unwrap(),
+ );
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(Int32Array::from_iter_values(0..num_vectors as i32)),
+ Arc::new(
+ FixedSizeListArray::try_new(
+ item,
+ dim as i32,
+ Arc::new(Float32Array::from(vectors.clone())),
+ None,
+ )
+ .unwrap(),
+ ),
+ ],
+ )
+ .unwrap();
+ let frag_id = memtable.insert(batch).await.unwrap();
+
+ let flusher = MemTableFlusher::new(
+ store.clone(),
+ base_path.clone(),
+ base_uri.clone(),
+ shard_id,
+ manifest_store.clone(),
+ );
+ let result = flusher
+ .flush_with_indexes(&memtable, epoch, &index_specs, 1, frag_id + 1)
+ .await
+ .unwrap();
+ let gen_uri = format!("{}/_mem_wal/{}/{}", base_uri, shard_id, result.sstable.path);
+ let dataset = Dataset::open(&gen_uri).await.unwrap();
+
+ let mut scan = dataset.scan();
+ scan.nearest("vector", &Float32Array::from(query.clone()), 10)
+ .unwrap();
+ scan.distance_metric(DistanceType::Cosine);
+ // A wide search list keeps graph recall out of what this measures.
+ scan.ef(500);
+ scan.fast_search();
+ let batch = scan.try_into_batch().await.unwrap();
+ let ids = batch["id"].as_any().downcast_ref::().unwrap();
+ let distances = batch["_distance"]
+ .as_any()
+ .downcast_ref::()
+ .unwrap();
+
+ let found = ids
+ .values()
+ .iter()
+ .filter(|id| nearest.contains(id))
+ .count();
+ assert!(found >= 9, "recall@10 was {found}/10");
+ // Lance's SQ index reports cosine as L2 over normalized vectors.
+ for (id, distance) in ids.values().iter().zip(distances.values()) {
+ let expected = 2.0 * cosine_distance(*id as usize);
+ assert!(
+ (distance - expected).abs() < 0.05,
+ "id={id} reported {distance}, expected {expected}"
+ );
+ }
+ }
+
#[tokio::test]
async fn test_flusher_with_fts_index() {
use super::super::super::index::IndexStore;
diff --git a/rust/lance/src/dataset/mem_wal/scanner/builder.rs b/rust/lance/src/dataset/mem_wal/scanner/builder.rs
index 82fbeb41ce9..ca4728587c1 100644
--- a/rust/lance/src/dataset/mem_wal/scanner/builder.rs
+++ b/rust/lance/src/dataset/mem_wal/scanner/builder.rs
@@ -48,9 +48,8 @@ struct LsmVectorQuery {
k: usize,
/// IVF partition probe bounds for indexed arms.
probe_bounds: ProbeBounds,
- /// Re-rank base candidates with exact distances when set (refine factor is
- /// treated as a boolean; the LSM merge needs exact base distances).
- refine: bool,
+ /// Refine factor for the base and SSTable arms; `None` leaves them unrefined.
+ refine_factor: Option,
/// Distance metric; `None` defaults to L2 (matching the unindexed memtable arm).
metric_type: Option,
}
@@ -490,7 +489,7 @@ impl LsmScanner {
key: key.slice(0, key.len()),
k,
probe_bounds: ProbeBounds::default(),
- refine: false,
+ refine_factor: None,
metric_type: None,
});
Ok(self)
@@ -528,12 +527,11 @@ impl LsmScanner {
self
}
- /// Re-rank base-table candidates with exact distances so they are directly
- /// comparable to the exact memtable distances in the cross-source merge.
- /// The factor is treated as a boolean today. No-op unless `nearest` was set.
+ /// Re-rank base and SSTable candidates by exact distance over `refine_factor`
+ /// times as many. `0` clears it; no-op unless `nearest` was set.
pub fn refine(mut self, refine_factor: u32) -> Self {
if let Some(q) = self.nearest.as_mut() {
- q.refine = refine_factor > 0;
+ q.refine_factor = (refine_factor > 0).then_some(refine_factor);
}
self
}
@@ -633,7 +631,8 @@ impl LsmScanner {
distance_type,
)
.with_identity_schema(Arc::clone(&self.identity_schema))
- .with_filter(self.filter.clone());
+ .with_filter(self.filter.clone())
+ .with_refine_factor(nearest.refine_factor);
if let BaseSource::Table(dataset) = &self.base {
planner = planner.with_dataset(dataset.clone());
}
@@ -660,7 +659,7 @@ impl LsmScanner {
per_source_k,
nearest.probe_bounds,
self.projection.as_deref(),
- nearest.refine,
+ false,
overfetch_factor,
)
.await?;
diff --git a/rust/lance/src/dataset/mem_wal/scanner/projection.rs b/rust/lance/src/dataset/mem_wal/scanner/projection.rs
index e0f6e470712..ca2bf7bc1aa 100644
--- a/rust/lance/src/dataset/mem_wal/scanner/projection.rs
+++ b/rust/lance/src/dataset/mem_wal/scanner/projection.rs
@@ -19,8 +19,9 @@ use std::collections::HashMap;
use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use datafusion::logical_expr::Operator;
use datafusion::physical_expr::PhysicalExpr;
-use datafusion::physical_expr::expressions::{CastExpr, Column, Literal};
+use datafusion::physical_expr::expressions::{BinaryExpr, CastExpr, Column, Literal};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::projection::ProjectionExec;
use datafusion::scalar::ScalarValue;
@@ -217,6 +218,40 @@ pub fn canonical_output_schema(
Ok(Arc::new(Schema::new(fields)))
}
+/// Multiply `plan`'s `_distance` by `factor`, forwarding other columns; no-op
+/// at `1.0`.
+pub fn scale_distance(plan: Arc, factor: f32) -> Result> {
+ if factor == 1.0 {
+ return Ok(plan);
+ }
+ let input_schema = plan.schema();
+ let project_exprs: Vec<(Arc, String)> = input_schema
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(idx, field)| {
+ let column: Arc = Arc::new(Column::new(field.name(), idx));
+ let expr = if field.name() == DISTANCE_COLUMN {
+ Arc::new(BinaryExpr::new(
+ column,
+ Operator::Multiply,
+ Arc::new(Literal::new(ScalarValue::Float32(Some(factor)))),
+ ))
+ } else {
+ column
+ };
+ (expr, field.name().clone())
+ })
+ .collect();
+ let projection_exec = ProjectionExec::try_new(project_exprs, plan).map_err(|e| {
+ lance_core::Error::internal(format!(
+ "Failed to build scale_distance ProjectionExec: {}",
+ e
+ ))
+ })?;
+ Ok(Arc::new(projection_exec))
+}
+
/// Wrap `plan` so the named columns become typed NULL literals; all
/// other columns are forwarded unchanged. Schema is preserved (same
/// fields, same dtypes). Useful for stripping the *value* of an
diff --git a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs
index 884c9b1067c..8b22118c33c 100644
--- a/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs
+++ b/rust/lance/src/dataset/mem_wal/scanner/vector_search.rs
@@ -31,7 +31,7 @@ use super::data_source::LsmDataSource;
use super::generation_read::{GenerationRead, memtable_matches_table};
use super::projection::{
DISTANCE_COLUMN, build_scanner_projection, canonical_output_schema, null_columns,
- project_to_canonical, validate_projection_names, wants_row_id,
+ project_to_canonical, scale_distance, validate_projection_names, wants_row_id,
};
use super::sstable_cache::{DatasetCache, SsTableWarmer, open_sstable};
use crate::session::Session;
@@ -71,12 +71,14 @@ impl ProbeBounds {
/// SortExec: order_by=[_distance ASC], fetch=k (per partition, parallel)
/// UnionExec
/// ProjectionExec (canonical output schema)
-/// SortExec(_distance, fetch=k)
+/// ProjectionExec (_distance rescaled, unrefined only; see `with_distance_scale`)
+/// SortExec(_distance, fetch=k)
/// MemTableBruteForceVectorExec or VectorIndexExec: active memtable KNN
/// ProjectionExec (canonical output schema)
/// ProjectionExec (null_columns _rowid)
/// PkBlockFilterExec: block-list (SSTable)
-/// KNNExec: SSTable gen N, fetch=ceil(k*overfetch) (fast_search)
+/// ProjectionExec (_distance rescaled, unrefined only)
+/// KNNExec: SSTable gen N, fetch=ceil(k*overfetch) (fast_search)[.refine()?]
/// … one per SSTable gen …
/// ProjectionExec (canonical output schema)
/// PkBlockFilterExec: block-list (base)
@@ -130,6 +132,19 @@ pub struct LsmVectorSearchPlanner {
/// generation's `IVF_HNSW_SQ` — and is ignored by an arm whose index is not
/// a graph. `None` leaves each arm at its own default.
ef: Option,
+ /// See [`Self::with_refine_factor`].
+ refine_factor: Option,
+ /// See [`Self::with_distance_scale`].
+ distance_scale: f32,
+}
+
+/// Multiple of the exact distance a flushed generation's `IVF_HNSW_SQ` index
+/// reports: 2 for cosine (L2 over normalized vectors), otherwise 1.
+fn sstable_distance_scale(distance_type: lance_linalg::distance::DistanceType) -> f32 {
+ match distance_type {
+ lance_linalg::distance::DistanceType::Cosine => 2.0,
+ _ => 1.0,
+ }
}
impl LsmVectorSearchPlanner {
@@ -164,6 +179,8 @@ impl LsmVectorSearchPlanner {
filter: None,
distance_range: (None, None),
ef: None,
+ refine_factor: None,
+ distance_scale: 1.0,
}
}
@@ -194,6 +211,30 @@ impl LsmVectorSearchPlanner {
self
}
+ /// Refine the base and SSTable arms by `rf` so every arm reports exact
+ /// distances; `None` leaves them unrefined.
+ pub fn with_refine_factor(mut self, refine_factor: Option) -> Self {
+ self.refine_factor = refine_factor;
+ self
+ }
+
+ /// Report unrefined distances, and take the distance range, as `scale`
+ /// times the exact distance. Cosine SQ/PQ indexes report `2 * (1 - cos)`,
+ /// so `2.0` matches such a base. Ignored when refining; defaults to `1.0`.
+ pub fn with_distance_scale(mut self, scale: f32) -> Self {
+ self.distance_scale = scale;
+ self
+ }
+
+ /// Factor from an arm's `native` multiple of the exact distance to the
+ /// plan's scale; `1.0` when refining.
+ fn arm_scale(&self, refine_factor: Option, native: f32) -> f32 {
+ match refine_factor {
+ Some(_) => 1.0,
+ None => self.distance_scale / native,
+ }
+ }
+
/// The table's schema carrying each field's id, which is what resolves a
/// generation's stored columns to the table's.
///
@@ -249,12 +290,9 @@ impl LsmVectorSearchPlanner {
/// * `k` - Number of nearest neighbors to return
/// * `nprobes` - Exact number of IVF partitions to search (for IVF-based indexes)
/// * `projection` - Columns to include in output (None = all columns)
- /// * `refine_base_table` - When true, the base-table arm re-ranks its
- /// candidates with exact distances (refine factor 1). Useful when the base
- /// table uses an approximate index (IVF-PQ) so cross-source distance
- /// comparison is exact. Memtable arms already use exact distances and
- /// never need refine. Auto-enabled whenever stale filtering is on (see
- /// below).
+ /// * `refine_base_table` - When true and no refine factor is set, refine
+ /// the base and SSTable arms by 1. Auto-enabled when a newer generation
+ /// blocks the base table.
/// * `overfetch_factor` - Controls how aggressively sources over-fetch to
/// backfill the rows dropped by stale-row filtering. Values below `1.0`
/// are rejected; stale filtering is always enabled.
@@ -353,11 +391,15 @@ impl LsmVectorSearchPlanner {
true, // include _distance — KNN always produces it
)?;
- // Refine the base table when explicitly requested, or whenever the base
- // is blocked (it then over-fetches its approximate-index candidates, so
- // distances must be re-ranked to exact before the cross-source merge).
+ // The merge needs every arm on one scale. A blocked base over-fetches
+ // approximate candidates, so it refines.
// `block_lists` is non-empty exactly when a newer generation exists.
- let refine_base = refine_base_table || !block_lists.is_empty();
+ let has_base = sources
+ .iter()
+ .any(|source| matches!(source, LsmDataSource::BaseTable { .. }));
+ let refine_factor = self
+ .refine_factor
+ .or_else(|| (refine_base_table || (has_base && !block_lists.is_empty())).then_some(1));
// Stage per-source over-fetch decisions, then build every KNN plan
// concurrently — the builds are independent and a sequential loop was
@@ -384,8 +426,8 @@ impl LsmVectorSearchPlanner {
// Type-erased, not merely boxed: the `Send` proof recurses through a
// boxed future's concrete type but stops at a trait object, and an arm
// resolves a generation's schema before it searches.
- let built = futures::future::try_join_all(arm_inputs.iter().map(
- |(source, is_base, _, _, fetch_k)| {
+ let built =
+ futures::future::try_join_all(arm_inputs.iter().map(|(source, _, _, _, fetch_k)| {
let arm: futures::future::BoxFuture<'_, Result>> =
Box::pin(self.build_knn_plan(
source,
@@ -393,12 +435,11 @@ impl LsmVectorSearchPlanner {
*fetch_k,
probe_bounds,
projection,
- *is_base && refine_base,
+ refine_factor,
));
arm
- },
- ))
- .await?;
+ }))
+ .await?;
let mut knn_plans = Vec::new();
for ((_, is_base, _, blocked, _), knn) in arm_inputs.iter().zip(built) {
@@ -516,7 +557,7 @@ impl LsmVectorSearchPlanner {
k: usize,
probe_bounds: ProbeBounds,
projection: Option<&[String]>,
- refine: bool,
+ refine_factor: Option,
) -> Result> {
match source {
LsmDataSource::BaseTable { dataset } => {
@@ -558,10 +599,8 @@ impl LsmVectorSearchPlanner {
}
// Memtables cover unindexed rows; only search indexed data here.
scanner.fast_search();
- // Re-rank base candidates with exact distances so they're
- // directly comparable to memtable distances in the merge.
- if refine {
- scanner.refine(1);
+ if let Some(refine_factor) = refine_factor {
+ scanner.refine(refine_factor);
}
scanner.create_plan().await
}
@@ -635,7 +674,12 @@ impl LsmVectorSearchPlanner {
}
};
scanner.nearest(&vector_column, query_arr.as_ref(), k)?;
- scanner.distance_range(self.distance_range.0, self.distance_range.1);
+ let scale =
+ self.arm_scale(refine_factor, sstable_distance_scale(self.distance_type));
+ scanner.distance_range(
+ self.distance_range.0.map(|lower| lower / scale),
+ self.distance_range.1.map(|upper| upper / scale),
+ );
if let Some(minimum_nprobes) = probe_bounds.minimum_nprobes {
scanner.minimum_nprobes(minimum_nprobes);
}
@@ -647,10 +691,14 @@ impl LsmVectorSearchPlanner {
scanner.ef(ef);
}
scanner.fast_search();
+ if let Some(refine_factor) = refine_factor {
+ scanner.refine(refine_factor);
+ }
// Boxed for the reason the scan planner's arm gives: a
// generation resolves its own schema before scanning, and
// the inlined future is too deep for the `Send` proof.
- generation.reconcile_above(Box::pin(scanner.create_plan()).await?, &above)
+ let plan = scale_distance(Box::pin(scanner.create_plan()).await?, scale)?;
+ generation.reconcile_above(plan, &above)
}
LsmDataSource::ActiveMemTable {
batch_store,
@@ -717,7 +765,12 @@ impl LsmVectorSearchPlanner {
}
};
scanner.nearest(&vector_column, query_vector, k)?;
- scanner.distance_range(self.distance_range.0, self.distance_range.1);
+ // This arm's distances are exact.
+ let scale = self.arm_scale(refine_factor, 1.0);
+ scanner.distance_range(
+ self.distance_range.0.map(|lower| lower / scale),
+ self.distance_range.1.map(|upper| upper / scale),
+ );
if let Some(minimum_nprobes) = probe_bounds.minimum_nprobes {
scanner.minimum_nprobes(minimum_nprobes);
}
@@ -728,7 +781,7 @@ impl LsmVectorSearchPlanner {
if let Some(ef) = self.ef {
scanner.ef(ef);
}
- let plan = Box::pin(scanner.create_plan()).await?;
+ let plan = scale_distance(Box::pin(scanner.create_plan()).await?, scale)?;
match generation {
Some(generation) => generation.reconcile_above(plan, &above),
None => Ok(plan),
@@ -3384,6 +3437,156 @@ mod tests {
);
}
+ /// SSTable `IVF_HNSW_SQ` reports cosine as `2 * (1 - cos)` and the memtable
+ /// reports it exactly; both must land on the requested scale.
+ #[tokio::test]
+ async fn test_vector_search_cosine_arms_share_one_scale() {
+ use crate::dataset::mem_wal::scanner::collector::{InMemoryMemTableRef, InMemoryMemTables};
+ use crate::dataset::mem_wal::scanner::data_source::ShardSnapshot;
+ use crate::dataset::mem_wal::write::{BatchStore, IndexStore};
+ use crate::index::DatasetIndexExt;
+ use crate::index::vector::VectorIndexParams;
+ use datafusion::prelude::SessionContext;
+ use futures::TryStreamExt;
+ use lance_index::IndexType;
+ use lance_index::vector::hnsw::builder::HnswBuildParams;
+ use lance_index::vector::ivf::IvfBuildParams;
+ use lance_index::vector::sq::builder::SQBuildParams;
+ use lance_linalg::distance::DistanceType;
+
+ // A vector of norm `norm` whose cosine distance to the query is `cos_distance`.
+ let at = |cos_distance: f32, norm: f32| {
+ let cos = 1.0 - cos_distance;
+ [norm * cos, norm * (1.0 - cos * cos).sqrt(), 0.0, 0.0]
+ };
+ let query = {
+ use arrow_array::builder::Float32Builder;
+ let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 4);
+ builder.values().append_slice(&[1.0, 0.0, 0.0, 0.0]);
+ builder.append(true);
+ builder.finish()
+ };
+ // Exact cosine distance by id; the arms alternate so a scale mismatch
+ // between them reorders the result.
+ let expected = [(1, 0.05), (3, 0.2), (2, 0.4), (4, 0.6)];
+
+ let schema = create_vector_schema();
+ let temp_dir = tempfile::tempdir().unwrap();
+ let base_uri = format!("{}/base", temp_dir.path().to_str().unwrap());
+ let shard_id = uuid::Uuid::new_v4();
+
+ // SSTable: ids 1 and 2, plus far rows so SQ has a spread to train on.
+ let mut sstable_rows = vec![(1, at(0.05, 3.0)), (2, at(0.4, 2.0))];
+ sstable_rows.extend((10..20).map(|id| (id, at(1.9, 1.0 + id as f32 * 0.1))));
+ let gen1_uri = format!("{}/_mem_wal/{}/gen_1", base_uri, shard_id);
+ let mut gen1 = create_dataset(&gen1_uri, vec![batch_rows(&schema, &sstable_rows)]).await;
+ let hnsw_sq = VectorIndexParams::with_ivf_hnsw_sq_params(
+ DistanceType::Cosine,
+ IvfBuildParams::new(1),
+ HnswBuildParams::default(),
+ SQBuildParams::default(),
+ );
+ gen1.create_index(&["vector"], IndexType::Vector, None, &hnsw_sq, true)
+ .await
+ .unwrap();
+
+ // Active memtable: ids 3 and 4.
+ let batch_store = Arc::new(BatchStore::with_capacity(16));
+ let mut index_store = IndexStore::new();
+ index_store.enable_pk_index(&[("id".to_string(), 0)]);
+ index_store.add_hnsw(
+ "vector_hnsw".to_string(),
+ 1,
+ "vector".to_string(),
+ DistanceType::Cosine,
+ 64,
+ 8,
+ );
+ let batch = batch_rows(&schema, &[(3, at(0.2, 0.5)), (4, at(0.6, 1.0))]);
+ batch_store.append(batch.clone()).unwrap();
+ index_store
+ .insert_with_batch_position(&batch, 0, Some(0))
+ .unwrap();
+ let active = InMemoryMemTableRef {
+ batch_store,
+ index_store: Arc::new(index_store),
+ schema: schema.clone(),
+ generation: 2,
+ };
+
+ let search = |scale: f32, refine: Option, range: (Option, Option)| {
+ let snapshot = ShardSnapshot::new(shard_id)
+ .with_current_generation(2)
+ .with_sstable(1, "gen_1".to_string());
+ let collector =
+ LsmDataSourceCollector::without_base_table(base_uri.clone(), vec![snapshot])
+ .with_in_memory_memtables(
+ shard_id,
+ InMemoryMemTables {
+ active: active.clone(),
+ frozen: vec![],
+ },
+ );
+ let planner = LsmVectorSearchPlanner::new(
+ collector,
+ vec!["id".to_string()],
+ schema.clone(),
+ "vector".to_string(),
+ DistanceType::Cosine,
+ )
+ .with_distance_scale(scale)
+ .with_refine_factor(refine)
+ .with_distance_range(range.0, range.1);
+ let query = query.clone();
+ async move {
+ let plan = planner
+ .plan_search(&query, 4, 1, None, false, 1.0)
+ .await
+ .unwrap();
+ let batches: Vec = plan
+ .execute(0, SessionContext::new().task_ctx())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ collect_id_dist(&batches)
+ }
+ };
+ let assert_on_scale = |rows: &[(i32, f32)], scale: f32, what: &str| {
+ let ids: Vec = rows.iter().map(|(id, _)| *id).collect();
+ assert_eq!(ids, vec![1, 3, 2, 4], "{what}: wrong order, got {rows:?}");
+ for ((id, distance), (_, exact)) in rows.iter().zip(expected) {
+ assert!(
+ (distance - exact * scale).abs() < 0.02 * scale,
+ "{what}: id={id} reported {distance}, expected {}",
+ exact * scale
+ );
+ }
+ };
+
+ assert_on_scale(&search(1.0, None, (None, None)).await, 1.0, "exact scale");
+ assert_on_scale(
+ &search(2.0, None, (None, None)).await,
+ 2.0,
+ "SQ/PQ base scale",
+ );
+ // Refining reports exact distances whatever the scale.
+ assert_on_scale(&search(2.0, Some(1), (None, None)).await, 1.0, "refined");
+
+ // On the 2x scale, [0.3, 1.0) is cosine distance [0.15, 0.5): one row
+ // from each arm.
+ let ids: Vec = search(2.0, None, (Some(0.3), Some(1.0)))
+ .await
+ .into_iter()
+ .map(|(id, _)| id)
+ .collect();
+ assert_eq!(
+ ids,
+ vec![3, 2],
+ "distance_range must apply on the scale it is given in"
+ );
+ }
+
#[tokio::test]
async fn test_vector_search_sstable_superseded_by_newer_sstable() {
// An older SSTable's stale row must be suppressed by a newer