diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index 69cf62cc867..a9e3f442f7c 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/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) 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 56cb421ae84..3f3881e6d4e 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/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 diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_cse_off.sql b/spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_cse_off.sql new file mode 100644 index 00000000000..c56529f2af6 --- /dev/null +++ b/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, 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 diff --git a/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit_cse_off.sql b/spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit_cse_off.sql new file mode 100644 index 00000000000..7894bc21be4 --- /dev/null +++ b/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, m2 map, 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 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 8d248669b0a..71f329eb17a 100644 --- a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala +++ b/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)) + } + } }