comphead commented on code in PR #5420:
URL: https://github.com/apache/datafusion-comet/pull/5420#discussion_r4277251013
##########
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:
Most of this test looks expressible as SQL fixtures. The `k = 1` and
`TRY_AVG` results could be plain `query` records in
`expressions/aggregate/avg.sql`, which already asserts native execution in
place of the `CometHashAggregateExec` size checks. The `k = 2` ANSI case could
be a `-- Config: spark.sql.ansi.enabled=true` file with `query
expect_error(NUMERIC_VALUE_OUT_OF_RANGE)`. `decimal AVG with object hash
aggregate falls back at maximum precision` (line 2540) could be `query
expect_fallback(Grouped decimal AVG at maximum precision)`, which other
fixtures in `expressions/aggregate` already use. I have not checked that
`expect_fallback` sees operator-level reasons. The global fallback and
shuffle-writer tests depend on partition layout and should stay in Scala.
##########
native/spark-expr/src/agg_funcs/avg_decimal.rs:
##########
@@ -671,31 +609,281 @@ impl GroupsAccumulator for AvgDecimalGroupsAccumulator {
}
/// Returns the `sum`/`count` as a i128 Decimal128 with
-/// target_scale and target_precision and return None if overflows.
+/// target_scale and target_precision, returning the rounded wide value on
overflow.
///
/// * sum: The total sum value stored as Decimal128 with sum_scale
/// * count: total count, stored as a i128 (*NOT* a Decimal128 value)
/// * target_min: The minimum output value possible to represent with the
target precision
/// * target_max: The maximum output value possible to represent with the
target precision
/// * scaler: scale factor for avg
#[inline(always)]
-fn avg(sum: i128, count: i128, target_min: i128, target_max: i128, scaler:
i128) -> Option<i128> {
- if let Some(value) = sum.checked_mul(scaler) {
- // `sum / count` with ROUND_HALF_UP
+fn avg(
+ sum: i128,
+ count: i128,
+ target_min: i128,
+ target_max: i128,
+ scaler: i128,
+) -> std::result::Result<i128, i256> {
+ // The scaled numerator can exceed i128 while both the sum and final
average fit.
+ // Use the inexpensive path when it fits, and widen only that intermediate
otherwise.
+ let rounded = if let Some(value) = sum.checked_mul(scaler) {
let (div, rem) = value.div_rem(&count);
let half = div_ceil(count, 2);
- let half_neg = half.neg_wrapping();
- let new_value = match value >= 0 {
+ let rounded = match value >= 0 {
true if rem >= half => div.add_wrapping(1),
- false if rem <= half_neg => div.sub_wrapping(1),
+ false if rem <= half.neg_wrapping() => div.sub_wrapping(1),
_ => div,
};
- if new_value >= target_min && new_value <= target_max {
- Some(new_value)
+ i256::from_i128(rounded)
+ } else {
+ // A Decimal128 numerator times the Decimal128 scale factor fits
within i256.
+ let value = i256::from_i128(sum) * i256::from_i128(scaler);
+ let divisor = i256::from_i128(count);
+ let (div, rem) = (value / divisor, value % divisor);
+ let half = i256::from_i128(div_ceil(count, 2));
+ if rem >= half {
+ div + i256::ONE
+ } else if rem <= -half {
+ div - i256::ONE
} else {
- None
+ div
}
+ };
+ if rounded >= i256::from_i128(target_min) && rounded <=
i256::from_i128(target_max) {
+ Ok(rounded.to_i128().expect("bounded Decimal128 average"))
} else {
- None
+ Err(rounded)
+ }
+}
+
+fn avg_result_overflow_error(value: i256, precision: u8, scale: i8) ->
crate::SparkError {
Review Comment:
The value in this message differs from Spark's.
`DecimalDivideWithOverflowCheck` calls `toPrecision` on the raw quotient from
`Decimal./`, which divides at scale `MAX_SCALE + 1` (39) with
`RoundingMode.DOWN`. On failure `toPrecision` passes that unrounded quotient to
`cannotChangeDecimalPrecisionError`, which renders `value.toPlainString`. For
the `DECIMAL(38,18)` row `100000000000000000` in `grouped decimal AVG widens
division and reports final result overflow`, I expect Spark to print the value
with 39 fractional digits, where this formats the rounded result at the result
scale. The test only checks for `NUMERIC_VALUE_OUT_OF_RANGE`, so it would not
catch the difference. I derived this from the Spark 4.0 source and did not run
it. Could the error format the quotient truncated at scale 39, or could a short
note go next to the known-divergence line for grouped `AVG`?
--
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]
