comphead commented on code in PR #6830:
URL: https://github.com/apache/datafusion-comet/pull/6830#discussion_r4265451115
##########
spark/src/test/resources/sql-tests/expressions/aggregate/corr.sql:
##########
@@ -37,3 +37,30 @@ INSERT INTO test_corr_nan VALUES (cast('NaN' as double),
cast('NaN' as double),
query tolerance=1e-6
SELECT grp, corr(x, y) FROM test_corr_nan GROUP BY grp ORDER BY grp
+
+-- A constant input makes Spark divide ck by sqrt(xMk * yMk) = 0
+-- https://github.com/apache/datafusion-comet/issues/6481
+statement
+CREATE TABLE test_corr_const(y double, x double, grp string) USING parquet
+
+statement
+INSERT INTO test_corr_const VALUES (0.0, 0.1, 'a'), (1.0, 0.1, 'a'), (2.0,
0.1, 'a'), (3.0, 0.1, 'a'), (4.0, 0.1, 'a'), (5.0, 0.1, 'a'), (1.0, 2.0, 'b'),
(NULL, 3.0, 'c')
+
+query
+SELECT corr(y, x), corr(x, y) FROM test_corr_const WHERE grp = 'a'
+
+query
+SELECT grp, corr(y, x), corr(x, y) FROM test_corr_const GROUP BY grp ORDER BY
grp
+
+statement
+SET spark.sql.ansi.enabled=true
+
+query expect_error(DIVIDE_BY_ZERO)
+SELECT corr(y, x) FROM test_corr_const WHERE grp = 'a'
+
+query expect_error(DIVIDE_BY_ZERO)
Review Comment:
Could these two use `expect_error_class(DIVIDE_BY_ZERO)` instead of
`expect_error(DIVIDE_BY_ZERO)`? `expect_error` only checks that both runs fail
with a message containing the pattern, so it still passes if `corr` falls back
to Spark, because Spark raises the same error. `expect_error_class` goes
through `checkSparkError`. That also checks that the operators are native and
that the exception class and SQLSTATE match Spark's. The non-ANSI `query` lines
above already pass the operator check on this table, so I expect the switch to
work (not run). With it, this file covers what the ANSI branch of the new Scala
test asserts for one partition.
##########
native/spark-expr/src/agg_funcs/correlation.rs:
##########
@@ -594,6 +640,199 @@ mod groups_tests {
}
}
+ fn assert_divide_by_zero<T: std::fmt::Debug>(result: Result<T>) {
+ let Err(DataFusionError::External(error)) = result else {
+ panic!("expected a structured Spark divide-by-zero error, got
{result:?}");
+ };
+ assert!(matches!(
+ error.downcast_ref::<SparkError>(),
+ Some(SparkError::DivideByZero)
+ ));
+ }
+
+ fn partial_state(input: &[ArrayRef]) -> Vec<ArrayRef> {
+ let mut partial =
+ CorrelationAccumulator::try_new(true,
zero_divisor(EvalMode::Legacy)).unwrap();
+ partial.update_batch(input).unwrap();
+ partial
+ .state()
+ .unwrap()
+ .iter()
+ .map(|s| s.to_array_of_size(1).unwrap())
+ .collect()
+ }
+
+ #[test]
+ fn correlation_zero_denominator_respects_eval_mode() {
+ // https://github.com/apache/datafusion-comet/issues/6481
+ // x is constant at 0.1, which binary floating point cannot represent
+ // exactly. Spark divides ck by sqrt(xMk * yMk) = 0, so its Divide
+ // returns null in LEGACY mode and raises DIVIDE_BY_ZERO in ANSI mode.
+ let y: Vec<f64> = (0..6).map(f64::from).collect();
+ let x = [0.1; 6];
+ let input = |range: std::ops::Range<usize>| -> Vec<ArrayRef> {
+ vec![
+ Arc::new(Float64Array::from(y[range.clone()].to_vec())),
+ Arc::new(Float64Array::from(x[range].to_vec())),
+ ]
+ };
+ for eval_mode in [EvalMode::Legacy, EvalMode::Ansi] {
+ for null_on_divide_by_zero in [false, true] {
Review Comment:
The `null_on_divide_by_zero` loop doubles this test without reaching a new
branch. Every accumulator here ends with `n` = 6, so evaluation never reaches
the `count == 1` branch, the only place the flag is read. The children run as
`StatsType::Population` and never read it either. Fixing it to one value keeps
the same coverage. `correlation_ansi_raises_only_on_zero_denominator` is the
test where the loop matters.
##########
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:
The non-ANSI branch and the `statsQuery` block repeat `statistical
aggregates merge fractional constants across partitions`. That test already
uses the same data (`y` from 0 to 5, `x` at 0.1, two files in both orders, with
and without `GROUP BY`). It asserts native Partial and Final aggregates, `corr`
= NULL and exact zeros for the other six aggregates. The single-partition cases
are in `corr.sql`. Could this test keep only the ANSI branch (error class,
SQLSTATE and the `checkQueryContextMatchesSpark` fragment)? Folding it as an
`ansi = true` pass into that existing test would also work.
--
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]
