From a94862eddd2b5cfb98d86a4031efe5f52fed17c8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 8 Jul 2026 08:09:07 -0600 Subject: [PATCH 1/6] refactor: make supportsMixedPartialFinal aggregate-function aware --- .../comet/serde/CometAggregateExpressionSerde.scala | 13 +++++++------ .../org/apache/comet/serde/QueryPlanSerde.scala | 2 +- .../scala/org/apache/comet/serde/aggregates.scala | 12 ++++++------ 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala b/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala index 9a83152168..a52d600821 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometAggregateExpressionSerde.scala @@ -82,15 +82,16 @@ trait CometAggregateExpressionSerde[T <: AggregateFunction] { def getSupportLevel(expr: T): SupportLevel = Compatible(None) /** - * Whether this aggregate's intermediate buffer format is compatible between Spark and Comet, - * making it safe to run the Partial in one engine and the Final in the other. Aggregates with - * simple single-value buffers (MIN, MAX, bitwise) are safe; those with complex or - * differently-encoded buffers (AVG, SUM with decimals, CollectSet, Variance) are not. COUNT is - * intentionally excluded: mixed COUNT partial/final regressed AQE's + * Whether this aggregate's intermediate buffer format is compatible between Spark and Comet for + * the given function instance, making it safe to run the Partial in one engine and the Final in + * the other. Aggregates with simple single-value buffers (MIN, MAX, bitwise) are always safe; + * SUM and non-decimal AVG match Spark's buffer and are safe except where noted per instance + * (e.g. TRY-mode SUM uses a Comet-internal flag column). COUNT is intentionally excluded + * despite a matching buffer: mixed COUNT partial/final regressed AQE's * PropagateEmptyRelationAfterAQE pattern (which matches BaseAggregateExec only) and the Spark * 4.0 count-bug decorrelation for correlated IN subqueries. */ - def supportsMixedPartialFinal: Boolean = false + def supportsMixedPartialFinal(fn: T): Boolean = false /** * Convert a Spark expression into a protocol buffer representation that can be passed into 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 91a93d633b..4d6126abce 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -432,7 +432,7 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { case Some(handler) => handler .asInstanceOf[CometAggregateExpressionSerde[AggregateFunction]] - .supportsMixedPartialFinal + .supportsMixedPartialFinal(fn) case None => false } } diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index 5710232cb4..dcc85f4a6a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -33,7 +33,7 @@ import org.apache.comet.shims.CometEvalModeUtil object CometMin extends CometAggregateExpressionSerde[Min] { - override def supportsMixedPartialFinal: Boolean = true + override def supportsMixedPartialFinal(fn: Min): Boolean = true override def getSupportLevel(expr: Min): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -70,7 +70,7 @@ object CometMin extends CometAggregateExpressionSerde[Min] { object CometMax extends CometAggregateExpressionSerde[Max] { - override def supportsMixedPartialFinal: Boolean = true + override def supportsMixedPartialFinal(fn: Max): Boolean = true override def getSupportLevel(expr: Max): SupportLevel = AggSerde.minMaxSupportLevel(expr.dataType) @@ -300,7 +300,7 @@ object CometLast extends CometAggregateExpressionSerde[Last] { } object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] { - override def supportsMixedPartialFinal: Boolean = true + override def supportsMixedPartialFinal(fn: BitAndAgg): Boolean = true override def getSupportLevel(expr: BitAndAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -339,7 +339,7 @@ object CometBitAndAgg extends CometAggregateExpressionSerde[BitAndAgg] { } object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] { - override def supportsMixedPartialFinal: Boolean = true + override def supportsMixedPartialFinal(fn: BitOrAgg): Boolean = true override def getSupportLevel(expr: BitOrAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -378,7 +378,7 @@ object CometBitOrAgg extends CometAggregateExpressionSerde[BitOrAgg] { } object CometBitXOrAgg extends CometAggregateExpressionSerde[BitXorAgg] { - override def supportsMixedPartialFinal: Boolean = true + override def supportsMixedPartialFinal(fn: BitXorAgg): Boolean = true override def getSupportLevel(expr: BitXorAgg): SupportLevel = if (AggSerde.bitwiseAggTypeSupported(expr.dataType)) { @@ -707,7 +707,7 @@ object CometCorr extends CometAggregateExpressionSerde[Corr] { object CometBloomFilterAggregate extends CometAggregateExpressionSerde[BloomFilterAggregate] { - override def supportsMixedPartialFinal: Boolean = true + override def supportsMixedPartialFinal(fn: BloomFilterAggregate): Boolean = true override def getSupportLevel(expr: BloomFilterAggregate): SupportLevel = expr.child.dataType match { From 7bb3981144dd2d68c53d9eb370976fe2d3e42caa Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 8 Jul 2026 08:18:37 -0600 Subject: [PATCH 2/6] feat: enable mixed partial/final execution for sum and non-decimal avg --- .../org/apache/comet/serde/aggregates.scala | 11 +++ .../comet/rules/CometExecRuleSuite.scala | 77 ++++++++++++++++++- 2 files changed, 84 insertions(+), 4 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index dcc85f4a6a..6c6544e93c 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.types.{ByteType, DecimalType, DoubleType, IntegerTyp import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT import org.apache.comet.CometSparkSessionExtensions.{isSpark41Plus, withFallbackReason} +import org.apache.comet.expressions.CometEvalMode import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProto, serializeDataType} import org.apache.comet.shims.CometEvalModeUtil @@ -130,6 +131,11 @@ object CometCount extends CometAggregateExpressionSerde[Count] { object CometAverage extends CometAggregateExpressionSerde[Average] { + override def supportsMixedPartialFinal(fn: Average): Boolean = + // Non-decimal AVG has a (sum: double, count: long) buffer matching Spark. Decimal AVG is + // deferred (overflow nulls count differently) and stays unsafe for mixed execution. + !fn.child.dataType.isInstanceOf[DecimalType] + override def getUnsupportedReasons(): Seq[String] = Seq( "YearMonthIntervalType and DayTimeIntervalType inputs are not supported") @@ -186,6 +192,11 @@ object CometAverage extends CometAggregateExpressionSerde[Average] { object CometSum extends CometAggregateExpressionSerde[Sum] { + override def supportsMixedPartialFinal(fn: Sum): Boolean = + // SUM's buffer matches Spark for Legacy/Ansi (decimal adds is_empty, also matching), but + // TRY-mode integer SUM carries a Comet-internal has_all_nulls column that Spark cannot read. + CometEvalModeUtil.fromSparkEvalMode(CometEvalModeUtil.sumEvalMode(fn)) != CometEvalMode.TRY + override def getSupportLevel(expr: Sum): SupportLevel = if (AggSerde.sumDataTypeSupported(expr.dataType)) { Compatible() diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 7fa06a26cc..1f76361ad7 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ShuffleEx import org.apache.spark.sql.types.{DataTypes, StructField, StructType} import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark42Plus} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark42Plus} import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator} /** @@ -152,8 +152,8 @@ class CometExecRuleSuite extends CometTestBase { CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val transformedPlan = applyCometExecRule(sparkPlan) - // SUM has incompatible intermediate buffers, so if the final aggregate cannot - // be converted to Comet, neither should be + // COUNT is intentionally excluded from mixed execution (AQE / count-bug reasons), so if + // the final aggregate cannot be converted to Comet, neither should the partial. assert( countOperators(transformedPlan, classOf[HashAggregateExec]) == originalHashAggCount) assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) @@ -177,7 +177,8 @@ class CometExecRuleSuite extends CometTestBase { CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { val transformedPlan = applyCometExecRule(sparkPlan) - // if the partial aggregate cannot be converted to Comet, then neither should be + // COUNT blocks mixed execution, so if the partial cannot be converted, neither should + // the final. assert( countOperators(transformedPlan, classOf[HashAggregateExec]) == originalHashAggCount) assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) @@ -231,6 +232,74 @@ class CometExecRuleSuite extends CometTestBase { } } + test("CometExecRule should allow SUM mixed Comet partial and Spark final") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT SUM(id) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + // SUM buffer matches Spark: partial converts to Comet, final stays Spark. + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) // final + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) // partial + } + } + } + + test("CometExecRule should allow SUM mixed Spark partial and Comet final") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT SUM(id) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) // partial + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) // final + } + } + } + + test("CometExecRule should allow AVG mixed Comet partial and Spark final") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) // final + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) // partial + } + } + } + + test("CometExecRule should not allow try_sum mixed execution") { + assume(isSpark35Plus, "try_sum was added in Spark 3.5") + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT try_sum(id) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + // TRY-mode SUM uses a Comet-internal buffer column, so mixing is unsafe: + // the partial must also fall back to Spark. + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 2) + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) + } + } + } + test("CometExecRule should allow BloomFilter mixed Comet partial and Spark final") { assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4142") val funcId = new FunctionIdentifier("bloom_filter_agg") From 2cd888cf1b7448013acd365e2092e18134285cf9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 8 Jul 2026 08:30:03 -0600 Subject: [PATCH 3/6] test: verify mixed-engine sum and avg match Spark --- .../comet/exec/CometAggregateSuite.scala | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index ae14c68207..f8c990fd09 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -183,6 +183,48 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("mixed engine sum/avg: Comet partial + Spark final matches Spark") { + val data = (0 until 100).map(i => (i, i.toLong, i.toDouble, i % 7)) + withParquetTable(data, "tbl") { + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + checkSparkAnswer( + "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_3) FROM tbl GROUP BY _4") + } + } + } + + test("mixed engine sum/avg: Spark partial + Comet final matches Spark") { + val data = (0 until 100).map(i => (i, i.toLong, i.toDouble, i % 7)) + withParquetTable(data, "tbl") { + withSQLConf( + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + checkSparkAnswer( + "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_3) FROM tbl GROUP BY _4") + } + } + } + + test("mixed engine decimal sum: both split directions match Spark") { + val data = (0 until 100).map(i => (BigDecimal(i), i % 7)) + withParquetTable(data, "tbl") { + Seq( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key, + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key).foreach { disabledConf => + withSQLConf( + disabledConf -> "false", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + checkSparkAnswer("SELECT _2, SUM(_1) FROM tbl GROUP BY _2") + } + } + } + } + test("Aggregation without aggregate expressions should use correct result expressions") { withSQLConf( CometConf.COMET_ENABLED.key -> "true", From 7c295ee7d5d720ec6e0264e594fbb4a64a83b0fe Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 8 Jul 2026 08:48:20 -0600 Subject: [PATCH 4/6] test: regenerate tpcds golden files for mixed sum/avg execution --- .../q70/extended.txt | 8 +- .../approved-plans-v1_4/q35/extended.txt | 2 +- .../approved-plans-v1_4/q45/extended.txt | 89 ++++++++++--------- .../q70a/extended.txt | 32 +++---- .../approved-plans-v2_7/q35/extended.txt | 2 +- 5 files changed, 67 insertions(+), 66 deletions(-) diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt index 08df1c04a0..f9fb64d9e4 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4-spark3_5/q70/extended.txt @@ -3,9 +3,9 @@ CometNativeColumnarToRow +- CometProject +- CometWindowExec +- CometSort - +- CometColumnarExchange - +- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - +- Exchange + +- CometExchange + +- CometHashAggregate + +- CometColumnarExchange +- HashAggregate +- Expand +- Project @@ -56,4 +56,4 @@ CometNativeColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 37 out of 53 eligible operators (69%). Final plan contains 4 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 39 out of 53 eligible operators (73%). Final plan contains 4 transitions between Spark and Comet. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q35/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q35/extended.txt index 50e9ebb965..fcb5e0003d 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q35/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q35/extended.txt @@ -1,5 +1,5 @@ TakeOrderedAndProject -+- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: avg, count] ++- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: count] +- Exchange +- HashAggregate +- Project diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q45/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q45/extended.txt index 2d463ef986..6f73e28782 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q45/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q45/extended.txt @@ -1,45 +1,46 @@ -TakeOrderedAndProject -+- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - +- Exchange - +- HashAggregate - +- Project - +- Filter - +- BroadcastHashJoin [COMET: Unsupported join type ExistenceJoin(exists#1)] - :- CometNativeColumnarToRow - : +- CometProject - : +- CometBroadcastHashJoin - : :- CometProject - : : +- CometBroadcastHashJoin - : : :- CometProject - : : : +- CometBroadcastHashJoin - : : : :- CometProject - : : : : +- CometBroadcastHashJoin - : : : : :- CometFilter - : : : : : +- CometNativeScan parquet spark_catalog.default.web_sales - : : : : : +- CometSubqueryBroadcast - : : : : : +- CometBroadcastExchange - : : : : : +- CometProject - : : : : : +- CometFilter - : : : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : : : +- CometBroadcastExchange - : : : : +- CometFilter - : : : : +- CometNativeScan parquet spark_catalog.default.customer - : : : +- CometBroadcastExchange - : : : +- CometProject - : : : +- CometFilter - : : : +- CometNativeScan parquet spark_catalog.default.customer_address - : : +- CometBroadcastExchange - : : +- CometProject - : : +- CometFilter - : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- CometBroadcastExchange - : +- CometProject - : +- CometFilter - : +- CometNativeScan parquet spark_catalog.default.item - +- BroadcastExchange - +- CometNativeColumnarToRow - +- CometProject - +- CometFilter - +- CometNativeScan parquet spark_catalog.default.item +CometNativeColumnarToRow ++- CometTakeOrderedAndProject + +- CometHashAggregate + +- CometColumnarExchange + +- HashAggregate + +- Project + +- Filter + +- BroadcastHashJoin [COMET: Unsupported join type ExistenceJoin(exists#1)] + :- CometNativeColumnarToRow + : +- CometProject + : +- CometBroadcastHashJoin + : :- CometProject + : : +- CometBroadcastHashJoin + : : :- CometProject + : : : +- CometBroadcastHashJoin + : : : :- CometProject + : : : : +- CometBroadcastHashJoin + : : : : :- CometFilter + : : : : : +- CometNativeScan parquet spark_catalog.default.web_sales + : : : : : +- CometSubqueryBroadcast + : : : : : +- CometBroadcastExchange + : : : : : +- CometProject + : : : : : +- CometFilter + : : : : : +- CometNativeScan parquet spark_catalog.default.date_dim + : : : : +- CometBroadcastExchange + : : : : +- CometFilter + : : : : +- CometNativeScan parquet spark_catalog.default.customer + : : : +- CometBroadcastExchange + : : : +- CometProject + : : : +- CometFilter + : : : +- CometNativeScan parquet spark_catalog.default.customer_address + : : +- CometBroadcastExchange + : : +- CometProject + : : +- CometFilter + : : +- CometNativeScan parquet spark_catalog.default.date_dim + : +- CometBroadcastExchange + : +- CometProject + : +- CometFilter + : +- CometNativeScan parquet spark_catalog.default.item + +- BroadcastExchange + +- CometNativeColumnarToRow + +- CometProject + +- CometFilter + +- CometNativeScan parquet spark_catalog.default.item -Comet accelerated 32 out of 41 eligible operators (78%). Final plan contains 2 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 35 out of 41 eligible operators (85%). Final plan contains 3 transitions between Spark and Comet. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt index e527d95f46..2aac2972c3 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt @@ -5,11 +5,11 @@ CometNativeColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange - +- HashAggregate - +- Union - :- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - : +- Exchange + +- CometExchange + +- CometHashAggregate + +- CometUnion + :- CometHashAggregate + : +- CometColumnarExchange : +- HashAggregate : +- Project : +- BroadcastHashJoin @@ -58,11 +58,11 @@ CometNativeColumnarToRow : +- CometProject : +- CometFilter : +- CometNativeScan parquet spark_catalog.default.date_dim - :- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - : +- Exchange - : +- HashAggregate - : +- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - : +- Exchange + :- CometHashAggregate + : +- CometExchange + : +- CometHashAggregate + : +- CometHashAggregate + : +- CometColumnarExchange : +- HashAggregate : +- Project : +- BroadcastHashJoin @@ -111,11 +111,11 @@ CometNativeColumnarToRow : +- CometProject : +- CometFilter : +- CometNativeScan parquet spark_catalog.default.date_dim - +- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - +- Exchange - +- HashAggregate - +- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: sum] - +- Exchange + +- CometHashAggregate + +- CometExchange + +- CometHashAggregate + +- CometHashAggregate + +- CometColumnarExchange +- HashAggregate +- Project +- BroadcastHashJoin @@ -165,4 +165,4 @@ CometNativeColumnarToRow +- CometFilter +- CometNativeScan parquet spark_catalog.default.date_dim -Comet accelerated 103 out of 156 eligible operators (66%). Final plan contains 10 transitions between Spark and Comet. \ No newline at end of file +Comet accelerated 117 out of 156 eligible operators (75%). Final plan contains 10 transitions between Spark and Comet. \ No newline at end of file diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q35/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q35/extended.txt index c340d4a42c..fcb5e0003d 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q35/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q35/extended.txt @@ -1,5 +1,5 @@ TakeOrderedAndProject -+- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: avg, count, sum] ++- HashAggregate [COMET: Spark Final aggregate without Comet Partial requires compatible intermediate buffer formats, but the following aggregate function(s) have incompatible buffers: count] +- Exchange +- HashAggregate +- Project From 8926d4e0340a8e0622fd363d54487b381272907e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 8 Jul 2026 09:03:39 -0600 Subject: [PATCH 5/6] test: add decimal-AVG-blocked guard and AVG parity/coverage tests --- .../comet/exec/CometAggregateSuite.scala | 4 +- .../comet/rules/CometExecRuleSuite.scala | 39 +++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index f8c990fd09..7e5eca4e99 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -191,7 +191,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { checkSparkAnswer( - "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_3) FROM tbl GROUP BY _4") + "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_2), AVG(_3) FROM tbl GROUP BY _4") } } } @@ -204,7 +204,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { checkSparkAnswer( - "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_3) FROM tbl GROUP BY _4") + "SELECT _4, SUM(_1), SUM(_2), SUM(_3), AVG(_1), AVG(_2), AVG(_3) FROM tbl GROUP BY _4") } } } diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 1f76361ad7..893405ff66 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -300,6 +300,45 @@ class CometExecRuleSuite extends CometTestBase { } } + test("CometExecRule should not allow decimal AVG mixed execution") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + // Precision must be large enough (prec + 4 > 15) that Spark's own DecimalAggregates + // optimizer rule does not rewrite AVG to operate on the unscaled Long value, which would + // sidestep the decimal buffer path this test is meant to exercise. + val sparkPlan = + createSparkPlan( + spark, + "SELECT AVG(CAST(id AS DECIMAL(20, 2))) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + // Decimal AVG is deferred (its overflow path nulls count differently from Spark), so + // mixed execution is unsafe and the partial must also fall back to Spark. + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 2) + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) + } + } + } + + test("CometExecRule should allow AVG mixed Spark partial and Comet final") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + val sparkPlan = + createSparkPlan(spark, "SELECT AVG(id) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 1) // partial + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 1) // final + } + } + } + test("CometExecRule should allow BloomFilter mixed Comet partial and Spark final") { assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4142") val funcId = new FunctionIdentifier("bloom_filter_agg") From 8baf10967444d45e1ec22a7d3041f4058464e9e7 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 8 Jul 2026 11:40:53 -0600 Subject: [PATCH 6/6] fix: resolve Spark SQL test failures from mixed SUM execution Enabling mixed Spark-partial / Comet-final SUM surfaced three Spark SQL test failures: - SPARK-28067 (decimal overflow): exclude DecimalType from CometSum.supportsMixedPartialFinal. Decimal overflow detection (ANSI throw / Legacy null) does not survive a Spark-partial / Comet-final split, so the required ArithmeticException was never raised. Mirrors the existing decimal AVG block. Adds a CometExecRule guard test and drops the now-invalid positive mixed-decimal-sum test. - SPARK-19471: relax the assertNoExceptions plan check in DataFrameAggregateSuite to also accept CometHashAggregateExec, since the final aggregate is now a Comet node. Comet's result is correct; only the plan shape changed. - SPARK-44040: ignore under Comet (#4412). A global SUM over an empty relation loses its mandatory single row because AQE's PropagateEmptyRelationAfterAQE only matches BaseAggregateExec, not CometHashAggregateExec. --- dev/diffs/3.4.3.diff | 33 ++++++++++++++--- dev/diffs/3.5.8.diff | 35 +++++++++++++++---- dev/diffs/4.0.2.diff | 35 +++++++++++++++---- dev/diffs/4.1.2.diff | 35 +++++++++++++++---- .../org/apache/comet/serde/aggregates.scala | 6 ++-- .../comet/exec/CometAggregateSuite.scala | 16 --------- .../comet/rules/CometExecRuleSuite.scala | 24 +++++++++++++ 7 files changed, 143 insertions(+), 41 deletions(-) diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 6faf9dab78..487c5d1482 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -1,5 +1,5 @@ diff --git a/pom.xml b/pom.xml -index d3544881af1..aae0ae3b27b 100644 +index d3544881af1..ff963395ec3 100644 --- a/pom.xml +++ b/pom.xml @@ -148,6 +148,8 @@ @@ -260,19 +260,32 @@ index cf40e944c09..bdd5be4f462 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 1cc09c3d7fc..f031fa45c33 100644 +index 1cc09c3d7fc..a84939c045b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -@@ -27,7 +27,7 @@ import org.apache.spark.SparkException +@@ -26,8 +26,9 @@ import org.scalatest.matchers.must.Matchers.the + import org.apache.spark.SparkException import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper ++import org.apache.spark.sql.comet.CometHashAggregateExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -755,7 +755,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -691,7 +692,9 @@ class DataFrameAggregateSuite extends QueryTest + case _ => false + }.isDefined) + } else { +- assert(stripAQEPlan(hashAggPlan).isInstanceOf[HashAggregateExec]) ++ val strippedPlan = stripAQEPlan(hashAggPlan) ++ assert(strippedPlan.isInstanceOf[HashAggregateExec] || ++ strippedPlan.exists(_.isInstanceOf[CometHashAggregateExec])) + } + + // test case for ObjectHashAggregate and SortAggregate +@@ -755,7 +758,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { @@ -1532,7 +1545,7 @@ index ac710c32296..2854b433dd3 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 593bd7bb4ba..32af28b0238 100644 +index 593bd7bb4ba..b327d84d5cc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -26,9 +26,11 @@ import org.scalatest.time.SpanSugar._ @@ -1923,6 +1936,16 @@ index 593bd7bb4ba..32af28b0238 100644 } assert(shuffles2.size == 4) val smj2 = findTopLevelSortMergeJoin(adaptive2) +@@ -2703,7 +2744,8 @@ class AdaptiveQueryExecSuite + } + } + +-test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec") { ++test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) { + val emptyDf = spark.range(1).where("false") + val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) + val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index bd9c79e5b96..2ada8c28842 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala diff --git a/dev/diffs/3.5.8.diff b/dev/diffs/3.5.8.diff index 7155e9a9b4..fe87568b96 100644 --- a/dev/diffs/3.5.8.diff +++ b/dev/diffs/3.5.8.diff @@ -1,5 +1,5 @@ diff --git a/pom.xml b/pom.xml -index edd2ad57880..a47b7dec672 100644 +index edd2ad57880..45f8fd01538 100644 --- a/pom.xml +++ b/pom.xml @@ -152,6 +152,8 @@ @@ -241,19 +241,32 @@ index e5494726695..00937f025c2 100644 test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 6f3090d8908..c08a60fb0c2 100644 +index 6f3090d8908..a0e9309888d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.plans.logical.Expand +@@ -27,8 +27,9 @@ import org.apache.spark.{SparkException, SparkThrowable} + import org.apache.spark.sql.catalyst.plans.logical.Expand import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper ++import org.apache.spark.sql.comet.CometHashAggregateExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -793,7 +793,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -729,7 +730,9 @@ class DataFrameAggregateSuite extends QueryTest + case _ => false + }.isDefined) + } else { +- assert(stripAQEPlan(hashAggPlan).isInstanceOf[HashAggregateExec]) ++ val strippedPlan = stripAQEPlan(hashAggPlan) ++ assert(strippedPlan.isInstanceOf[HashAggregateExec] || ++ strippedPlan.exists(_.isInstanceOf[CometHashAggregateExec])) + } + + // test case for ObjectHashAggregate and SortAggregate +@@ -793,7 +796,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { @@ -1498,7 +1511,7 @@ index 5a413c77754..207b66e1d7b 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 2f8e401e743..dbcf3171946 100644 +index 2f8e401e743..7849c685b19 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,9 +27,11 @@ import org.scalatest.time.SpanSugar._ @@ -1904,7 +1917,17 @@ index 2f8e401e743..dbcf3171946 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2980,7 +3023,9 @@ class AdaptiveQueryExecSuite +@@ -2938,7 +2981,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec") { ++ test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) { + val emptyDf = spark.range(1).where("false") + val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) + val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) +@@ -2980,7 +3024,9 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) diff --git a/dev/diffs/4.0.2.diff b/dev/diffs/4.0.2.diff index 28cac668c8..d8e506b631 100644 --- a/dev/diffs/4.0.2.diff +++ b/dev/diffs/4.0.2.diff @@ -39,7 +39,7 @@ index 6c51bd4ff2e..e72ec1d26e2 100644 withSpark(sc) { sc => TestUtils.waitUntilExecutorsUp(sc, 2, 60000) diff --git a/pom.xml b/pom.xml -index 252cfdf9073..60cb9dcb7cf 100644 +index 252cfdf9073..50ec9d6314e 100644 --- a/pom.xml +++ b/pom.xml @@ -148,6 +148,8 @@ @@ -378,19 +378,32 @@ index 0f42502f1d9..e9ff802141f 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 9db406ff12f..245e4caa319 100644 +index 9db406ff12f..19b4bb8e39d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -@@ -30,7 +30,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId +@@ -29,8 +29,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS + import org.apache.spark.sql.errors.DataTypeErrors.toSQLId import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper ++import org.apache.spark.sql.comet.CometHashAggregateExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -855,7 +855,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -791,7 +792,9 @@ class DataFrameAggregateSuite extends QueryTest + case _ => false + }.isDefined) + } else { +- assert(stripAQEPlan(hashAggPlan).isInstanceOf[HashAggregateExec]) ++ val strippedPlan = stripAQEPlan(hashAggPlan) ++ assert(strippedPlan.isInstanceOf[HashAggregateExec] || ++ strippedPlan.exists(_.isInstanceOf[CometHashAggregateExec])) + } + + // test case for ObjectHashAggregate and SortAggregate +@@ -855,7 +858,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { @@ -2122,7 +2135,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 272be70f9fe..12daa1f5932 100644 +index 272be70f9fe..4a175083adf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -28,12 +28,14 @@ import org.apache.spark.SparkException @@ -2539,7 +2552,17 @@ index 272be70f9fe..12daa1f5932 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3129,7 +3173,8 @@ class AdaptiveQueryExecSuite +@@ -3039,7 +3083,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec") { ++ test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) { + val emptyDf = spark.range(1).where("false") + val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) + val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) +@@ -3129,7 +3174,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) diff --git a/dev/diffs/4.1.2.diff b/dev/diffs/4.1.2.diff index 946dae3949..9a1b30b010 100644 --- a/dev/diffs/4.1.2.diff +++ b/dev/diffs/4.1.2.diff @@ -39,7 +39,7 @@ index 6df8bc85b51..dabb75e2b75 100644 withSpark(sc) { sc => TestUtils.waitUntilExecutorsUp(sc, 2, 60000) diff --git a/pom.xml b/pom.xml -index dc201151999..d5c08f11ded 100644 +index dc201151999..20ee0e7482a 100644 --- a/pom.xml +++ b/pom.xml @@ -152,6 +152,8 @@ @@ -392,19 +392,32 @@ index 0d807aeae4d..6d7744e771b 100644 withTempView("t0", "t1", "t2") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index bfe15b33768..55c23a38ccc 100644 +index bfe15b33768..31f8ba4b876 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -@@ -31,7 +31,7 @@ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId +@@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS + import org.apache.spark.sql.errors.DataTypeErrors.toSQLId import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper ++import org.apache.spark.sql.comet.CometHashAggregateExec import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -856,7 +856,7 @@ class DataFrameAggregateSuite extends QueryTest +@@ -792,7 +793,9 @@ class DataFrameAggregateSuite extends QueryTest + case _ => false + }.isDefined) + } else { +- assert(stripAQEPlan(hashAggPlan).isInstanceOf[HashAggregateExec]) ++ val strippedPlan = stripAQEPlan(hashAggPlan) ++ assert(strippedPlan.isInstanceOf[HashAggregateExec] || ++ strippedPlan.exists(_.isInstanceOf[CometHashAggregateExec])) + } + + // test case for ObjectHashAggregate and SortAggregate +@@ -856,7 +859,7 @@ class DataFrameAggregateSuite extends QueryTest assert(objHashAggPlans.nonEmpty) val exchangePlans = collect(aggPlan) { @@ -2245,7 +2258,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 3e7d26f74bd..7e70e72fa3e 100644 +index 3e7d26f74bd..79232dc3664 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,12 +27,14 @@ import org.apache.spark.SparkException @@ -2662,7 +2675,17 @@ index 3e7d26f74bd..7e70e72fa3e 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3310,7 +3354,8 @@ class AdaptiveQueryExecSuite +@@ -3220,7 +3264,8 @@ class AdaptiveQueryExecSuite + } + } + +- test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec") { ++ test("SPARK-44040: Fix compute stats when AggregateExec nodes above QueryStageExec", ++ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4412")) { + val emptyDf = spark.range(1).where("false") + val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) + val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) +@@ -3310,7 +3355,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index 6c6544e93c..510efd3c9f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -193,9 +193,11 @@ object CometAverage extends CometAggregateExpressionSerde[Average] { object CometSum extends CometAggregateExpressionSerde[Sum] { override def supportsMixedPartialFinal(fn: Sum): Boolean = - // SUM's buffer matches Spark for Legacy/Ansi (decimal adds is_empty, also matching), but + // Decimal SUM is excluded: overflow detection (ANSI throw / Legacy null) does not survive a + // Spark-partial / Comet-final split, so the required ArithmeticException is never raised. // TRY-mode integer SUM carries a Comet-internal has_all_nulls column that Spark cannot read. - CometEvalModeUtil.fromSparkEvalMode(CometEvalModeUtil.sumEvalMode(fn)) != CometEvalMode.TRY + !fn.child.dataType.isInstanceOf[DecimalType] && + CometEvalModeUtil.fromSparkEvalMode(CometEvalModeUtil.sumEvalMode(fn)) != CometEvalMode.TRY override def getSupportLevel(expr: Sum): SupportLevel = if (AggSerde.sumDataTypeSupported(expr.dataType)) { diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 7e5eca4e99..e597e6e6e4 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -209,22 +209,6 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - test("mixed engine decimal sum: both split directions match Spark") { - val data = (0 until 100).map(i => (BigDecimal(i), i % 7)) - withParquetTable(data, "tbl") { - Seq( - CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key, - CometConf.COMET_ENABLE_PARTIAL_HASH_AGGREGATE.key).foreach { disabledConf => - withSQLConf( - disabledConf -> "false", - CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", - CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { - checkSparkAnswer("SELECT _2, SUM(_1) FROM tbl GROUP BY _2") - } - } - } - } - test("Aggregation without aggregate expressions should use correct result expressions") { withSQLConf( CometConf.COMET_ENABLED.key -> "true", diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 893405ff66..5135d8483c 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -323,6 +323,30 @@ class CometExecRuleSuite extends CometTestBase { } } + test("CometExecRule should not allow decimal SUM mixed execution") { + withTempView("test_data") { + createTestDataFrame.createOrReplaceTempView("test_data") + // Precision must be large enough (prec + 4 > 15) that Spark's own DecimalAggregates + // optimizer rule does not rewrite SUM to operate on the unscaled Long value, which would + // sidestep the decimal buffer path this test is meant to exercise. + val sparkPlan = + createSparkPlan( + spark, + "SELECT SUM(CAST(id AS DECIMAL(20, 2))) FROM test_data GROUP BY (id % 3)") + assert(countOperators(sparkPlan, classOf[HashAggregateExec]) == 2) + withSQLConf( + CometConf.COMET_ENABLE_FINAL_HASH_AGGREGATE.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val transformedPlan = applyCometExecRule(sparkPlan) + // Decimal SUM overflow detection (ANSI throw / Legacy null) does not survive a + // Spark-partial / Comet-final split, so mixed execution is unsafe and the partial + // must also fall back to Spark. + assert(countOperators(transformedPlan, classOf[HashAggregateExec]) == 2) + assert(countOperators(transformedPlan, classOf[CometHashAggregateExec]) == 0) + } + } + } + test("CometExecRule should allow AVG mixed Spark partial and Comet final") { withTempView("test_data") { createTestDataFrame.createOrReplaceTempView("test_data")