sunchao commented on code in PR #5420:
URL: https://github.com/apache/datafusion-comet/pull/5420#discussion_r4316319464
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -2372,6 +2409,228 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("high-precision global decimal AVG and TRY_AVG fall back to Spark") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native") {
+ withTempDir { dir =>
+ Seq((1, "0.6"), (2, "0.6"))
+ .toDF("ord", "raw_v")
+ .selectExpr(
+ "ord",
+ "CAST(raw_v AS DECIMAL(27,27)) AS v27",
+ "CAST(raw_v AS DECIMAL(28,28)) AS v28",
+ "CAST(raw_v AS DECIMAL(38,38)) AS v38")
+ .write
+ .mode("overwrite")
+ .parquet(dir.toString)
+ withParquetTable(spark.read.parquet(dir.toString).coalesce(1),
"global_avg") {
+ val expectedAverage = new java.math.BigDecimal("0.6").setScale(38)
+ // Precision 27 stays native; 28 starts fallback. At 38, the
intermediate 1.2
+ // overflows in Comet although Spark can divide it and return 0.6.
+ for ((ansi, aggregates, nativeCount, expected) <- Seq(
+ (false, "AVG(v27)", 2, Row(expectedAverage.setScale(31))),
+ (false, "AVG(v28)", 0, Row(expectedAverage.setScale(32))),
+ (true, "AVG(v38)", 0, Row(expectedAverage)),
+ (true, "TRY_AVG(v38)", 0, Row(expectedAverage)),
+ (false, "AVG(v38), COUNT(DISTINCT ord)", 0, Row(expectedAverage,
2L)),
+ // Object aggregation overflows when materializing its partial
buffer.
+ (false, "AVG(v38), sort_array(collect_list(ord))", 0, Row(null,
Seq(1, 2))))) {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) {
+ val df = sql(s"SELECT $aggregates FROM global_avg")
+ checkAnswer(df, Seq(expected))
+ assert(getNumCometHashAggregate(df) == nativeCount)
+ }
+ }
+ // Merging individually valid partials can overflow too; keep the
shuffled final safe.
+ spark.read
+ .parquet(dir.toString)
+ .repartition(2, col("ord"))
+ .createOrReplaceTempView("global_avg")
+ val shuffled = sql("SELECT AVG(v38), TRY_AVG(v38) FROM global_avg")
+ checkAnswer(shuffled, Seq(Row(expectedAverage, expectedAverage)))
+ assert(getNumCometHashAggregate(shuffled) == 0)
+ }
+ }
+ }
+ }
+
+ test("high-precision DISTINCT decimal AVG preserves Spark shuffle
semantics") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ // Native wide-decimal hash routing is tracked separately in #6005.
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
+ withTempDir { dir =>
+ Seq((1, "0.6"), (2, "0.6"), (3, "0.2"), (4, "0.3"))
+ .toDF("ord", "raw_v")
+ .selectExpr("ord", "CAST(raw_v AS DECIMAL(38,38)) AS v")
+ .write
+ .mode("overwrite")
+ .parquet(dir.toString)
+ withParquetTable(
+ spark.read.parquet(dir.toString).repartition(2, col("ord")),
+ "distinct_avg") {
+ // Spark hashes all three distinct values to one partition.
Materializing its
+ // partial sum overflows, unlike the adjacent scalar stages
exercised above.
+ val df = sql("SELECT AVG(DISTINCT v) FROM distinct_avg")
+ checkSparkAnswer(df)
+ checkAnswer(df, Seq(Row(null)))
+ }
+ }
+ }
+ }
+
+ test("grouped decimal AVG widens division and reports final result
overflow") {
Review Comment:
Addressed in `8c1658a9a4f4459b470ebec7848c73f3eed484be`.
Moved the grouped wide-division, TRY_AVG, final-result overflow, and focused
ObjectHash fallback cases into expressions/aggregate/avg.sql and avg_ansi.sql.
The ANSI fixture uses expect_error(NUMERIC_VALUE_OUT_OF_RANGE); the object
cases use expect_fallback with main's current operator reason. That matcher
uses the same operator-tag helper as the Scala assertion. Partition-sensitive
global/shuffle tests and main's broader object-buffer cases remain in Scala.
Default-native precision controls also replace main's temporary
rescaling-fallback test, since this PR fixes that overflow. Both SQL fixtures
pass on Spark 4.0.4 as part of the 326-test run. Spark 4.1 and cross-version
hosted checks are pending.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
