fix: honor subexpression elimination off in the native null short-circuit

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -1162,8 +1162,12 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
 
       val builder = OperatorOuterClass.Operator.newBuilder().setPlanId(op.id)
       // Which subexpressions Spark evaluates for every row decides where native execution may
-      // skip a function's arguments (see `QueryPlanSerde.canShortCircuitNulls`)
-      val nativeOp = QueryPlanSerde.withSharedSubexpressions(op.expressions) {
+      // skip a function's arguments (see `QueryPlanSerde.canShortCircuitNulls`). Spark's
+      // generated code for an aggregate shares them whatever
+      // spark.sql.subexpressionElimination.enabled says, and other operators only when it is on.
+      val nativeOp = QueryPlanSerde.withSharedSubexpressions(
+        op.expressions,
+        alwaysShared = op.isInstanceOf[BaseAggregateExec]) {
         if (op.children.nonEmpty && op.children.forall(_.isInstanceOf[CometNativeExec])) {
           val childOp = op.children.map(_.asInstanceOf[CometNativeExec].nativeOp)
           childOp.foreach(builder.addChildren)

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -1300,38 +1300,47 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
 
   /**
    * The subexpressions that occur more than once in the expressions of the operator being
-   * serialized, as Spark's subexpression elimination finds them. Spark's generated code for a
-   * projection or an aggregation evaluates each of them for every row, before the expressions
-   * that use them. They are treated that way for every operator, and whatever
-   * `spark.sql.subexpressionElimination.enabled` says, which only keeps them evaluated for every
-   * row where Spark might not. They are found on the first `canShortCircuitNulls` call for the
-   * operator, so an operator without an argument to skip never runs `EquivalentExpressions`.
+   * serialized, as Spark's subexpression elimination finds them. Where Spark eliminates them, it
+   * evaluates each for every row, before the expressions that use them: in the generated code of
+   * an aggregation always, and in a projection, such as a project or a shuffle's partitioning
+   * keys, only when `spark.sql.subexpressionElimination.enabled` is on. Other operators, such as
+   * a filter or a sort, are treated like a projection, which at worst keeps them evaluated for
+   * every row where Spark does not. They are found on the first `canShortCircuitNulls` call for
+   * the operator, so an operator without an argument to skip never runs `EquivalentExpressions`.
    * `None` outside `withSharedSubexpressions`, where they are unknown.
    */
   private val sharedSubexpressions = new DynamicVariable[Option[SharedSubexpressions]](None)
 
-  private final class SharedSubexpressions(expressions: Seq[Expression]) {
-    lazy val common: Seq[Expression] = {
-      val equivalence = new EquivalentExpressions
-      expressions.foreach(equivalence.addExprTree(_))
-      equivalence.getCommonSubexpressions
-    }
+  private final class SharedSubexpressions(expressions: Seq[Expression], alwaysShared: Boolean) {
+    lazy val common: Seq[Expression] =
+      if (alwaysShared || SQLConf.get.subexpressionEliminationEnabled) {
+        val equivalence = new EquivalentExpressions
+        expressions.foreach(equivalence.addExprTree(_))
+        equivalence.getCommonSubexpressions
+      } else {
+        Nil
+      }
   }
 
   /**
    * Runs `f`, which serializes an operator whose expressions are `expressions`. Code that
    * serializes expressions outside an operator's conversion, such as when the plan runs, calls it
-   * with the expressions Spark evaluates together, or no argument is skipped.
+   * with the expressions Spark evaluates together, or no argument is skipped. `alwaysShared` is
+   * for an aggregation, whose shared subexpressions Spark's generated code evaluates for every
+   * row whatever `spark.sql.subexpressionElimination.enabled` says. An aggregation that Spark
+   * runs without generated code shares none when it is off, and stays eager in Comet.
    */
-  def withSharedSubexpressions[T](expressions: Seq[Expression])(f: => T): T =
-    sharedSubexpressions.withValue(Some(new SharedSubexpressions(expressions)))(f)
+  def withSharedSubexpressions[T](expressions: Seq[Expression], alwaysShared: Boolean = false)(
+      f: => T): T =
+    sharedSubexpressions.withValue(Some(new SharedSubexpressions(expressions, alwaysShared)))(f)
 
   /**
    * Whether native execution may skip the arguments of `expr` after the first on the rows where
    * an earlier argument is NULL, as Spark's evaluation of a null-intolerant `BinaryExpression` or
    * `TernaryExpression` does. Not if one of them holds a subexpression that the operator's
-   * expressions share: Spark evaluates that for every row, and raises if it fails on a row whose
-   * array is NULL. Nor outside `withSharedSubexpressions`, where what they share is unknown.
+   * expressions share and Spark eliminates: Spark evaluates that for every row, and raises if it
+   * fails on a row whose array is NULL. Nor outside `withSharedSubexpressions`, where what they
+   * share is unknown.
    */
   def canShortCircuitNulls(expr: Expression): Boolean = {
     val later = expr.children.tail

spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_cse_off.sql

@@ -0,0 +1,74 @@
+-- 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.
+
+-- With spark.sql.subexpressionElimination.enabled off, Spark does not evaluate a subexpression
+-- that an operator's expressions share ahead of them, so a null-intolerant function skips a later
+-- argument holding one on a row whose array is NULL, and Comet skips it too. Spark's generated
+-- code for an aggregate shares subexpressions whatever the setting says, so there the cast still
+-- raises. The same queries with the setting on are in array_null_short_circuit.sql.
+-- https://github.com/apache/datafusion-comet/issues/6972
+
+-- Config: spark.sql.ansi.enabled=true
+-- Config: spark.sql.subexpressionElimination.enabled=false
+-- Config: spark.sql.codegen.wholeStage=true
+
+statement
+CREATE TABLE array_null_short_circuit_cse_off(ai array<int>, s string, i string) USING parquet
+
+-- `s` is malformed only where the array is NULL. COALESCE(1) writes a single file, so the NULL
+-- arrays share a batch with the others.
+statement
+INSERT INTO array_null_short_circuit_cse_off SELECT /*+ COALESCE(1) */ * FROM VALUES
+  (NULL, 'bad', 'bad'),
+  (array(1, 2), '2', '1'),
+  (array(3), '7', '0'),
+  (NULL, 'worse', 'worse'),
+  (array(4, NULL, 5), '5', '2')
+  AS v(ai, s, i)
+
+-- A projection's expressions share the cast
+query expect_native(array_contains, array_position)
+SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))
+FROM array_null_short_circuit_cse_off
+
+-- So do a shuffle's partitioning keys, for hash and range partitioning, and a
+-- TakeOrderedAndProject's projection list
+query
+SELECT ai, s FROM array_null_short_circuit_cse_off
+DISTRIBUTE BY array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))
+
+query
+SELECT ai, s FROM array_null_short_circuit_cse_off
+ORDER BY array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT)), s
+
+query expect_native(array_contains, array_position)
+SELECT array_contains(ai, CAST(s AS INT)), array_position(ai, CAST(s AS INT))
+FROM (SELECT * FROM array_null_short_circuit_cse_off ORDER BY i LIMIT 10)
+
+-- An aggregate's update expressions still share the cast, with or without grouping keys
+query expect_error_class(CAST_INVALID_INPUT)
+SELECT sum(array_position(ai, CAST(s AS INT))), max(array_contains(ai, CAST(s AS INT)))
+FROM array_null_short_circuit_cse_off
+
+query expect_error_class(CAST_INVALID_INPUT)
+SELECT i, sum(array_position(ai, CAST(s AS INT))), max(array_contains(ai, CAST(s AS INT)))
+FROM array_null_short_circuit_cse_off GROUP BY i
+
+-- Where the whole call is shared rather than the cast, the call skips the cast on its own
+query expect_native(array_position)
+SELECT sum(array_position(ai, CAST(s AS INT))), max(array_position(ai, CAST(s AS INT)))
+FROM array_null_short_circuit_cse_off

spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit_cse_off.sql

@@ -0,0 +1,48 @@
+-- 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.
+
+-- With spark.sql.subexpressionElimination.enabled off, Spark's projection evaluates `1 div i`
+-- only where the map is not NULL, though both lookups share it. Its generated code for an
+-- aggregate shares it whatever the setting says, so there the division by zero still raises. The
+-- array counterparts are in array/array_null_short_circuit_cse_off.sql.
+-- https://github.com/apache/datafusion-comet/issues/6972
+
+-- Config: spark.sql.ansi.enabled=true
+-- Config: spark.sql.subexpressionElimination.enabled=false
+-- Config: spark.sql.codegen.wholeStage=true
+
+statement
+CREATE TABLE map_null_short_circuit_cse_off(m1 map<int, int>, m2 map<int, int>, i int)
+USING parquet
+
+-- `i` is 0 only where the maps are NULL
+statement
+INSERT INTO map_null_short_circuit_cse_off SELECT /*+ COALESCE(1) */ * FROM VALUES
+  (NULL, NULL, 0),
+  (map(1, 5), map(1, 6), 1),
+  (map(2, 7), map(1, 8), 1)
+  AS v(m1, m2, i)
+
+query expect_native(getmapvalue)
+SELECT m1[1 div i], m2[1 div i] FROM map_null_short_circuit_cse_off
+
+query expect_error_class(DIVIDE_BY_ZERO)
+SELECT max(m1[1 div i]), min(m2[1 div i]) FROM map_null_short_circuit_cse_off
+
+-- Where the whole lookup is shared rather than the division, the lookup skips it on its own
+query expect_native(getmapvalue)
+SELECT max(m1[1 div i]), min(m1[1 div i]) FROM map_null_short_circuit_cse_off

spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala

@@ -22,6 +22,7 @@ package org.apache.comet.serde
 import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.spark.sql.catalyst.expressions.{ArrayContains, ArrayPosition, AttributeReference, Cast}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
 import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus}
@@ -121,4 +122,31 @@ class QueryPlanSerdeSuite extends AnyFunSuite {
       assert(!canShortCircuitNulls(position))
     }
   }
+
+  test("canShortCircuitNulls follows subexpression elimination except in aggregates") {
+    val cast = Cast(AttributeReference("s", StringType)(), IntegerType)
+    val array = AttributeReference("ai", ArrayType(IntegerType))()
+    val contains = ArrayContains(array, cast)
+    val shared = Seq(contains, ArrayPosition(array, cast))
+
+    def withSubexpressionElimination(enabled: Boolean)(f: => Unit): Unit = {
+      val conf = new SQLConf
+      conf.setConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED, enabled)
+      SQLConf.withExistingConf(conf) {
+        assert(SQLConf.get.subexpressionEliminationEnabled == enabled)
+        f
+      }
+    }
+
+    withSubexpressionElimination(enabled = false) {
+      withSharedSubexpressions(shared) {
+        assert(canShortCircuitNulls(contains))
+      }
+      withSharedSubexpressions(shared, alwaysShared = true) {
+        assert(!canShortCircuitNulls(contains))
+      }
+      // Outside an operator's conversion what they share is still unknown
+      assert(!canShortCircuitNulls(contains))
+    }
+  }
 }