feat: indexed nearest-by join over Lance (Spark 4.2 SQL) - #796
Draft
sezruby wants to merge 9 commits into
Draft
Conversation
Accelerate Spark 4.2's `APPROX NEAREST k BY DISTANCE ...` join
(SPARK-56395) over Lance tables by intercepting the `NearestByJoin`
operator and running a per-row native ANN/exact probe against the
Lance vector index instead of the default O(|L|x|R|) cross-product +
MaxMinByK lowering.
New opt-in module `lance-spark-knn-4.2_2.13`:
- Catalyst postHocResolutionRule (`IndexedNearestByJoinRule`) that
recognizes `NearestByJoin` and rewrites it to a single no-shuffle
logical node; must run in postHoc (not the optimizer) because
Spark's `RewriteNearestByJoin` fires first in `FinishAnalysis`.
- Planner strategy + exec (`LanceKnnJoinStrategy`/`LanceKnnJoinExec`)
lowering to one `mapPartitions` -> per-left-row native probe ->
bounded top-K heap -> late materialize by row address. No shuffle,
no broadcast.
- Probe core (`LanceProbe`/`LanceKnnJoinStage`/`TopKHeap`/`Metric`)
supporting L2, cosine, and dot metrics with `nprobes`/`refineFactor`
tuning.
- Right-side WHERE pushdown as a Lance prefilter, including nested
struct-field predicates rendered as dotted paths (`col.field`).
Conservative: if any sub-expression is unsupported the whole
rewrite is refused and Spark's brute-force lowering runs unchanged.
Off by default (`spark.lance.knn.indexedNearestByJoin.enabled=false`);
wired via `LanceKnnSparkSessionExtensions`. Part of the effort tracked
in lance-format#541.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
sezruby
marked this pull request as draft
August 26, 2026 16:33
Fixes the six findings raised on the PR:
1. pom parent version -> 0.8.0-beta.1 (deps already ${project.version}).
2. Carry the full pinned Lance read context through the join: the rule
captures the base read options AND the DataSourceV2 relation options
(where a DataFrame read hides branch/version/storage credentials) plus
the namespace context; LanceKnnJoinExec merges and version-pins that
context once on the driver so every executor probe opens one snapshot.
3. Materialize right-side payloads schema-aware via coerceToSpark
(Map -> Row, recursive for array/map/struct); LanceProbe stays
Spark-agnostic.
4. Quote/escape non-simple identifiers in the Lance prefilter string
(bare for [A-Za-z_][A-Za-z0-9_]*, double-quoted otherwise, per
dotted-path segment).
5. Stream join output through a lazy iterator and close the probe on task
completion (eager drain-and-close fallback when there is no
TaskContext), so the native handle can't leak or close early.
6. Add docs: a DQL operations page for the APPROX NEAREST join and a
config.md subsection for the KNN extension and its tuning configs.
Adds regression tests: relation context capture, identifier quoting,
payload coercion, and iterator laziness.
Verified locally: lance-spark-knn-4.2_2.13 builds and all 40 tests pass,
including the native e2e suite (oracle equivalence, WHERE-prefilter
pushdown, IVF-PQ recall, and rule-disabled fallthrough).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Covers finding #2 at the runtime level (the earlier test only asserted capture into the stage Conf): a `version` read option supplied on the DataFrame read (`spark.read.format("lance").option("version", "1")`) lands only on the DataSourceV2 relation options, so it must be captured, merged, and pinned on the driver for the probe to scan that snapshot. The test writes v1, appends v2 rows that duplicate each left query vector (distance 0), then asserts: pinned to v1 the appended rows are invisible and every left row still gets k v1 hits; on latest each left row's own duplicate is a nearest hit — so the pin is what excluded them (the assertion is discriminating, not vacuous). Verified: IndexedNearestByJoinSqlTest 7/7 pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
sezruby
marked this pull request as ready for review
August 26, 2026 17:52
…T join
Fixes the six findings from the re-review, each with a regression test
matching the reviewer's reproducer:
1. LanceProbe.openDataset no longer calls runtimeNamespace unconditionally
(which force-loaded the namespace impl class). It now applies the exact
worker-open namespace policy of LanceFragmentScanner.create: touch the
namespace only when a namespace impl is configured AND executor credential
refresh is enabled, then rebuild or clear it accordingly.
2. LanceKnnJoinStage.mergeReadOptions now preserves the connector's
incompatible-ref guard: when the base carries a pinned ref and the relation
also pins version/branch, a same-named branch keeps the table ref while an
incompatible combination is rejected, rather than letting the relation
silently win and read a different snapshot.
3. coerceToSpark rebuilds a real Spark map from Arrow's map representation
(a LIST of {key,value} entry structs surfaces as a Seq, not a Scala Map),
so a MapType payload column encodes correctly instead of failing/garbling.
4. quoteIdentifier always-quotes column identifiers so a name colliding with a
SQL keyword/literal (e.g. `true`) can't collapse into a tautology. The quote
character is a BACKTICK: Lance's filter dialect treats a double-quoted token
as a string literal, so `"category" = 'A'` prefilters to zero rows — verified
against a real Lance dataset via the WHERE-pushdown oracle test.
5. The Catalyst rewrite now declines when a standard analysis guard would fire
(spark.sql.crossJoin.enabled=false or a streaming child), leaving Spark's own
path to reject with the expected error instead of silently rewriting past it.
6. Docs gain an Installation section with the module Maven coordinate
(org.lance:lance-spark-knn-4.2_2.13) and a --packages spark-submit example.
All 47 tests in lance-spark-knn-4.2_2.13 pass; spotless clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…yload materialization Address two gatekeeper findings on PR lance-format#796: Finding A — require the connector's canonical fixed-size-vector metadata on the right attribute before rewriting. A plain variable-length List<Float> and a searchable Lance fixed-size-list vector both map to Spark ArrayType(FloatType); only the `arrow.fixed-size-list.size` field metadata distinguishes them. The rule now gates recognizeRanking on VectorUtils.isVectorField (the connector's canonical marker) and leaves the NearestByJoin unchanged otherwise, so a variable list falls through to Spark's brute-force path instead of handing Lance a column it cannot search. Regression: variableListVectorFallsBackInsteadOfFailing. Finding B — materialize projected right-side fields through the canonical connector Arrow-to-Spark adapter (LanceArrowColumnVector -> ColumnarBatch row -> CatalystTypeConverters external values) instead of raw Arrow getObject, so non-numeric payload columns (DateType, TimestampType) round-trip as the external java.sql.Date / java.sql.Timestamp the join's ExpressionEncoder expects. `_rowid` (not in the Spark schema) keeps the getObject fallback. Regression: dateAndTimestampPayloadMatchesOrdinarySparkTypes (payload parity vs an ordinary Spark Lance read). All 49 knn-4.2 module tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Blob columns are late-materialized: the connector's canonical reader threads the dataset URI, column name, and row addresses into BlobStructAccessor.setBlobReferenceContext so a blob descriptor resolves to its payload. The no-shuffle probe path fetches only _rowid and wraps vectors without that context, so a non-null legacy (v1) blob would resolve to an empty payload. Conservatively decline the rewrite when the Lance relation output carries any blob column (v1 lance-encoding:blob, or v2 lance.blob.v2 extension), detected via the connector's own BlobUtils.isBlobReadColumn. The query then falls through to Spark's brute-force cross-product, whose canonical blob-aware scan returns the true payload — so payload parity is owned by the fallback. Adds rule-level v1/v2 decline regression tests, each with a positive control (same schema minus the blob column rewrites) proving the blob column — not a missing fixed-size vector — is the discriminating cause. Module suite: 51 tests, 0 failures. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
On the SQL path the rule sets internalK == k, so every probed row is kept and the separate late-materialization scan just re-fetches the exact rows the search already found — a second Lance scan plus a k-element `_rowid IN (arrow_cast(...))` filter parse per query, for nothing. Add LanceProbe.probeRows, which runs the nearest search AND projects the payload columns in a single scan, returning MaterializedHit(rowAddr, score, row). LanceKnnJoinStage.processRow uses it when internalK <= k and keeps the existing probe -> trim -> materialize path for a future over-fetching caller (internalK > k), where deferring the payload fetch to survivors is the win. Factor the shared nearest-query builder and the _rowid / score vector readers so probe and probeRows search identically. Add a probeRows-vs- (probe+materialize) parity test asserting identical (rowAddr, score) hits and identical projected payload. Document the fixed-size-vector and blob right-side fallback triggers. 52 tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Contributor
Author
|
For reviewability, this feature is being split into a stack of smaller PRs. This PR stays open as the full-feature reference; the slices land independently:
|
This was referenced Aug 26, 2026
…rved-name routing Carry the probe-core review fixes into the umbrella revision and close the fold's headline finding at the stage layer: - Metric: cosine/dot are DISTANCES in Lance (1 - similarity), so smaller is better for all three metrics; flip smallerIsBetter and fix the doc. - TopKHeap: admit by the metric Ordering (ord.lt), not a raw float compare, so a NaN score sorts as worst-and-evictable instead of silently winning. - LanceProbe.probeRows: preserve projected payload columns that lack a Spark type via the generic Arrow fallback; guard + fusesCleanly reject a projection that collides with the injected _rowid / _distance / _score columns. - LanceKnnJoinStage: route reserved-name projections away from the fold to the split probe -> materialize path (materialize injects only _rowid), via the new unit-testable foldsInOneScan predicate. Tests: +3 fold-routing regressions in LanceKnnJoinStageTest, +metric-direction and reserved-name regressions carried from probe-core. Full module: 59 pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Make the nearest eligibility contract schema-level, not projection-level,
end to end (rule + stage + probe).
Lance's nearest scan always injects _rowid and the _distance/_score
metadata. If the right table's own schema already has a column by one of
those names, the injected metadata shadows it and NO indexed route
recovers the physical column: an all-columns fused scan reads the user's
_distance out-of-band as the ranking score and silently drops it from the
payload — data loss, not an error. Fold-vs-split routing does not help;
both scans inject _rowid.
- IndexedNearestByJoinRule: add hasReservedColumn, a decline guard
modeled on the existing hasBlobColumn — when the right schema owns a
reserved-named column the rule leaves the NearestByJoin in place so
Spark's brute-force nearest-by (whose ordinary scan returns the true
payload) handles it.
- LanceProbe: schemaSupportsNearest / reservedSchemaColumns eligibility
primitive + requireNearestCompatibleSchema() backstop called at the
top of probe()/probeRows(), superseding the projection-only
fusesCleanly guard (which missed the empty/all-columns form).
- LanceKnnJoinStage.foldsInOneScan: drop the removed fusesCleanly term
— reserved-column tables are declined upstream, so the fold decision
is now purely about over-fetch (internalK <= k).
Tests:
- rule: reservedColumnDeclinesRewrite — a right schema owning each
reserved name declines, with a vector-only control that rewrites.
- SQL e2e (real Lance): sqlUserDistanceColumnDeclinesAndReturnsStoredPayload
— a table with a physical _distance column declines and Spark's
fallback returns the true stored value, not the search score.
- probe: real _distance-column dataset regression — probe()/all-columns
probeRows() fail fast naming _distance instead of dropping it; pure
schemaSupportsNearest / reservedSchemaColumns contract test.
- stage routing tests reworked to the over-fetch-only decision.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The reserved-name collision is closed: indexed interception now declines any right relation whose schema owns _rowid, _distance, or _score, leaving Spark’s ordinary scan to preserve the user payload. The probe enforces the same invariant defensively, while compatible schemas retain the no-overfetch fused path.
sezruby
marked this pull request as draft
August 27, 2026 05:38
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
Accelerates Spark 4.2's native
APPROX NEAREST k BY DISTANCE ...join (SPARK-56395) over Lance tables. The default lowering turnsNearestByJoininto an O(|L|×|R|) cross-product +MaxMinByK; this rule intercepts the operator and runs a per-left-row native ANN/exact probe against the Lance vector index instead — no shuffle, no broadcast.Part of the effort tracked in #541. This is the narrower Spark 4.2 SQL-interception slice only; the manual
df.kNearestJoinDataFrame API is intentionally deferred and can be added later as a version-agnostic base module.How
New opt-in module
lance-spark-knn-4.2_2.13:IndexedNearestByJoinRule(apostHocResolutionRule) recognizesNearestByJoinand rewrites it to a single no-shuffle logical node. It must run in postHoc, not the optimizer — Spark'sRewriteNearestByJoinfires first in theFinishAnalysisbatch, so an injected optimizer rule would only ever see the already-lowered cross-product.LanceKnnJoinStrategy/LanceKnnJoinExeclower to onemapPartitions→ per-left-row native probe → bounded top-K heap → late-materialize by row address.LanceProbe/LanceKnnJoinStage/TopKHeap/Metric) supports L2, cosine, and dot metrics withnprobes/refineFactortuning.WHEREis translated to a Lance filter string, including nested struct-field predicates rendered as dotted paths (col.field). Conservative by design — if any sub-expression is unsupported the whole rewrite is refused and Spark's brute-force lowering runs unchanged (never a silently dropped residual).Off by default (
spark.lance.knn.indexedNearestByJoin.enabled=false); wired viaLanceKnnSparkSessionExtensions.Config
spark.lance.knn.indexedNearestByJoin.enabledfalsespark.lance.knn.nprobesspark.lance.knn.refineFactorTests
IndexedNearestByJoinRuleTest— pure Catalyst unit tests (rule matching, filter translation accept/refuse, nested struct-field paths); no Lance execution.IndexedNearestByJoinSqlTest— end-to-end SQL over real Lance tables: exact-path oracle equivalence,WHEREprefilter pushdown, rule-off fallthrough, and IVF-PQ recall (default / clustered / refineFactor) checks.TopKHeapTest,LanceProbeValidationTest— probe-core units.🤖 Generated with Claude Code