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 d3eacddf5e0..3a60be3f788 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 @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ -import org.apache.spark.sql.catalyst.plans.physical.Partitioning +import org.apache.spark.sql.catalyst.plans.physical.{IdentityBroadcastMode, Partitioning} import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CodegenSupport, ColumnarToRowTransition, SparkPlan, SQLExecution} @@ -137,7 +137,13 @@ case class CometColumnarToRowExec(child: SparkPlan) numOutputRows) val mode = cometBroadcastExchange.get.mode - val relation = mode.transform(rows, Some(numOutputRows.value)) + // Identity broadcasts retain rows, but both JVM conversion paths reuse the returned row. + // Hashed relations copy row bytes into their own storage while consuming the iterator. + val broadcastRows = mode match { + case IdentityBroadcastMode => rows.map(_.copy()) + case _ => rows + } + val relation = mode.transform(broadcastRows, Some(numOutputRows.value)) val broadcasted = sparkContext.broadcastInternal(relation, serializedOnly = true) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq) diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index bf0bacf0d9e..e67ec19c7b3 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -28,13 +28,13 @@ import org.scalatest.PrivateMethodTester._ import org.apache.logging.log4j.Level import org.apache.spark.sql._ -import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow} import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Descending, Expression, ExpressionInfo, GreaterThan, In, InSet, KnownFloatingPointNormalized, Literal, Not, RuntimeReplaceable, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate, Final, Min, Partial, PartialMerge} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} -import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RoundRobinPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, IdentityBroadcastMode, RoundRobinPartitioning} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.CometExec.nativeLeaves @@ -98,6 +98,31 @@ class CometExecRuleSuite extends CometTestBase { CometExecRule(spark).apply(stripAQEPlan(plan)) } + test("columnar-to-row identity broadcasts retain every build row") { + for (direct <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true", + CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_ENABLED.key -> direct.toString, + CometConf.COMET_DIRECT_COLUMNAR_TO_ROW_MIN_BATCH_SIZE.key -> "1") { + val build = applyCometExecRule( + createSparkPlan(spark, "SELECT CAST(id AS INT) AS y FROM range(2024, 2028, 1, 2)")) + assert(build.isInstanceOf[CometNativeExec], build.treeString) + val mode = IdentityBroadcastMode + val original = BroadcastExchangeExec(mode, build) + val exchange = CometBroadcastExchangeExec(original, build.output, mode, build) + val rows = CometColumnarToRowExec(exchange) + .executeBroadcast[Array[InternalRow]]() + .value + // Two rows per partition expose row reuse in both JVM conversion paths. + withClue(s"direct=$direct: ") { + assert(rows.map(_.getInt(0)).sorted.toSeq == Seq(2024, 2025, 2026, 2027)) + } + } + } + } + /** Create a test data frame that is used in all tests */ private def createTestDataFrame = { val testSchema = new StructType(