peter-toth commented on code in PR #56101:
URL: https://github.com/apache/spark/pull/56101#discussion_r3652531387
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala:
##########
@@ -72,7 +73,14 @@ 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 !NearestByJoin.canBroadcastRight(j, SQLConf.get) =>
Review Comment:
**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.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala:
##########
@@ -0,0 +1,172 @@
+/*
+ * 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.{Comparator, 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.catalyst.util.TypeUtils
+import org.apache.spark.sql.execution.{ExplainUtils, SparkPlan}
+import org.apache.spark.sql.execution.metric.SQLMetrics
+
+/**
+ * Heap entry storing an index into the broadcast array alongside its ranking
value.
+ * Using a case class with primitive `Int` field avoids boxing that `(Int,
Any)` tuples incur.
+ */
+private[joins] case class HeapEntry(index: Int, rankingValue: Any)
+
+/**
+ * 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_BROADCASTJOIN_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 BaseJoinExec {
+
+ override def condition: Option[Expression] = None
+ override def leftKeys: Seq[Expression] = Seq.empty
+ override def rightKeys: Seq[Expression] = Seq.empty
+
+ override def simpleStringWithNodeId(): String = {
+ val opId = ExplainUtils.getOpId(this)
+ s"$nodeName $joinType k=$numResults $direction ($opId)".trim
+ }
+
+ override def output: Seq[Attribute] = joinType match {
+ case _: InnerLike | LeftOuter =>
+ left.output.map(_.withNullability(true)) ++
right.output.map(_.withNullability(true))
+ 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 streamedRowsMetric = longMetric("streamedRows")
+ val localJoinType = joinType
+ val k = numResults
+ val isDistance = direction == NearestByDistance
+ val leftOutput = left.output
+ val rightOutput = right.output
+ val rankExpr = rankingExpression
+ val allOutput = output
+ val ordering = TypeUtils.getInterpretedOrdering(rankExpr.dataType)
+
+ left.execute().mapPartitionsInternal { leftIter =>
+ val rightRows = broadcastedRight.value
+ if (rightRows.isEmpty && localJoinType != LeftOuter) {
+ Iterator.empty
+ } else {
+ val joinedRow = new JoinedRow
+ val rankingProj = UnsafeProjection.create(
Review Comment:
**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):
```scala
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.
--
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]
