Skip to content

feat: indexed nearest-neighbor join over Lance (no-shuffle) - #13

Closed
sezruby wants to merge 1 commit into
mainfrom
feat/indexed-nearest-join
Closed

feat: indexed nearest-neighbor join over Lance (no-shuffle)#13
sezruby wants to merge 1 commit into
mainfrom
feat/indexed-nearest-join

Conversation

@sezruby

@sezruby sezruby commented Aug 18, 2026

Copy link
Copy Markdown
Owner

What

Adds a Lance-native, no-shuffle indexed nearest-neighbor join, in three new modules:

  • lance-spark-knn_2.12 / _2.13 — Spark-version-agnostic core plus the DataFrame API (df.kNearestJoin). The _2.13 module cross-builds the _2.12 sources.
  • lance-spark-knn-4.2_2.13 — Catalyst integration for the SQL APPROX NEAREST k BY DISTANCE ... syntax (the NearestByJoin operator from SPARK-56395, present in Spark 4.2).

How it works

Per partition, a single mapPartitions drives one LanceProbe per left row — a native ANN/exact vector search against the right-side Lance dataset — keeps a bounded top-K heap, then late-materializes the matched right rows by row address. There is no shuffle, no broadcast, and no requiredChildDistribution.

The DataFrame path (df.kNearestJoin) and the SQL path share the exact same LanceKnnJoinStage.runPartition core:

  • DataFrame: left.rdd.mapPartitions(...)spark.createDataFrame(...).
  • SQL: a post-hoc resolution rule rewrites NearestByJoin to a LanceKnnJoinLogicalPlan, and a planner strategy lowers it to LanceKnnJoinExec. Opt-in behind spark.lance.knn.indexedNearestByJoin.enabled (off by default); when disabled, the query falls through to Spark's built-in brute-force rewrite. A right-side WHERE is pushed down to Lance as a prefilter, or the rewrite is refused when the predicate can't be pushed in full.

Metrics supported: L2, cosine, dot.

Injection point

The rule is registered via injectPostHocResolutionRule, not injectOptimizerRule: Spark's RewriteNearestByJoin runs in FinishAnalysis, before the optimizer batch, so by the time an injected optimizer rule fires the NearestByJoin operator has already been replaced. Post-hoc resolution is the only injection point that still sees the unrewritten operator.

Tests

  • IVF-PQ recall vs. a brute-force oracle.
  • LanceProbe, top-K heap, and size-gate unit tests.
  • SQL end-to-end (parser → rule → strategy → exec against a real Lance dataset, oracle-checked) plus a rule-unit suite covering metric/direction matching, opt-in gating, alias/filter unwrapping, and the prefilter translator.

Notes for review

  • I could not compile or run these locally in this environment — CI on the fork and your review are the gate. The code is written against lance 11.0.0-beta.10 / lance-spark 0.7.1 / Spark 4.2 APIs, but has not been locally built.
  • The dataset size gate (LanceKnnSizeGate) is wired on the DataFrame path only. It is intentionally not invoked from the SQL rule, because the rule runs at plan time and opening the dataset there would break rule-unit tests that use fake URIs.
  • Cross-module dependencies use ${project.version}; junit-jupiter and lance-core are inherited from the root pom's global dependencies.

🤖 Generated with Claude Code

Add a Lance-native, no-shuffle indexed nearest-neighbor join in three modules.

- lance-spark-knn_2.12 / _2.13: Spark-version-agnostic core plus the DataFrame
  API (df.kNearestJoin). A per-partition mapPartitions drives one LanceProbe per
  left row (native ANN/exact vector search on the right-side Lance dataset), keeps
  a bounded top-K heap, then late-materializes the matched right rows by row
  address. No shuffle, no broadcast, no requiredChildDistribution.
- lance-spark-knn-4.2_2.13: Catalyst integration that rewrites Spark 4.2's
  "APPROX NEAREST k BY DISTANCE ..." (the NearestByJoin operator from SPARK-56395)
  onto the same no-shuffle path, via a post-hoc resolution rule plus a planner
  strategy. Opt-in behind spark.lance.knn.indexedNearestByJoin.enabled; when
  disabled the query falls through to Spark's built-in brute-force rewrite. The
  rule also pushes a right-side WHERE down to Lance as a prefilter, or refuses the
  rewrite when the predicate cannot be pushed in full.

Metrics: L2, cosine, dot. Tests: IVF-PQ recall vs. brute-force oracle, probe /
top-K heap / size-gate unit tests, and a SQL end-to-end plus rule-unit suite for
the Catalyst path.

The 2.13 module cross-builds the 2.12 sources. The size gate is wired on the
DataFrame path only.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@sezruby

sezruby commented Aug 20, 2026

Copy link
Copy Markdown
Owner Author

Superseded by #14 — collapsed to a single Spark 4.2 SQL Catalyst module (lance-spark-knn-4.2_2.13); the manual DataFrame API is deferred and can be re-added later. Leaving this branch intact/reopenable.

@sezruby sezruby closed this Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant