feat(knn): indexed Lance vector-join execution stage (Spark 4.2) [2/3] - #799
Open
sezruby wants to merge 3 commits into
Open
feat(knn): indexed Lance vector-join execution stage (Spark 4.2) [2/3]#799sezruby wants to merge 3 commits into
sezruby wants to merge 3 commits into
Conversation
First slice of the indexed APPROX NEAREST join over Lance (SPARK-56395), split out of the umbrella PR for reviewability. This slice is the JVM primitive layer only — no Spark plan integration: - LanceProbe: opens a Lance dataset once and serves per-query nearest searches. `probe` returns row refs + scores (payload fetched later); `probeRows` folds the search and payload projection into one scan for the no-overfetch path. 64-bit-unsigned row-id handling and the canonical Arrow -> Spark payload adapter live here. - TopKHeap: bounded best-first merge, metric-direction aware. - Metric / ScoredRowRef / MaterializedHit: value types. Tests run standalone against a real Lance dataset (no vector index needed — the brute-force scan is a recall=1.0 oracle): LanceProbeValidationTest (probe shape, brute-force equivalence, probeRows == probe+materialize parity, handle reuse, fragment restriction, namespace policy) and TopKHeapTest. 13 tests pass. New opt-in module `lance-spark-knn-4.2_2.13` (Spark 4.2 / Scala 2.13). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…reserved-name collision Addresses the gatekeeper findings on the probe-core slice: - Metric: Lance returns a distance for every metric, including cosine (1 - cosine_similarity) and dot (1 - dot_product), so smaller is better for all three. Cosine/Dot were flagged larger-is-better, which made the merge heap retain the farthest neighbor. - TopKHeap: derive admission from the heap's own ordering (ord.lt) instead of a raw float comparison, so a NaN worst-survivor is evictable rather than pinning a slot forever. - LanceProbe.probeRows: preserve projected payload columns that lack a supplied Spark type through the generic Arrow conversion (matching materialize/readRows) instead of silently dropping them. - LanceProbe.probeRows: reject a projection that names a column the nearest scan injects (_rowid / _distance / _score) so it cannot collide inside the fused scan; expose fusesCleanly/ReservedProjectionColumns so the join stage can route such schemas to the split probe + materialize path. - .bumpversion.toml: register lance-spark-knn-4.2_2.13/pom.xml so release version bumps reach the new module. Regressions added: metric direction through the size-1 heap (all three metrics), NaN eviction, unmapped-field preservation, and reserved-name rejection. Module test phase: 17 tests, 0 failures. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add the per-partition execution stage that drives the indexed nearest-by join, layered on the probe-core primitives from [1/3]: - LanceKnnJoinStage.runPartition: open one LanceProbe per partition, stream the join output lazily (no per-partition buffering), and close the native handle on task completion. - Fused vs split routing (foldsInOneScan): when there is no over-fetch and the payload projection does not collide with an injected column, search + project in ONE native scan; otherwise probe -> trim -> late-materialize by _rowid. - resolveReadContext / mergeReadOptions: pin the Lance snapshot once on the driver and mirror the connector's incompatible-pinned-ref guard. - coerceToSpark: schema-aware payload coercion (structs, arrays, and the Arrow entry-list shape a MapType cell actually arrives in) to the encoder's shape. Tests: LanceKnnJoinStageTest covers lazy output, fold-vs-split routing for reserved-name projections, read-option merge conflicts, and payload coercion — backend-free. Full module: 30 pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
![lance-gatekeeper[bot]](https://avatars.githubusercontent.com/in/4433141?s=60&v=4){kind=small}
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The per-partition native probe is a worthwhile replacement for Spark’s all-pairs rewrite, but this stack is not safe to merge yet. The verified reserved-metadata discussion shows that both fused and split nearest scans fail on a legal user _distance column. Fix that shared eligibility contract in prerequisite #797, merge it first, then rebase this slice.
This execution slice also needs a real-dataset stage regression that exercises the task-scoped probe and both routing paths, so the shipped stage—not only its helpers—is verified.
| * the Spark-type coercion of materialized right-side payloads | ||
| * ([[LanceKnnJoinStage.coerceToSpark]]). Both are `private[knn]`, so this test lives in the | ||
| * `org.lance.spark.knn.internal` package to reach them. The full probe → trim → materialize | ||
| * pipeline is covered by the e2e test in this module against a real Lance dataset. |
There was a problem hiding this comment.
This comment claims end-to-end coverage for the new execution stage, but no submitted test invokes runPartition or processRow against a real Lance dataset; this suite exercises only a stubbed iterator, coercion helpers, option merging, and routing predicates. That leaves the probe-to-output path, outer/null behavior, task-scoped cleanup, and fused/split routing unverified, contrary to the repository requirement that new functionality include integration tests. Add a real-dataset happy-path stage test plus the key error cases here, including empty/all-columns and reserved-metadata projections.
lance-gatekeeper
Bot
added
the
K-changes
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
1 participant
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.
Part [2/3] of the indexed nearest-by vector join over Lance, split by architectural layer from the umbrella #796.
This slice adds the per-partition execution stage (
internal/LanceKnnJoinStage) that drives the join, layered on the probe-core primitives from #797 ([1/3]).Stacked on #797
This PR is stacked on #797. Until #797 merges to
main, this diff also shows the [1/3] probe-core files (Metric,TopKHeap,ScoredRowRef,LanceProbe) becausemaindoes not yet contain them; once #797 merges, this diff reduces to the join-stage layer alone. Review target here:internal/LanceKnnJoinStage.scalaandinternal/LanceKnnJoinStageTest.scala.What this stage does
runPartition: open oneLanceProbeper partition, stream the join output lazily (no per-partition buffering), close the native handle on task completion (success or failure).foldsInOneScan): with no over-fetch and a payload projection that does not collide with an injected column (_rowid/_distance/_score), search and project in ONE native scan; otherwise probe → trim → late-materialize survivors by_rowid.resolveReadContext/mergeReadOptions: pin the Lance snapshot once on the driver so every task probes one consistent version; mirror the connector's incompatible-pinned-ref guard.coerceToSpark: schema-aware payload coercion (structs, arrays, and the Arrow entry-list shape aMapTypecell arrives in) to the shape the join's encoder accepts.Tests
LanceKnnJoinStageTest(backend-free) covers lazy output, fold-vs-split routing for reserved-name projections, read-option merge conflicts, and payload coercion. Full moduletestphase: 30 pass.The SQL Catalyst rule, physical/logical plan, strategy, and session extension — plus their tests and docs — follow in [3/3].
🤖 Generated with Claude Code