From 3ca36ab7b3cfb7cb332c7a5376600c0dbd141f52 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 14 Oct 2026 21:42:10 +0700 Subject: [PATCH 1/2] fix: honor subexpression elimination off in the native null short-circuit canShortCircuitNulls kept a later argument eager whenever the operator's expressions shared a subexpression it held, whatever spark.sql.subexpressionElimination.enabled says. With the setting off, Spark's projection, shuffle partition keys and TakeOrderedAndProject share nothing and skip that argument on a row whose first argument is NULL, so Comet raised where Spark returns rows. withSharedSubexpressions now takes an alwaysShared flag. With the setting off an operator shares nothing unless it asks for that flag, which CometExecRule does for aggregates, since Spark's whole-stage codegen aggregates share the subexpressions of their update expressions regardless of the setting. Nothing changes with the setting on. --- .../apache/comet/rules/CometExecRule.scala | 8 +- .../apache/comet/serde/QueryPlanSerde.scala | 43 ++++++----- .../array_null_short_circuit_cse_off.sql | 74 +++++++++++++++++++ .../map/map_null_short_circuit_cse_off.sql | 48 ++++++++++++ .../comet/serde/QueryPlanSerdeSuite.scala | 39 ++++++++++ 5 files changed, 193 insertions(+), 19 deletions(-) create mode 100644 spark/src/test/resources/sql-tests/expressions/array/array_null_short_circuit_cse_off.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/map/map_null_short_circuit_cse_off.sql 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 5a09275a16c..1c3b3035fa8 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -1135,8 +1135,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..7d12bcd7b9a 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,42 @@ 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)) + } + withSubexpressionElimination(enabled = true) { + Seq(false, true).foreach { alwaysShared => + withSharedSubexpressions(shared, alwaysShared) { + assert(!canShortCircuitNulls(contains)) + } + withSharedSubexpressions(Seq(contains), alwaysShared) { + assert(canShortCircuitNulls(contains)) + } + } + assert(!canShortCircuitNulls(contains)) + } + } } From dfc91afc048edc37f1b74517d93576771ab49a73 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Thu, 15 Oct 2026 08:48:31 +0700 Subject: [PATCH 2/2] test: drop the subexpression elimination on case that repeats an existing test With elimination on, alwaysShared does not change the shared set, and the remaining assertions are the ones the default-conf test already makes. --- .../org/apache/comet/serde/QueryPlanSerdeSuite.scala | 11 ----------- 1 file changed, 11 deletions(-) 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 7d12bcd7b9a..71f329eb17a 100644 --- a/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/QueryPlanSerdeSuite.scala @@ -148,16 +148,5 @@ class QueryPlanSerdeSuite extends AnyFunSuite { // Outside an operator's conversion what they share is still unknown assert(!canShortCircuitNulls(contains)) } - withSubexpressionElimination(enabled = true) { - Seq(false, true).foreach { alwaysShared => - withSharedSubexpressions(shared, alwaysShared) { - assert(!canShortCircuitNulls(contains)) - } - withSharedSubexpressions(Seq(contains), alwaysShared) { - assert(canShortCircuitNulls(contains)) - } - } - assert(!canShortCircuitNulls(contains)) - } } }