peter-toth commented on code in PR #56101:
URL: https://github.com/apache/spark/pull/56101#discussion_r3763410679
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -2660,6 +2660,19 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val NEAREST_BY_BROADCAST_ENABLED =
+ buildConf("spark.sql.join.nearestBy.broadcast.enabled")
+ .internal()
+ .doc("When true, NearestByJoin uses a streaming heap operator instead of
the " +
+ "cross-product + aggregate rewrite. The right side is always
broadcast, regardless " +
+ "of its size and of spark.sql.autoBroadcastJoinThreshold, so a right
side too large " +
+ "to broadcast fails the query instead of falling back to the rewrite.
Because no " +
+ "Join node is built, spark.sql.crossJoin.enabled does not apply on
this path.")
Review Comment:
**Finding 15.** This last sentence is now false for the case the new guard
carves out, and the `CheckAnalysis` gate that mirrors it is wrong in the same
way.
`CheckAnalysis.scala:713` waives `NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED`
whenever `conf.nearestByBroadcastEnabled`, on the premise that no `Join` node
gets built. `containsCrossChildPythonUDF` breaks that premise: those nodes
*are* rewritten, a `Join(taggedLeft, right, joinType, None, JoinHint.NONE)`
*is* built, and the `Check Cartesian Products` batch (`Optimizer.scala:290`)
runs after `Finish Analysis` (`:185`), so it sees it.
Measured on this head — flag on, `spark.sql.crossJoin.enabled=false`, a
two-sided `PythonUDF` ranking (the same fixture as the new suite test), through
`spark.sessionState.optimizer.execute`:
org.apache.spark.sql.AnalysisException: Detected implicit cartesian
product for INNER join between logical plans
Either: use the CROSS JOIN syntax to allow cartesian products between
these relations, or:
enable implicit cartesian products by setting the configuration variable
spark.sql.crossJoin.enabled=true.
at
org.apache.spark.sql.catalyst.optimizer.CheckCartesianProducts$.apply(Optimizer.scala:2665)
The same query with `crossJoin.enabled=true` optimizes fine, so this is the
gate, not the fixture. So the user who opted out of cross products still gets
rejected — which is arguably the right outcome — but with a cartesian-product
message about a join they never wrote, after `CheckAnalysis` deliberately
decided not to reject them, and against a conf doc that says the setting
doesn't apply here.
Cheapest fix is to make the two sites agree by construction, the way the
rewrite guard and the strategy were made to agree at round 3: lift the
predicate onto the `NearestByJoin` companion and use it in both places, so the
dedicated error class is what fires.
```scala
// NearestByJoin companion
def hasCrossChildPythonUDF(j: NearestByJoin): Boolean = ...
// CheckAnalysis
case j: NearestByJoin if !conf.crossJoinEnabled &&
(!conf.nearestByBroadcastEnabled ||
NearestByJoin.hasCrossChildPythonUDF(j)) =>
```
and then the sentence here becomes true again as "...does not apply to nodes
that reach the operator". This is the behaviour behind @cloud-fan's wording nit
at `CheckAnalysis.scala:717` — replying there rather than here if you'd prefer
to keep it in one thread.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ValidateSparkPlan.scala:
##########
@@ -58,6 +58,11 @@ object ValidateSparkPlan extends Rule[SparkPlan] {
validate(buildPlan)
}
validate(probePlan)
+ case b: BroadcastNearestByJoinExec =>
Review Comment:
**Finding 16.** This registration is right, but it isn't the only AQE site
that enumerates broadcast joins by concrete type. `CoalesceShufflePartitions`
has two, and this operator takes the catch-all in both:
```scala
// CoalesceShufflePartitions.scala:183-200
private def childrenNeedCompatiblePartitioning(p: SparkPlan): Boolean = p
match {
...
case _: BroadcastHashJoinExec => false
case _: BroadcastNestedLoopJoinExec => false
case _ => true // <- BroadcastNearestByJoinExec
}
private def isExplodingJoin(p: SparkPlan): Boolean = p match {
case _: BroadcastNestedLoopJoinExec => true
case _: CartesianProductExec => true
case _ => false // <- BroadcastNearestByJoinExec
}
```
`isExplodingJoin` is the one that matters. It feeds
`CoalesceGroup.hasExplodingJoin`, and `advisoryPartitionSize` (`:134-137`)
returns `COALESCE_PARTITIONS_MIN_PARTITION_SIZE` instead of
`ADVISORY_PARTITION_SIZE_IN_BYTES` when it is set — precisely so a coalesced
input partition doesn't blow up after the join multiplies it. This operator
emits up to `numResults` rows per left row, and `numResults` is validated up to
`NearestByJoinValidation.MaxNumResults = 100000`, so it is exactly as exploding
as `BroadcastNestedLoopJoinExec`: AQE will coalesce the streamed side to the
full advisory size (64 MB by default) and then each partition fans out k-fold
downstream.
`p.exists(isExplodingJoin)` at `:176` means the one-liner is enough:
```scala
private def isExplodingJoin(p: SparkPlan): Boolean = p match {
case _: BroadcastNearestByJoinExec => true
case _: BroadcastNestedLoopJoinExec => true
case _: CartesianProductExec => true
case _ => false
}
```
`childrenNeedCompatiblePartitioning` is worth the matching line for
consistency with the other two broadcast joins, though `true` there is
conservative rather than wrong — the children genuinely don't need compatible
partitioning when one side is broadcast.
Failure mode is silent degradation (oversized downstream tasks), never an
exception, so nothing in the suite or CI can catch it — same shape as the
`ValidateSparkPlan` gap from last round.
##########
sql/core/benchmarks/NearestByJoinBenchmark-results.txt:
##########
@@ -0,0 +1,12 @@
+================================================================================================
+NearestByJoin Benchmark
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.260-264.1054.amzn2int.x86_64
Review Comment:
**Finding 17.** These numbers weren't produced by the standard path, so they
can't be read against any other file in this directory.
Every other committed result file comes from the
`.github/workflows/benchmark.yml` runners. Across all of
`sql/core/benchmarks/*-results.txt`:
$ grep -h "^OpenJDK" sql/core/benchmarks/*.txt | sort | uniq -c | sort
-rn
935 OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure
935 OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure
935 OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure
40 OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
40 OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux
6.17.0-1010-azure
40 OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
6.17.0-1010-azure
1 OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.260-264.1054.amzn2int.x86_64 <- this file
The equal counts across three JDKs are the tell: the workflow emits a
17/21/25 block per case, and this file has a single 17 block from a different
machine. Please regenerate it through the benchmark workflow
(`.github/workflows/benchmark.yml`, targeting
`org.apache.spark.sql.execution.benchmark.NearestByJoinBenchmark`) and commit
that output. The 28x ratio the PR description leads with will almost certainly
survive; the point is that a future reader can compare it with
`JoinBenchmark-results.txt` next door.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExecSuite.scala:
##########
@@ -0,0 +1,823 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.joins
+
+import java.sql.Date
+
+import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec,
AdaptiveSparkPlanHelper,
+ BroadcastQueryStageExec}
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+class BroadcastNearestByJoinExecSuite extends QueryTest with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+
+ import testImplicits._
+
+ private def withStreamingHeap(f: => Unit): Unit = {
+ withSQLConf(
+ SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true") {
+ f
+ }
+ }
+
+ test("empty right table - INNER returns nothing") {
+ withStreamingHeap {
+ val left = spark.range(5).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = spark.range(0).toDF("rid").withColumn("y", lit(0.0))
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ assert(result.count() == 0)
+ }
+ }
+
+ test("empty right table - LEFT OUTER returns left with nulls") {
+ withStreamingHeap {
+ val left = spark.range(3).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = spark.range(0).toDF("rid").withColumn("y", lit(0.0))
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance", joinType =
"left_outer")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ assert(result.count() == 3)
+ result.collect().foreach { row =>
+ assert(row.isNullAt(2)) // rid is null
+ assert(row.isNullAt(3)) // y is null
+ }
+ }
+ }
+
+ test("k=1 returns single nearest") {
+ withStreamingHeap {
+ val left = Seq((1, 10.0), (2, 20.0)).toDF("id", "x")
+ val right = Seq((10, 9.0), (11, 15.0), (12, 21.0)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 1, mode = "exact", direction = "distance")
+ .orderBy("id")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ checkAnswer(result, Seq(
+ Row(1, 10.0, 10, 9.0), // nearest to 10.0 is 9.0
+ Row(2, 20.0, 12, 21.0) // nearest to 20.0 is 21.0
+ ))
+ }
+ }
+
+ test("k > right table size returns all right rows per left row") {
+ withStreamingHeap {
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ val right = Seq((10, 1.0), (11, 2.0)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 10, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // Only 2 right rows exist, so we get 2 results
+ assert(result.count() == 2)
+ checkAnswer(result.orderBy("rid"), Seq(
+ Row(1, 5.0, 10, 1.0),
+ Row(1, 5.0, 11, 2.0)
+ ))
+ }
+ }
+
+ test("NaN ranking values participate in ordering") {
+ withStreamingHeap {
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ val right = Seq((10, Double.NaN), (11, 3.0), (12, 7.0)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // NaN participates in natural ordering (sorts after all non-NaN for
distance)
+ assert(result.count() == 3)
+ checkAnswer(result.orderBy("rid"), Seq(
+ Row(1, 5.0, 10, Double.NaN),
+ Row(1, 5.0, 11, 3.0),
+ Row(1, 5.0, 12, 7.0)
+ ))
+ }
+ }
+
+ test("null ranking values are excluded, not treated as 0.0") {
+ withStreamingHeap {
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ // y=null will produce null ranking expression (abs(5.0 - null) = null)
+ val right = Seq((10, Some(3.0)), (11, None), (12,
Some(7.0))).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // null row excluded, only 2 results
+ assert(result.count() == 2)
+ checkAnswer(result.orderBy("rid"), Seq(
+ Row(1, 5.0, 10, 3.0),
+ Row(1, 5.0, 12, 7.0)
+ ))
+ }
+ }
+
+ test("asymmetric - right small left big, operator fires") {
+ withStreamingHeap {
+ val left = spark.range(1000).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = spark.range(50).toDF("rid").withColumn("y",
col("rid").cast("double") * 20)
+ val df = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val plan = df.queryExecution.executedPlan.toString()
+ assert(plan.contains("BroadcastNearestByJoin"),
+ s"Expected BroadcastNearestByJoinExec in plan: $plan")
+ assert(df.count() == 1000 * 3)
+ // Spot check: id=0 (x=0.0) nearest to y values 0,20,40 -> rids 0,1,2
+ val row0 = df.filter(col("id") === 0).orderBy(abs(col("x") -
col("y"))).collect()
+ assert(row0.length == 3)
+ assert(row0(0).getAs[Long]("rid") == 0L) // y=0, distance=0
+ }
+ }
+
+ test("asymmetric - right exceeds broadcast threshold, broadcast still used
(flag ON)") {
+ // When the broadcast flag is ON, BroadcastNearestByJoinExec is always
used regardless
+ // of right-side size. There is no size decision and no fallback; an
oversized right
+ // side will fail the query at runtime (SPARK-57091).
+ withSQLConf(
+ SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") {
+ val left = spark.range(10).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = spark.range(50).toDF("rid").withColumn("y",
col("rid").cast("double"))
+ val df = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ val plan = df.queryExecution.executedPlan.toString()
+ assert(plan.contains("BroadcastNearestByJoin"),
+ s"Expected BroadcastNearestByJoinExec in plan (flag ON always
broadcasts): $plan")
+ // Results still correct
+ assert(df.count() == 10 * 2)
+ val row0 = df.filter(col("id") === 0).orderBy(abs(col("x") -
col("y"))).collect()
+ assert(row0(0).getAs[Long]("rid") == 0L)
+ }
+ }
+
+ test("asymmetric - left small right big, operator fires") {
+ withStreamingHeap {
+ val left = spark.range(10).toDF("id").withColumn("x",
col("id").cast("double") * 50)
+ val right = spark.range(500).toDF("rid").withColumn("y",
col("rid").cast("double"))
+ val df = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 5, mode = "exact", direction = "distance")
+ val plan = df.queryExecution.executedPlan.toString()
+ assert(plan.contains("BroadcastNearestByJoin"),
+ s"Expected BroadcastNearestByJoinExec in plan: $plan")
+ assert(df.count() == 10 * 5)
+ // id=0 (x=0.0): nearest are y=0,1,2,3,4
+ val row0 = df.filter(col("id") === 0).orderBy(abs(col("x") -
col("y"))).collect()
+ assert(row0(0).getAs[Long]("rid") == 0L)
+ assert(row0(4).getAs[Long]("rid") == 4L)
+ }
+ }
+
+ test("asymmetric - both sides moderate, operator fires") {
+ withStreamingHeap {
+ val left = spark.range(200).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = spark.range(200).toDF("rid").withColumn("y",
col("rid").cast("double") + 0.5)
+ val df = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val plan = df.queryExecution.executedPlan.toString()
+ assert(plan.contains("BroadcastNearestByJoin"),
+ s"Expected BroadcastNearestByJoinExec in plan: $plan")
+ assert(df.count() == 200 * 3)
+ // id=100 (x=100.0): nearest y values are 99.5(rid=99), 100.5(rid=100),
101.5(rid=101)
+ val row100 = df.filter(col("id") === 100).orderBy(abs(col("x") -
col("y"))).collect()
+ assert(row100.length == 3)
+ assert(row100(0).getAs[Long]("rid") == 99L) // y=99.5, distance=0.5
+ }
+ }
+
+ test("basic correctness - small dataset") {
+ withStreamingHeap {
+ val left = Seq((1, 0.0), (2, 10.0), (3, 20.0)).toDF("id", "x")
+ val right = Seq((100, 1.0), (101, 9.0), (102, 11.0), (103, 19.0), (104,
25.0))
+ .toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ .orderBy("id", "y")
+ // id=1 (x=0.0): nearest are y=1.0 (d=1), y=9.0 (d=9)
+ // id=2 (x=10.0): nearest are y=9.0 (d=1), y=11.0 (d=1)
+ // id=3 (x=20.0): nearest are y=19.0 (d=1), y=25.0 (d=5)
+ checkAnswer(result, Seq(
+ Row(1, 0.0, 100, 1.0),
+ Row(1, 0.0, 101, 9.0),
+ Row(2, 10.0, 101, 9.0),
+ Row(2, 10.0, 102, 11.0),
+ Row(3, 20.0, 103, 19.0),
+ Row(3, 20.0, 104, 25.0)
+ ))
+ }
+ }
+
+ test("similarity direction - keeps largest ranking values") {
+ withStreamingHeap {
+ // Higher ranking value = more similar; top-k should be the largest
values
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ val right = Seq((10, 1.0), (11, 3.0), (12, 8.0), (13, 10.0)).toDF("rid",
"y")
+ // Use y directly as ranking: higher y = more similar
+ val result = left.nearestByJoin(right, col("y"),
+ numResults = 2, mode = "exact", direction = "similarity")
+ .orderBy(col("y").desc)
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // Top-2 by largest y: rid=13 (y=10.0), rid=12 (y=8.0)
+ checkAnswer(result, Seq(
+ Row(1, 5.0, 13, 10.0),
+ Row(1, 5.0, 12, 8.0)
+ ))
+ }
+ }
+
+ test("integer ranking expression - does not corrupt values") {
+ withStreamingHeap {
+ // x and y are IntegerType, so (x - y) produces IntegerType ranking.
+ // Use negative ranking values to expose getDouble corruption:
+ // int -1 stored as 0x00000000FFFFFFFF reads as a tiny positive double,
not -1.0.
+ val left = Seq((1, 5)).toDF("id", "x")
+ val right = Seq((10, 6), (11, 3), (12, 100)).toDF("rid", "y")
+ // ranking = x - y: (5-6)=-1, (5-3)=2, (5-100)=-95
+ // direction=similarity means largest ranking wins, so top-2 =
rid=11(2), rid=10(-1)
+ val result = left.nearestByJoin(right, col("x") - col("y"),
+ numResults = 2, mode = "exact", direction = "similarity")
+ .orderBy((col("x") - col("y")).desc)
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ checkAnswer(result, Seq(
+ Row(1, 5, 11, 3), // ranking = 2 (largest)
+ Row(1, 5, 10, 6) // ranking = -1 (second largest)
+ ))
+ }
+ }
+
+ test("tie-breaking - equal distances produce correct count") {
+ withStreamingHeap {
+ // 4 right rows all at distance 1.0 from x=5.0; k=2
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ val right = Seq((10, 4.0), (11, 6.0), (12, 4.0), (13, 6.0)).toDF("rid",
"y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // All 4 are tied at distance 1.0; we should get exactly k=2 results
+ assert(result.count() == 2)
+ // Each result should have distance 1.0
+ result.collect().foreach { row =>
+ val y = row.getAs[Double]("y")
+ assert(math.abs(5.0 - y) == 1.0)
+ }
+ }
+ }
+
+ // ==========================================================================
+ // SPARK-57091: Tests for ranking comparison fix
(TypeUtils.getInterpretedOrdering)
+ // ==========================================================================
+
+ test("SPARK-57091: DateType ranking - nearest by earliest date (distance)") {
+ withStreamingHeap {
+ // Date ranking: the old Cast(_, DoubleType) cannot cast dates to double,
+ // producing null rankings and empty results for INNER join.
+ val left = Seq((1, Date.valueOf("2024-06-15"))).toDF("id", "ref_date")
+ val right = Seq(
+ (10, Date.valueOf("2024-06-10")),
+ (11, Date.valueOf("2024-06-20")),
+ (12, Date.valueOf("2024-12-01"))
+ ).toDF("rid", "event_date")
+ // Use event_date as ranking; direction=distance means earliest dates win
+ val result = left.nearestByJoin(right, col("event_date"),
+ numResults = 2, mode = "exact", direction = "distance")
+ .orderBy("event_date")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ checkAnswer(result, Seq(
+ Row(1, Date.valueOf("2024-06-15"), 10, Date.valueOf("2024-06-10")),
+ Row(1, Date.valueOf("2024-06-15"), 11, Date.valueOf("2024-06-20"))
+ ))
+ }
+ }
+
+ test("SPARK-57091: DateType ranking - similarity picks latest date") {
+ withStreamingHeap {
+ // Same DateType scenario but direction=similarity (largest value wins =
latest date).
+ // The old Cast(_, DoubleType) cannot cast dates, yielding empty results.
+ val left = Seq((1, Date.valueOf("2024-06-15"))).toDF("id", "ref_date")
+ val right = Seq(
+ (10, Date.valueOf("2024-01-01")),
+ (11, Date.valueOf("2024-06-20")),
+ (12, Date.valueOf("2024-12-01"))
+ ).toDF("rid", "event_date")
+ val result = left.nearestByJoin(right, col("event_date"),
+ numResults = 2, mode = "exact", direction = "similarity")
+ .orderBy(col("event_date").desc)
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // Largest (latest) dates: 2024-12-01, 2024-06-20
+ checkAnswer(result, Seq(
+ Row(1, Date.valueOf("2024-06-15"), 12, Date.valueOf("2024-12-01")),
+ Row(1, Date.valueOf("2024-06-15"), 11, Date.valueOf("2024-06-20"))
+ ))
+ }
+ }
+
+ test("SPARK-57091: Long ranking past 2^53 - distinguishes values equal as
Double") {
+ withStreamingHeap {
+ // Two Long values that differ by 1 but are equal when cast to Double:
+ // Long.MAX_VALUE - 1 and Long.MAX_VALUE both cast to the same Double
(9.223372036854776E18)
+ // The old Cast(_, DoubleType) approach would see them as equal and pick
arbitrarily.
+ val v1 = Long.MaxValue - 1 // 9223372036854775806
+ val v2 = Long.MaxValue // 9223372036854775807
+ // Verify they are indeed equal as Double (precondition)
+ assert(v1.toDouble == v2.toDouble,
+ "precondition: these Longs must be equal as Double to test the fix")
+
+ val left = Seq((1, 0L)).toDF("id", "x")
+ val right = Seq((10, v1), (11, v2)).toDF("rid", "y")
+ // direction=distance: smallest y wins. v1 < v2 as Long but equal as
Double.
+ val result = left.nearestByJoin(right, col("y"),
+ numResults = 1, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // With proper Long ordering, v1 (smaller) is chosen deterministically
+ checkAnswer(result, Seq(Row(1, 0L, 10, v1)))
+ }
+ }
+
+ test("SPARK-57091: Decimal ranking preserves precision beyond Double") {
+ withStreamingHeap {
+ // Decimals with 18 digits of precision: these are identical when cast
to Double
+ // but differ in the last digit as Decimal.
+ val d1 = new java.math.BigDecimal("1.000000000000000001")
+ val d2 = new java.math.BigDecimal("1.000000000000000002")
+ // Verify they are equal as Double (precondition)
+ assert(d1.doubleValue() == d2.doubleValue(),
+ "precondition: these Decimals must be equal as Double to test the fix")
+
+ val left = spark.createDataFrame(
+ Seq((1, new java.math.BigDecimal("0")))).toDF("id", "x")
+ val right = spark.createDataFrame(
+ Seq((10, d1), (11, d2))).toDF("rid", "y")
+ // direction=distance: smallest y wins. d1 < d2 as Decimal but equal as
Double.
+ val result = left.nearestByJoin(right, col("y"),
+ numResults = 1, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ checkAnswer(result, Seq(Row(1, new java.math.BigDecimal("0"), 10, d1)))
+ }
+ }
+
+ test("SPARK-57091: output schema nullability - INNER join has all-nullable
columns") {
+ withStreamingHeap {
+ // The fix ensures both left and right output attributes are nullable
for INNER join.
+ // With the old approach, left side would retain original nullability
(non-nullable for
+ // spark.range), causing schema mismatch with the logical plan.
+ val left = spark.range(3).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = Seq((10, 1.0), (11, 2.0)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 1, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.toString.contains("BroadcastNearestByJoin"),
+ "expected BroadcastNearestByJoinExec in plan")
+ // All output columns must be nullable
+ plan.output.foreach { attr =>
+ assert(attr.nullable,
+ s"column '${attr.name}' should be nullable in INNER join output but
was not")
+ }
+ }
+ }
+
+ test("SPARK-57091: gate off - broadcast operator does not fire, rewrite path
used") {
+ // When spark.sql.join.nearestBy.broadcast.enabled is false (default),
+ // the BroadcastNearestByJoinExec must NOT appear and the rewrite path
must be used.
+ withSQLConf(
+ SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "false",
+ SQLConf.CROSS_JOINS_ENABLED.key -> "true") {
+ val left = Seq((1, 10.0), (2, 20.0)).toDF("id", "x")
+ val right = Seq((10, 9.0), (11, 21.0)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 1, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan.toString()
+ assert(!plan.contains("BroadcastNearestByJoin"),
+ "BroadcastNearestByJoinExec must NOT appear when gate is off")
+ // Results are still correct via rewrite path
+ checkAnswer(result.orderBy("id"), Seq(
+ Row(1, 10.0, 10, 9.0),
+ Row(2, 20.0, 11, 21.0)
+ ))
+ }
+ }
+
+ test("SPARK-57091: Long ranking - similarity direction orders correctly past
2^53") {
+ withStreamingHeap {
+ // v1 = 2^53 - 1, v2 = 2^53, v3 = 2^53 + 1. v2 and v3 are equal as Double
+ // but v1 < v2 < v3 as Long. With k=1 and similarity (largest wins),
proper
+ // Long ordering deterministically keeps v3 because v3 > v2 strictly.
+ // With a Double cast v2==v3, and the result would be non-deterministic.
+ val v2 = (1L << 53) // 9007199254740992
+ val v3 = (1L << 53) + 1 // 9007199254740993
+ assert(v2.toDouble == v3.toDouble,
+ "precondition: v2 and v3 must be equal as Double")
+
+ val left = Seq((1, 0L)).toDF("id", "x")
+ // v3 is inserted first; with correct Long comparison the heap retains it
+ // because v3 > v2 strictly, so v2 fails the retention check and is
never offered.
+ val right = Seq((12, v3), (11, v2)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, col("y"),
+ numResults = 1, mode = "exact", direction = "similarity")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // With proper Long ordering, v3 (larger) must be kept
+ checkAnswer(result, Seq(Row(1, 0L, 12, v3)))
+ }
+ }
+
+ test("SPARK-57091: AQE ValidateSparkPlan accepts
BroadcastNearestByJoinExec") {
+ // Directly verify that ValidateSparkPlan does not reject a
+ // BroadcastNearestByJoinExec whose right child is a
BroadcastQueryStageExec.
+ // Without the explicit `case b: BroadcastNearestByJoinExec` in
ValidateSparkPlan,
+ // the default case recurses into the BroadcastQueryStageExec child and
throws
+ // InvalidAQEPlanException.
+ import org.apache.spark.sql.execution.adaptive.ValidateSparkPlan
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true") {
+ val left = Seq((1, 10.0), (2, 20.0), (3, 30.0)).toDF("id", "x")
+ val right = Seq((10, 9.0), (11, 15.0), (12, 21.0), (13,
29.0)).toDF("rid", "y")
+ val df = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ val executedPlan = df.queryExecution.executedPlan
+ // After execution, the plan tree has BroadcastNearestByJoinExec with
+ // BroadcastQueryStageExec as its right child
+ val adaptivePlan = executedPlan match {
+ case aqe: AdaptiveSparkPlanExec =>
+ df.collect() // force execution
+ aqe.executedPlan
+ case other => other
+ }
+ // Find the BroadcastNearestByJoinExec with its BroadcastQueryStageExec
right child
+ val nearestByOps = collect(adaptivePlan) {
+ case b: BroadcastNearestByJoinExec => b
+ }
+ assert(nearestByOps.nonEmpty,
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
adaptivePlan.treeString)
+ val nb = nearestByOps.head
+ assert(nb.right.isInstanceOf[BroadcastQueryStageExec],
+ s"Expected BroadcastQueryStageExec as right child but got:
${nb.right.getClass.getName}")
+ // Directly invoke ValidateSparkPlan on the subtree rooted at
+ // BroadcastNearestByJoinExec. This must not throw
InvalidAQEPlanException.
+ // If the case in ValidateSparkPlan for our operator is missing, this
throws.
+ ValidateSparkPlan.apply(nb)
+ }
+ }
+
+ test("SPARK-57091: StringType ranking - buffer retention with heap
eviction") {
+ withStreamingHeap {
+ // Use StringType as the ranking column. UnsafeProjection reuses its
output buffer,
+ // so UTF8String values point into the mutable buffer. Without .copy(),
earlier
+ // heap entries get corrupted when the buffer is overwritten on
subsequent iterations.
+ // Data is ordered so that LATER rows have BETTER (smaller) ranking
values and
+ // DISPLACE earlier retained entries, exercising the variable-length
.copy() path.
+ val left = Seq((1, "ref")).toDF("id", "x")
+ val right = Seq(
+ (10, "hhh"), (11, "ggg"), (12, "fff"), (13, "eee"),
+ (14, "ddd"), (15, "ccc"), (16, "bbb"), (17, "aaa")
+ ).toDF("rid", "label")
+ // direction=distance: smallest string wins (lexicographic). k=2 ->
"aaa", "bbb"
+ // The first entries in the heap are "hhh","ggg" which get displaced by
later
+ // better values, forcing .copy() of the retained ranking values.
+ val result = left.nearestByJoin(right, col("label"),
+ numResults = 2, mode = "exact", direction = "distance")
+ .orderBy("label")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ checkAnswer(result, Seq(
+ Row(1, "ref", 17, "aaa"),
+ Row(1, "ref", 16, "bbb")
+ ))
+ }
+ }
+
+ test("NaN ranking values under similarity direction") {
+ withStreamingHeap {
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ val right = Seq((10, Double.NaN), (11, 3.0), (12, 8.0), (13,
10.0)).toDF("rid", "y")
+ // direction=similarity: largest ranking value wins. NaN is largest in
Java ordering.
+ val result = left.nearestByJoin(right, col("y"),
+ numResults = 2, mode = "exact", direction = "similarity")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // NaN is greatest per Java Double ordering, so top-2 similarity = NaN,
10.0
+ checkAnswer(result.orderBy(col("y").desc_nulls_last), Seq(
+ Row(1, 5.0, 10, Double.NaN),
+ Row(1, 5.0, 13, 10.0)
+ ))
+ }
+ }
+
+ test("null in non-ranking right column propagates correctly") {
+ withStreamingHeap {
+ val left = Seq((1, 5.0)).toDF("id", "x")
+ // "label" column has nulls but is NOT the ranking column
+ val right = Seq(
+ (10, 3.0, Some("a")),
+ (11, 7.0, None),
+ (12, 100.0, Some("c"))
+ ).toDF("rid", "y", "label")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ .orderBy("rid")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // Top-2 nearest: rid=10 (d=2.0, label="a"), rid=11 (d=2.0, label=null)
+ checkAnswer(result, Seq(
+ Row(1, 5.0, 10, 3.0, "a"),
+ Row(1, 5.0, 11, 7.0, null)
+ ))
+ }
+ }
+
+ test("SPARK-57091: non-deterministic ranking expression (rand()) does not
throw") {
+ withStreamingHeap {
+ // A non-deterministic ranking expression must have its projection
initialized
+ // with the partition index before evaluation. Without initialization,
rand() throws
+ // "Nondeterministic expression ... has not been initialized".
+ val left = spark.range(10).toDF("id").withColumn("x",
col("id").cast("double"))
+ val right = Seq((10, 1.0), (11, 2.0), (12, 3.0)).toDF("rid", "y")
+ // Use rand() as ranking: non-deterministic, returns Double
+ val result = left.nearestByJoin(right, rand(),
+ numResults = 2, mode = "exact", direction = "similarity")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // Each of 10 left rows should get 2 right rows (k=2, right has 3 rows)
+ assert(result.count() == 20)
+ }
+ }
+
+ test("SPARK-57091: DSv2 right side does not throw INTERNAL_ERROR with flag
ON") {
+ // Regression: before the stats-free fix, reading right.stats.sizeInBytes
in the
+ // optimizer's FinishAnalysis batch triggered computeStats on a DSv2
source before
+ // filter/partition pushdown completed, throwing [INTERNAL_ERROR]. With
the fix,
+ // the optimizer no longer reads stats -- the NearestByJoin node is left
intact for
+ // the planner, which unconditionally plans BroadcastNearestByJoinExec.
+ withSQLConf(
+ "spark.sql.catalog.testcat" ->
+
classOf[org.apache.spark.sql.connector.catalog.InMemoryTableCatalog].getName,
+ SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true") {
+ spark.sql(
+ """CREATE TABLE testcat.right_tbl (rid INT, y DOUBLE)
+ |USING foo""".stripMargin)
+ try {
+ spark.sql("INSERT INTO testcat.right_tbl VALUES (10, 1.0), (11, 5.0),
(12, 9.0)")
+ val left = Seq((1, 3.0), (2, 7.0)).toDF("id", "x")
+ val right = spark.table("testcat.right_tbl")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ // Verify correct results:
+ // id=1(x=3.0) nearest y=1.0(d=2),5.0(d=2)
+ // id=2(x=7.0) nearest y=5.0(d=2),9.0(d=2)
+ assert(result.count() == 4)
+ val row1 = result.filter(col("id") === 1).orderBy(abs(col("x") -
col("y"))).collect()
+ assert(row1.length == 2)
+ } finally {
+ spark.sql("DROP TABLE IF EXISTS testcat.right_tbl")
+ }
+ }
+ }
+
+ test("SPARK-57091: partitioned right table with flag ON uses
BroadcastNearestByJoin") {
+ // Regression: before the stats-free fix, partitioned file tables reported
inflated
+ // sizeInBytes before filter pushdown. With the unconditional broadcast
contract,
+ // the planner always plans BroadcastNearestByJoinExec when the flag is ON.
+ withStreamingHeap {
+ withTempPath { dir =>
+ // Create partitioned parquet data
+ val data = (1 to 100).map(i => (i % 5, i, i.toDouble))
+ spark.createDataFrame(data).toDF("part", "rid", "y")
+ .write.partitionBy("part").parquet(dir.getAbsolutePath)
+
+ val right = spark.read.parquet(dir.getAbsolutePath).filter(col("part")
=== 0)
+ val left = Seq((1, 10.0), (2, 50.0)).toDF("id", "x")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec for partitioned right but
got:\n" +
+ plan.treeString)
+ // part=0 has values 5,10,15,...,100 (20 rows). Each left row gets 3
nearest.
+ assert(result.count() == 6)
+ }
+ }
+ }
+
+ // ==========================================================================
+ // F7: crossJoin.enabled divergence -- operator path does not require it
+ // ==========================================================================
+
+ test("SPARK-57091: flag ON succeeds without crossJoin.enabled") {
+ // When the broadcast flag is ON, no Join node is built so
CheckCartesianProducts
+ // does not apply. NEAREST BY should succeed even with crossJoin.enabled =
false.
+ withSQLConf(
+ SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true",
+ SQLConf.CROSS_JOINS_ENABLED.key -> "false") {
+ val left = Seq((1, 10.0), (2, 20.0)).toDF("id", "x")
+ val right = Seq((10, 9.0), (11, 15.0), (12, 21.0)).toDF("rid", "y")
+ val result = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 2, mode = "exact", direction = "distance")
+ val plan = result.queryExecution.executedPlan
+ assert(plan.treeString.contains("BroadcastNearestByJoin"),
+ "Expected BroadcastNearestByJoinExec in plan but got:\n" +
plan.treeString)
+ checkAnswer(result.filter(col("id") === 1).orderBy(abs(col("x") -
col("y"))), Seq(
+ Row(1, 10.0, 10, 9.0),
+ Row(1, 10.0, 11, 15.0)
+ ))
+ checkAnswer(result.filter(col("id") === 2).orderBy(abs(col("x") -
col("y"))), Seq(
+ Row(2, 20.0, 12, 21.0),
+ Row(2, 20.0, 11, 15.0)
+ ))
+ }
+ }
+
+ // ==========================================================================
+ // EXPLAIN FORMATTED output
+ // ==========================================================================
+
+ test("SPARK-57091: EXPLAIN FORMATTED shows ranking, k, and direction") {
+ withStreamingHeap {
+ val left = Seq((1, 10.0)).toDF("id", "x")
+ val right = Seq((10, 9.0)).toDF("rid", "y")
+ val df = left.nearestByJoin(right, abs(col("x") - col("y")),
+ numResults = 3, mode = "exact", direction = "distance")
+ val explain = df.queryExecution.explainString(
+ org.apache.spark.sql.execution.FormattedMode)
+ assert(explain.contains("NumResults: 3"),
+ s"EXPLAIN FORMATTED should contain 'NumResults: 3'\n$explain")
+ assert(explain.contains("Direction: NearestByDistance"),
+ s"EXPLAIN FORMATTED should contain 'Direction:
NearestByDistance'\n$explain")
+ assert(explain.contains("Ranking:"),
+ s"EXPLAIN FORMATTED should contain 'Ranking:'\n$explain")
+ assert(explain.contains("Join type: Inner"),
+ s"EXPLAIN FORMATTED should contain 'Join type: Inner'\n$explain")
+ }
+ }
+
+ // ==========================================================================
+ // Cross-child Python UDF fallback to rewrite path
+ // ==========================================================================
+
+ test("SPARK-57091: cross-child Python UDF in ranking routes through rewrite
path") {
Review Comment:
**Finding 18.** This test and its `RewriteNearestByJoinSuite` twin assert
which plan shape the rule picks, and nothing else —
`PythonUDF("cross_child_distance", null, ...)` can't be evaluated, as the
comment says. That leaves both Python-UDF paths untested at execution:
1. **The fallback path.** @cloud-fan asked for "a two-sided Python UDF
parity test" at `RewriteNearestByJoin.scala:80`. A rewrite-vs-operator parity
assertion is exactly what a mock UDF can't give you: the whole risk in routing
these back through the rewrite is that the results still match, and that needs
a UDF that runs.
2. **The path the guard deliberately admits.** A single-child ranking
(`udf(right.y)`) stays on the binary `NearestByJoin`, and `ExtractPythonUDFs`
then rewrites the *right child* to `BatchEvalPython(udf, pythonUDF0, right)`
and repoints the ranking at `pythonUDF0`. So the broadcast side gains a column,
the operator's `output`/`resultProj` widen with it, and a trimming
`Project(plan.output, ...)` is stacked on top. I traced that and believe it
works, but it is a non-obvious interaction between UDF extraction and a
broadcast build side, and the second new test asserts only that the node was
*left* intact — it never plans or runs it.
`python/pyspark/sql/tests/test_nearest_by_join.py` already exists (12 tests,
no UDF ranking among them) and is the cheap home for both, e.g.:
```python
def test_broadcast_flag_with_python_udf_ranking(self):
dist = udf(lambda a, b: float(abs(a - b)), DoubleType())
for enabled in ["true", "false"]:
with self.sql_conf({"spark.sql.join.nearestBy.broadcast.enabled":
enabled}):
# two-sided: routed through the rewrite either way -- parity
...nearestByJoin(right, dist(left.x, right.y), numResults=2, ...)
# right-side only: reaches BroadcastNearestByJoinExec when
enabled
...nearestByJoin(right, dist(lit(5.0), right.y), numResults=2,
...)
```
with the two collected results compared across the flag. The connect parity
file next door picks it up for free.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
