From 3adfb1eccbbd1aebd4185f634296b74777d09541 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 08:42:10 +0100 Subject: [PATCH] fix: keep a columnar-to-row over a Comet broadcast out of whole-stage codegen A Spark BroadcastHashJoinExec can read a Comet broadcast through CometColumnarToRowExec, whose doExecuteBroadcast turns the Arrow batches into the join's relation. This happens under AQE when a re-optimization leaves the join in Spark after its build side was already planned as a Comet broadcast stage, for example when OptimizeSkewedJoin puts an AQEShuffleRead under the sort-merge join below it. When the join itself is outside whole-stage codegen (more output fields than spark.sql.codegen.maxFields), CollapseCodegenStages wrapped the CometColumnarToRowExec into a WholeStageCodegenExec of its own, and the join failed with "WholeStageCodegen (n) does not implement doExecuteBroadcast". CometColumnarToRowExec over a Comet broadcast exchange, a broadcast query stage or a reused exchange of one now reports supportCodegen = false, so it stays the join's build plan, or sits behind an InputAdapter that forwards executeBroadcast when the join is codegen'd. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../sql/comet/CometColumnarToRowExec.scala | 9 +++ .../apache/comet/exec/CometJoinSuite.scala | 57 ++++++++++++++++++- 2 files changed, 63 insertions(+), 3 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala index 98c1105edc1..de3ecb31206 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowExec.scala @@ -65,6 +65,15 @@ case class CometColumnarToRowExec(child: SparkPlan) override def outputOrdering: Seq[SortOrder] = child.outputOrdering + override def supportCodegen: Boolean = !readsBroadcast(child) + + private def readsBroadcast(plan: SparkPlan): Boolean = plan match { + case _: CometBroadcastExchangeExec => true + case stage: BroadcastQueryStageExec => readsBroadcast(stage.plan) + case reused: ReusedExchangeExec => readsBroadcast(reused.child) + case _ => false + } + // `ColumnarToRowExec` processes the input RDD directly, which is kind of a leaf node in the // codegen stage and needs to do the limit check. protected override def canCheckLimitNotReached: Boolean = true diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 8bb8d5ddf70..392ad516a00 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,11 +31,12 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} -import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometColumnarToRowExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} +import org.apache.spark.sql.execution.{InputAdapter, LocalTableScanExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} @@ -1479,6 +1480,56 @@ class CometJoinSuite extends CometTestBase { } } + for (costBased <- Seq("true", "false"); maxFields <- Seq("3", "100")) { + test( + "a Spark broadcast hash join over a skewed join reads a Comet broadcast stage " + + s"(costBasedEngines=$costBased, codegen.maxFields=$maxFields)") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "4", + SQLConf.SKEW_JOIN_ENABLED.key -> "true", + SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_FACTOR.key -> "2", + SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100", + SQLConf.WHOLESTAGE_MAX_NUM_FIELDS.key -> maxFields, + CometConf.COMET_EXEC_COST_BASED_ENGINES_ENABLED.key -> costBased) { + withParquetTable( + (0 until 4000).map(i => (if (i % 10 < 8) 0 else i % 50, i % 10, i)), + "skew_f") { + withParquetTable((0 until 500).map(i => (i % 50, i)), "skew_g") { + withParquetTable((0 until 10).map(i => (i, s"name_$i")), "skew_d") { + val query = + """SELECT /*+ BROADCAST(d), MERGE(f, g) */ f._1, f._2, f._3, g._2, d._2 + |FROM skew_f f JOIN skew_g g ON f._1 = g._1 + |LEFT JOIN skew_d d ON f._2 = d._1 + |DISTRIBUTE BY f._3""".stripMargin + val (_, plan) = checkSparkAnswer(sql(query)) + val skewed = collect(plan) { case j: SortMergeJoinExec if j.isSkewJoin => j } + assert(skewed.nonEmpty, s"expected a skewed Spark sort-merge join:\n$plan") + val sparkBroadcastJoins = collect(plan) { case j: BroadcastHashJoinExec => j } + assert(sparkBroadcastJoins.size == 1, s"plan:\n$plan") + val build = sparkBroadcastJoins.head.right + val readsCometBroadcast = build.collectFirst { case c: CometColumnarToRowExec => + c.child match { + case InputAdapter(stage) => stage + case stage => stage + } + } + assert( + readsCometBroadcast.exists(_.isInstanceOf[BroadcastQueryStageExec]), + s"expected the join to read a Comet broadcast stage:\n$plan") + assert(!build.isInstanceOf[WholeStageCodegenExec], s"plan:\n$plan") + } + } + } + } + } + } + test("BroadcastNestedLoopJoin with unequal filter") { withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { // Include NULL keys: predicate `_1 > _1` returns NULL with a NULL operand, so