sunchao commented on code in PR #5420:
URL: https://github.com/apache/datafusion-comet/pull/5420#discussion_r4178657719
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -2205,6 +2243,140 @@ 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 preserves overflow across JVM and native shuffle")
{
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true") {
+ withTempDir { dir =>
+ Seq((1, 0, "0.6"), (1, 0, "0.6"), (2, 2, "0.1"), (2, 2, "0.2"), (3, 4,
null))
Review Comment:
Addressed in `f99c3c8406fbb8b882734c34a353defbda342676`: group 1 now
includes `(1, 0, "-0.6")`. The native payload returns within range while its
null marker must stay latched, so the native and JVM shuffle legs both depend
on preserving that marker. The full Spark 4.1 reactor test compilation passes;
runtime CI for this head is pending.
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -1879,6 +1879,28 @@ trait CometBaseAggregate {
case _ => false
})
+ protected def aggregateSupportLevel(op: BaseAggregateExec): SupportLevel = {
+ val unsupportedAverage = op.groupingExpressions.isEmpty &&
Review Comment:
Addressed in `f99c3c8406fbb8b882734c34a353defbda342676` with
`hasMaxPrecisionDecimalAvg` and the grouped ObjectHashAggregate guard alongside
the SUM guard. Added an explicit regression for `0.6, 0.6, -0.4` plus
`collect_list`, checking Spark's `0.26666666666666666666666666666666666667`
result and no native aggregate in both ANSI modes. Updated the compatibility
documentation and linked #5509 for closure. Native core/test and full Spark 4.1
reactor test compilation pass; runtime CI is 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]
