diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index eda40fb597c..b1e3216161d 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -149,8 +149,8 @@ so users hunting an unexpected value have a single place to check: when a partial is written out or the sum is evaluated, which matches Spark's whole-stage codegen path. Without codegen Spark buffers the aggregate in an `UnsafeRow` and latches as soon as a running sum leaves the precision. Comet falls back at precision 38 when codegen is disabled by - `spark.sql.codegen.wholeStage`, by `spark.sql.codegen.factoryMode=NO_CODEGEN`, by an - imperative sibling aggregate, by an expression Spark cannot compile such as a lambda function, + `spark.sql.codegen.wholeStage`, by `spark.sql.codegen.factoryMode=NO_CODEGEN` on Spark 3.5+, + by an imperative sibling aggregate, by an expression Spark cannot compile such as a lambda function, or by the `spark.sql.codegen.maxFields` limit on the aggregate's own output and inputs, but not when Spark abandons codegen at runtime, after a compile failure under `spark.sql.codegen.fallback` or when the generated code exceeds diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index a43e34ffebe..7bf6e845485 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -56,7 +56,7 @@ import com.google.common.base.Objects import com.google.protobuf.CodedOutputStream import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, ConfigEntry, ContribServices} -import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, withFallbackReason} +import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, isSpark35Plus, withFallbackReason} import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} @@ -2234,11 +2234,12 @@ object CometHashAggregateExec // Without codegen Spark buffers an ungrouped aggregate in an UnsafeRow, which latches a // decimal sum that leaves the precision, while the native accumulator keeps it unbounded. // Spark turns codegen off by config, for an imperative aggregate, for a non-leaf - // CodegenFallback expression, or when the output or an input exceeds its field limit. + // CodegenFallback expression, or when the output or an input exceeds its field limit. The + // NO_CODEGEN factory mode turns whole-stage codegen off only from Spark 3.5 (SPARK-44236). val codegenOff = !op.conf.wholeStageEnabled || - op.conf + (isSpark35Plus && op.conf .getConfString(SQLConf.CODEGEN_FACTORY_MODE.key) - .equalsIgnoreCase(CodegenObjectFactoryMode.NO_CODEGEN.toString) || + .equalsIgnoreCase(CodegenObjectFactoryMode.NO_CODEGEN.toString)) || op.aggregateExpressions.exists(_.aggregateFunction.isInstanceOf[ImperativeAggregate]) || WholeStageCodegenExec.isTooManyFields(op.conf, op.schema) || op.children.exists(child => WholeStageCodegenExec.isTooManyFields(op.conf, child.schema)) || diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 5e80bee03cc..b4ef6f3cc8d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -43,7 +43,7 @@ import org.apache.spark.sql.types.{ArrayType, DataTypes, StructField, StructType import org.apache.comet.CometConf import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT -import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus} import org.apache.comet.rules.CometExecRule import org.apache.comet.serde.RegrSparkVersions import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} @@ -1963,11 +1963,18 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { Row(null)) } withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN") { - assertDecimalSumFallsBackLikeSpark( - sql("SELECT SUM(v) FROM dec_no_codegen"), - reason, - ansiEnabled, - Row(null)) + val df = sql("SELECT SUM(v) FROM dec_no_codegen") + if (isSpark35Plus) { + assertDecimalSumFallsBackLikeSpark(df, reason, ansiEnabled, Row(null)) + } else { + // Before SPARK-44236 (3.5) the factory mode leaves whole-stage codegen on, so + // Spark recovers the sum and the aggregate stays native. + checkSparkAnswerAndOperator(df) + val answer = df.collect().toSeq + assert( + answer == Seq(Row(recoveredSum.bigDecimal.setScale(38))), + s"NO_CODEGEN on Spark 3.4 returned $answer, expected $recoveredSum") + } } } } diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala index f3e59e75101..20caba36b1f 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -19,6 +19,9 @@ package org.apache.comet.parquet +import java.util.concurrent.atomic.AtomicReference + +import org.apache.spark.CometListenerBusUtils import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometNativeWriteExec, CometWriteFilesExec} import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} @@ -61,12 +64,12 @@ abstract class CometParquetWriterTestBase extends CometTestBase { /** As above, for a write that names its own target (an `INSERT INTO`, for example). */ protected def captureWritePlan(writeOp: => Unit): SparkPlan = { - var capturedPlan: Option[QueryExecution] = None + val capturedPlan = new AtomicReference[QueryExecution]() val listener = new org.apache.spark.sql.util.QueryExecutionListener { override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { if (funcName == "save" || funcName.contains("command")) { - capturedPlan = Some(qe) + capturedPlan.set(qe) } } @@ -76,27 +79,18 @@ abstract class CometParquetWriterTestBase extends CometTestBase { exception: Exception): Unit = {} } + // Listener events are delivered asynchronously, so drain the bus before registering: an + // earlier write's event still in flight would otherwise be captured in place of this one. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { writeOp + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - // Wait for listener to be called with timeout - val maxWaitTimeMs = 15000 - val checkIntervalMs = 100 - val maxIterations = maxWaitTimeMs / checkIntervalMs - var iterations = 0 - - while (capturedPlan.isEmpty && iterations < maxIterations) { - Thread.sleep(checkIntervalMs) - iterations += 1 - } - - assert( - capturedPlan.isDefined, - s"Listener was not called within ${maxWaitTimeMs}ms - no execution plan captured") - - stripAQEPlan(capturedPlan.get.executedPlan) + val plan = capturedPlan.get() + assert(plan != null, "Listener was not called - no execution plan captured") + stripAQEPlan(plan.executedPlan) } finally { spark.listenerManager.unregister(listener) }