Indexed nearest-neighbor join via distributed per-query Lance index probes

[sezruby]

Summary

Add a KNN-join optimization for the case where the right side is a Lance dataset with a vector index (IVF-PQ / HNSW). Instead of the default O(|L| × |R|) cross-product + top-K, distribute the left rows across tasks and have each task run a native nearest search against the Lance index — one mapPartitions, one native call per query row, no Spark-side shuffle. Falls back to the default plan when the right side isn't an indexed Lance dataset.

Deliberately narrow: the indexed, per-query path only (benchmark name SPARK_DIST), exposed via a df.kNearestJoin DataFrame extension on released Spark (3.5 / 4.0 / 4.1).

Why

• The default rewrite scans the whole right side per query — quadratic and memory-heavy at scale.
• A single native Dataset.nearest() call is already the complete distributed search: Lance probes the IVF index, scans candidate fragments in parallel across its own threads, merges with an in-process heap, and returns the final top-K per query. So the fastest shape is to get out of Lance's way — Spark distributes the queries, Lance does the search.
• Vector indexes are fragment-local, and native search already merges across fragments internally, so per-task probing is trivially parallel with no cross-task recombine.
• The per-query probe scales sub-linearly with corpus size, so it holds where the cross product OOMs.

Approach (TL;DR)

distribute L across tasks (mapPartitions)
   → each task: open the Lance index, run a native nearest search per query row
       (Lance internally probes IVF + scans candidate fragments in parallel and
        returns the final top-K rowIds + distances)
   → return results — no shuffle, no separate merge stage, no second scan

Entry point: df.kNearestJoin(right, expr, k, ...) (Spark 3.5 / 4.0 / 4.1). Recall/latency knobs: nprobes, refineFactor.

Not a staged/shuffle pipeline. An earlier probe → exchange → merge → materialize design was benchmarked at cluster scale and dropped: it added a shuffle, a redundant merge stage, and a second materialize scan to reproduce work a single native call already does in-process, netting out ~break-even with the naive driver loop and ~13× slower than the plain mapPartitions shape. The fewer Spark stages between the query and Lance, the faster.

Scope

In scope

• per-query distributed native index probe (indexed R only), mapPartitions → native call, no shuffle
• recall/latency knobs (nprobes, refineFactor) + an IVF-PQ recall test
• df.kNearestJoin DataFrame entry point + graceful fallback when R isn't an indexed Lance dataset

Out of scope (deferred)

• Fragment/segment-scoped distributed probe. For datasets whose index is too large to hold per task, the native fix is per-task index-segment scoping — deferred to upstream lance#7169 (segment-based distributed vector search), which scopes each task to a segment by UUID and structurally removes the whole-dataset resident cost. Not hand-rolled in the connector: doing the recombine here re-introduces the cross-task per-query merge that benchmarked slower than the plain per-task probe (fragment-grouping with probeParallelism > 1 was measured and lost at every scale — native search already threads across fragments internally).
• Batched multi-query probe (setKeys) on the indexed path — depends on lance#7640 (share IVF partition scans across batch queries), still unmerged; without it the indexed batch path regresses. Follow-up once #7640 lands and is wired end-to-end.
• No-index brute-force path.
• Native APPROX NEAREST … BY DISTANCE interception (SPARK-56395 / Spark 4.2). Same engine, different trigger. Now unblocked — Spark 4.2.0 is on Maven Central and lance-spark ships the lance-spark-4.2_2.13 module — so this is the next follow-up there.

Memory & size gate

Each task opens the whole-dataset index (fragmentIds = None), so per-executor resident memory
grows with |R|, not with query count. Benchmarking this on a standalone Spark cluster
(synthetic uniform dim=128 vectors, IVF-PQ, 4 concurrent probes per executor, -Xmx fixed)
gave a clear picture — and it is not the hard-OOM cliff we expected:

| |R| (rows) | per-executor peak RSS | JVM heap |
|-------------|-----------------------|----------|
| 20M | ~1.8 GiB | ~0.14 GiB |
| 50M | ~4.3 GiB | ~0.15 GiB |
| 100M | ~6.9 GiB | ~0.16 GiB |

Two things stand out:

1. The JVM heap stays flat (~0.15 GiB) regardless of |R|. The Spark side of this path is thin —
   broadcast the query set, mapPartitions a native probe, collect small (id, rowId, distance)
   rows. All the growth is off-heap.

2. The growth does not translate into OOM. Peak RSS here is VmHWM. Re-running the same probe
   over the same index makes RSS climb monotonically across repetitions while heap stays flat —
   the signature of index data becoming resident as more of it is touched (OS page cache of the
   mmap'd index and/or a bounded internal index cache), not of unbounded anonymous allocation.
   Empirically this never crossed an OOM cliff: across every configuration tried — a 14× reduction
   in configured executor memory, peak RSS running well past the container's nominal limit, and
   corpora up to hundreds of millions of high-dimensional vectors — no executor was OOM-killed, and
   end-to-end latency stayed flat/sub-linear.

So the "gate" is not an OOM guardrail — the unevictable floor (heap + native working buffers)
is small and roughly |R|-independent. It is a page-cache / working-set guardrail: past the
point where the touched index working set no longer fits in an executor's available RAM, index
pages fault in and out of storage per probe and latency degrades. The degradation is gradual (the
measured query-latency curve is sub-linear in |R|, with no wall), not a crash.

Heuristic. Model the resident index footprint as a + b·|R| per executor and fall back to the
plain Spark read when it would exceed a safe fraction of executor RAM:

residentFootprint(|R|)  ≈  a + b · |R|          # per executor, at C concurrent probes
apply native probe   iff  residentFootprint(|R|)  ≤  SAFETY · executorRAM

Least-squares over the three points above gives a ≈ 0.70 GiB, b ≈ 64 MiB per 1M rows
(dim=128, C=4). Notes on the constants:

• Growth is sub-linear — the per-1M slope falls from ~82 MiB (20M→50M) to ~56 MiB (50M→100M) —
  because a fixed query set touches a shrinking fraction of a larger index, so a linear a + b·|R|
  slightly over-predicts at the extremes, i.e. it errs conservative (falls back a little early).
• b tracks the resident index (memory-mapped index pages are shared across an executor's
  concurrent probes, so it is largely independent of C) and scales with the per-row index size —
  PQ code bytes plus the fraction of full vectors pulled during refine — so it grows with vector
  dimension and with nprobes/refineFactor. a (JVM + native runtime + per-task buffers) grows
  with C. Treat both as deployment-calibrated, not universal.

Worked thresholds (|R|_max = (SAFETY·executorRAM − a) / b, dim=128):

executor RAM	SAFETY=0.7	SAFETY=0.8
8 GiB	~78M rows	~91M rows
16 GiB	~167M rows	~192M rows
32 GiB	~345M rows	~396M rows

Where it ships. Because there is no hard OOM cliff — the native path degrades gracefully rather
than crashing — the size gate is a latency optimization, not a safety requirement, so nothing
blocks on it landing early. The natural home is the SQL operator path (APPROX NEAREST … BY DISTANCE, SPARK-56395): the estimate gates at plan time, before any task is scheduled, off the
right dataset's row count and vector dimension. So it's sequenced as a follow-up in the
lance-spark-4.2_2.13 module — now available, since Spark 4.2.0 is on Maven Central — alongside the
native-interception work above; the measured constants here are exactly what that plan-time gate
consumes. Longer term
the per-task resident cost is removed structurally by native segment-scoping (lance#7169), at which
point each task holds only a segment and the gate relaxes to per-segment cost.

Delivery plan

Small, independently-reviewable PRs:

1. LanceProbe primitive — per-task native nearest search + metric/result types; returns (leftId, rowId, distance) per query row. Standalone, unit-tested.
2. df.kNearestJoin + recall — mapPartitions → native probe per task (no shuffle), gating + graceful fallback when R isn't an indexed Lance dataset, nprobes / refineFactor, IVF-PQ recall test.

Critical path 1 → 2; I'd open PR 1 first once the direction looks good. (Size-threshold gate is sequenced as a follow-up in the Spark-4.2 SQL module — see Memory & size gate — since the benchmark has fixed its constants and shown it's a soft latency guardrail, not an OOM requirement; nothing here blocks on it.)

---

Narrower, indexed-per-query re-scope of lance-format#541.