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

[hamersaw]

Problem

LsmVectorSearchPlanner merges its arms by _distance, but with cosine the arms reach the merge on different scales:

• A flushed generation is indexed as IVF_HNSW_SQ, which ranks cosine as L2 over normalized vectors: 2 * (1 - cos).
• The memtable ranks cosine exactly: 1 - cos.
• A caller merging its own base-table search (sophon's query node does) gets whatever that index reports, which is 2 * (1 - cos) for SQ and PQ but 1 - cos for FLAT and RQ.

So SSTable rows always lose to memtable rows of equal or worse similarity, and distance_range cuts each arm on a different scale.

Separately, flush wrote the memtable's unnormalized vectors into the cosine SQ index. The index's search normalizes the query and ranks by L2, so a generation of non-unit embeddings ranked by ‖q̂ − x‖², not cosine: recall@10 was 2/10 against exact cosine.

Fix

• write_sq_storage normalizes vectors before SQ training when the index is cosine, as a regular cosine build does. The graph is unchanged: cosine over raw vectors orders neighbors the same as L2 over normalized ones.
• Two planner builders:
  • with_refine_factor(Option<u32>): refines the base and SSTable arms, so every arm reports exact distances.
  • with_distance_scale(f32): when unrefined, every arm reports this multiple of the exact distance. The memtable scales its exact distances, SSTables convert from their SQ index's multiple, and the distance range is converted per arm. Defaults to 1.0.
• plan_search keeps its signature (the Python and Java bindings expose it). refine_base_table now refines SSTable arms too, so the arms stay on one scale.
• LsmScanner::refine passes its factor through instead of collapsing it to a bool.

Tests

• test_flusher_cosine_index_ranks_unnormalized_vectors: non-unit cosine vectors flushed to a generation keep recall@10 ≥ 9/10 and report 2 * (1 - cos). Without the fix: 2/10.
• test_vector_search_cosine_arms_share_one_scale: an SSTable and a memtable come back in true order at scale 1, at scale 2, and refined, and distance_range admits one row from each arm on the 2× scale.
• cargo test -p lance --lib dataset::mem_wal: 856 passed. cargo fmt is clean.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

Keep candidate-selection bounds and exact-rerank bounds in their respective distance units. The new SSTable refinement path still drops valid cosine neighbors when a distance range is supplied.

[lance-gatekeeper]

Refined cosine searches silently lose in-range SSTable rows. When refining, arm_scale returns 1.0, so these bounds remain in exact cosine units. However, Scanner::ann forwards them to HNSW/SQ, which filters on 2 * (1 - cos), before Scanner::flat_knn applies the same bounds to exact distances. In this PR's fixture, id 2 has exact distance 0.4 but SQ distance about 0.763; an upper bound of 0.5 discards it before refinement. Automatic refinement with a base table reaches this path too, and retrying the query cannot recover the missing neighbor.

Use separate bounds for SQ candidate selection and exact reranking, or an exact bounded search path, and cover the refined-range case. Rescaling the final distances cannot restore candidates already discarded.

Reproducer

Add this at the end of test_vector_search_cosine_arms_share_one_scale:

let unrefined = search(1.0, None, (Some(0.15), Some(0.5))).await;
assert_eq!(
    unrefined.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
    vec![3, 2],
);
let refined = search(1.0, Some(1), (Some(0.15), Some(0.5))).await;
assert_eq!(
    refined.iter().map(|(id, _)| *id).collect::<Vec<_>>(),
    vec![3, 2],
    "refined bounded rows: {refined:?}",
);

Run cargo test -p lance --lib test_vector_search_cosine_arms_share_one_scale -- --nocapture.

On 8c504101aaf8380ede4f07d253fc23e8d60ed9b7, the unrefined assertion passes. The refined assertion fails with left: [3], right: [3, 2]; its only row is (3, 0.19999999).

[hamersaw]

❌ Gate recommendation: request changes.

Keep candidate-selection bounds and exact-rerank bounds in their respective distance units. The new SSTable refinement path still drops valid cosine neighbors when a distance range is supplied.

I think we're fixing this in #9851 as part of the follow up. The plan was to merge this fix afterwards because it applies to current non-WAL tables as well. So this is partitioning WAL-fixes with general vector fixes.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

Keep the shared scanner fix in #9851, as you proposed, and land it before #9849. Then rebase this PR and verify the refined SSTable distance-range regression.

The current head still silently drops valid cosine neighbors before exact reranking. #9851 targets this path by converting candidate-selection bounds to the index’s distance units while preserving exact bounds for refinement.