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