andygrove commented on code in PR #5420:
URL: https://github.com/apache/datafusion-comet/pull/5420#discussion_r4112605650
##########
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:
Could group 1 get a third row, `(1, 0, "-0.6")`? As written, the native leg
passes even with the `partial_sums.is_null(idx)` check removed from
`AvgDecimalGroupsAccumulator::merge_batch`. Native shuffle keeps the payload of
the null sum, which here is the overflowed 1.2, so the final merge overflows
again without needing the null bit. Only the JVM leg's ANSI case catches that
mutation. With the extra row the payload comes back to 0.6, and the native leg
fails under the same mutation. Unmutated, the test still passes on both legs
with the same expected results, because Spark's `UnsafeRow` buffer has already
latched group 1 to null by then.
##########
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:
Grouped max-precision AVG under `ObjectHashAggregateExec` still diverges,
which is #5509. With main merged in, `SELECT g, AVG(v),
sort_array(collect_list(ord)) FROM t GROUP BY g` over `DECIMAL(38,38)` values
`0.6, 0.6, -0.4` returns NULL natively, or `ARITHMETIC_OVERFLOW` under ANSI,
where Spark returns `0.26666666666666666666666666666666666667`. The `SUM`
version of the same query falls back and matches, because #6041 declines it in
`CometObjectHashAggregateExec.getSupportLevel`, and its comment there already
says decimal AVG has the same gap. Extending this check to
`ObjectHashAggregateExec` with grouping keys fixed both ANSI modes for me, and
all of `CometAggregateSuite` still passed. A `hasMaxPrecisionDecimalAvg` next
to `hasMaxPrecisionDecimalSum` might be the tidiest way to write it. Could this
PR decline that case too and close #5509? If you'd rather keep it separate,
could #5509 go in the known result-value divergences list in
`compatibility/index.md`, since it returns NUL
L without an error in legacy mode?
--
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]
