fix: correct two nightly test failures on Spark 3.4 and 4.2

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

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)) ||

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")
+            }
           }
         }
       }

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)
     }