Indexed nearest-neighbor join optimization for lance-spark (tracking issue, 7-PR split)

[sezruby]

Summary

Add an indexed nearest-neighbor join optimization to lance-spark: route SQL of the shape

SELECT *
FROM   queries q
LEFT OUTER JOIN documents d
  APPROX NEAREST 10 BY DISTANCE vector_l2_distance(q.vec, d.vec)

through a per-fragment Lance index probe + Spark-side merge instead of the default
O(|L| × |R|) cross product. Targets workloads where the right side is a Lance
dataset with a vector index (IVF-PQ / HNSW); falls back to the existing Spark plan when
those conditions don't hold.

PoC PR (against my fork's main, ~1 week behind upstream — same diff a draft PR
against lance-format/lance-spark:main would show):
sezruby#1

Headline numbers from the PoC (full results in BENCHMARK_RESULTS.md):

• DataFrame API vs. naive crossJoin + UDF + window: ~600× (small scale, oracle-validated)
• SQL with Spark 4.2 NearestByJoin rule ON vs. OFF (apples-to-apples): ~17×
• Cohere wiki, 1.8M rows × dim 1024, on a real cluster: ~160× (median across 7 iterations, range 100–200×)

Design (TL;DR)

Lance vector indexes are fragment-local, so per-fragment probes are independent
and trivially parallel. The runtime path is:

LanceProbeExec        ← per-task: open Lance, run nearest search, emit (leftId, rowId, dist)
   │
   ▼ ShuffleExchangeExec (Catalyst-inserted, hash-partitioned on leftId)
   │
LanceMergeExec        ← per-leftId top-K heap merge across probe outputs
   │
   ▼
LanceMaterializeExec  ← point-fetch the chosen rowIds back from Lance

All three are real Catalyst SparkPlan nodes — visible in df.explain(),
AQE engages on the merge shuffle, ColumnPruning / projection fusion all work.
No public API change for end users once the rule lands: registering one
Spark session extension is the entire opt-in.

Full design (problem, alternatives, why each piece is shaped the way it is, the
ColumnPruning gotcha, the Catalyst-vs-Kryo payload analysis, fragment-grouping
math): DESIGN.md.

Relationship to SPARK-56395

SPARK-56395 lands a
first-class NearestByJoin operator in Spark 4.2:

• SQL: ... APPROX NEAREST K BY DISTANCE expr
• DataFrame (Scala/Java/PySpark) and Spark Connect: df.nearestByJoin(other, expr, numResults, mode, direction, joinType)
  (apache/spark#55629 catalyst,
  apache/spark#55682 DataFrame API,
  both targeted at master / 4.2)
• Default rewrite (RewriteNearestByJoinRule): cross product + min_by(struct, expr, K) — correct, but slow on large right sides

This is a tailwind for our work, not a conflict. Once Spark 4.2.0 publishes,
the lance-spark-knn-4.2_2.13 module (deferred from the split below) ships
a Catalyst rule that intercepts NearestByJoin when the right side is a
Lance dataset with a vector index, and substitutes the same 3-exec staged
pipeline. End-user code is unchanged — it's just APPROX NEAREST SQL or
df.nearestByJoin(...), and the optimization engages automatically. The
df.kNearestJoin extension (PR #5 below) covers Spark 3.5 / 4.0 / 4.1
users until 4.2 reaches them.

Also worth linking here is an ANN-proposal doc
(NEARESTBYJOIN_ANN_PROPOSAL.md)
written during the SPARK-56395 conversation — describes how the per-format
indexed path generalizes (Lance today; the same shape works for any format
that has a vector index).

PR split (7 PRs, ~13.3k LoC total)

The PoC branch is too large to land as one PR. Split into 7 self-contained,
independently-reviewable PRs (full rationale in
UPSTREAM_DELIVERY_PLAN.md).
Phase ordering is preserved so commits read chronologically.

#	Title	Size	Tests	Review effort
1	Phase 0 — LanceProbe primitive (per-task nearest search)	~700 LoC	4	2–3 h
2	Phase 0/1 — IndexedNearestJoin.apply + staged RDD pipeline + TopKHeap	~1.5k LoC	~10	4–6 h
3	Phase 1.5 — fragment-grouped probe (probeParallelism > 1, LPT bin-packing)	~600 LoC	4	2–3 h
4	Phase 2 — 3-exec staged Catalyst operators + AQE visibility	~1.3k LoC	~24	6–10 h
5	df.kNearestJoin DataFrame extension (sugar on top of #4)	~250 LoC	6	1 h
6	Phase 3 — refineFactor, ef, IVF-PQ recall test	~600 LoC	5	2 h
7	Spark 4.0 compat (reflection bridge for Dataset.ofRows move)	~60 LoC	(existing)	30 min

Critical path is #1 → #2 → #4. Once #4 is in, #3 / #5 / #6 / #7 are
parallel-reviewable.

Each PR is self-contained: it adds tests that exercise only the feature it
ships, and doesn't depend on later PRs to compile. PR #4 is the heaviest —
it carries a "3-exec staged split — root cause and fix" post-mortem in
the description (the ColumnPruning guard + Project(Nil) insertion +
0-field UnsafeRow crash + references = child.outputSet fix), since
that's the most subtle part of the feature.

Out of scope for this issue / kept on the fork

• All benchmarks (benchmark/ tree) — internal tooling, pulls in
  HTTP-fetch / local-FS logic that doesn't belong in the connector.
  External BENCHMARK_RESULTS.md post is the artifact instead.
• lance-spark-knn-4.2_2.13 module — the SPARK-56395 interception path,
  ready against 4.2-SNAPSHOT today, gated on Spark 4.2.0 publishing to
  Maven Central. Will land as a follow-up PR (PR Discuss the version of lance and lance-catalog to depend on #8) at that point. Once
  in, APPROX NEAREST SQL and df.nearestByJoin(...) on a Lance right
  side automatically engage the staged pipeline — no per-user opt-in.
• Deployment-specific build config (Linux-x86_64 shade filter for
  managed-Spark distributions with volume-upload ingress timeouts).

Asks before opening PRs

1. Any concerns with the overall direction or the 7-PR split before I start
   sending the first one?
2. Is there a preferred maintainer to tag for review on each phase, or
   should I just open the PRs and let assignment happen normally?
3. Anything in flight (zonemap, distributed index builds, full-text search
   SQL ext, vector-index roadmap) that this should coordinate with?

I'll start with PR #1 (LanceProbe primitive — smallest, no Spark integration
surface) once there's a nod here.

[sezruby]

Superseded by #798.

The 7-PR staged-shuffle design tracked here is outdated. The delivered approach is narrower: a single Spark 4.2 SQL-only module (lance-spark-knn-4.2_2.13) that intercepts APPROX NEAREST / NearestByJoin with a native no-shuffle per-partition index probe — there is no staged Probe → Shuffle → Merge → Materialize pipeline, and the df.kNearestJoin DataFrame API and multi-version (3.5 / 4.0 / 4.1) support are deferred.

Implementation: umbrella PR #796, with reviewable slices starting at #797. Tracking continues in #798.