peter-toth commented on code in PR #56101: URL: https://github.com/apache/spark/pull/56101#discussion_r3881698062
########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,264 @@ +/* + * 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. + * Mutable so that a fixed pool of entries can be reused across all left rows in a + * partition: the pool is sized to min(k, rightRows.length) and allocated once per + * partition, bounding total HeapEntry allocations to min(k, rightRows.length) per task + * regardless of the number of left rows. + */ +private[joins] class HeapEntry(var index: Int, var rankingValue: Any) { + HeapEntry.allocationCount += 1 Review Comment: **Finding 19.** `HeapEntry.allocationCount` is production state that exists only for a test. The pooling itself is right, and this is not a re-open of @cloud-fan's thread - it is about the hook added to prove it. Three separate problems: 1. **No production purpose, and no precedent for this shape.** I checked every other test-only member under `sql/core/.../execution/`: each one is a read-only *accessor* over state that already exists for a real reason - `CacheManager.numCachedEntries` over `cachedData`, `WindowSegmentTree.peekBlockBytes` / `peekBlockCount` over `blockBytes` / `numBlocks`. None introduces new mutable state, and none writes on an allocation path. 2. **The counter cannot fail the assertion it exists for.** `allocationCount += 1` on a `@volatile var` is a non-atomic read-modify-write, so concurrent tasks lose increments. Lost increments only shrink the count, so `assert(allocations <= 2)` can never fail *because of* the race - the race can only hide a real regression. 3. **The assertion is local-mode-only and does not say so.** `resetAllocationCount()` runs on the driver; the increments happen on executors. In any non-local deployment the driver's counter stays `0` and the test passes vacuously. The invariant is structural: `entryPool` is allocated outside `leftIter.flatMap`, so it cannot be per-left-row. I would delete the increment, the companion object, and `"SPARK-57091: HeapEntry allocations bounded by min(k, rightRows) per partition"`. ```scala private[joins] class HeapEntry(var index: Int, var rankingValue: Any) ``` ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,264 @@ +/* + * 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. + * Mutable so that a fixed pool of entries can be reused across all left rows in a + * partition: the pool is sized to min(k, rightRows.length) and allocated once per + * partition, bounding total HeapEntry allocations to min(k, rightRows.length) per task + * regardless of the number of left rows. + */ +private[joins] class HeapEntry(var index: Int, var rankingValue: Any) { + HeapEntry.allocationCount += 1 +} + +private[joins] object HeapEntry { + /** + * Test-only allocation counter. Incremented by the HeapEntry constructor. + * Used to verify the per-partition pooling invariant: total allocations should be + * bounded by min(k, rightRows.length) per partition, NOT by leftRows * k. + */ + @volatile var allocationCount: Long = 0L + + def resetAllocationCount(): Unit = { allocationCount = 0L } +} + +/** + * 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 unconditionally when + * `spark.sql.join.nearestBy.broadcast.enabled` is on -- except when the ranking + * expression contains a scalar Python UDF referencing both children, in which case + * [[org.apache.spark.sql.catalyst.optimizer.RewriteNearestByJoin]] falls back to the + * aggregate rewrite for correctness (ExtractPythonUDFs cannot split a two-sided UDF + * across a binary operator). + * For nodes that reach this operator, there is no size test and no fallback. + * A right side too large to broadcast will fail the query. Tie-breaking among equal + * ranking values is non-deterministic (matches the rewrite). + * + * Because no `Join` node is built on the operator path, `CheckCartesianProducts` does not + * apply and `spark.sql.crossJoin.enabled = false` does not reject NEAREST BY queries. + * This is intentional: the operator produces at most k rows per left row (bounded), not a + * true cross product. + */ +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 verboseStringWithOperatorId(): String = { + s""" + |$formattedNodeName + |${ExplainUtils.generateFieldString("Ranking", rankingExpression.sql)} + |${ExplainUtils.generateFieldString("NumResults", numResults.toString)} + |${ExplainUtils.generateFieldString("Direction", direction.toString)} + |${ExplainUtils.generateFieldString("Join type", joinType.toString)} + |""".stripMargin + } + + 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 Review Comment: **Finding 25.** Both declarations here are load-bearing, and neither has a test. I deleted lines 116-121 (falling back to `SparkPlan.outputPartitioning` = `UnknownPartitioning(0)` and `outputOrdering` = `Nil`) and ran `BroadcastNearestByJoinExecSuite`: **33/33 still pass**. The suite cannot tell the difference. They do real work. Two probes on this head with AQE off: - `left.repartition(5, $"id")` -> `nearestByJoin` -> `groupBy("id").count()`: **1** `ShuffleExchangeExec` with `outputPartitioning`, **2** without. - `left.repartition(3, $"id").sortWithinPartitions("id")` -> `nearestByJoin` -> `.sortWithinPartitions("id")`: **1** `SortExec` with `outputOrdering` (`RemoveRedundantSorts` drops the outer one), **2** without. `outputPartitioning` is the one that worries me. It is a promise `EnsureRequirements` acts on, so if it were ever wrong a downstream operator would silently skip a shuffle it needs - wrong results, not a lost optimization - and nothing in the suite would notice. `outputOrdering` was my finding 11, so the coverage gap is mine to have missed at the time. One test covers both. I ran this in two arms on this head: it passes as written, and fails on the first assertion once lines 116-121 are gone. Needs `org.apache.spark.sql.execution.SortExec` and `org.apache.spark.sql.execution.exchange.ShuffleExchangeExec` imported; `collect` already comes from `AdaptiveSparkPlanHelper`. ```scala test("SPARK-57091: output partitioning and ordering are preserved from the left side") { withSQLConf( SQLConf.NEAREST_BY_BROADCAST_ENABLED.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val left = spark.range(20).toDF("id").withColumn("x", col("id").cast("double")) .repartition(3, col("id")).sortWithinPartitions("id") val right = Seq((10, 1.0), (11, 2.0)).toDF("rid", "y") val df = left.nearestByJoin(right, abs(col("x") - col("y")), numResults = 2, mode = "exact", direction = "distance") // left.outputPartitioning survives: the group-by needs no extra shuffle. val grouped = df.groupBy("id").count() assert(collect(grouped.queryExecution.executedPlan) { case s: ShuffleExchangeExec => s }.size == 1) // left.outputOrdering survives: RemoveRedundantSorts drops the outer sort. val sorted = df.sortWithinPartitions("id") assert(collect(sorted.queryExecution.executedPlan) { case s: SortExec => s }.size == 1) } } ``` ########## 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 " + Review Comment: **Finding 21.** This is the fifth site with the scoping problem @cloud-fan raised, and the only one still unfixed. `CheckAnalysis:717`, `SparkStrategies:429`, the operator scaladoc and the rewrite parenthetical were all scoped to the operator path in `338726e2`. `SQLConf.scala` is byte-identical to `3e6e08a`, so two claims here are still false for a ranking whose scalar Python UDF references both children: - "When true, NearestByJoin uses a streaming heap operator instead of the cross-product + aggregate rewrite" - that ranking still uses the rewrite. - "The right side is always broadcast" - it is not; nothing is broadcast on that path. This is separate from my thread on line 2670. The closing `crossJoin.enabled` sentence reads correctly now that "this path" means the operator path. ```scala .doc("When true, NearestByJoin is planned as a streaming heap operator instead of " + "the cross-product + aggregate rewrite, except when the ranking expression " + "contains a scalar Python UDF referencing both children, which always takes the " + "rewrite. On the operator path 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 that path.") ``` ########## sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala: ########## @@ -953,6 +954,34 @@ class AnalysisErrorSuite extends AnalysisTest with DataTypeErrorsBase { } } + test("NearestByJoin with cross-child Python UDF is rejected when crossJoin disabled " + + "even with broadcast flag ON") { + // When the broadcast flag is ON but the ranking contains a cross-child scalar Python UDF, + // the rewrite path is forced (creating a Join node). CheckAnalysis must emit the dedicated + // NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED error rather than letting the generic + // CheckCartesianProducts "Detected implicit cartesian product" message surface. + // This test FAILS before the CheckAnalysis fix (would get the cartesian-product error) Review Comment: **Finding 23.** This comment names a failure the test cannot produce. `assertAnalysisErrorCondition` is `analyzer.checkAnalysis(analyzer.execute(inputPlan))` (`AnalysisTest.scala:196-198`) - analyzer only, no optimizer - so `CheckCartesianProducts` never runs from here. Without the `CheckAnalysis` change the test fails at `intercept[AnalysisException]` on *no exception at all*, not on the cartesian-product error. I hit exactly that when checking your fix: dropping the `hasCrossChildPythonUDF` disjunct makes the analyzer raise nothing, because `NON_ORDERABLE_RANKING_EXPRESSION` does not fire on an `IntegerType` UDF. ```scala // Without the CheckAnalysis fix the analyzer raises nothing at all here, and the query // would instead die later in the optimizer's CheckCartesianProducts with the generic // "Detected implicit cartesian product" message. ``` ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala: ########## @@ -71,8 +71,30 @@ import org.apache.spark.sql.catalyst.rules._ object RewriteNearestByJoin extends Rule[LogicalPlan] { private lazy val random = new scala.util.Random() + /** + * Returns true when the ranking expression of a [[NearestByJoin]] contains a scalar + * Python UDF whose attribute references span both the left and right children. + * + * Delegates to [[NearestByJoin.hasCrossChildPythonUDF]] -- a single source of truth shared + * with `CheckAnalysis` (which needs the same predicate to decide whether to waive the + * cross-join error). + */ + private def containsCrossChildPythonUDF(j: NearestByJoin): Boolean = Review Comment: **Finding 24.** Now that the predicate lives on the companion, this delegate and its duplicated scaladoc can go. Line 97 can call `NearestByJoin.hasCrossChildPythonUDF(j)` directly and still fit in 100 columns (82 chars): ```scala if !conf.nearestByBroadcastEnabled || NearestByJoin.hasCrossChildPythonUDF(j) => ``` ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNearestByJoinExec.scala: ########## @@ -0,0 +1,264 @@ +/* + * 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. + * Mutable so that a fixed pool of entries can be reused across all left rows in a + * partition: the pool is sized to min(k, rightRows.length) and allocated once per + * partition, bounding total HeapEntry allocations to min(k, rightRows.length) per task + * regardless of the number of left rows. + */ +private[joins] class HeapEntry(var index: Int, var rankingValue: Any) { + HeapEntry.allocationCount += 1 +} + +private[joins] object HeapEntry { + /** + * Test-only allocation counter. Incremented by the HeapEntry constructor. + * Used to verify the per-partition pooling invariant: total allocations should be + * bounded by min(k, rightRows.length) per partition, NOT by leftRows * k. + */ + @volatile var allocationCount: Long = 0L + + def resetAllocationCount(): Unit = { allocationCount = 0L } +} + +/** + * 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 unconditionally when + * `spark.sql.join.nearestBy.broadcast.enabled` is on -- except when the ranking + * expression contains a scalar Python UDF referencing both children, in which case + * [[org.apache.spark.sql.catalyst.optimizer.RewriteNearestByJoin]] falls back to the + * aggregate rewrite for correctness (ExtractPythonUDFs cannot split a two-sided UDF + * across a binary operator). + * For nodes that reach this operator, there is no size test and no fallback. + * A right side too large to broadcast will fail the query. Tie-breaking among equal + * ranking values is non-deterministic (matches the rewrite). + * + * Because no `Join` node is built on the operator path, `CheckCartesianProducts` does not + * apply and `spark.sql.crossJoin.enabled = false` does not reject NEAREST BY queries. + * This is intentional: the operator produces at most k rows per left row (bounded), not a + * true cross product. + */ +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 verboseStringWithOperatorId(): String = { + s""" + |$formattedNodeName + |${ExplainUtils.generateFieldString("Ranking", rankingExpression.sql)} + |${ExplainUtils.generateFieldString("NumResults", numResults.toString)} + |${ExplainUtils.generateFieldString("Direction", direction.toString)} + |${ExplainUtils.generateFieldString("Join type", joinType.toString)} + |""".stripMargin + } + + 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 + + // Matches BroadcastNestedLoopJoinExec: for BuildRight with InnerLike or LeftOuter, + // the streamed (left) side's ordering is preserved because we iterate left rows in + // order and emit each row's matches contiguously. + override def outputOrdering: Seq[SortOrder] = left.outputOrdering + + 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().mapPartitionsWithIndexInternal { (index, leftIter) => + val rightRows = broadcastedRight.value + if (rightRows.isEmpty && localJoinType != LeftOuter) { + Iterator.empty + } else { + val joinedRow = new JoinedRow + val rankingProj = UnsafeProjection.create( + Seq(rankExpr), leftOutput ++ rightOutput) + rankingProj.initialize(index) + val resultProj = UnsafeProjection.create(allOutput, allOutput) + val rankingNeedsCopy = !UnsafeRow.isFixedLength(rankExpr.dataType) + + // Pre-allocate the all-null right row for LEFT OUTER unmatched left rows. + // Hoisted here so it is built once per partition rather than once per left row. + val nullRight: InternalRow = new GenericInternalRow(rightOutput.size) + + // Pool of HeapEntry objects allocated ONCE per partition and reused across all + // left rows. The pool size is min(k, rightRows.length) -- the maximum number of + // entries that can simultaneously live in the heap. This bounds total HeapEntry + // allocations to min(k, rightRows.length) per partition (per task), regardless + // of the number of left rows processed. + val poolSize = math.min(k, rightRows.length) + val entryPool = new Array[HeapEntry](poolSize) + var pi = 0 + while (pi < poolSize) { + entryPool(pi) = new HeapEntry(0, null) + pi += 1 + } + + // Hoist heap outside flatMap -- created once per partition, cleared per left row. + val heapCapacity = poolSize + 1 + val heap = if (isDistance) { + new JPriorityQueue[HeapEntry](heapCapacity, + new Comparator[HeapEntry] { + override def compare(a: HeapEntry, b: HeapEntry): Int = + ordering.compare(b.rankingValue, a.rankingValue) + }) + } else { + new JPriorityQueue[HeapEntry](heapCapacity, + new Comparator[HeapEntry] { + override def compare(a: HeapEntry, b: HeapEntry): Int = + ordering.compare(a.rankingValue, b.rankingValue) + }) + } + + // Index into entryPool for the next unused pooled entry during heap filling. + var nextPoolIdx = 0 + + leftIter.flatMap { leftRow => + streamedRowsMetric += 1 + heap.clear() + nextPoolIdx = 0 + + var i = 0 + while (i < rightRows.length) { + val rightRow = rightRows(i) + joinedRow(leftRow, rightRow) + val rankingRow = rankingProj(joinedRow) + if (!rankingRow.isNullAt(0)) { + val rawValue = rankingRow.get(0, rankExpr.dataType) + // Only insert if the heap has room or the candidate beats the current worst. + // This avoids unnecessary .copy() allocations for rows that would be + // immediately evicted, and reduces PriorityQueue churn. + // For distance (isDistance=true): smaller is better, worst=largest on peek. + // For similarity: larger is better, worst=smallest on peek. + val shouldRetain = heap.size() < k || (if (isDistance) { + ordering.compare(rawValue, heap.peek().rankingValue) < 0 + } else { + ordering.compare(rawValue, heap.peek().rankingValue) > 0 + }) + if (shouldRetain) { + val rankingValue = if (rankingNeedsCopy) { + // Deep-copy variable-width values (UTF8String, BinaryView, structs, + // arrays, maps) without allocating the enclosing one-column UnsafeRow. + InternalRow.copyValue(rawValue) + } else { + rawValue + } + if (heap.size() < k) { + // Heap has room -- use the next pooled entry. + val entry = entryPool(nextPoolIdx) + nextPoolIdx += 1 + entry.index = i + entry.rankingValue = rankingValue + heap.offer(entry) + } else { + // Heap is at capacity -- reuse the evicted (worst) entry. + val evicted = heap.poll() + evicted.index = i + evicted.rankingValue = rankingValue + heap.offer(evicted) + } + } + } + i += 1 + } + + // CORRECTNESS: results must be fully materialized (output rows copied) BEFORE + // returning from this flatMap closure, because the pooled HeapEntry objects + // will be reused for the next left row. The resultProj(joinedRow).copy() below + // produces a standalone UnsafeRow that does not alias the pool or the heap. + // The Array[InternalRow] is eagerly built and its iterator consumed lazily, + // but since flatMap fully drains each inner iterator before advancing leftIter, + // the pool is not touched until all output rows for this left row are consumed. + if (heap.isEmpty && localJoinType == LeftOuter) { + joinedRow(leftRow, nullRight) + numOutput += 1 + Iterator.single(resultProj(joinedRow).copy()) + } else { + val results = new Array[InternalRow](heap.size()) Review Comment: **Finding 20.** The pool made every output row a copy; draining indices instead avoids that. The `CORRECTNESS` comment above argues both sides. Its first sentence says rows *must* be copied before the closure returns; its last sentence explains why they need not be, and the last sentence is the correct one - `Iterator.flatMap` only calls `self.next()` once the current inner iterator reports `hasNext == false`. Nothing downstream needs a `HeapEntry`, only `entry.index`. Draining into an `Array[Int]` releases the pool before the iterator is returned, so the rows can stay lazy over the reused projection buffer. That is the normal `doExecute` contract - `ProjectExec` is `iter.map(project)`, and `BroadcastNestedLoopJoinExec.innerJoin` is `buildRows.iterator.map(r => joinedRow(streamedRow, r))`. It saves one `UnsafeRow` allocation plus memcpy per *output* row, up to `leftRows * k`, and drops the live footprint per left row from k rows to one. The LEFT OUTER branch loses its `.copy()` too. I applied exactly this on `ea8f0c7` and ran `BroadcastNearestByJoinExecSuite` + `DataFrameNearestByJoinSuite`: 54/54 green. ```scala // Drain the heap into a plain Array[Int] of broadcast indices. That releases the // pooled HeapEntry objects before the inner iterator is returned, so the output // rows can stay lazy over the projection's reused buffer. if (heap.isEmpty && localJoinType == LeftOuter) { joinedRow(leftRow, nullRight) numOutput += 1 Iterator.single(resultProj(joinedRow)) } else { val indices = new Array[Int](heap.size()) var idx = heap.size() - 1 while (!heap.isEmpty) { indices(idx) = heap.poll().index idx -= 1 } indices.iterator.map { i => joinedRow(leftRow, rightRows(i)) numOutput += 1 resultProj(joinedRow) } } ``` The per-row `.copy()` predates this commit, so that half is a late catch on my side. The eager `Array[InternalRow]` and the comment defending it are new here. ########## python/pyspark/sql/tests/test_nearest_by_join.py: ########## @@ -252,6 +252,80 @@ def test_streaming_inputs_rejected(self): messageParameters={}, ) + def test_python_udf_two_sided_parity(self): + """A two-sided scalar Python UDF ranking (referencing both left and right columns) + is routed through the rewrite path regardless of the broadcast flag. Results must + be identical across flag ON and flag OFF.""" + from pyspark.sql.functions import udf + from pyspark.sql.types import DoubleType + + dist_udf = udf(lambda a, b: float(abs(a - b)), DoubleType()) + + users = self.spark.createDataFrame([(1, 10.0), (2, 20.0), (3, 30.0)], ["user_id", "score"]) + products = self.spark.createDataFrame( + [("A", 11.0), ("B", 22.0), ("C", 5.0)], ["product", "pscore"] + ) + + def run_with_flag(flag_value): + with self.sql_conf( + {"spark.sql.join.nearestBy.broadcast.enabled": str(flag_value).lower()} + ): + return ( + users.nearestByJoin( + products, + dist_udf(users.score, products.pscore), + numResults=2, + mode="approx", + direction="distance", + ) + .select("user_id", "product") + .orderBy("user_id", "product") + .collect() + ) + + result_off = run_with_flag(False) + result_on = run_with_flag(True) + self.assertEqual(result_off, result_on) + # Sanity: each user gets 2 products + self.assertEqual(len(result_off), 6) + + def test_python_udf_right_side_only_parity(self): + """A right-side-only scalar Python UDF ranking reaches BroadcastNearestByJoinExec Review Comment: **Finding 22.** Both new tests assert only flag-on == flag-off, which a fallback to the rewrite also satisfies. This is the coverage I asked for in [my earlier thread](https://github.com/apache/spark/pull/56101#discussion_r3763410696), and getting both paths to actually execute is the important half. But the docstring here claims the right-side-only ranking "reaches BroadcastNearestByJoinExec", and neither test can observe that. `crossJoin.enabled=false` discriminates: on the operator path no `Join` is built and the query succeeds, while the rewrite fallback builds one and `CheckAnalysis` raises `NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED`. The two-sided arm is already pinned in Scala by the new `AnalysisErrorSuite` test, so what is missing here is the right-side-only arm succeeding under the same conf. I measured this in two arms on `ea8f0c7`, with the assembly built and `python/run-tests`: - **As written:** the test below passes. - **With `hasCrossChildPythonUDF`'s `&&` widened to `||`** (so a right-side-only ranking wrongly falls back to the rewrite): the test below fails with `[NEAREST_BY_JOIN.CROSS_JOIN_NOT_ENABLED]`, while `test_python_udf_right_side_only_parity` and `test_python_udf_two_sided_parity` both still pass. ```python def test_python_udf_right_side_only_reaches_operator(self): """With the broadcast flag ON a right-side-only scalar Python UDF ranking takes the operator path, which builds no Join, so crossJoin.enabled=False does not reject it.""" from pyspark.sql.functions import udf, lit from pyspark.sql.types import DoubleType dist_udf = udf(lambda a, b: float(abs(a - b)), DoubleType()) users = self.spark.createDataFrame([(1, 10.0)], ["user_id", "score"]) products = self.spark.createDataFrame([("A", 11.0), ("B", 22.0)], ["product", "pscore"]) with self.sql_conf( { "spark.sql.join.nearestBy.broadcast.enabled": "true", "spark.sql.crossJoin.enabled": "false", } ): self.assertEqual( users.nearestByJoin( products, dist_udf(lit(5.0), products.pscore), numResults=2, mode="approx", direction="distance", ).count(), 2, ) ``` A plan assertion is not an option here because `connect/test_parity_nearest_by_join.py` inherits this mixin. -- 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]
