fix: let shared subexpressions that cannot raise be skipped by AtLeastNNonNulls

docs/source/user-guide/latest/expressions.md

@@ -281,7 +281,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci
 | `hour` | ✅ | Native |  |
 | `last_day` | ✅ | Native |  |
 | `localtimestamp` | ✅ | — |  |
-| `make_date` | ✅ | Native | Computed arguments preserve year → month → day NULL checks. Expressions fall back when shared subexpressions or a parent could change which arguments Spark evaluates. IN probes stay native for literal/column candidates and array/struct constructors of those values; other computed candidates require Spark evaluation. Operator-level safeguards and their opt-out are described in [make_date evaluation](#make_date-evaluation). |
+| `make_date` | ✅ | Native | Computed arguments preserve year → month → day NULL checks. Expressions fall back when shared subexpressions that can raise, or a parent, could change which arguments Spark evaluates. IN probes stay native for literal/column candidates and array/struct constructors of those values; other computed candidates require Spark evaluation. Operator-level safeguards and their opt-out are described in [make_date evaluation](#make_date-evaluation). |
 | `make_dt_interval` | ✅ | Codegen dispatch |  |
 | `make_interval` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; intervals outside Arrow's nanosecond range are tracked by [#5279](https://github.com/apache/datafusion-comet/issues/5279); the native path is opt-in via allowIncompatible ([details](compatibility/expressions/datetime.md)) |
 | `make_time` | ✅ | — | Spark 4.1+; requires `spark.sql.timeType.enabled=true`, which Spark leaves off by default. Runs natively; remaining TIME type work is tracked by [#4288](https://github.com/apache/datafusion-comet/issues/4288) |

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

@@ -1339,22 +1339,56 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
    * 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.
+   * `None` outside `withSharedSubexpressions`, where they are unknown. Only those that can raise
+   * are kept, see `cannotRaise`.
    */
   private val sharedSubexpressions = new DynamicVariable[Option[SharedSubexpressions]](None)
 
   private final class SharedSubexpressions(expressions: Seq[Expression], alwaysShared: Boolean) {
-    // Find shared nodes in Spark's original trees, then compare them after decimal promotion.
-    lazy val common: Seq[Expression] =
+    // The shared subexpressions that can raise. Find shared nodes in Spark's original trees, then
+    // compare them after decimal promotion. Promotion wraps decimal arithmetic in CheckOverflow,
+    // which `cannotRaise` does not list, so that arithmetic is always kept.
+    lazy val raising: Seq[Expression] =
       if (alwaysShared || SQLConf.get.subexpressionEliminationEnabled) {
         val equivalence = new EquivalentExpressions
         expressions.foreach(equivalence.addExprTree(_))
-        equivalence.getCommonSubexpressions.map(DecimalPrecision.promote)
+        equivalence.getCommonSubexpressions.map(DecimalPrecision.promote).filterNot(cannotRaise)
       } else {
         Nil
       }
   }
 
+  /**
+   * Whether evaluating `expr` on a row cannot raise, whatever the row holds. This fails closed:
+   * an expression is assumed to raise unless it is listed. Casts and numeric arithmetic raise
+   * only in the ANSI eval mode they carry, but some expressions raise in any mode, such as
+   * `element_at` with the index 0 and the arithmetic of intervals, so those stay out.
+   */
+  private def cannotRaise(expr: Expression): Boolean =
+    !expr.exists(!nodeCannotRaise(_))
+
+  private def nodeCannotRaise(node: Expression): Boolean = {
+    def plainArithmetic(a: BinaryExpression, mode: EvalMode.Value): Boolean =
+      mode != EvalMode.ANSI && a.dataType.isInstanceOf[NumericType]
+    node match {
+      case _: Literal | _: AttributeReference | _: BoundReference | _: GetStructField => true
+      case _: BinaryComparison | _: IsNull | _: IsNotNull | _: Not | _: And | _: Or | _: In |
+          _: InSet | _: IsNaN | _: Coalesce | _: If | _: CaseWhen =>
+        true
+      case _: Length | _: Upper | _: Lower | _: Substring | _: StartsWith | _: EndsWith |
+          _: Contains =>
+        true
+      case c: Cast => evalMode(c) != CometEvalMode.ANSI
+      case a: Add => plainArithmetic(a, a.evalMode)
+      case a: Subtract => plainArithmetic(a, a.evalMode)
+      case a: Multiply => plainArithmetic(a, a.evalMode)
+      case a: Divide => plainArithmetic(a, a.evalMode)
+      case a: IntegralDivide => plainArithmetic(a, a.evalMode)
+      case a: Remainder => plainArithmetic(a, a.evalMode)
+      case _ => false
+    }
+  }
+
   /**
    * 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
@@ -1370,26 +1404,26 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
   /**
    * 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 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.
+   * `TernaryExpression` does. Not if one of them holds a subexpression that can raise and that
+   * the operator's 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 = canSkipEvaluation(expr.children.tail)
 
   /**
    * Whether native execution may leave `skippable`, expressions that Spark evaluates only for
    * some rows, unevaluated on the rows Spark skips. Not if one of them holds a subexpression that
-   * the operator's expressions share and Spark eliminates, nor outside
+   * can raise and that the operator's expressions share and Spark eliminates, nor outside
    * `withSharedSubexpressions`; see `canShortCircuitNulls`. Normalize both sides with decimal
    * promotion so callers may pass either original or already-promoted trees.
    */
   def canSkipEvaluation(skippable: Seq[Expression]): Boolean = {
     // Subexpression elimination never shares a leaf, such as a column or a literal
     skippable.forall(_.isInstanceOf[LeafExpression]) || sharedSubexpressions.value.exists { s =>
       // Avoid decimal promotion when Spark shares nothing.
-      s.common.isEmpty || !skippable.exists { expr =>
-        DecimalPrecision.promote(expr).exists(e => s.common.exists(_.semanticEquals(e)))
+      s.raising.isEmpty || !skippable.exists { expr =>
+        DecimalPrecision.promote(expr).exists(e => s.raising.exists(_.semanticEquals(e)))
       }
     }
   }

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -24,7 +24,7 @@ import scala.util.Random
 import org.apache.hadoop.fs.Path
 import org.apache.spark.SparkThrowable
 import org.apache.spark.sql.{Column, CometTestBase, DataFrame, Row}
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AtLeastNNonNulls, AttributeReference, Cast, Divide, FromUnixTime, In, InSet, IntegralDivide, Literal, Multiply, Remainder, StructsToJson, Subtract, TruncDate, TruncTimestamp}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AtLeastNNonNulls, AttributeReference, Cast, Divide, FromUnixTime, GetStructField, In, InSet, IntegralDivide, Literal, Lower, Multiply, Remainder, StructsToJson, Subtract, TruncDate, TruncTimestamp, Upper}
 import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, OptimizeIn, SimplifyExtractValueOps}
 import org.apache.spark.sql.comet.{CometFilterExec, CometProjectExec, CometSortExec, CometTakeOrderedAndProjectExec}
 import org.apache.spark.sql.execution.{LocalTableScanExec, ProjectExec, SparkPlan}
@@ -199,6 +199,39 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("AtLeastNNonNulls: a shared subexpression that cannot raise does not force a fallback") {
+    withParquetTable(Seq((Some(1), "Bad"), (None, "7"), (None, null)), "nonnulls") {
+      withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+        // Upper(_2) is shared by the last two children, which are skipped on a row that reaches n=1
+        // at its first child. Spark evaluates it for every row, but it cannot fail, so skipping it
+        // shows nothing.
+        val df = spark.table("nonnulls")
+        val shared = Upper(df.col("_2").expr)
+        checkSparkAnswerAndImpl(
+          df.select(
+            getColumnFromExpression(
+              AtLeastNNonNulls(1, Seq(df.col("_1").expr, shared, Lower(shared))))),
+          native = Seq("atleastnnonnulls"))
+      }
+    }
+  }
+
+  test("AtLeastNNonNulls: a shared field of a struct does not force a fallback") {
+    withParquetTable(Seq((Some((1, "Bad")), 1), (None, 2), (Some((2, null)), 3)), "nonnulls") {
+      withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") {
+        // Both the field and the Upper of it are shared by the last two children, as when the structs
+        // an explode produces are cleaned up with na.drop. Reading a field never fails.
+        val df = spark.table("nonnulls")
+        val name = GetStructField(df.col("_1").expr, 1)
+        checkSparkAnswerAndImpl(
+          df.select(
+            getColumnFromExpression(
+              AtLeastNNonNulls(1, Seq(df.col("_2").expr, Upper(name), Lower(Upper(name)))))),
+          native = Seq("atleastnnonnulls"))
+      }
+    }
+  }
+
   test("AtLeastNNonNulls: shared decimal subexpressions preserve divide-by-zero errors") {
     val rows = Seq((Some(1), BigDecimal(1), BigDecimal(0)), (None, BigDecimal(7), BigDecimal(1)))
     withParquetTable(rows, "nonnulls") {

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

@@ -23,7 +23,7 @@ import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.SimpleAnalyzer
-import org.apache.spark.sql.catalyst.expressions.{Alias, ArrayContains, ArrayPosition, AttributeReference, AttributeSeq, BindReferences, Cast, CreateNamedStruct, Divide, Expression, In, Literal, MakeDate, Not, UnsafeProjection}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ArrayContains, ArrayPosition, AttributeReference, AttributeSeq, BindReferences, Cast, CreateNamedStruct, Divide, ElementAt, Expression, GetStructField, In, Length, Literal, MakeDate, Not, UnsafeProjection}
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project}
 import org.apache.spark.sql.comet.DecimalPrecision
 import org.apache.spark.sql.comet.util.Utils.equalsIgnoreCompatibleNullability
@@ -208,13 +208,14 @@ class QueryPlanSerdeSuite extends AnyFunSuite {
   test("canShortCircuitNulls skips an argument only where the shared subexpressions are known") {
     val array = AttributeReference("ai", ArrayType(IntegerType))()
     val value = AttributeReference("i", IntegerType)()
-    val cast = Cast(AttributeReference("s", StringType)(), IntegerType)
-    val contains = ArrayContains(array, cast)
-    val position = ArrayPosition(array, cast)
+    // element_at with the index 0 raises in every mode, so it is a shared subexpression to keep
+    val risky = ElementAt(AttributeReference("other", ArrayType(IntegerType))(), Literal(0))
+    val contains = ArrayContains(array, risky)
+    val position = ArrayPosition(array, risky)
 
     // A column is never shared, so it can be skipped anywhere
     assert(canShortCircuitNulls(ArrayContains(array, value)))
-    // Something serialized outside an operator's conversion might share the cast
+    // Something serialized outside an operator's conversion might share the element_at
     assert(!canShortCircuitNulls(contains))
     withSharedSubexpressions(Seq(contains)) {
       assert(canShortCircuitNulls(contains))
@@ -226,10 +227,10 @@ class QueryPlanSerdeSuite extends AnyFunSuite {
   }
 
   test("canShortCircuitNulls follows subexpression elimination except in aggregates") {
-    val cast = Cast(AttributeReference("s", StringType)(), IntegerType)
+    val risky = ElementAt(AttributeReference("other", ArrayType(IntegerType))(), Literal(0))
     val array = AttributeReference("ai", ArrayType(IntegerType))()
-    val contains = ArrayContains(array, cast)
-    val shared = Seq(contains, ArrayPosition(array, cast))
+    val contains = ArrayContains(array, risky)
+    val shared = Seq(contains, ArrayPosition(array, risky))
 
     def withSubexpressionElimination(enabled: Boolean)(f: => Unit): Unit = {
       val conf = new SQLConf
@@ -252,6 +253,39 @@ class QueryPlanSerdeSuite extends AnyFunSuite {
     }
   }
 
+  test("canSkipEvaluation ignores shared subexpressions that cannot raise") {
+    val array = AttributeReference("ai", ArrayType(IntegerType))()
+    val text = AttributeReference("s", StringType)()
+    val struct = AttributeReference("st", StructType(Seq(StructField("f", IntegerType))))()
+    val months = AttributeReference("m", YearMonthIntervalType())()
+    // element_at with the index 0 raises in every mode
+    val risky = ElementAt(AttributeReference("other", ArrayType(IntegerType))(), Literal(0))
+
+    // Whether `e` can be skipped when the operator shares it between two expressions
+    def skippable(e: Expression): Boolean =
+      withSharedSubexpressions(Seq(e, e), alwaysShared = true)(canSkipEvaluation(Seq(e)))
+
+    assert(skippable(Length(text)))
+    // A field of a struct is NULL for a NULL struct and never fails
+    assert(skippable(GetStructField(struct, 0)))
+    assert(!skippable(risky))
+    // Interval arithmetic raises on overflow in every mode, unlike numeric arithmetic
+    assert(!skippable(Add(months, months)))
+    // An expression that cannot raise but holds one that does is kept as well
+    assert(!skippable(Length(Cast(risky, StringType))))
+
+    // A shared subexpression that raises does not keep a consumer of another one from being skipped
+    val safe = Length(text)
+    val riskyContains = ArrayContains(array, risky)
+    val safeContains = ArrayContains(array, safe)
+    withSharedSubexpressions(
+      Seq(riskyContains, ArrayPosition(array, risky), safeContains, ArrayPosition(array, safe)),
+      alwaysShared = true) {
+      assert(!canShortCircuitNulls(riskyContains))
+      assert(canShortCircuitNulls(safeContains))
+    }
+  }
+
   test("canSkipEvaluation recognizes shared subexpressions across decimal promotion") {
     val left = AttributeReference("a", DecimalType(10, 0))()
     val right = AttributeReference("b", DecimalType(10, 0))()