fix(mem_wal): rank every vector search arm on one distance scale

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}:
    *     <ul>
@@ -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(

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

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<dyn Array> = 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<f64> = sq.update_bounds::<Float32Type>(vectors)?;

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<f32> = (0..dim).map(|_| next()).collect();
+            let norm = row.iter().map(|x| x * x).sum::<f32>().sqrt();
+            let target = 0.2 + 4.8 * (next() * 0.5 + 0.5);
+            vectors.extend(row.iter().map(|x| x * target / norm));
+        }
+        let query: Vec<f32> = (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<usize> = (0..num_vectors).collect();
+        by_distance.sort_by(|a, b| cosine_distance(*a).total_cmp(&cosine_distance(*b)));
+        let nearest: std::collections::HashSet<i32> =
+            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::<Int32Array>().unwrap();
+        let distances = batch["_distance"]
+            .as_any()
+            .downcast_ref::<Float32Array>()
+            .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;

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<u32>,
     /// Distance metric; `None` defaults to L2 (matching the unindexed memtable arm).
     metric_type: Option<DistanceType>,
 }
@@ -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?;

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<dyn ExecutionPlan>, factor: f32) -> Result<Arc<dyn ExecutionPlan>> {
+    if factor == 1.0 {
+        return Ok(plan);
+    }
+    let input_schema = plan.schema();
+    let project_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = input_schema
+        .fields()
+        .iter()
+        .enumerate()
+        .map(|(idx, field)| {
+            let column: Arc<dyn PhysicalExpr> = 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