[SPARK-57091][SQL] Add BroadcastNearestByJoinExec to avoid cross-product materialization

[yadavay-amzn]

What changes were proposed in this pull request?

Add BroadcastNearestByJoinExec, a dedicated physical operator for NearestByJoin that replaces the existing cross-product rewrite (RewriteNearestByJoin) when spark.sql.join.nearestBy.broadcast.enabled is set to true (default false).

The operator broadcasts the right side and performs a single-pass iteration per left row with a bounded priority queue of size k. It exploits the asymmetric pattern (small reference table, large fact table) by never materializing the NxM cross product.

Changes:

• New physical operator: BroadcastNearestByJoinExec (extends BaseJoinExec)
• Ranking comparison via TypeUtils.getInterpretedOrdering (supports all orderable types: String, Date, Decimal, Long, etc.; no Cast-to-double)
• New planner strategy: NearestByJoinSelection in SparkStrategies
• Flag-gated optimizer skip: when spark.sql.join.nearestBy.broadcast.enabled is on, RewriteNearestByJoin leaves the NearestByJoin node intact (stats-free gate, no size threshold) for the planner's NearestByJoinSelection strategy, which plans BroadcastNearestByJoinExec. Exception: rankings whose expression contains a scalar Python UDF referencing both children fall back to the cross-product rewrite even when the flag is on (the operator cannot evaluate a cross-child Python UDF), so flag-on is not synonymous with the operator path. The predicate is NearestByJoin.hasCrossChildPythonUDF, shared by the rewrite guard and CheckAnalysis so the two sites agree by construction
• The right side is broadcast unconditionally (no size decision, no autoBroadcastJoinThreshold check). An oversized right side that cannot be broadcast will fail the query; there is no fallback to the rewrite
• spark.sql.crossJoin.enabled = false does not reject NEAREST BY queries that reach the operator when the flag is on (the operator is a bounded top-k per left row, not an unconditioned cross product, and builds no Join node). Rankings that fall back to the rewrite (cross-child scalar Python UDF) do build a Join, so CheckAnalysis still emits NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED for them when crossJoin.enabled = false rather than letting them fail later with a generic cartesian-product error
• Variable-length ranking values (String, Decimal, etc.) are copied only when needed (!UnsafeRow.isFixedLength)
• Non-deterministic ranking expressions are initialized with the partition index via mapPartitionsWithIndexInternal
• AQE support: registered in ValidateSparkPlan, and in CoalesceShufflePartitions (isExplodingJoin = true, since the operator fans out up to k rows per left row; childrenNeedCompatiblePartitioning = false, since the right side is broadcast)
• New SQLConf: spark.sql.join.nearestBy.broadcast.enabled (internal, default false, since 4.4.0)
• GC-friendly heap: HeapEntry is a small mutable class (not a case class); a pool of min(k, rightRows.length) entries is allocated once per partition and reused across all left rows, bounding total HeapEntry allocations per task while avoiding boxing of the index field

Why are the changes needed?

The current NearestByJoin implementation rewrites to a cross-join + aggregate. This materializes all NxM row pairs, shuffles them by synthetic UUID, and then the aggregate discards the vast majority. At moderate scale this becomes the bottleneck.

Benchmark (10,000x10,000, k=5): the streaming-heap operator is roughly 24-30x faster than the cross-product rewrite (depending on JDK). Results are generated by the standard benchmark workflow (.github/workflows/benchmark.yml, class org.apache.spark.sql.execution.benchmark.NearestByJoinBenchmark) across JDK 17/21/25 and committed under sql/core/benchmarks/.

Why a dedicated operator instead of optimizing the existing rewrite?

The cross-join materializes all NxM rows before the aggregate can bound them. There is no "streaming top-k aggregate" that can short-circuit the cross-join mid-execution. The operator avoids this by never materializing the full product -- for each left row, only k heap entries exist at any time.

Does this PR introduce any user-facing change?

No. The feature is opt-in via a new SQLConf that defaults to false. When disabled, the existing RewriteNearestByJoin path is used unchanged.

How was this patch tested?

• BroadcastNearestByJoinExecSuite (33 tests): correctness for distance/similarity directions, all orderable ranking types (Double, Int, Long, Decimal, Date, String), NULL/NaN handling, tie-breaking, empty tables and k boundary cases, output nullability, non-deterministic ranking initialization, DSv2 and partitioned sources, the unconditional-broadcast contract (no threshold, no fallback), the crossJoin.enabled divergence (a query that reaches the operator succeeds without it), StringType buffer retention with heap eviction, per-partition HeapEntry pool allocation bound, EXPLAIN output, AQE re-optimization, rewrite-vs-operator parity, and the cross-child scalar Python UDF fallback routing (rule-decision level)
• RewriteNearestByJoinSuite (15 tests): includes new coverage that flag-on leaves the NearestByJoin node intact and that cross-child scalar Python UDF rankings still take the rewrite path even with the flag on; passes with default conf
• AnalysisErrorSuite: a cross-child scalar Python UDF NearestByJoin with the broadcast flag on and crossJoin.enabled = false is rejected with the dedicated NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED error (not a generic cartesian-product error)
• python/pyspark/sql/tests/test_nearest_by_join.py: PySpark parity tests including two-sided and right-side-only scalar Python UDF rankings, each executed and compared across the broadcast flag (on vs off)
• NearestByJoinBenchmark (SqlBasedBenchmark): comparative benchmark, not run in CI; results committed under sql/core/benchmarks/
• Existing DataFrameNearestByJoinSuite (21 tests): passes with default conf (no regression)

Was this patch authored or co-authored using generative AI tooling?

Yes.

---

JIRA: SPARK-57091

[yadavay-amzn]

@dilipbiswal @cloud-fan Could you please take a look at this change when you get a chance?
This PR adds a dedicated BroadcastNearestByJoinExec operator that avoids the N * M cross-product materialization in RewriteNearestByJoin (noted as an out of scope improvement for SPARK-56395).

This improvement is inspired by @sarutak's SortMergeAsOfJoinExec for AS-OF join (#55912) and is similar in spirit to it by using a dedicated physical operator to replace an expensive rewrite for a specialized join type.

The operator broadcasts the right side and maintains a bounded k-heap per left row to avoid the shuffling cost and full materialization when the right table fits within autoBroadcastJoinThreshold. Benchmark numbers have been promising (see description). I've also added several unit tests covering correctness and edge cases.

Does this approach align with the planned evolution of the feature? Any concerns about adding a dedicated physical operator vs. optimizing the existing rewrite?
Happy to collaborate and adjust the approach based on your feedback. Thanks!

[cloud-fan]

3 blocking, 4 non-blocking, 3 nits.
Solid operator skeleton and a well-motivated optimization, but the broadcast path silently diverges from the rewrite it replaces for a real slice of inputs, narrows the output schema contract, and ships a heavyweight benchmark as a CI test. Off-by-default contains the blast radius, but the divergences contradict the "no user-facing change" claim once the flag is on. A separate operator is the right structural call (top-k is not a BNLJ-style boolean filter), but it hand-reimplements BNLJ's broadcast scaffolding — which is where the two correctness bugs entered.

Correctness (2)

• BroadcastNearestByJoinExec.scala:88: ranking coerced via Cast(_, DoubleType) — wrong/empty for non-numeric orderable ranking types, lossy for big Long/Decimal, and NaN diverges under similarity — see inline
• BroadcastNearestByJoinExec.scala:54: operator output narrows nullability the logical NearestByJoin.output widens (both sides); convention copied from BNLJ — codegen null-elision hazard + inter-path schema divergence — see inline

Design / architecture (3)

• NearestByJoinBenchmark.scala:25: benchmark is a QueryTest suite, so the 200K x 200K / 30K x 30K cases run in CI — should be a SqlBasedBenchmark — see inline
• RewriteNearestByJoin.scala:83: broadcast-threshold predicate duplicated in the optimizer guard and the planner strategy with no shared source of truth — see inline
• BroadcastNearestByJoinExec.scala:49: operator hand-reimplements BNLJ's broadcast-iterate scaffolding instead of reusing a shared join base — a separate operator is right, but the plumbing should be shared (this is where the ranking + nullability bugs entered) — see inline

Suggestions (2)

• BroadcastNearestByJoinExec.scala:93: heap stores boxed (Int, Double) tuples — O(N*M) allocations per partition — see inline
• Test coverage gap (general): no rewrite-vs-operator parity test for non-numeric ranking, NaN-under-similarity, or null in a non-ranking right column on INNER — the exact cases the two correctness findings break on.

Nits: 3 minor items (see inline comments) — one of them (a non-ASCII arrow in a comment) fails scalastyle and will break CI.

Verification

Traced the rewrite-vs-operator equivalence (the operator must produce identical top-k to RewriteNearestByJoin, since the flag only chooses the path). Equivalent for: numeric ranking, null ranking (both exclude), empty right (Inner empties, LeftOuter null-right), and NaN under distance. NOT equivalent for: (1) non-numeric orderable ranking types — Cast(_, DoubleType) yields null so every row is dropped, while the rewrite ranks by natural ordering; (2) high-magnitude Long / high-precision Decimal — double cast is lossy; (3) NaN under similarity — the rewrite's MAX_BY ranks NaN as the top match (Spark NaN ordering) while the operator drops it. NearestByJoinSelection applies no ranking-type guard before emitting the operator, so these reach it unguarded.

[yadavay-amzn]

@cloud-fan Thank you for the detailed review. All items addressed:

Blocking (3/3):

1. Ranking: Cast(_, DoubleType) → TypeUtils.getInterpretedOrdering(rankExpr.dataType). Also added .copy() on ranking rows to prevent buffer corruption for variable-length types.
2. Output nullability: matches logical NearestByJoin.output (both sides nullable).
3. Benchmark: converted to SqlBasedBenchmark object.

Non-blocking (4/4):

1. Extends BaseJoinExec with condition=None, leftKeys/rightKeys=Seq.empty, overrides simpleStringWithNodeId.
2. Shared NearestByJoin.canBroadcastRight(j, conf) with sizeInBytes >= 0 guard.
3. Boxed tuples → HeapEntry case class (unboxed Int).
4. Added tests: NaN-under-similarity, null-in-non-ranking-column, DateType, Long past 2^53, Decimal, StringType.

Nits (3/3): All fixed.

25 tests total, all passing.

[cloud-fan]

can you resolve the merge conflicts?

[yadavay-amzn]

Resolved conflict and rebased on master.

The conflict was in SparkPlanner.scala: master added AsOfJoinSelection to the strategy list at the same spot we added NearestByJoinSelection. Kept both (both confirmed defined as Strategy objects)

[yadavay-amzn]

@cloud-fan ready for another round of reviews when you get a chance, thanks!

[cloud-fan]

10 addressed, 0 remaining, 2 new. (2 new = 1 newly introduced, 1 late catch.)

0 blocking, 2 non-blocking, 0 nits. All prior findings verified resolved; the operator produces results equivalent to the rewrite with a matching output-nullability contract.

Correctness (1)

• SQLConf.scala:2360: NEAREST_BY_BROADCAST_ENABLED uses .version("5.0.0") but the feature is backportable (NearestByJoin already on branch-4.x, 4.3.0-SNAPSHOT); should be 4.3.0 — see inline

Suggestions (1)

• BroadcastNearestByJoinExec.scala:131: per-pair .copy() allocates a byte[] for every (left, right) pair; only needed for variable-length ranking types — see inline

Verification

Traced operator-vs-rewrite equivalence (the flag only selects the execution path, so results must be identical): equivalent for numeric, non-numeric-orderable (String/Date), null, and NaN ranking under both directions — TypeUtils.getInterpretedOrdering resolves to the same PhysicalDataType.ordering that MaxMinByK uses, so NaN ranks largest in both (top match under SIMILARITY, dropped last under DISTANCE); equivalent for empty-right INNER (no rows) and LEFT OUTER (null-right), and k > right-size; ties are non-deterministic in both by contract. Output metadata now widens both sides to nullable (BroadcastNearestByJoinExec.scala:68-69), matching the logical NearestByJoin.output — the prior INNER REF-NARROWER (copied from BNLJ's generic convention) is fixed.

[yadavay-amzn]

Thanks @cloud-fan - both addressed in 36f02e3: the conf version is now 4.3.0, and the ranking-row copy is gated on !UnsafeRow.isFixedLength(...) so fixed-width ranking avoids the per-pair allocation while variable-length ranking stays retain-safe (String/Decimal tests cover it). PTAL.

[yadavay-amzn]

@peter-toth Could you also review this please when you get a chance.
cc @cloud-fan

[yadavay-amzn]

@cloud-fan Could you please take another look when you get a chance?

Apache Spark 4.3.0 whose Feature Freeze is scheduled on 1st August.

Wondering if you think this PR is in good shape to make it to this code freeze date

[cloud-fan]

2 addressed, 0 remaining, 1 new. (0 newly introduced, 1 late catch.)
0 blocking, 0 non-blocking, 1 nit.

Already raised in existing discussion (1)

• This Scaladoc link still cannot resolve because SQLConf is neither imported nor fully qualified in this file. Please add org.apache.spark.sql.internal.SQLConf to the imports (or fully qualify the link) so the generated API documentation links to AUTO_BROADCASTJOIN_THRESHOLD instead of rendering an unresolved reference. -- existing discussion

Verification

Traced both routing outcomes through the shared eligibility predicate: disabled or oversized inputs still reach RewriteNearestByJoin, while eligible inputs are planned as BroadcastNearestByJoinExec. Compared heap ordering and null exclusion with MaxMinByK, checked empty-right and left-outer behavior, and verified that physical output nullability matches NearestByJoin.output. No tests were run as part of this review.

[peter-toth]

Thanks for the PR, @yadavay-amzn!

I did an independent read of the operator, the optimizer skip, the planner strategy, and the AQE validation. The shape is sound: a dedicated broadcast top-k operator with a bounded k-heap per left row, ordering via TypeUtils.getInterpretedOrdering, output nullability matched to the logical NearestByJoin, and -- nice detail -- the eager heap-drain into results before returning the inner iterator, which keeps the reused heap from corrupting across left rows. cloud-fan's earlier rounds covered the big ones (ranking type, nullability, benchmark), so this is mostly in good shape. Two things I'd still flag, one a correctness gap.

The main one: a non-deterministic ranking expression -- which NearestByJoin explicitly supports and the rewrite handles -- crashes on the broadcast path, because the ranking UnsafeProjection is never initialized with a partition index. Details inline. The second is a non-blocking note on where the rewrite-skip decision is made.

Not re-opening cloud-fan's still-open Scaladoc-link nit on [[SQLConf.AUTO_BROADCASTJOIN_THRESHOLD]] (SQLConf isn't imported in that file) -- just +1 to it.

Blocking

• 1. Non-deterministic ranking crashes the broadcast operator: the ranking UnsafeProjection is built inside mapPartitionsInternal and never initialize(partitionIndex)d, so a rand() / non-deterministic-UDF ranking (allowed by NearestByJoin, handled by the rewrite) throws at runtime once the conf is on and the right is broadcast-sized. Untested -- the suite only uses deterministic rankings. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala:104]

Non-blocking

• 2. Skip decided on the pre-optimization size estimate: RewriteNearestByJoin runs in the first (FinishAnalysis) optimizer batch, so canBroadcastRight reads the right's least-refined sizeInBytes, while the planner re-reads it on the fully-optimized right. The realistic effect is silent under-firing (not the planning failure in the other thread -- that needs the estimate to grow, which optimization doesn't do). A one-line comment noting the decision uses the pre-optimization estimate would help. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala:83]

[peter-toth]

Finding 1. rankingProj can hold a non-deterministic ranking expression, but it's created inside mapPartitionsInternal (line 98) and its initialize(partitionIndex) is never called. NearestByJoin explicitly allows a non-deterministic ranking (allowNonDeterministicExpression = true, and its scaladoc calls out rand() for randomized tie-breaking / scoring UDFs), and canBroadcastRight doesn't gate on determinism -- so left.nearestByJoin(right, rand(), ...) with spark.sql.join.nearestBy.broadcast.enabled=true and a broadcast-sized right routes here and throws the moment the ranking is evaluated (require(initialized ...) in Nondeterministic.eval for the interpreted path; NPE on the null RNG in codegen). The rewrite path handles this -- it materializes __ranking__ in a Project so the standard projection machinery runs initialize -- so this is a silent regression on the broadcast path.

Fix: thread the partition index in and initialize the projection, matching ProjectExec / BroadcastNestedLoopJoinExec (both call .initialize(index) for exactly this):

left.execute().mapPartitionsWithIndexInternal { (index, leftIter) =>
  ...
  val rankingProj = UnsafeProjection.create(Seq(rankExpr), leftOutput ++ rightOutput)
  rankingProj.initialize(index)
  ...

If you'd rather not support a non-deterministic ranking on the broadcast path yet, the simpler alternative is to route those queries back to the rewrite: add && j.rankingExpression.deterministic to NearestByJoin.canBroadcastRight. Either way, please add a rand()-ranked test on the broadcast path -- the current suite only exercises deterministic rankings, so this crash isn't caught.

[yadavay-amzn]

You're right, thanks. Switched to mapPartitionsWithIndexInternal and call rankingProj.initialize(partitionIndex) before eval, so a rand() ranking no longer throws. Added a test for it.

[peter-toth]

Finding 2. RewriteNearestByJoin runs in the FinishAnalysis batch, which is the first optimizer batch -- before the size-refining rules (filter inference, empty-relation propagation, subquery/CTE optimization, CBO). So canBroadcastRight here reads the right's earliest, least-refined stats.sizeInBytes, whereas NearestByJoinSelection re-reads it at planning time on the fully-optimized right. Since the estimate generally only shrinks through optimization, this isn't the planning failure raised in the other thread (that would need the estimate to grow), but it does mean the operator can silently under-fire: a right whose optimized estimate fits autoBroadcastJoinThreshold still gets rewritten to the cross-product if its FinishAnalysis estimate was over the threshold. Non-blocking (opt-in, correct either way) -- the two-pass approach you've deferred would resolve it -- but a comment noting the skip uses the pre-optimization estimate would set expectations.

[yadavay-amzn]

The gate is needed in the rewrite: without it a large-right NearestByJoin reaches NearestByJoinSelection, which returns Nil (a planning failure). And since estimates only shrink through optimization, reading the early size is conservative: it can miss a broadcast but never picks one wrongly, and the planner re-checks with refined stats. Kept it as-is, but happy to revisit if the missed-opportunity case matters.

[peter-toth]

Re-checked through 11ba350 — finding 1 resolved (rankingProj.initialize(index) is now called; I verified the new rand() test genuinely fails without that line, so it's covering the fix, not just the path).

Finding 2 is still open and I'm escalating it to Blocking, because your reply — that reading the early size "is conservative: it can miss a broadcast but never picks one wrongly" — holds for the direction of the estimate but not for the legality of reading it there. RewriteNearestByJoin runs in FinishAnalysis, the optimizer's first batch, and Optimizer.scala:239-242 states the invariant directly: "Before this batch, the logical plan may contain nodes that do not report stats. Anything that uses stats must run after this batch." Two consequences I reproduced on this head, neither of which is a missed optimization:

• DSv2 right side fails outright. With the flag on, left.nearestByJoin(v2Table, ...) throws [INTERNAL_ERROR] BUG: computeStats called before pushdown on DSv2 relation: testcat.rightTbl. The identical query passes with the flag off. DataSourceV2Relation.computeStats throws by design when Utils.isTesting, and in production returns full-table/all-columns stats after building a throwaway scan.
• Partitioned file tables never fire. Their pre-pushdown sizeInBytes is defaultSizeInBytes (Long.MaxValue), so the predicate is false regardless of the real size — I measured a right side that is 9 KB after partition pruning still being rewritten to the cross product.

I owe you a correction here: the DSv2 half was present in round 1 and I under-diagnosed it as merely "under-fires", so it's a late catch on my side, not new breakage. Non-partitioned v1 tables and in-memory DataFrames do work — which is exactly and only what the 26 tests exercise, so nothing in the suite catches this.

Everything else I checked holds: heap ordering and result-array direction are right for both directions, the eager drain into results keeps the reused heap safe across left rows, the !UnsafeRow.isFixedLength copy gate is correct (it can only over-copy, never under-copy), the physical output correctly mirrors the logical node's deliberate widen-both-sides contract, and 4.3.0 is the right conf version for a branch-4.x feature. Nothing to add on cloud-fan's resolved threads.

On the 1 Aug feature freeze: the fix for finding 2 is small if you take the conf-only gate — the blocker is the DSv2 failure, not the design.

Blocking

• 2. Broadcast eligibility reads stats in the optimizer's first batch (round 1): the predicate is legal at planning time but not inside FinishAnalysis, where DSv2 relations throw and partitioned relations report no size. Make the rewrite-skip stats-free and keep the size decision in the planner. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/NearestByJoin.scala:32]

Non-blocking

• 3. Heap still allocates and offers per candidate pair (late catch): HeapEntry fixed the boxing cloud-fan raised but not the count — every non-null pair still allocates and does offer+poll. Guarding on heap.peek() before offering drops that to O(k log(M/k)) expected insertions and also skips the variable-length .copy() for evicted rows, which completes cloud-fan's two hot-loop threads rather than re-opening them.

Minor

• 4. NearestByJoinSelection was inserted between JoinSelection's scaladoc and JoinSelection (late catch): the block describing equi-joins and shuffle-replicate nested loop join now documents the new strategy, and JoinSelection has none. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala:182]

[peter-toth]

Finding 2. (Carried from round 1, escalated to Blocking — continues the round-1 thread on RewriteNearestByJoin.scala:83.)

This predicate is fine when NearestByJoinSelection calls it, but RewriteNearestByJoin calls it from the FinishAnalysis batch — the optimizer's first batch. Optimizer.scala:239-242 states the invariant:

This batch pushes filters and projections into scan nodes. Before this batch, the logical plan may contain nodes that do not report stats. Anything that uses stats must run after this batch.

Your round-1 reply argued the early read is conservative because estimates only shrink. That's true of the direction, but the problem is that some relations cannot answer at all this early:

1. DSv2 right side throws. DataSourceV2Relation.computeStats (DataSourceV2Relation.scala:86-103) throws when Utils.isTesting, precisely to catch this:

org.apache.spark.SparkException: [INTERNAL_ERROR] BUG: computeStats called before
pushdown on DSv2 relation: testcat.rightTbl SQLSTATE: XX000

Reproduced on this head — the same query passes with the flag off:

sql("CREATE TABLE testcat.rightTbl (rid INT, y DOUBLE) USING foo")
sql("INSERT INTO testcat.rightTbl VALUES (10, 1.0), (11, 2.0), (12, 3.0)")
Seq((1, 2.5)).toDF("id", "x").nearestByJoin(
  spark.table("testcat.rightTbl"), abs($"x" - $"y"),
  numResults = 2, mode = "exact", direction = "distance").collect()
// flag on  -> INTERNAL_ERROR above
// flag off -> [1,2.5,12,3.0], [1,2.5,11,2.0]

Outside testing it takes the else branch, which builds a throwaway scan (table.asReadable.newScanBuilder(options).build()) inside the optimizer's first batch and returns full-table/all-columns stats — the comment there says "bad stats are better than failing a query". So on DSv2 this is either a hard failure or a decision made on a number the file itself calls bad, plus per-query scan-build cost.

2. Partitioned file tables never fire. Their pre-pushdown sizeInBytes is defaultSizeInBytes (Long.MaxValue), so the predicate is false at any threshold:

right = spark.table("rightPart").filter($"p" === 0)   // partitioned parquet
analyzed  sizeInBytes = 9223372036854775807
optimized sizeInBytes = 9398
BroadcastNearestByJoin fired = false

A 9 KB right side takes the cross-product path. (Non-partitioned v1 tables and in-memory DataFrames are unaffected, which is why the suite is green.)

Fix. Broadcast sizing belongs where JoinSelection does it — the planner. Keep this predicate as-is for the strategy and make the rewrite gate stats-free:

// RewriteNearestByJoin
case j @ NearestByJoin(...) if !SQLConf.get.nearestByBroadcastEnabled =>

That leaves the planner obligated to handle every NearestByJoin when the flag is on, which is the "no strategy handles it" hazard your code comment calls out. Two ways to discharge it:

• (a) Drop the size test from NearestByJoinSelection too, so the flag means "always use the operator". For an .internal(), default-off flag that's a reasonable contract, and it makes the two call sites agree by construction. It gives up the automatic large-right fallback.
• (b) Keep the size test in the strategy and, when it fails, plan the rewrite from there — planLater(RewriteNearestByJoin.rewriteOne(j)) after factoring the rewrite body out of the rule. This preserves today's semantics with valid stats; the cost is that the rewritten subtree misses the operator-optimization batches. This is essentially the two-pass approach you deferred.

Either way, please add a DSv2 right-side test with the flag on — the repro above is about 10 lines with InMemoryTableCatalog, and it's the shape most users on Iceberg/Delta will hit first.

[yadavay-amzn]

You're right, and thanks for the repro. I took option (a): RewriteNearestByJoin now gates only on the flag (!nearestByBroadcastEnabled) and reads no stats, so nothing touches stats in FinishAnalysis anymore. NearestByJoinSelection drops the size test and always plans the operator when the flag is on, so the two sites agree by construction and the DSv2/partitioned cases work. The contract is now what you described: for this internal, default-off flag, on means always use the operator (giving up the automatic large-right fallback). Added a flag-on DSv2 test with InMemoryTableCatalog (your repro, which threw INTERNAL_ERROR before) and a partitioned-table test that now fires the operator.

Separately on Finding 3 (non-blocking): the heap only offers when it isn't full or the candidate beats peek(), and the variable-length .copy() now runs only for rows that are actually kept, so evicted rows cost nothing. Same top-k output.

[peter-toth]

Finding 4. The new strategy was inserted between the strategy-selection scaladoc block (ending */ at line 180) and JoinSelection, so that block — which documents equi-join/non-equi-join support, broadcast hash join build-side preferences, and "Shuffle-and-replicate nested loop join ... Supports only inner like joins" — now attaches to NearestByJoinSelection, which does none of those things. JoinSelection is left undocumented.

Moving the object below JoinSelection restores both (and keeps the join-strategy doc adjacent to the code it describes):

  object JoinSelection extends Strategy with JoinSelectionHelper {
    ...
  }

  object NearestByJoinSelection extends Strategy {
    ...
  }

Placing it above the scaladoc block would work equally well.

[yadavay-amzn]

Done, moved NearestByJoinSelection below JoinSelection so the strategy-selection scaladoc attaches back to JoinSelection.

[uros-b]

Please undo this unrelated change

[uros-b]

dominated is true when the candidate should be kept, which is the opposite of what the name suggests?

[cloud-fan]

1 addressed, 0 remaining, 9 new. (2 newly introduced, 7 late catches, 0 previously raised.)
1 blocking, 7 non-blocking, 1 nit.
The execution design is sound, but the current head still has one misleading contract plus several focused maintainability and test gaps that should be resolved or explicitly accepted.

Already raised in existing discussion (7)

• The class documentation still promises a broadcast-threshold check and aggregate fallback, but enabling the flag now broadcasts every right side unconditionally. Update the class and config documentation to state that there is no size fallback and oversized inputs can fail during broadcast. -- existing discussion
• Remove NearestByJoin.canBroadcastRight and its SQLConf import: neither the rewrite nor the strategy calls it after switching to unconditional flag-on broadcasting, and leaving the early-stats helper invites reintroducing the DSv2 failure. -- existing discussion
• Document and test that flag-on execution bypasses spark.sql.crossJoin.enabled=false; the rewrite path deliberately rejects this query while the dedicated operator no longer creates a Join for CheckCartesianProducts to inspect. -- existing discussion
• Expose the ranking expression in formatted explain output; the current custom simple string and inherited BaseJoinExec details omit the expression that determines which rows are selected. -- existing discussion
• Bound the initial priority-queue capacity by the broadcast row count as well as k; k can be 100000 even when the right side has only a few rows, causing a large backing array per task. -- existing discussion
• Add a direct rewrite-vs-operator parity test that runs the same inputs with the feature flag on and off; hand-authored expected rows do not pin the core equivalence contract across both execution paths. -- existing discussion
• Rename dominated: the value is true when a candidate should be retained, which is the opposite of the term's usual meaning and makes the hot-path condition easy to misread. -- existing discussion

Suggestions (2)

• sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExecSuite.scala:493: This test says its ascending String inputs force heap evictions, but after aaa and bbb fill the distance heap every later value is worse and is rejected before insertion. Reverse or otherwise reorder the labels so the test actually exercises retained variable-width values across evictions. -- see inline
• sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/NearestByJoinBenchmark.scala:44: Commit the benchmark results file so the performance claims are reproducible from the tree, and avoid keeping both inputs cached during the comparison unless that retained cache is part of the benchmark contract. -- see inline

Verification

Traced both configuration branches from RewriteNearestByJoin through planner selection and AQE validation. Compared heap admission and drain ordering for distance and similarity, verified null exclusion and left-outer null padding, and checked that natural ordering and variable-width copies align with the existing aggregate path. No tests were run as part of this review.

PR description suggestions

• Update the broadcast eligibility and test-coverage bullets: flag-on execution now broadcasts unconditionally and has no threshold fallback.

[cloud-fan]

These ascending labels never force an eviction: after aaa and bbb fill the distance heap, every later label is worse and is rejected before insertion. Please reorder the input so later rows displace retained variable-width values; otherwise this test does not exercise the buffer-retention behavior its comment claims.

[cloud-fan]

Please commit the generated benchmark results so the performance claims are reproducible from this tree. Also either unpersist the inputs between cases or explain why retaining both caches is part of the intended benchmark setup, since that retained memory can bias the cross-product comparison.

[cloud-fan]

9 addressed, 0 remaining, 5 new. (3 newly introduced, 2 late catches, 0 previously raised.)
0 blocking, 1 non-blocking, 4 nits.
The core execution design and previously raised contract issues are resolved; only one small allocation cleanup and four documentation/test-comment nits remain.

Nits: 4 minor items (see inline comments).

Suggestions (1)

• sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala:178: The LEFT OUTER no-match branch allocates an identical null-right row for every unmatched left row. -- see inline

Verification

Traced both flag modes from RewriteNearestByJoin through planner selection and AQE validation. Compared the operator heap with the aggregate rewrite for natural ordering, null exclusion, empty/all-null right sides, LEFT OUTER padding, nondeterministic initialization, and output nullability; the direct parity cases and focused source tests cover those contracts. No tests were run as part of this review.

[cloud-fan]

This failure explanation does not match the heap path: once k=1 is full, an equal candidate fails the strict retention comparison and never reaches offer/poll; PriorityQueue also does not guarantee insertion order for ties. Please keep the useful Long-ordering assertion but remove the claim that this ordering forces the old Double implementation to select v2.

[cloud-fan]

Both links in this paragraph are unresolved because these logical-plan symbols are outside the joins package and are not imported. Please fully qualify RewriteNearestByJoin and NearestByJoin (or import them) so generated Scaladoc links to the intended declarations.

[cloud-fan]

This allocates the same all-null right row for every unmatched left row. Hoist it alongside the per-partition projections and reuse it before the immediate projection/copy.

[cloud-fan]

Please remove the stray uros-b :62: reviewer artifact from this heading.

[cloud-fan]

Please remove the stray :47 fragment from this heading.

[yadavay-amzn]

@cloud-fan @uros-b @peter-toth Thanks for the thorough review! Addressed all comments in latest commit, please take another look when you get a chance. Thanks again!

[cloud-fan]

5 addressed, 0 remaining, 1 new. (0 newly introduced, 1 late catch, 0 previously raised.)
0 blocking, 1 non-blocking, 0 nits.
The execution design and prior review concerns are resolved; one small hot-path allocation cleanup remains worth making.

Suggestions (1)

• sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala:172: This copies the entire one-column UnsafeRow for every retained variable-width ranking value. InternalRow.copyValue(rawValue) already deep-copies UTF8String, BinaryView, structs, arrays, and maps, so using it here preserves buffer safety without allocating and copying the enclosing row on this nested scan's hot path. -- see inline

Verification

Traced both configuration branches from RewriteNearestByJoin through planner selection and AQE validation. Compared natural ordering, null exclusion, empty/all-null right sides, LEFT OUTER padding, nondeterministic initialization, output nullability, and direct operator/rewrite parity. No tests were run as part of this review.

[cloud-fan]

InternalRow.copyValue(rawValue) already deep-copies the variable-width values handled here, so it avoids copying the enclosing one-column UnsafeRow for every retained candidate.

Suggested change

	rankingRow.copy().get(0, rankExpr.dataType)
	InternalRow.copyValue(rawValue)

[peter-toth]

Re-checked through 97808f0 — findings 5, 6, 7 resolved: the class scaladoc, the strategy scaladoc, the rewrite-guard comment and the conf doc all state the unconditional-broadcast contract now, canBroadcastRight and its SQLConf import are gone, and the crossJoin divergence is documented and pinned by a test that leaves crossJoin.enabled at false. Nothing regressed, and cloud-fan's and uros-b's threads all look addressed on this head.

Two new blocking items. One is my own earlier call going stale: I confirmed 4.3.0 as the conf version in round 2 and it was right then, but branch-4.x has moved on since. The other is the AQE registration — it does real work, but I measured that nothing in the suite would notice if it disappeared.

Blocking

• 8. Conf .version is now stale (new): branch-4.3 has been cut and branch-4.x is 4.4.0-SNAPSHOT, so a conf added by a PR merging today first ships in 4.4.0. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:2534]
• 9. AQE test doesn't cover the ValidateSparkPlan registration (new): I deleted the new case and this test still passed, so the registration has no coverage — and without it reOptimize silently rejects every re-optimization attempt for queries using this operator. Repro and a working test in the inline. [inline: sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExecSuite.scala:463]

Non-blocking

• 10. Parity case 4 doesn't test the ordering it names (new): the .orderBy(...) erases the emission order the case exists to check, so it's case 1 again with one left row. Both paths do promise best-first per left row, so it's worth pinning for real. [inline: sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExecSuite.scala:737]
• 11. outputOrdering = Nil throws away the left ordering (new): BroadcastNestedLoopJoinExec returns streamed.outputOrdering for exactly the (join type, build side) pairs this operator supports, and the same argument holds here. Nil costs a downstream SortExec that isn't needed. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala:100]

Minor

• 12. EXPLAIN FORMATTED renames a standard field (new): BaseJoinExec prints Join type; this override prints JoinType for this one operator. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala:79]
• 13. SQLConf.get instead of the rule's own conf (new): Rule extends SQLConfHelper, and the matching check in CheckAnalysis uses conf.nearestByBroadcastEnabled. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala:81]

[peter-toth]

Finding 8. This needs to be 4.4.0 now.

When cloud-fan asked for 4.3.0 (here), and when I confirmed it in round 2, branch-4.x was 4.3.0-SNAPSHOT. Since then branch-4.3 has been cut and branch-4.x has moved on, so the next open feature release — which is what cloud-fan's rule points at — is 4.4.0:

$ dev/next_version_candidates.py
master       5.0.0
branch-4.x   4.4.0

$ git show apache/branch-4.x:pom.xml | grep -m1 -A1 spark-parent
    <version>4.4.0-SNAPSHOT</version>

Suggested change

	.version("4.3.0")
	.version("4.4.0")

[peter-toth]

Finding 9. I deleted the case b: BroadcastNearestByJoinExec block from ValidateSparkPlan on this head and re-ran this test — it still passes, in 2.2s. So the registration this PR adds has no coverage.

Why it can't be seen from here. ValidateSparkPlan sits in AdaptiveSparkPlanExec.queryStagePreparationRules (AdaptiveSparkPlanExec.scala:129), so it runs twice. On the initial plan the right child is a plain BroadcastExchangeExec, so the new case is entered but takes validate(b.right) — the same walk the old catch-all did. It runs again inside reOptimize, after the broadcast stage has materialized and LogicalQueryStageStrategy has put a BroadcastQueryStageExec under the operator. That second run is the one the registration exists for, and reOptimize catches InvalidAQEPlanException and returns None (AdaptiveSparkPlanExec.scala:854-859), after which the loop at :388 just keeps currentPhysicalPlan. No exception, no wrong answer — AQE simply stops re-planning for the whole query. Neither assertion here can observe that: plan.contains("BroadcastNearestByJoin") reads the pre-execution AdaptiveSparkPlanExec (still the initial plan), and checkAnswer runs a separate query that succeeds either way.

Here is a test that does see it. It puts a sort-merge join on the left that AQE demotes to a broadcast hash join once the shuffle stats arrive, which only happens if reOptimize succeeds. I ran it both ways on this head: it passes as-is, and fails with the case deleted.

  test("SPARK-57091: AQE re-optimization is not rejected by ValidateSparkPlan") {
    withSQLConf(
      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
      SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true",
      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
      SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") {
      val a = spark.range(0, 200).toDF("k").withColumn("x", col("k").cast("double"))
      val b = spark.range(0, 200).toDF("k2")
      val right = Seq((10, 9.0), (11, 15.0)).toDF("rid", "y")
      val df = a.join(b, col("k") === col("k2"))
        .nearestByJoin(right, abs(col("x") - col("y")),
          numResults = 2, mode = "exact", direction = "distance")

      val initialPlan = df.queryExecution.executedPlan
      assert(collect(initialPlan) { case j: SortMergeJoinExec => j }.size == 1,
        "precondition: the left side must start as a sort-merge join\n" + initialPlan)
      df.collect()
      val finalPlan = initialPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan
      assert(collect(finalPlan) { case j: SortMergeJoinExec => j }.isEmpty,
        "AQE re-optimization was rejected; the sort-merge join survived:\n" + finalPlan)
      assert(collect(finalPlan) { case j: BroadcastHashJoinExec => j }.size == 1,
        "AQE re-optimization did not demote the sort-merge join:\n" + finalPlan)
    }
  }

It needs with AdaptiveSparkPlanHelper on the suite (that is where collect comes from — plain SparkPlan.collect won't do, since AdaptiveSparkPlanExec is a LeafExecNode) and one import:

import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper}

With the case deleted, this is what it reports — the sort-merge join is still there in the final plan:

AQE re-optimization was rejected; the sort-merge join survived:
ResultQueryStage 3
+- BroadcastNearestByJoin Inner, 2, abs((x#47 - y#62)), NearestByDistance
   :- *(6) SortMergeJoin [k#45L], [k2#50L], Inner
   :  :- *(4) Sort [k#45L ASC NULLS FIRST], false, 0
   :  :  +- AQEShuffleRead coalesced
   :  :     +- ShuffleQueryStage 0

Putting it in AdaptiveQueryExecSuite instead, next to the other re-optimization tests, would be just as good by me — that suite already has runAdaptiveAndVerifyResult and findTopLevelBroadcastHashJoin.

[peter-toth]

Finding 10. The case is labelled "Per-left-row result ordering (best-first)", but this orderBy sorts both sides before the comparison, so the assertion says nothing about the order rows come out in. As written it's case 1 again with a different sort key and a single left row.

Both paths do promise best-first per left row — RewriteNearestByJoin's scaladoc says Inline preserves MaxMinByK's array order, and the operator drains the heap back-to-front into results for the same reason — so this is worth pinning. Drop the sort and compare the raw order:

Suggested change

	.orderBy(abs(col("x") - col("y")), col("rid")).collect()
	.collect()

With one left row, no shuffle above the join and no ties in this data (distances 1.0, 5.0, 9.5, 11.0, 90.0 for k=3), collect() order is deterministic and best-first on both paths, so it won't be flaky.

[peter-toth]

Finding 11. BroadcastNestedLoopJoinExec keeps the streamed side's ordering for exactly the (join type, build side) combinations this operator supports:

// BroadcastNestedLoopJoinExec.scala:72-77
override def outputOrdering: Seq[SortOrder] = (joinType, buildSide) match {
  case (_: InnerLike, _) | (LeftOuter, BuildRight) | (RightOuter, BuildLeft) |
       (LeftSingle, BuildRight) | (LeftSemi, BuildRight) | (LeftAnti, BuildRight) =>
    streamed.outputOrdering
  case _ => Nil
}

and the argument carries over here: doExecute walks leftIter in order and emits each left row's matches contiguously (including the single null-padded row for an unmatched LEFT OUTER row), so any ordering on left columns still holds on the output — repeated adjacent keys are fine for a SortOrder. The nullability difference between left.output and this node's widened output doesn't matter either, since AttributeReference.canonicalized drops nullability and SortOrder matching in EnsureRequirements goes through semanticEquals.

Nil isn't wrong, it just makes any downstream operator that wants the left ordering (a sort-merge join, a window, an orderBy on a left column) pay for a SortExec it doesn't need.

Suggested change

	override def outputOrdering: Seq[SortOrder] = Nil
	override def outputOrdering: Seq[SortOrder] = left.outputOrdering

[peter-toth]

Finding 12. BaseJoinExec.verboseStringWithOperatorId labels this field Join type (BaseJoinExec.scala:47 and :53), and so does every other join operator's EXPLAIN FORMATTED output. Ranking, NumResults and Direction are new fields and can be named however reads best, but this one already exists, so renaming it just for this operator makes EXPLAIN inconsistent across joins.

Suggested change

	|${ExplainUtils.generateFieldString("JoinType", joinType.toString)}
	|${ExplainUtils.generateFieldString("Join type", joinType.toString)}

The assert(explain.contains("JoinType: Inner")) in BroadcastNearestByJoinExecSuite needs the same update.

[peter-toth]

Finding 13. Rule extends SQLConfHelper (Rule.scala:24), which already gives this object a conf, and the check this pairs with in CheckAnalysis reads conf.nearestByBroadcastEnabled. Going through SQLConf.get directly is the only such read in the rule, and it's the sole reason the SQLConf import was added in this PR.

Suggested change

	if !SQLConf.get.nearestByBroadcastEnabled =>
	if !conf.nearestByBroadcastEnabled =>

With that, import org.apache.spark.sql.internal.SQLConf can come back out.

[cloud-fan]

1 addressed, 0 remaining, 6 new. (0 newly introduced, 6 late catches, 0 previously raised.)
0 blocking, 4 non-blocking, 2 nits.
The execution design remains sound, but six already-raised cleanup and coverage issues are still current, including a stale release annotation and two tests that do not assert their named contracts.

Already raised in existing discussion (6)

• Update this to .version("4.4.0"). The latest branch-4.x now identifies itself as 4.4.0-SNAPSHOT, so this normally backported config first ships in 4.4.0 rather than 4.3.0. -- existing discussion
• This test does not cover the new ValidateSparkPlan branch: it inspects the pre-execution adaptive plan, while the branch matters after a broadcast query stage is installed during re-optimization. Add the focused adaptive join-demotion assertion from the existing thread so deleting the registration makes the test fail. -- existing discussion
• Drop this orderBy before collecting case 4. Sorting both paths erases the best-first emission order the case is named to verify, so the current assertion duplicates the value-parity coverage instead of testing ordering. -- existing discussion
• Return left.outputOrdering here. The operator walks the left iterator in order and emits each left row's matches contiguously, so discarding that ordering forces avoidable downstream sorts compared with the established build-right broadcast join behavior. -- existing discussion
• Use Join type for this field, matching BaseJoinExec and the other join operators' formatted explain output; update the assertion with it. -- existing discussion
• Use the rule's inherited conf.nearestByBroadcastEnabled accessor here and remove the dedicated SQLConf import, matching the paired CheckAnalysis read and normal Rule convention. -- existing discussion

Verification

Traced both configuration modes from CheckAnalysis through the non-excludable FinishAnalysis rewrite, planner selection, identity broadcast, and AQE validation. Rechecked natural ordering, null exclusion, left-outer padding, projection-buffer retention, nondeterministic initialization, output metadata, and direct operator/rewrite parity. No tests were run as part of this review.

PR metadata suggestions

• Update the config's stated since-version from 4.3.0 to 4.4.0.
• Qualify the AQE test claim: the current test exercises execution with AQE enabled but does not prove post-stage re-optimization succeeds.

[cloud-fan]

6 addressed, 0 remaining, 2 new to this AI review. (0 newly introduced, 1 late catch, 1 previously raised.)
1 blocking, 1 non-blocking, 0 nits.
The JVM execution path is substantially covered, but a supported two-sided Python UDF ranking fails only when the new flag is enabled.

Already raised in existing discussion (1)

• The hot loop still allocates a new HeapEntry for every retained candidate. In descending-best input, every right row can replace the current worst, so this remains O(leftRows * rightRows) escaping objects even though only k entries stay live. Reuse the polled entry (or use a primitive/mutable bounded heap) so the operator's allocation count is bounded by k per task. -- existing discussion

Correctness (1)

• sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala:80: A Python UDF ranking that references both sides works on the rewrite path but fails when the broadcast flag is enabled. The enabled path leaves the UDF on the binary NearestByJoin; ExtractPythonUDFs can evaluate a UDF only under one child and throws when its references span both children. Route Python-UDF rankings through a joined evaluation node (or keep them on the rewrite path), and add a two-sided Python UDF parity test. -- see inline

Verification

Traced both configuration branches from analysis through optimizer and physical planning. Compared the operator with the rewrite for ordering, null exclusion, outer-row padding, output metadata, nondeterministic initialization, and AQE validation. Also traced Python UDF extraction: the disabled rewrite first creates a joined child, while the enabled binary NearestByJoin leaves a two-sided UDF with no single child containing all references, which ExtractPythonUDFs rejects.

[cloud-fan]

Please keep rankings containing a Python UDF on a plan shape where the UDF can be evaluated after both inputs are joined. With this gate enabled, a UDF such as udf((x, y) => ...)(left.x, right.y) remains on the binary NearestByJoin; ExtractPythonUDFs cannot place it under either child and throws Invalid PythonUDF ... requires attributes from more than one child. The flag-off rewrite works because it creates the Join before UDF extraction. Please route this case through the rewrite (or introduce a joined evaluation node) and add a two-sided Python UDF parity test.

[cloud-fan]

1 addressed, 1 remaining, 4 new to this AI review. (4 newly introduced, 0 late catches, 0 previously raised.)
0 blocking, 1 non-blocking, 4 nits.
The Python-UDF fallback is implemented, but four comments still describe the pre-fallback universal flag behavior; one previously raised allocation concern also remains.

Remaining from prior review (1)

• Non-blocking: This still allocates up to k HeapEntry objects for every left row, even though the inner result iterator is exhausted before flatMap advances. Please pool min(k, rightRows.length) mutable entries per partition and reset them when refilling the heap; that makes the allocation bound match the operator's GC-focused design across the whole partition, not once per left row. -- existing thread

Nits: 4 minor items (see inline comments).

Verification

Traced both configuration branches through CheckAnalysis, RewriteNearestByJoin, planner selection, identity broadcast, and AQE validation. Verified that the new cross-child scalar Python UDF guard routes those rankings through the established joined rewrite, while other enabled plans reach the dedicated operator. Rechecked heap admission and reuse, ordering, null/outer behavior, and the direct parity coverage. No tests were run as part of this review.

PR metadata suggestions

• Document that enabled queries with cross-child scalar Python UDF rankings fall back to the aggregate rewrite instead of reaching BroadcastNearestByJoinExec.

[peter-toth]

Re-checked through 3e6e08a — findings 8, 9, 10, 11, 12, 13 all resolved, nothing regressed. Specifically: the conf is .version("4.4.0") and I re-ran dev/next_version_candidates.py on this head to confirm that's still the right value (master 5.0.0 / branch-4.x 4.4.0); the new AQE test asserts the right child really is a BroadcastQueryStageExec and then calls ValidateSparkPlan.apply on that subtree, so deleting the case b: BroadcastNearestByJoinExec now makes it throw via the catch-all rather than passing silently; parity case 4 dropped its orderBy; outputOrdering is left.outputOrdering; the EXPLAIN field is Join type again; and the rule reads conf with the SQLConf import gone.

The new commit changes direction slightly: flag-ON is no longer synonymous with the operator path, because a cross-child scalar Python UDF ranking is now routed back through the rewrite. That is the right call, but the gate it interacts with in CheckAnalysis wasn't updated with it — finding 15, which I measured. Not re-opening @cloud-fan's four wording nits on that same fallback, the HeapEntry pooling thread, or the metadata suggestion about documenting the fallback (+1 to all of them).

Blocking

• 14. PR description no longer matches the code (new): HeapEntry is described as a "case class" but it is a mutable plain class — and the mutability is the point of this commit; BroadcastNearestByJoinExecSuite has 32 tests, not 31; and RewriteNearestByJoinSuite is listed under "Existing ... (13) ... pass with default conf (no regression)" although this PR adds two tests to it (now 15).

Non-blocking

• 15. Flag ON waives the crossJoin check for queries that still build the Join (new): with the flag on, CheckAnalysis skips NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED, but the new Python-UDF fallback does build the synthetic Join, so those queries fail later in CheckCartesianProducts with a generic cartesian-product error instead. I reproduced it; the conf doc sentence below is false for that case. [inline: sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:2670]
• 16. Operator not registered in CoalesceShufflePartitions (late catch): isExplodingJoin lists BroadcastNestedLoopJoinExec and CartesianProductExec; this operator multiplies each left row by up to k (validated to 100000) yet falls through to case _ => false, so AQE coalesces the left side to the full advisory partition size instead of the min size. Silent — no test or CI job can fail on it. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ValidateSparkPlan.scala:61]
• 17. Committed benchmark results were not produced by the benchmark workflow (late catch): this is the only header among all sql/core/benchmarks/*-results.txt that is not from the GitHub Actions runner, and it carries one JDK block where .github/workflows/benchmark.yml emits three (17/21/25), so the numbers aren't comparable with any neighbouring file. [inline: sql/core/benchmarks/NearestByJoinBenchmark-results.txt:5]
• 18. Neither Python-UDF path is ever executed (new): both new tests build PythonUDF(..., null, ...) and assert only the rule's decision, so @cloud-fan's "two-sided Python UDF parity test" ask is still open, and the single-child ranking the guard deliberately admits onto the operator path has no coverage at all. python/pyspark/sql/tests/test_nearest_by_join.py can run both for real. [inline: sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExecSuite.scala:706]

[yadavay-amzn]

@cloud-fan @peter-toth Thanks again for the thorough review!
Ready for another look when you get the chance