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>
|
@geruh @yanghua @LuciferYang — could I get a maintainer review on this when you have a moment? Why this exists: the current path to KNN entries in Spark is a JVM cross join + distance compute. That's only viable on small datasets, and even there it's very slow; at larger scale it doesn't complete at all — brute-force cross join isn't a feasible way to get nearest neighbors at any real size. This series replaces it with Lance's native indexed vector search: The offload is gated on index availability, with a dataset-size gate planned as well. This PR: slice 1 of 3 of the indexed Status: CI is green across Spark 3.4–4.2, and the Lance Gatekeeper's recommendation is approve with a non-blocking risk (the |
|
I’m missing some context. Is this a feature that was already discussed and agreed to add? |
There was a problem hiding this comment.
🟡 Gate recommendation: maintainer decision required.
The open context question exposes a product-scope decision the prior Review did not resolve. The available tracking record—#541, superseded by #798—documents the proposal and split, but not a verified maintainer agreement to add the Spark 4.2 indexed APPROX NEAREST path.
The probe-core implementation remains technically supported, with only the bounded conservative _score fallback risk previously noted. Before accepting this foundation, maintainers should confirm that the feature direction and staged landing are wanted.
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