feat: indexed nearest-neighbor join over Lance (no-shuffle) - #13
Closed
sezruby wants to merge 1 commit into
Closed
Conversation
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>
Owner
Author
|
Superseded by #14 — collapsed to a single Spark 4.2 SQL Catalyst module ( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.13module cross-builds the_2.12sources.lance-spark-knn-4.2_2.13— Catalyst integration for the SQLAPPROX NEAREST k BY DISTANCE ...syntax (theNearestByJoinoperator from SPARK-56395, present in Spark 4.2).How it works
Per partition, a single
mapPartitionsdrives oneLanceProbeper 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 norequiredChildDistribution.The DataFrame path (
df.kNearestJoin) and the SQL path share the exact sameLanceKnnJoinStage.runPartitioncore:left.rdd.mapPartitions(...)→spark.createDataFrame(...).NearestByJointo aLanceKnnJoinLogicalPlan, and a planner strategy lowers it toLanceKnnJoinExec. Opt-in behindspark.lance.knn.indexedNearestByJoin.enabled(off by default); when disabled, the query falls through to Spark's built-in brute-force rewrite. A right-sideWHEREis 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, notinjectOptimizerRule: Spark'sRewriteNearestByJoinruns inFinishAnalysis, before the optimizer batch, so by the time an injected optimizer rule fires theNearestByJoinoperator has already been replaced. Post-hoc resolution is the only injection point that still sees the unrewritten operator.Tests
LanceProbe, top-K heap, and size-gate unit tests.Notes for review
lance11.0.0-beta.10/lance-spark0.7.1/ Spark 4.2 APIs, but has not been locally built.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.${project.version};junit-jupiterandlance-coreare inherited from the root pom's global dependencies.🤖 Generated with Claude Code