feat(knn): probe-core primitives for indexed Lance vector join (Spark 4.2) [1/3] - #797
feat(knn): probe-core primitives for indexed Lance vector join (Spark 4.2) [1/3]#797sezruby 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>
|
ACTION NEEDED The PR title and description are used as the merge commit message. Please update your PR title and description to match the specification. For details on the error please inspect the "PR Title Check" action. |
…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>
Make the nearest-probe eligibility contract schema-level instead of
projection-level. Lance's nearest scan always injects _rowid and the
_distance/_score metadata; if the dataset's own schema already has a
column by one of those names, the injected metadata shadows it and NO
probe route recovers the physical column. Empirically an all-columns
probeRows scan reads a physical _distance out-of-band as the ranking
score and silently drops it from the payload — data loss, not an error.
Replace the projection-only fusesCleanly guard (which only caught an
explicit reserved column in the projection list, missing the empty /
all-columns form) with:
- LanceProbe.schemaSupportsNearest / reservedSchemaColumns: the pure
eligibility primitive the Catalyst rule consults to DECLINE such a
table up front and fall back to default nearest-by execution.
- requireNearestCompatibleSchema(): a defensive backstop reading the
dataset schema once, called at the top of probe() and probeRows(),
throwing a clear error naming the offending column.
Regression: write a real dataset WITH a _distance column and assert the
schema is reported non-nearest-compatible and BOTH probe entry points
fail fast naming _distance instead of returning a lossy payload; plus a
pure schemaSupportsNearest / reservedSchemaColumns contract test.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
There was a problem hiding this comment.
The remaining reserved-name failure is fixed: nearest probing now checks the opened dataset schema before scanner creation and declines when Lance-injected metadata would collide with a user column. The focused real-dataset tests cover both probe entry points and the empty/all-columns case.
The guard also treats _score as reserved even though the pinned vector-search path currently emits _distance. That can conservatively exclude some schemas from this optimization, but the planner fallback preserves correctness. No further change is requested for this pull request.
What
The probe-core layer of the indexed
APPROX NEARESTjoin over Lance (SPARK-56395) — a new opt-in modulelance-spark-knn-4.2_2.13. This is the JVM primitive layer only, with no Spark plan integration, so it can be reviewed on its own.This is slice 1 of 3, split out of the umbrella PR #796 for reviewability. The full end-to-end feature (with the Catalyst rewrite and SQL tests) lives in #796; this PR carries only the foundation.
Contents
LanceProbe— opens a Lance dataset once and serves per-query nearest searches against a fixed fragment set.probereturns row references + scores (payload deferred);probeRowsfolds the search and payload projection into a single scan for the no-overfetch path. Handles 64-bit-unsigned row ids and materializes payloads through the connector's canonical Arrow→Spark adapter.TopKHeap— bounded, best-first top-K merge, metric-direction aware.Metric/ScoredRowRef/MaterializedHit— small value types.Tests
Run standalone against a real Lance dataset written by Spark. No vector index is required — Lance's brute-force scan is an exact recall=1.0 oracle, which isolates probe-core correctness from index quality:
LanceProbeValidationTest— probe result shape, brute-force-oracle equivalence,probeRows==probe+materializeparity, dataset-handle reuse across probes, fragment-id restriction, and the executor namespace policy.TopKHeapTest— bounded merge / ordering.13 tests pass.
Follow-ups
LanceKnnJoinStage).APPROX NEARESTinterception rule, physical operator, session extension, and end-to-end SQL/recall tests + docs.🤖 Generated with Claude Code