cloud-fan commented on code in PR #56101: URL: https://github.com/apache/spark/pull/56101#discussion_r3333254172
########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,149 @@ +/* + * 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.util.{PriorityQueue => JPriorityQueue} + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.{InnerLike, JoinType, LeftOuter, NearestByDirection, NearestByDistance} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.types.DoubleType + +/** + * Physical operator for NearestByJoin that avoids materializing the full cross product. + * For each left row, iterates all broadcast right rows maintaining a bounded priority + * queue of size k, then emits the top-k matches directly. + * + * The right side is fully broadcast to all partitions. This operator only fires when + * the right side fits within [[SQLConf.AUTO_BROADCAST_JOIN_THRESHOLD]]. For right tables + * exceeding this threshold, the existing cross-product + aggregate rewrite is used as + * fallback. Tie-breaking among equal ranking values is non-deterministic (matches the + * existing rewrite behavior). + */ +case class BroadcastNearestByJoinExec( + left: SparkPlan, + right: SparkPlan, + joinType: JoinType, + numResults: Int, + rankingExpression: Expression, + direction: NearestByDirection) extends BinaryExecNode { + + override def output: Seq[Attribute] = joinType match { + case LeftOuter => + left.output ++ right.output.map(_.withNullability(true)) + case _: InnerLike => + left.output ++ right.output + case other => + throw SparkException.internalError( + s"$nodeName does not support join type: $other") + } + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), + "streamedRows" -> SQLMetrics.createMetric(sparkContext, "number of left rows processed")) + + override def requiredChildDistribution: Seq[Distribution] = + UnspecifiedDistribution :: BroadcastDistribution(IdentityBroadcastMode) :: Nil + + override def outputPartitioning: Partitioning = left.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = Nil + + protected override def doExecute(): RDD[InternalRow] = { + val broadcastedRight = right.executeBroadcast[Array[InternalRow]]() + val numOutput = longMetric("numOutputRows") + val k = numResults + val isDistance = direction == NearestByDistance + val leftOutput = left.output + val rightOutput = right.output + val rankExpr = rankingExpression + val allOutput = output + + left.execute().mapPartitionsInternal { leftIter => + val rightRows = broadcastedRight.value + if (rightRows.isEmpty && joinType != LeftOuter) { + Iterator.empty + } else { + val joinedRow = new JoinedRow + val rankingProj = UnsafeProjection.create( Review Comment: **[blocking]** The ranking is computed as `Cast(rankExpr, DoubleType)` and compared as doubles, but `CheckAnalysis` admits any *orderable* ranking type (`RowOrdering.isOrderable`), and the rewrite path ranks via `MaxMinByK` using the type's natural ordering. So the operator diverges from the rewrite for: - **non-numeric orderable types** (StringType, DateType, ...) — `Cast` to Double yields `null`, so every row is excluded: INNER returns empty, LEFT OUTER returns all-null right side; - **high-magnitude Long / high-precision Decimal** — the double cast loses precision past ~2^53, causing mis-ranking / spurious ties; - **NaN under `direction=similarity`** — this code drops NaN (line 116), but the rewrite's `MAX_BY` ranks NaN as the *largest* value, selecting NaN rows as the top matches. This silently produces different results with the flag on, contradicting the PR's "no user-facing change" / "matches the existing rewrite behavior". Either rank by the ranking expression's natural ordering (as the rewrite does), or — minimally — have `NearestByJoinSelection` fall back to the rewrite when `rankingExpression.dataType` is non-numeric, and decide the Long/Decimal/NaN policy explicitly. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/NearestByJoinBenchmark.scala: ########## @@ -0,0 +1,203 @@ +/* + * 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 org.apache.spark.sql.QueryTest +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +class NearestByJoinBenchmark extends QueryTest with SharedSparkSession { Review Comment: **[blocking]** This extends `QueryTest with SharedSparkSession` with `test(...)` cases, so the ScalaTest runner executes it in normal CI — including a 200K x 200K join (~40e9 inner iterations, ~62 min by the PR's own table) and a 30K x 30K memory test that drives a 900M-row cross product on the fallback path. Spark benchmarks are `object X extends SqlBasedBenchmark` with `runBenchmarkSuite`, live under `.../execution/benchmark/`, and are never discovered as ScalaTest cases (cf. `JoinBenchmark`, `AnsiIntervalSortBenchmark`). As written this will severely slow / time out / OOM CI. Please convert to a `SqlBasedBenchmark`, or drop the benchmark from this PR. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,149 @@ +/* + * 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.util.{PriorityQueue => JPriorityQueue} + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.{InnerLike, JoinType, LeftOuter, NearestByDirection, NearestByDistance} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.types.DoubleType + +/** + * Physical operator for NearestByJoin that avoids materializing the full cross product. + * For each left row, iterates all broadcast right rows maintaining a bounded priority + * queue of size k, then emits the top-k matches directly. + * + * The right side is fully broadcast to all partitions. This operator only fires when + * the right side fits within [[SQLConf.AUTO_BROADCAST_JOIN_THRESHOLD]]. For right tables + * exceeding this threshold, the existing cross-product + aggregate rewrite is used as + * fallback. Tie-breaking among equal ranking values is non-deterministic (matches the + * existing rewrite behavior). + */ +case class BroadcastNearestByJoinExec( + left: SparkPlan, + right: SparkPlan, + joinType: JoinType, + numResults: Int, + rankingExpression: Expression, + direction: NearestByDirection) extends BinaryExecNode { + + override def output: Seq[Attribute] = joinType match { + case LeftOuter => + left.output ++ right.output.map(_.withNullability(true)) + case _: InnerLike => Review Comment: **[blocking]** The logical `NearestByJoin.output` deliberately widens *both* sides to nullable (with a comment explaining it must match the rewrite's analyzed schema and "what users see in cached or written outputs"). This operator uses the generic join convention instead — the same one as `BroadcastNestedLoopJoinExec.output` (Inner: neither side widened; LeftOuter: only right widened). That convention is correct for a generic join but wrong here. Two consequences: - the two execution paths declare different schemas for the same logical node; - for INNER, a right data column that actually contains nulls is declared non-nullable — a parent codegen operator bound to this output may elide null checks (NPE / wrong results). The operator's `output` should match the logical node's nullability contract (widen both sides to nullable). This is the concrete symptom of re-deriving BNLJ's `output` by hand (see the class-level comment below). ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala: ########## @@ -72,7 +73,15 @@ object RewriteNearestByJoin extends Rule[LogicalPlan] { private lazy val random = new scala.util.Random() def apply(plan: LogicalPlan): LogicalPlan = plan.transformUp { - case j @ NearestByJoin(left, right, joinType, _, numResults, rankingExpression, direction) => + case j @ NearestByJoin(left, right, joinType, _, numResults, rankingExpression, direction) + // The optimizer checks both the config flag AND the broadcast threshold to decide + // whether to skip the rewrite. This mixes optimizer/planner concerns but is necessary: + // if we only checked the config flag, a NearestByJoin node with right side exceeding + // the broadcast threshold would reach the planner unrewritten, and no strategy would + // handle it (NearestByJoinSelection returns Nil for large right sides), causing a + // planning failure. The alternative (two-pass approach) is deferred to future work. + if !(SQLConf.get.nearestByBroadcastEnabled && Review Comment: **[non-blocking]** This `nearestByBroadcastEnabled && right.stats.sizeInBytes <= autoBroadcastJoinThreshold` predicate is duplicated verbatim in `NearestByJoinSelection` (SparkStrategies). They are two independent sources of truth for the same decision: if they ever drift, a `NearestByJoin` with a large right reaches the planner unrewritten and the strategy returns `Nil` -> planning failure. The layering tension you note (optimizer consulting a physical broadcast-size decision) is acknowledged and the two-pass refactor reasonably deferred — but the *duplication* is independently fixable now: extract one shared predicate (e.g. a method on `NearestByJoin`, or a shared helper) that both the rule guard and the strategy call. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,149 @@ +/* + * 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.util.{PriorityQueue => JPriorityQueue} + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.{InnerLike, JoinType, LeftOuter, NearestByDirection, NearestByDistance} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.types.DoubleType + +/** + * Physical operator for NearestByJoin that avoids materializing the full cross product. + * For each left row, iterates all broadcast right rows maintaining a bounded priority + * queue of size k, then emits the top-k matches directly. + * + * The right side is fully broadcast to all partitions. This operator only fires when + * the right side fits within [[SQLConf.AUTO_BROADCAST_JOIN_THRESHOLD]]. For right tables + * exceeding this threshold, the existing cross-product + aggregate rewrite is used as + * fallback. Tie-breaking among equal ranking values is non-deterministic (matches the + * existing rewrite behavior). + */ +case class BroadcastNearestByJoinExec( + left: SparkPlan, + right: SparkPlan, + joinType: JoinType, + numResults: Int, + rankingExpression: Expression, + direction: NearestByDirection) extends BinaryExecNode { + + override def output: Seq[Attribute] = joinType match { + case LeftOuter => + left.output ++ right.output.map(_.withNullability(true)) + case _: InnerLike => + left.output ++ right.output + case other => + throw SparkException.internalError( + s"$nodeName does not support join type: $other") + } + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), + "streamedRows" -> SQLMetrics.createMetric(sparkContext, "number of left rows processed")) + + override def requiredChildDistribution: Seq[Distribution] = + UnspecifiedDistribution :: BroadcastDistribution(IdentityBroadcastMode) :: Nil + + override def outputPartitioning: Partitioning = left.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = Nil + + protected override def doExecute(): RDD[InternalRow] = { + val broadcastedRight = right.executeBroadcast[Array[InternalRow]]() + val numOutput = longMetric("numOutputRows") + val k = numResults + val isDistance = direction == NearestByDistance + val leftOutput = left.output + val rightOutput = right.output + val rankExpr = rankingExpression + val allOutput = output + + left.execute().mapPartitionsInternal { leftIter => + val rightRows = broadcastedRight.value + if (rightRows.isEmpty && joinType != LeftOuter) { + Iterator.empty + } else { + val joinedRow = new JoinedRow + val rankingProj = UnsafeProjection.create( + Seq(Cast(rankExpr, DoubleType)), leftOutput ++ rightOutput) + val resultProj = UnsafeProjection.create(allOutput, allOutput) + val streamedRowsMetric = longMetric("streamedRows") + + // Hoist heap outside flatMap to reduce GC pressure Review Comment: **[suggestion]** The heap stores boxed `(Int, Double)` tuples, so each `heap.offer((i, rankingValue))` allocates a `Tuple2` and boxes both fields — O(N*M) allocations per partition, which works against the GC-pressure win this operator is built to deliver. Consider a primitive-keyed bounded structure (parallel `int`/`double` arrays, or a small primitive heap) to avoid the boxing on the hot path. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/NearestByJoinBenchmark.scala: ########## @@ -0,0 +1,203 @@ +/* + * 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 org.apache.spark.sql.QueryTest +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +class NearestByJoinBenchmark extends QueryTest with SharedSparkSession { + + private val leftSize = 30000 + private val rightSize = 30000 + private val k = 5 + + override def sparkConf = super.sparkConf + .set("spark.driver.memory", "4g") + .set("spark.executor.memory", "4g") + + test("correctness: streaming heap produces same results as rewrite") { + // Use data with no ties to avoid tie-breaking differences + val left = spark.range(0, 20).toDF("id") + .withColumn("x", col("id").cast("double") * 7.3) + val right = spark.range(0, 15).toDF("rid") + .withColumn("y", col("rid").cast("double") * 11.1 + 0.5) + + // Get results with current rewrite + val rewriteResult = withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "false", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin( + right, + abs(col("x") - col("y")), + numResults = 3, + mode = "exact", + direction = "distance") + .orderBy("id", "rid") + .collect() + } + + // Get results with streaming heap + val heapResult = withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin( + right, + abs(col("x") - col("y")), + numResults = 3, + mode = "exact", + direction = "distance") + .orderBy("id", "rid") + .collect() + } + + assert(rewriteResult.length == heapResult.length, + s"Row count mismatch: rewrite=${rewriteResult.length}, heap=${heapResult.length}") + rewriteResult.zip(heapResult).zipWithIndex.foreach { case ((r, h), i) => + assert(r == h, s"Row $i mismatch: rewrite=$r, heap=$h") + } + // scalastyle:off println + println(s"Correctness check passed: ${rewriteResult.length} rows match") + // scalastyle:on println + } + + private def measureMemoryMB(block: => Long): (Long, Double, Double) = { + val runtime = Runtime.getRuntime + runtime.gc() + Thread.sleep(100) + val before = runtime.totalMemory() - runtime.freeMemory() + val result = block + val after = runtime.totalMemory() - runtime.freeMemory() + val deltaMB = Math.max(0, after - before) / (1024.0 * 1024.0) + val peakMB = after / (1024.0 * 1024.0) + (result, deltaMB, peakMB) + } + + test("benchmark: streaming heap at 200K x 200K") { + val size = 200000 + val left = spark.range(0, size).toDF("id").withColumn("x", rand(42) * 1000.0) + val right = spark.range(0, size).toDF("rid").withColumn("y", rand(43) * 1000.0) + left.cache().count() + right.cache().count() + + val start = System.nanoTime() + val count = withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin(right, abs(col("x") - col("y")), + numResults = k, mode = "exact", direction = "distance").count() + } + val secs = (System.nanoTime() - start) / 1e9 + // scalastyle:off println + println("=" * 70) + println(s"200K x 200K streaming heap: ${secs}s, rows=$count") + println("=" * 70) + // scalastyle:on println + assert(count == size.toLong * k) + } + + test("memory benchmark: streaming heap vs cross-product at 30K x 30K") { + val left = spark.range(0, leftSize).toDF("id") + .withColumn("x", rand(42) * 1000.0) + val right = spark.range(0, rightSize).toDF("rid") + .withColumn("y", rand(43) * 1000.0) + + left.cache().count() + right.cache().count() + + // Streaming heap: should complete with minimal memory (M rows + k heap) + val heapStart = System.nanoTime() + val (heapCount, heapDeltaMB, heapPeakMB) = measureMemoryMB { + withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin(right, abs(col("x") - col("y")), + numResults = k, mode = "exact", direction = "distance").count() + } + } + val heapMs = (System.nanoTime() - heapStart) / 1e6 + + // Cross-product: tries to materialize N*M = 900M rows → likely OOM at 1GB + var rewriteMs = -1.0 + var rewriteCount = -1L + var rewriteDeltaMB = 0.0 + var rewritePeakMB = 0.0 + var rewriteOOM = false + try { + val rewriteStart = System.nanoTime() + val (count, delta, peak) = measureMemoryMB { + withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "false", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin(right, abs(col("x") - col("y")), + numResults = k, mode = "exact", direction = "distance").count() + } + } + rewriteMs = (System.nanoTime() - rewriteStart) / 1e6 + rewriteCount = count + rewriteDeltaMB = delta + rewritePeakMB = peak + } catch { + case _: OutOfMemoryError => + rewriteOOM = true + case e: Exception if e.getCause != null && + e.getCause.isInstanceOf[OutOfMemoryError] => + rewriteOOM = true + case e: org.apache.spark.SparkException + if e.getCause != null && e.getCause.isInstanceOf[OutOfMemoryError] => + rewriteOOM = true + } + + // scalastyle:off println + println("=" * 70) + println("MEMORY BENCHMARK: NearestByJoin") + println(s" Config: left=$leftSize, right=$rightSize, k=$k") + println(s" Memory constraint: spark.driver.memory=1g, spark.executor.memory=1g") Review Comment: **[nit]** This prints `spark.driver.memory=1g, spark.executor.memory=1g`, but `sparkConf` (lines 32-33) sets both to `4g` — the benchmark output misreports its own configuration. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,149 @@ +/* + * 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.util.{PriorityQueue => JPriorityQueue} + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.{InnerLike, JoinType, LeftOuter, NearestByDirection, NearestByDistance} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.types.DoubleType + +/** + * Physical operator for NearestByJoin that avoids materializing the full cross product. + * For each left row, iterates all broadcast right rows maintaining a bounded priority + * queue of size k, then emits the top-k matches directly. + * + * The right side is fully broadcast to all partitions. This operator only fires when + * the right side fits within [[SQLConf.AUTO_BROADCAST_JOIN_THRESHOLD]]. For right tables + * exceeding this threshold, the existing cross-product + aggregate rewrite is used as + * fallback. Tie-breaking among equal ranking values is non-deterministic (matches the + * existing rewrite behavior). + */ +case class BroadcastNearestByJoinExec( + left: SparkPlan, + right: SparkPlan, + joinType: JoinType, + numResults: Int, + rankingExpression: Expression, + direction: NearestByDirection) extends BinaryExecNode { Review Comment: **[non-blocking, design question]** This extends `BinaryExecNode` directly and rebuilds the broadcast-iterate plumbing from scratch — `requiredChildDistribution`, `executeBroadcast[Array[InternalRow]]`, the `JoinedRow` loop, and `output` — all essentially identical to `BroadcastNestedLoopJoinExec` (BuildRight). To be clear, a *separate* operator looks like the right call and I wouldn't fold this into BNLJ: top-k isn't expressible as BNLJ's boolean `condition`, so BNLJ can't recognize/optimize the pattern; and top-k only fits `Inner`/`LeftOuter` (2 of BNLJ's 7 join-type/buildSide branches), so a shared class would need a `ranking`/`numResults`/`direction` mode flag that's meaningless elsewhere. That matches the precedent of giving `SortMergeAsOfJoinExec` its own operator. The narrower ask: reuse the shared join base (`BaseJoinExec` / the join boilerplate BNLJ already builds on) for the plumbing, and own *only* the genuinely new part — the k-heap + ranking. The hand-reimplemented scaffolding is exactly where the ranking divergence and the `output` nullability bug entered (the latter copied from BNLJ's convention); sharing the base would prevent the nullability one by construction. Could you say why a from-scratch `BinaryExecNode` over a shared join base? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,149 @@ +/* + * 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.util.{PriorityQueue => JPriorityQueue} + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.{InnerLike, JoinType, LeftOuter, NearestByDirection, NearestByDistance} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.sql.types.DoubleType + +/** + * Physical operator for NearestByJoin that avoids materializing the full cross product. + * For each left row, iterates all broadcast right rows maintaining a bounded priority + * queue of size k, then emits the top-k matches directly. + * + * The right side is fully broadcast to all partitions. This operator only fires when + * the right side fits within [[SQLConf.AUTO_BROADCAST_JOIN_THRESHOLD]]. For right tables Review Comment: **[nit]** Broken Scaladoc link: the config val is `AUTO_BROADCASTJOIN_THRESHOLD` (no underscore between `BROADCAST` and `JOIN`); `AUTO_BROADCAST_JOIN_THRESHOLD` doesn't exist. Note `SQLConf` isn't imported in this file either, so the link won't resolve as Scaladoc even with the right name — add the import, or reference `SQLConf.autoBroadcastJoinThreshold`. ```suggestion * the right side fits within [[SQLConf.AUTO_BROADCASTJOIN_THRESHOLD]]. For right tables ``` ########## sql/core/src/test/scala/org/apache/spark/sql/execution/joins/NearestByJoinBenchmark.scala: ########## @@ -0,0 +1,203 @@ +/* + * 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 org.apache.spark.sql.QueryTest +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +class NearestByJoinBenchmark extends QueryTest with SharedSparkSession { + + private val leftSize = 30000 + private val rightSize = 30000 + private val k = 5 + + override def sparkConf = super.sparkConf + .set("spark.driver.memory", "4g") + .set("spark.executor.memory", "4g") + + test("correctness: streaming heap produces same results as rewrite") { + // Use data with no ties to avoid tie-breaking differences + val left = spark.range(0, 20).toDF("id") + .withColumn("x", col("id").cast("double") * 7.3) + val right = spark.range(0, 15).toDF("rid") + .withColumn("y", col("rid").cast("double") * 11.1 + 0.5) + + // Get results with current rewrite + val rewriteResult = withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "false", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin( + right, + abs(col("x") - col("y")), + numResults = 3, + mode = "exact", + direction = "distance") + .orderBy("id", "rid") + .collect() + } + + // Get results with streaming heap + val heapResult = withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin( + right, + abs(col("x") - col("y")), + numResults = 3, + mode = "exact", + direction = "distance") + .orderBy("id", "rid") + .collect() + } + + assert(rewriteResult.length == heapResult.length, + s"Row count mismatch: rewrite=${rewriteResult.length}, heap=${heapResult.length}") + rewriteResult.zip(heapResult).zipWithIndex.foreach { case ((r, h), i) => + assert(r == h, s"Row $i mismatch: rewrite=$r, heap=$h") + } + // scalastyle:off println + println(s"Correctness check passed: ${rewriteResult.length} rows match") + // scalastyle:on println + } + + private def measureMemoryMB(block: => Long): (Long, Double, Double) = { + val runtime = Runtime.getRuntime + runtime.gc() + Thread.sleep(100) + val before = runtime.totalMemory() - runtime.freeMemory() + val result = block + val after = runtime.totalMemory() - runtime.freeMemory() + val deltaMB = Math.max(0, after - before) / (1024.0 * 1024.0) + val peakMB = after / (1024.0 * 1024.0) + (result, deltaMB, peakMB) + } + + test("benchmark: streaming heap at 200K x 200K") { + val size = 200000 + val left = spark.range(0, size).toDF("id").withColumn("x", rand(42) * 1000.0) + val right = spark.range(0, size).toDF("rid").withColumn("y", rand(43) * 1000.0) + left.cache().count() + right.cache().count() + + val start = System.nanoTime() + val count = withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin(right, abs(col("x") - col("y")), + numResults = k, mode = "exact", direction = "distance").count() + } + val secs = (System.nanoTime() - start) / 1e9 + // scalastyle:off println + println("=" * 70) + println(s"200K x 200K streaming heap: ${secs}s, rows=$count") + println("=" * 70) + // scalastyle:on println + assert(count == size.toLong * k) + } + + test("memory benchmark: streaming heap vs cross-product at 30K x 30K") { + val left = spark.range(0, leftSize).toDF("id") + .withColumn("x", rand(42) * 1000.0) + val right = spark.range(0, rightSize).toDF("rid") + .withColumn("y", rand(43) * 1000.0) + + left.cache().count() + right.cache().count() + + // Streaming heap: should complete with minimal memory (M rows + k heap) + val heapStart = System.nanoTime() + val (heapCount, heapDeltaMB, heapPeakMB) = measureMemoryMB { + withSQLConf( + SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", + SQLConf.CROSS_JOINS_ENABLED.key -> "true") { + left.nearestByJoin(right, abs(col("x") - col("y")), + numResults = k, mode = "exact", direction = "distance").count() + } + } + val heapMs = (System.nanoTime() - heapStart) / 1e6 + + // Cross-product: tries to materialize N*M = 900M rows → likely OOM at 1GB Review Comment: **[nit]** The `→` arrow is non-ASCII and trips scalastyle's `nonascii` rule, which fails the build. Also "OOM at 1GB" is stale — `sparkConf` sets driver/executor memory to `4g` (lines 32-33), and the println at line 171 likewise still claims `1g`. Suggested fix for the arrow (and dropping the wrong size): ```suggestion // Cross-product: tries to materialize N*M = 900M rows -> likely OOM ``` -- 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]
