From ce1feae285b0a21e319d688097351c6be8695b01 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Thu, 5 Nov 2026 20:49:43 -0700 Subject: [PATCH] fix: let shared subexpressions that cannot raise be skipped by AtLeastNNonNulls AtLeastNNonNulls, and the null short-circuit of function arguments, refuse to skip an argument that holds a subexpression Spark shares between the operator's expressions. Spark evaluates a shared subexpression for every row, so skipping one could hide an error that Spark raises. The check applied to every shared subexpression, so a filter whose columns merely read the same struct field or take the same substring, such as na.drop over columns derived from one input, fell back to Spark, with a transition on each side. Only a shared subexpression that can raise needs to be kept. canSkipEvaluation now leaves out those made of expressions that cannot fail whatever the row holds: literals and columns, struct field reads, comparisons and boolean operators, string functions such as substr and upper, and casts and numeric arithmetic outside ANSI mode. The list fails closed, so anything it does not name, such as element_at, an ANSI cast, interval arithmetic or a decimal expression after promotion, still counts as able to raise. --- docs/source/user-guide/latest/expressions.md | 2 +- .../apache/comet/serde/QueryPlanSerde.scala | 56 +++++++++++++++---- .../apache/comet/CometExpressionSuite.scala | 35 +++++++++++- .../comet/serde/QueryPlanSerdeSuite.scala | 50 ++++++++++++++--- 4 files changed, 122 insertions(+), 21 deletions(-) diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index cf8baac011d..2bebe49c59f 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/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) | diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index c8b89b0742c..84dcb3f0460 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/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,17 +1404,17 @@ 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. */ @@ -1388,8 +1422,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { // 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))) } } } diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 3a4f8140506..89cc140268b 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/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") { diff --git a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala index d7e595d4a83..b79acedca58 100644 --- a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala +++ b/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))()