comphead commented on code in PR #6069:
URL: https://github.com/apache/datafusion-comet/pull/6069#discussion_r4354855411
##########
native/spark-expr/src/agg_funcs/sum_int.rs:
##########
@@ -914,6 +1085,199 @@ mod tests {
use arrow::array::Int64Array;
use datafusion::logical_expr::{EmitTo, GroupsAccumulator};
+ #[test]
+ fn sliding_sum_matches_ordered_checked_addition() {
+ let choices = [
+ None,
+ Some(0),
+ Some(1),
+ Some(-1),
+ Some(i64::MIN),
+ Some(i64::MAX),
+ ];
+ // Exhaust every 5-row sequence of nulls, boundaries, and cancellation.
+ for mut seed in 0..choices.len().pow(5) {
+ let values: Vec<_> = (0..5)
+ .map(|_| {
+ let value = choices[seed % choices.len()];
+ seed /= choices.len();
+ value
+ })
+ .collect();
+ let array: ArrayRef = Arc::new(Int64Array::from(values.clone()));
+ for width in [1, 2, 3, 5] {
+ for mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] {
+ let mut acc = SlidingSumIntegerAccumulator::new(mode);
+ for end in 1..=values.len() {
+ // Deliberately update first, as DataFusion does. The
+ // transient union may overflow while both frames fit.
+ acc.update_batch(&[array.slice(end - 1, 1)]).unwrap();
+ let start = end.saturating_sub(width);
+ if end > width {
+ acc.retract_batch(&[array.slice(start - 1,
1)]).unwrap();
+ }
+ let mut expected = None;
+ let mut overflow = false;
+ for &v in values[start..end].iter().flatten() {
+ let sum = expected.unwrap_or(0_i64);
+ overflow |= sum.checked_add(v).is_none();
+ expected = Some(sum.wrapping_add(v));
+ }
+ let actual = acc.evaluate();
+ if overflow && mode == EvalMode::Ansi {
+ assert!(actual.is_err(), "{values:?},
{start}..{end}");
+ } else {
+ if overflow && mode == EvalMode::Try {
+ expected = None;
+ }
+ assert_eq!(
+ actual.unwrap(),
+ ScalarValue::Int64(expected),
+ "{values:?}, {start}..{end}, {mode:?}"
+ );
+ }
+ }
+
acc.retract_batch(&[array.slice(values.len().saturating_sub(width), width)])
+ .unwrap();
+ assert_eq!(acc.evaluate().unwrap(),
ScalarValue::Int64(None));
+ assert!(acc.minima.is_empty() && acc.maxima.is_empty());
+ // Empty -> non-empty, with absolute prefix counters
retained.
+ acc.update_batch(&[Arc::new(Int64Array::from(vec![7]))])
+ .unwrap();
+ assert_eq!(acc.evaluate().unwrap(),
ScalarValue::Int64(Some(7)));
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn sliding_sum_matches_advancing_frames() {
+ // Cross the bounds by accumulating ordinary values as well as MIN/MAX.
+ let half = 1_i64 << 62;
+ let choices = [
+ None,
+ Some(half),
+ Some(-half),
+ Some(i64::MAX),
+ Some(i64::MIN),
+ ];
+ for mut seed in 0..choices.len().pow(5) {
+ let values: Vec<_> = (0..5)
+ .map(|_| {
+ let value = choices[seed % choices.len()];
+ seed /= choices.len();
+ value
+ })
+ .collect();
+ let array: ArrayRef = Arc::new(Int64Array::from(values.clone()));
+ // Growing prefixes, shrinking suffixes, batch updates/retractions,
+ // and disjoint frames.
+ for frames in [
+ vec![0..1, 0..2, 0..3, 0..4, 0..5],
+ vec![0..5, 1..5, 2..5, 3..5, 4..5, 5..5],
+ vec![0..2, 1..4, 2..5, 4..5, 5..5],
+ vec![0..1, 3..4, 4..5, 5..5],
+ ] {
+ for mode in [EvalMode::Ansi, EvalMode::Try] {
+ let mut acc = SlidingSumIntegerAccumulator::new(mode);
+ let mut previous = 0..0;
+ for frame in &frames {
+ acc.update_batch(&[array.slice(previous.end, frame.end
- previous.end)])
+ .unwrap();
+ acc.retract_batch(&[
+ array.slice(previous.start, frame.start -
previous.start)
+ ])
+ .unwrap();
+ let expected = values[frame.clone()]
+ .iter()
+ .flatten()
+ .try_fold(None, |sum, &value| {
+
sum.unwrap_or(0_i64).checked_add(value).map(Some)
+ });
+ match (expected, mode) {
+ (None, EvalMode::Ansi) =>
assert!(acc.evaluate().is_err()),
+ _ => assert_eq!(
+ acc.evaluate().unwrap(),
+ ScalarValue::Int64(expected.flatten()),
+ "{values:?}, {frame:?}, {mode:?}"
+ ),
+ }
+ previous = frame.clone();
+ }
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn sliding_sum_ordinary_suffixes_need_no_queue_allocation() {
+ for value in [Some(8), Some(-8), Some(0), None] {
+ let array: ArrayRef = Arc::new(Int64Array::from(vec![value;
65536]));
+ let mut acc = SlidingSumIntegerAccumulator::new(EvalMode::Try);
+ acc.update_batch(&[Arc::clone(&array)]).unwrap();
+ for start in (0..65536).step_by(1024) {
+ assert_eq!(
+ acc.evaluate().unwrap(),
+ ScalarValue::Int64(value.map(|v| v * (65536 - start) as
i64))
+ );
+ assert_eq!(acc.size(), std::mem::size_of_val(&acc));
+ acc.retract_batch(&[array.slice(start, 1024)]).unwrap();
+ }
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(None));
+ assert_eq!(acc.size(), std::mem::size_of_val(&acc));
+ }
+ }
+
+ #[test]
+ fn sliding_sum_exact_bounds_need_no_queue_allocation() {
+ for value in [i64::MIN, i64::MAX] {
+ let array: ArrayRef = Arc::new(Int64Array::from(vec![Some(value),
None, Some(0)]));
+ let mut acc = SlidingSumIntegerAccumulator::new(EvalMode::Ansi);
+ acc.update_batch(&[Arc::clone(&array)]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(),
ScalarValue::Int64(Some(value)));
+ assert_eq!(acc.size(), std::mem::size_of_val(&acc));
+ acc.retract_batch(&[array.slice(0, 1)]).unwrap();
+ assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(Some(0)));
+ assert_eq!(acc.size(), std::mem::size_of_val(&acc));
+ }
+ }
+
+ #[test]
+ fn sliding_sum_checks_intermediate_overflow_and_recovers() {
Review Comment:
As far as I can tell `sliding_sum_matches_ordered_checked_addition` already
covers this. Its sequences include `[i64::MAX, 1, -1, NULL, NULL]` and
`[i64::MIN, -1, 1, NULL, NULL]` at width 3 in TRY mode, which go through the
overflowing frame and recover at the next one, with the update before retract
order. Could this test be dropped?
##########
native/core/src/execution/planner/window.rs:
##########
@@ -591,7 +590,8 @@ impl PhysicalPlanner {
Ok((udaf(func), vec![child]))
}
DataType::Int8 | DataType::Int16 | DataType::Int32 |
DataType::Int64
- if is_ever_expanding =>
+ if is_ever_expanding
+ || from_protobuf_eval_mode(expr.eval_mode)? !=
EvalMode::Legacy =>
{
let eval_mode =
from_protobuf_eval_mode(expr.eval_mode)?;
let func = SumInteger::try_new(arrow_type, eval_mode)?;
Review Comment:
`from_protobuf_eval_mode(expr.eval_mode)?` runs in the guard and again in
the body, and the Decimal arm above decodes it too. Could it be decoded once
above `match arrow_type`? I have not checked whether that would make an unknown
eval-mode value an error for float and other sums, which the serde probably
always sets.
##########
spark/src/test/scala/org/apache/spark/sql/benchmark/CometSlidingSumBenchmark.scala:
##########
@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.comet.CometWindowExec
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+
+/**
+ * Compare DataFusion's legacy SUM, checked native SUM/TRY_SUM, and Spark
windows. Arguments: row
+ * count (default 65536), frames (all, bounded, or suffix). For example,
`4000000 suffix` measures
+ * a whole-partition initial frame.
+ */
+object CometSlidingSumBenchmark extends CometBenchmarkBase {
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ val rows = mainArgs.headOption.map(_.toLong).getOrElse(65536L)
+ val frames = mainArgs.lift(1).getOrElse("all")
+ require(rows > 0 && Set("all", "bounded", "suffix").contains(frames))
+ val boundedFrames =
+ Seq(16, 1024).map(width => (s"rows=$width", s"${width - 1} PRECEDING AND
CURRENT ROW"))
+ val suffixFrame = Seq(("suffix", "CURRENT ROW AND UNBOUNDED FOLLOWING"))
+ val selectedFrames = frames match {
+ case "bounded" => boundedFrames
+ case "suffix" => suffixFrame
+ case _ => boundedFrames ++ suffixFrame
+ }
+ val nativeCases = Seq(
+ ("DataFusion SUM (legacy)", "sum", false, true),
+ ("Comet SUM (ANSI)", "sum", true, true),
+ ("Comet TRY_SUM", "try_sum", true, true))
+ val sparkCases =
+ Seq(("Spark SUM (ANSI)", "sum", true, false), ("Spark TRY_SUM",
"try_sum", true, false))
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ "spark.comet.operator.WindowExec.allowIncompatible" -> "true") {
Review Comment:
`CometWindowExec` has no `getSupportLevel` override, so it is always
`Compatible` and `spark.comet.operator.WindowExec.allowIncompatible` has no
effect. andygrove asked for the same removal on the SQL fixture and it was
dropped there. Could it go here too so it does not get copied into the next
benchmark?
##########
spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala:
##########
@@ -607,17 +611,83 @@ class CometWindowExecSuite extends CometTestBase {
s"SELECT $function OVER() FROM t1",
s"SELECT $function OVER(order by _2) FROM t1",
s"SELECT $function OVER(order by _2 desc) FROM t1",
- s"SELECT $function OVER(partition by _2 order by _2) FROM t1",
+ s"SELECT $function OVER(partition by _2 order by _2) FROM t1")
+ queries.foreach { query =>
+ checkSparkAnswerAndOperator(query)
+ }
+
+ val slidingQueries = Seq(
s"SELECT $function OVER(rows between 1 preceding and 1
following) FROM t1",
s"SELECT $function OVER(order by _2 rows between 1 preceding and
current row) FROM t1",
s"SELECT $function OVER(order by _2 rows between current row and
1 following) FROM t1")
- queries.foreach { query =>
+ slidingQueries.foreach { query =>
checkSparkAnswerAndOperator(query)
}
Review Comment:
Now that sliding sums run natively in both modes, `queries` and
`slidingQueries` go through the same `checkSparkAnswerAndOperator` loop, so the
split no longer does anything. On main these three queries sit in the single
`queries` list. Could they go back there so the only change to this test is the
ANSI dimension?
--
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]