Skip to content

feat: indexed nearest-by join over Lance (Spark 4.2 SQL) - #796

Draft
sezruby wants to merge 9 commits into
lance-format:mainfrom
sezruby:knn-4.2-sql
Draft

feat: indexed nearest-by join over Lance (Spark 4.2 SQL)#796
sezruby wants to merge 9 commits into
lance-format:mainfrom
sezruby:knn-4.2-sql

Conversation

@sezruby

@sezruby sezruby commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

What

Accelerates Spark 4.2's native APPROX NEAREST k BY DISTANCE ... join (SPARK-56395) over Lance tables. The default lowering turns NearestByJoin into 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.kNearestJoin DataFrame 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:

  • Catalyst IndexedNearestByJoinRule (a postHocResolutionRule) recognizes NearestByJoin and rewrites it to a single no-shuffle logical node. It must run in postHoc, not the optimizer — Spark's RewriteNearestByJoin fires first in the FinishAnalysis batch, so an injected optimizer rule would only ever see the already-lowered cross-product.
  • LanceKnnJoinStrategy / LanceKnnJoinExec lower to one mapPartitions → per-left-row native probe → bounded top-K heap → late-materialize by row address.
  • Probe core (LanceProbe / LanceKnnJoinStage / TopKHeap / Metric) supports L2, cosine, and dot metrics with nprobes / refineFactor tuning.
  • Prefilter pushdown: right-side WHERE is 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 via LanceKnnSparkSessionExtensions.

Config

key default meaning
spark.lance.knn.indexedNearestByJoin.enabled false master switch for the rule
spark.lance.knn.nprobes (index default) IVF probe count
spark.lance.knn.refineFactor (unset) over-fetch factor for re-ranking

Tests

  • 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, WHERE prefilter pushdown, rule-off fallthrough, and IVF-PQ recall (default / clustered / refineFactor) checks.
  • TopKHeapTest, LanceProbeValidationTest — probe-core units.

Note: the default CI matrix builds connector modules with -am, which pulls upstream deps only; this module is downstream of the connectors, so a reactor build (e.g. ./mvnw test -pl lance-spark-knn-4.2_2.13 -am) is needed to exercise it.

🤖 Generated with Claude Code

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>
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
@github-actions github-actions Bot added the enhancement New feature or request label Aug 26, 2026
@sezruby
sezruby marked this pull request as draft August 26, 2026 16:33
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
sezruby and others added 2 commits August 26, 2026 10:16
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
sezruby marked this pull request as ready for review August 26, 2026 17:52
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
…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>
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
…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>
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

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>
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 26, 2026
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>
@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 26, 2026
@sezruby

sezruby commented Aug 26, 2026

Copy link
Copy Markdown
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:

lance-gatekeeper[bot]

This comment was marked as outdated.

…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>
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 26, 2026
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>
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 27, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 27, 2026
@sezruby
sezruby marked this pull request as draft August 27, 2026 05:38
@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant