Skip to content

DATABRICKS 2026-10-05 Tier 3

Read original ↗

NEAREST BY Join: Scaling Vector Search in Databricks Runtime

Summary

Classical vector search is a serving problem: one query embedding arrives and the system returns top-k in tens of milliseconds. But a large share of real vector-search work is batch — entity resolution, deduplication, semantic tagging, classification, record enrichment, batch recommendations — where millions of queries run against millions-to-billions of base vectors on a schedule, and success is measured by whether the whole job finishes within its SLA at reasonable cost, not by single-lookup latency. Databricks argues these batch workloads deserve a fundamentally different architecture, and built vector search as a native feature of the Databricks Runtime (on Spark + Photon) rather than federating to an external real-time serving endpoint. The result is a "deliberately small yet deep stack": a new SQL join syntax (NEAREST BY) that makes the top-k ranking join a first-class relational operation; a rewrite that lowers it onto a cross join + a scalar distance function + a bounded top-k aggregate; a fused Photon operator backed by a blocked GEMM kernel that collapses the plan's middle and crosses the roofline ridge point; and an optional IVF index built as an ordinary liquid-clustered Delta table so APPROX queries score a fraction of the base vectors with the same kernels.

Key takeaways

  1. Batch vector search is structurally a top-k ranking join, not a million small lookups. "A batch vector search is not a million small searches. It is one large query: for each row on the left, find the k nearest rows on the right — a top-k ranking join. Executing enormous joins is exactly what the runtime engine excels at." The first-version VECTOR_SEARCH SQL function federated to an external real-time endpoint as a Generate node streaming one query row at a time — every row a network request + deserialize + possible retry — so throughput was capped by endpoint sizing, not cluster size, with the runtime reduced to a dispatcher. (Source: this article)

  2. Building search into the engine that already holds the data pays off two ways. Single copy of data — embeddings stay in Delta tables on the Lakehouse, so there is no separate vector store, no sync pipeline, no second system to operate and pay for. Single engine for execution — search runs in one elastically-scaling, fault-tolerant engine with kernels built for batch query shapes; tasks retry automatically and memory pressure spills to disk, so there is no client-side concurrency control, rate limiting, or retry loop to build.

  3. The requirements are a textbook batch-engine spec, deliberately trading per-request latency for throughput. Aggregate throughput over per-request latency; scale on both sides of the join (hundreds of millions of query vectors × billions of base vectors); elastic parallelism that partitions cleanly across hundreds-to-thousands of cores and scales to zero after the run; peak arithmetic per core (distance scoring is expensive, so inner loops must run near the hardware's FLOPs/s and bytes/s ceilings); and fault tolerance (survive worker loss, task failure, and memory pressure via spilling) — "properties of an execution engine, not features we can wrap around a real-time serving endpoint."

  4. The syntax encodes the relational structure directly — and makes APPROX/EXACT a semantic contract. NEAREST BY is an asymmetric join (like LATERAL): the left side drives, the right is searched; ranking direction is explicit (BY SIMILARITY descending, BY DISTANCE ascending); LEFT OUTER keeps query rows with no candidates; and the BY expression is a pluggable orderable scalar over both sides. Crucially, EXACT guarantees the true top-k by exhaustive evaluation while APPROX lets the optimizer substitute an ANN strategy — so creating or dropping an index can never silently change results; only queries that opt into APPROX consent to approximation. Contrast: Postgres+pgvector and Snowflake compose distance operators with ORDER BY … LIMIT and need a LATERAL subquery per driving row (fragile fast-path recognition); BigQuery exposes a table-valued function where column references are unvalidated strings.

  5. The rewrite isolates exactly two primitives all runtime flows through. NEAREST BY parses to a logical join node the optimizer lowers onto standard relational operators: tag each query row with a generated id, score every (query, base) pair, keep the k best per id with a grouped top-k, inline the kept rows back out. Because every operator is ordinary, "the plan distributes, spills, and retries like any other — correctness and fault tolerance come for free." What remains performance-critical is the distance function that scores a pair and the aggregate that keeps each group's k best.

  6. The distance kernels are native SIMD with four deliberate design choices. All seven vector SQL functions over ARRAY<FLOAT> run as native Photon SIMD kernels where a single fused multiply-add (FMA) computes a·b + c across a whole vector register per issue: (a) portable by construction — each kernel is compiled into several ISA-specific clones (AVX-512 on Intel, SVE2 on Graviton) and the best supported one is chosen at runtime; (b) zero-copy input — an ARRAY<FLOAT> is contiguous within a row, read as a raw pointer into the column buffer with no copy and no per-element indexing in the hot loop; (c) relaxed float semantics — dropping strict IEEE ordering lets the compiler fuse FMAs and split the reduction into parallel accumulator chains so the loop isn't serialized on one running sum; (d) nothing scalar in the loop — sqrt/divide done once outside. (These compose with concepts/memory-bandwidth-bound / concepts/matrix-multiplication-accumulate.)

  7. The bounded top-k aggregate replaces a full sort with O(k) state. In plain SQL, grouped top-k is a window function (ROW_NUMBER() OVER (PARTITION BY query ORDER BY score)) that sorts every partition in full then discards all but k. Databricks instead extended max_by/min_by with a third K overload built on four properties — one-comparison selection (per-group state is a bounded heap whose root is both eviction candidate and admission threshold; once full, each candidate is accepted/rejected with a single comparison, no sort over the candidate stream); O(k) state (memory per group is independent of input size; a query against a billion rows carries only k rows of state, and k goes to 100,000); late materialization (the heap stores indices, not copies — a candidate stays a pointer into the live columnar batch and is copied into aggregate state only if it survives when the batch is recycled); and distributable (partial/merge contract — each partition emits its local top-k, the shuffle moves k-element arrays instead of raw pairs, the merge re-inserts them; result is exact because the global top-k is always a subset of the union of partials). A canonical custom hot-path data structure.

  8. The roofline model explains why a fully Photonized naive plan is still far from optimal — and what to do. Scoring two aligned columns yields one score per row with no data reuse. But the NEAREST BY shape yields nq × nb scores from only nq + nb distinct vectors — each base vector is needed by every query. Under the pairwise cross join every operand is loaded fresh and used once: bytes moved = 2 × 4 × d × nq × nb, arithmetic intensity AI = 1/4, pinned to the DRAM slope at 0.8% of peak, scaling O(1). Under a fused GEMM that buffers Q and streams B: bytes moved = 4 × d × nb, AI = nq/2, scaling O(nq) — crossing the ridge at nq = 64 and riding the compute roof for larger batches. FLOPs are identical under both (2 × d × nq × nb); only the bytes change. "The vector distance and similarity kernels are near-optimal for their input shape, but the naive plan shape has no data reuse to exploit… The fix is a fused operator plan." (On a reference m6i.2xlarge, 64 executors × 4 cores × 204.8 GFLOP/s ≈ 52.5 TFLOP/s of fp32 FMA; L1/L2 roofs sit 30–60× above the DRAM fair share but the base table dwarfs any cache, so every base byte crosses the DRAM boundary at least once.)

  9. The fused operator + blocked GEMM turn the roofline analysis into arithmetic intensity. A single Photon execution node replaces the cross join, the distance/similarity projection, and optionally the partial top-k: it buffers the smaller input, streams the other in batches, scores query-by-base tiles with a custom GEMM kernel, and maintains per-query top-k state across tiles when k is small. With top-k fused, output is at most nq × k candidate rows rather than nq × nb pairs. Two levels of blocking improve reuse: packed base panels reused across query tiles for cache hits (panel small enough to fit cache, large enough to limit score- buffer round-trips at each embedding-dimension boundary) and register blocking so each loaded value feeds multiple FMAs (tile large enough to overlap FMA latency, small enough to avoid register spilling). The full nq × nb score matrix is never written to DRAM — finished tiles land in a bounded, recycled score buffer that streaming top-k consumes in place; once a query holds k entries, its worst score becomes the admission threshold, letting the merge reject multiple scores per vectorized compare. For large k, the operator runs the blocked GEMM alone and hands scored tiles to the max_by/min_by partial — even then, emitting one fp32 score per pair costs 4 bytes against 2d FLOPs, so AI = d/2, still far past the ridge for any realistic dimension. Distribution follows standard join strategies: broadcast the smaller side if it fits, otherwise partition both and run block-cartesian.

  10. Exhaustive KNN is O(nq × nb) no matter how good the kernel — that is what APPROX + the IVF index are for. "A million queries against a billion rows is 10¹⁵ dot products, and no register tile amortizes away an exponent." The index is a classic IVF (inverted file): k-means over a corpus sample produces centroids, each base row is assigned to its nearest centroid, and at query time each query scores only the vectors in its nprobe nearest clusters — at billions scale a query probes ≤ 0.1% of the corpus. Databricks chose IVF over graph indexes because independent cluster scans parallelize across executors and the layout is naturally columnar, whereas graph traversal is a chain of serial lookups (concepts/sequential-vs-random-io). Physically the index is an ordinary Delta table liquid-clustered by centroid id: assignment rows carry the vector next to its centroid id; each cluster's candidates sit contiguous on blob storage; files whose clusters no query probed are pruned before read; refresh is transactional + incremental. An APPROX query rewrites onto the same primitives — probe centroids with a top-k NEAREST BY join, equi-join on centroid id to restrict each query to its probed clusters, score + top-k, merge — and files added since the last refresh are brute-forced in a compensation branch unioned into the same merge, so a stale index prunes less but never degrades quality.

  11. Query execution is shaped by query and base cardinalities (four cases), and liquid clustering pays off hardest in the large-large case. Small index / small or large query → broadcast index (BNLJ), everything in memory or queries stream + index broadcast to all. Large index / small query → broadcast queries. Large index / large query → equi-join by centroid id with in-place shuffle of the index: only the query side moves — each probed query is shuffled to its clusters' partitions while the index side directly scans files for matching centroid ids, so liquid clustering avoids a full shuffle of the (large) index. Within a partition each cluster's candidates arrive as a dense contiguous block, so the same GEMM kernel applies — the exact path runs it once globally; the approximate path runs it once per cluster; local top-k per query, regroup, merge the nprobe partials via the aggregate's partial/merge contract.

  12. Measured results on canonical customer workloads (indexed ANN, ≥96% recall@K; base 100K–5B vectors, query batches up to 10M). Semantic deduplication: self-join of a 10M-record table completed in minutes. Classification/tagging: 10M queries × 100K reference vectors in under a minute. Recommendation refreshes: 1M queries × 10M catalog vectors in minutes. Entity resolution/enrichment: 1M queries × 1B reference vectors in minutes. Thesis: "Batch vector search did not need a new system — it needed to become a first-class citizen of the one that already holds the data."

Systems / concepts / patterns extracted

Operational numbers / specifics

  • Reference per-core roofline (m6i.2xlarge): 204.8 GFLOP/s fp32 FMA per core; a 64-executor × 4-core cluster peaks at ≈ 52.5 TFLOP/s; L1/L2 cache roofs sit 30–60× above the DRAM fair share.
  • Pairwise cross join: bytes = 2 × 4 × d × nq × nb, AI = 1/4, 0.8% of peak, O(1) scaling (DRAM-pinned).
  • Fused GEMM: bytes = 4 × d × nb, AI = nq/2, O(nq) scaling; crosses the ridge at nq = 64. Large-k fallback: AI = d/2.
  • FLOPs (both plans): 2 × d × nq × nb (identical — fusion changes bytes, not FLOPs).
  • Top-k aggregate: O(k) state per group; k up to 100,000; heap stores indices (late materialization), partial/merge-exact.
  • IVF pruning: at billions scale a query probes ≤ 0.1% of the corpus.
  • Measured (≥96% recall@K): 10M self-dedup → minutes; 10M × 100K classify → <1 min; 1M × 10M recommend → minutes; 1M × 1B entity-resolution → minutes.

Caveats / scope

  • Tier-3 Databricks blog with a product/launch framing (closes with docs + hiring CTAs), but the architecture content — the top-k-ranking-join rewrite, the SIMD distance kernels, the roofline analysis, the fused blocked-GEMM operator, the bounded top-k aggregate, and the IVF-as-liquid-clustered-Delta index — is the overwhelming majority of the body and well above the 20% bar: an include.
  • Benchmark results are vendor-reported ("minutes" / "under a minute"), with no absolute cost figures, no per-stage breakdowns, and no independent baseline; the "canonical customer workloads" are described, not named with reproducible configs (concepts/benchmark-methodology-bias).
  • Roofline constants (204.8 GFLOP/s, 52.5 TFLOP/s, nq = 64 ridge) are explicitly illustrative for the m6i.2xlarge reference and shift with hardware.
  • IVF specifics undisclosed: nprobe defaults, centroid count / hierarchy depth, rerank shortlist size, and filtered-recall degradation under APPROX with SQL predicates.
  • The partial-top-k fusion path is k-dependent (small-k fused in-operator; large-k falls back to the max_by/min_by partial) — the exact crossover k is not given.

Source

Last updated · 854 distilled / 2,448 read