kazuyukitanimura commented on code in PR #6830:
URL: https://github.com/apache/datafusion-comet/pull/6830#discussion_r4272950219
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -3357,6 +3370,90 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("statistical aggregates corr preserves its captured evaluation mode") {
+ for (ansi <- Seq(false, true)) {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) {
+ val expr = Corr(Literal(1.0), Literal(2.0))
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> (!ansi).toString) {
+ val serialized = CometCorr.convert(
+ expr.toAggregateExpression(),
+ expr,
+ Seq.empty,
+ binding = false,
+ conf = SQLConf.get)
+ assert(serialized.isDefined)
+ val expected =
+ if (ansi) ExprOuterClass.EvalMode.ANSI else
ExprOuterClass.EvalMode.LEGACY
+ assert(serialized.get.getCorrelation.getEvalMode == expected)
+ }
+ }
+ }
+ }
+
+ test("statistical aggregates corr zero denominator respects ANSI mode") {
+ // https://github.com/apache/datafusion-comet/issues/6481
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1048576",
+ SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "1048576",
+ "spark.sql.files.minPartitionNum" -> "1",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "native") {
+ for (numPartitions <- Seq(1, 2)) {
+ withTempPath { path =>
+ // x is constant at 0.1, which binary floating point cannot
represent exactly.
+ // Two files also exercise merging the partials in the final
aggregate.
+ val files = if (numPartitions == 1) Seq(0 until 6) else Seq(0 until
3, 3 until 6)
+ files.foreach { ys =>
+ ys.map(y => (0, y.toDouble, 0.1))
+ .toDF("g", "y", "x")
+ .coalesce(1)
+ .write
+ .mode("append")
+ .parquet(path.getCanonicalPath)
+ }
+ withParquetTable(path.getCanonicalPath, "corr_constant") {
+ assert(spark.table("corr_constant").rdd.getNumPartitions ==
numPartitions)
+ for {
+ ansi <- Seq(false, true)
+ groupBy <- Seq("", " GROUP BY g")
+ } {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) {
+ for (args <- Seq("y, x", "x, y", "x, x")) {
+ val query = s"SELECT corr($args) FROM corr_constant" +
groupBy
+ val df = sql(query)
+ val aggregates =
stripAQEPlan(df.queryExecution.executedPlan).collect {
+ case a: CometHashAggregateExec => a
+ }
+ assert(aggregates.size == 2)
+ assert(aggregates.exists(_.modes.contains(Partial)))
+ assert(aggregates.exists(_.modes.contains(Final)))
+ // Spark divides ck by sqrt(xMk * yMk) = 0.
+ if (ansi) {
+ val error = checkSparkError(df, "DIVIDE_BY_ZERO")
+ assert(error.getSqlState == "22012")
+ checkQueryContextMatchesSpark(error, query)
+ } else {
+ checkSparkAnswerAndOperator(df)
+ checkCometAnswer(df, Seq(Row(null)))
+ }
+ }
+
+ // The other statistical aggregates have no zero divisor here.
+ val statsQuery = "SELECT covar_pop(y, x), covar_samp(y, x),
var_pop(x), " +
+ "var_samp(x), stddev_pop(x), stddev_samp(x), corr(y, y) " +
+ "FROM corr_constant" + groupBy
+ checkSparkAnswerAndOperator(statsQuery)
+ checkAnswer(sql(statsQuery), Seq(Row(0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 1.0)))
Review Comment:
Done in 0cc2c7f75. The test keeps only the ANSI branch: error class,
SQLSTATE and the checkQueryContextMatchesSpark fragment. It has a comment
pointing to the fractional-constants test for the LEGACY results.
--
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]
