diff --git a/docs/source/contributor-guide/expression-audits/window_funcs.md b/docs/source/contributor-guide/expression-audits/window_funcs.md index c692a42a555..7b6111f1fa9 100644 --- a/docs/source/contributor-guide/expression-audits/window_funcs.md +++ b/docs/source/contributor-guide/expression-audits/window_funcs.md @@ -35,6 +35,7 @@ ## sum (window) -- 4.1.1, audited 2026-06-25: `SUM()` over a sliding window frame (lower bound not `UNBOUNDED PRECEDING`) falls back to Spark. The native sliding path would use DataFusion's built-in `sum`, which wraps on overflow rather than returning Spark's NULL, and overflow cannot be detected at plan time, so the whole sliding decimal `SUM` case falls back. Ever-expanding frames use Comet's overflow-aware `SumDecimal` UDAF and run natively, matching Spark including overflow-to-NULL. Bigint sliding `SUM` overflow matches Spark (both wrap) and stays native. +- 4.1.1, audited 2026-06-25: `SUM()` over a sliding window frame (lower bound not `UNBOUNDED PRECEDING`) falls back to Spark. The native sliding path would use DataFusion's built-in `sum`, which wraps on overflow rather than returning Spark's NULL, and overflow cannot be detected at plan time, so the whole sliding decimal `SUM` case falls back. Ever-expanding frames use Comet's overflow-aware `SumDecimal` UDAF and run natively, matching Spark including overflow-to-NULL. Bigint sliding `SUM` overflow matches Spark in legacy mode (both wrap) and stays native. +- 3.5.9 and 4.1.3, audited 2026-09-30: Integral sliding `SUM` uses a native retractable accumulator in ANSI and TRY mode. It checks ordered prefix sums, so intermediate overflow throws or returns NULL even when later values cancel it, and results recover once the offending values leave the frame. Legacy sliding sums retain the DataFusion path. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 2f992c4843f..d7e511ee004 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -2480,10 +2480,9 @@ impl PhysicalPlanner { // `PERCENT_RANK` / `CUME_DIST` / `NTILE` // (`!uses_bounded_memory()` — "Can not execute X in a streaming // fashion") and keeps the Spark-compatible Comet UDAFs - // (`SumDecimal` / `SumInteger` / `AvgDecimal` / `Avg`) on the - // non-streaming path since they don't implement `retract_batch`. - // Because `process_agg_func` already picks DataFusion's - // retract-capable built-ins for sliding aggregate frames, + // (`SumDecimal` / `AvgDecimal` / `Avg`) on the non-streaming path + // where needed. `process_agg_func` picks retract-capable UDAFs + // for sliding aggregate frames; // ever-expanding aggregate frames (all that route to // `BoundedWindowAggExec` as `PlainAggregateWindowExpr`) never // trigger a retract call. @@ -3321,9 +3320,9 @@ impl PhysicalPlanner { // DataFusion uses `PlainAggregateWindowExpr` which does not call // `retract_batch`, so we can safely use Comet's Spark-compatible // UDAFs (SumDecimal/SumInteger/AvgDecimal/Avg). Otherwise it uses - // `SlidingAggregateWindowExpr` which requires retract — Comet's UDAFs - // don't implement it, so the caller must fall back to DataFusion's - // built-ins (which do). + // `SlidingAggregateWindowExpr` which requires retract. process_agg_func + // selects Comet's checked integer accumulator or a DataFusion built-in + // that supports retraction. let is_ever_expanding = spark_expr .spec .as_ref() @@ -3539,9 +3538,9 @@ impl PhysicalPlanner { Some(AggExprStruct::Sum(expr)) => { // For ever-expanding frames, use Comet's Spark-compatible Sum UDAFs // (SumDecimal / SumInteger) which enforce Spark overflow semantics. - // For sliding frames, those UDAFs can't be used (no retract_batch), - // so delegate to DataFusion's built-in `sum`, which supports retract - // but doesn't enforce Spark's decimal precision overflow-to-NULL. + // Checked integral sliding sums also use Comet's retractable + // accumulator. Legacy sliding sums keep DataFusion's fast path; + // decimal sliding frames are rejected by the JVM planner. let child = self.create_expr(expr.child.as_ref().unwrap(), Arc::clone(&schema))?; let arrow_type = to_arrow_datatype(expr.datatype.as_ref().unwrap()); match arrow_type { @@ -3556,7 +3555,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)?; diff --git a/native/spark-expr/src/agg_funcs/sum_int.rs b/native/spark-expr/src/agg_funcs/sum_int.rs index 56a8c122c41..b7c8dcf411c 100644 --- a/native/spark-expr/src/agg_funcs/sum_int.rs +++ b/native/spark-expr/src/agg_funcs/sum_int.rs @@ -29,6 +29,7 @@ use datafusion::logical_expr::Volatility::Immutable; use datafusion::logical_expr::{ Accumulator, AggregateUDFImpl, EmitTo, GroupsAccumulator, ReversedUDAF, Signature, }; +use std::collections::VecDeque; use std::sync::Arc; #[derive(Debug, PartialEq, Eq, Hash)] @@ -72,6 +73,10 @@ impl AggregateUDFImpl for SumInteger { } } + fn create_sliding_accumulator(&self, _args: AccumulatorArgs) -> DFResult> { + Ok(Box::new(SlidingSumIntegerAccumulator::new(self.eval_mode))) + } + fn state_fields(&self, _args: StateFieldsArgs) -> DFResult> { if self.eval_mode == EvalMode::Try { Ok(vec![ @@ -99,7 +104,173 @@ impl AggregateUDFImpl for SumInteger { } fn reverse_expr(&self) -> ReversedUDAF { - ReversedUDAF::Identical + // Checked addition depends on input order: [MAX, 1, -1] overflows, + // while [-1, 1, MAX] does not. Reversing a frame must not change that. + if self.eval_mode == EvalMode::Legacy { + ReversedUDAF::Identical + } else { + ReversedUDAF::NotSupported + } + } +} + +/// Spark recomputes each sliding frame in input order. Checking only its final +/// sum misses an intermediate overflow followed by cancellation. Track the min +/// and max prefix sums instead: subtracting the prefix before the frame gives +/// every partial sum Spark would visit. Monotonic queues retain just the extrema +/// candidates; each entry is pushed/popped once, for amortized O(1) work per row. +/// +/// A partial sum is bounded above by the positive sum and below by the negative +/// sum of retained values, including update-before-retract overlap. +/// Only enqueue a prefix when that bound exceeds the corresponding i64 limit. +/// A skipped prefix is safe for every later frame too: moving the left boundary +/// forward only removes values from those bounds. +/// Ordinary values therefore need no queue allocation. +/// +/// Worst-case queue space is still linear in the largest non-null frame, +/// including update-before-retract overlap. A suffix frame can span the entire +/// partition. Each entry occupies 32 bytes on a 64-bit host, and VecDeque retains +/// its grown capacity after retraction. size() reports this capacity, but +/// DataFusion 55.1's window operators do not account for accumulator sizes. +/// +/// i128 holds the exact sum of at most usize::MAX i64 values on a 64-bit host. +/// Nulls need no entries because retraction receives the original input values. +#[derive(Debug)] +struct SlidingSumIntegerAccumulator { + eval_mode: EvalMode, + end: i128, + start: i128, + positive_sum: i128, + added: usize, + removed: usize, + minima: VecDeque<(usize, i128)>, + maxima: VecDeque<(usize, i128)>, +} + +impl SlidingSumIntegerAccumulator { + fn new(eval_mode: EvalMode) -> Self { + Self { + eval_mode, + end: 0, + start: 0, + positive_sum: 0, + added: 0, + removed: 0, + minima: VecDeque::new(), + maxima: VecDeque::new(), + } + } + + fn add(&mut self, value: i64) { + self.end += i128::from(value); + self.positive_sum += i128::from(value.max(0)); + self.added += 1; + while self.minima.back().is_some_and(|&(_, v)| v >= self.end) { + self.minima.pop_back(); + } + while self.maxima.back().is_some_and(|&(_, v)| v <= self.end) { + self.maxima.pop_back(); + } + // Keep the dominance pops above even for a skipped prefix. It also + // bounds any older entries it dominates in every future frame. + let negative_sum = self.end - self.start - self.positive_sum; + if negative_sum < i128::from(i64::MIN) { + self.minima.push_back((self.added, self.end)); + } + if self.positive_sum > i128::from(i64::MAX) { + self.maxima.push_back((self.added, self.end)); + } + } + + fn remove(&mut self, value: i64) { + self.start += i128::from(value); + self.positive_sum -= i128::from(value.max(0)); + self.removed += 1; + } + + fn visit(&mut self, values: &ArrayRef, retract: bool) -> DFResult<()> { + fn visit( + acc: &mut SlidingSumIntegerAccumulator, + values: &PrimitiveArray, + retract: bool, + ) -> DFResult<()> { + for value in values.iter().flatten() { + let value = value.to_i64().ok_or_else(|| { + DataFusionError::Internal("Expected an integral SUM input".into()) + })?; + if retract { + acc.remove(value); + } else { + acc.add(value); + } + } + Ok(()) + } + match values.data_type() { + DataType::Int8 => visit(self, as_primitive_array::(values), retract), + DataType::Int16 => visit(self, as_primitive_array::(values), retract), + DataType::Int32 => visit(self, as_primitive_array::(values), retract), + DataType::Int64 => visit(self, as_primitive_array::(values), retract), + other => not_impl_err!("Sliding integer SUM does not support {other}"), + } + } +} + +impl Accumulator for SlidingSumIntegerAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> { + // DataFusion updates before retracting. Defer overflow checks until + // evaluate(), when the accumulator represents the actual output frame. + self.visit(&values[0], false) + } + + fn retract_batch(&mut self, values: &[ArrayRef]) -> DFResult<()> { + self.visit(&values[0], true)?; + while self.minima.front().is_some_and(|&(i, _)| i <= self.removed) { + self.minima.pop_front(); + } + while self.maxima.front().is_some_and(|&(i, _)| i <= self.removed) { + self.maxima.pop_front(); + } + Ok(()) + } + + fn supports_retract_batch(&self) -> bool { + true + } + + fn evaluate(&mut self) -> DFResult { + if self.added == self.removed { + return Ok(ScalarValue::Int64(None)); + } + let overflow = self + .minima + .front() + .is_some_and(|&(_, v)| v - self.start < i128::from(i64::MIN)) + || self + .maxima + .front() + .is_some_and(|&(_, v)| v - self.start > i128::from(i64::MAX)); + if overflow && self.eval_mode != EvalMode::Legacy { + return match self.eval_mode { + EvalMode::Ansi => Err(arithmetic_overflow_error("integer").into()), + _ => Ok(ScalarValue::Int64(None)), + }; + } + Ok(ScalarValue::Int64(Some((self.end - self.start) as i64))) + } + + fn size(&self) -> usize { + std::mem::size_of_val(self) + + (self.minima.capacity() + self.maxima.capacity()) + * std::mem::size_of::<(usize, i128)>() + } + + fn state(&mut self) -> DFResult> { + not_impl_err!("Sliding integer SUM is a window accumulator") + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> DFResult<()> { + not_impl_err!("Sliding integer SUM does not merge partial aggregates") } } @@ -903,6 +1074,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() { + for (values, expected) in [(vec![i64::MAX, 1, -1], 0), (vec![i64::MIN, -1, 1], 0)] { + let array: ArrayRef = Arc::new(Int64Array::from(values)); + let mut acc = SlidingSumIntegerAccumulator::new(EvalMode::Try); + acc.update_batch(&[Arc::clone(&array)]).unwrap(); + assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(None)); + acc.retract_batch(&[array.slice(0, 1)]).unwrap(); + assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(Some(expected))); + } + } + + #[test] + fn sliding_sum_accepts_all_integral_types() { + for datatype in [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + ] { + let array = arrow::compute::cast( + &Int64Array::from(vec![Some(100), None, Some(-100), Some(7)]), + &datatype, + ) + .unwrap(); + let mut acc = SlidingSumIntegerAccumulator::new(EvalMode::Ansi); + acc.update_batch(&[Arc::clone(&array)]).unwrap(); + acc.retract_batch(&[array.slice(0, 2)]).unwrap(); + assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(Some(-93))); + } + for mode in [EvalMode::Ansi, EvalMode::Try] { + let udf = SumInteger::try_new(DataType::Int64, mode).unwrap(); + assert!(matches!(udf.reverse_expr(), ReversedUDAF::NotSupported)); + } + } + fn run_update_batch_with_filter( acc: &mut dyn GroupsAccumulator, values: Vec, diff --git a/spark/src/test/resources/sql-tests/windows/sliding_integer_sum.sql b/spark/src/test/resources/sql-tests/windows/sliding_integer_sum.sql new file mode 100644 index 00000000000..2070fdae71f --- /dev/null +++ b/spark/src/test/resources/sql-tests/windows/sliding_integer_sum.sql @@ -0,0 +1,106 @@ +-- 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. + +-- Config: spark.sql.adaptive.enabled=false +-- Config: spark.sql.ansi.enabled=false + +-- ROWS ... N PRECEDING is covered through the DataFrame API in CometWindowExecSuite +-- so this fixture can keep the harness's ConstantFolding exclusion. + +statement +CREATE TABLE sliding_integer_sum(g INT, id INT, v BIGINT) USING parquet + +statement +INSERT INTO sliding_integer_sum VALUES + (1, 1, 9223372036854775807), (1, 2, 1), (1, 3, -1), (1, 4, NULL), (1, 5, NULL), + (2, 1, -9223372036854775808), (2, 2, -1), (2, 3, 1), (2, 4, NULL), (2, 5, NULL), + (3, 1, NULL), (3, 2, NULL), (3, 3, NULL) + +-- TRY mode runs natively even when ANSI is disabled. Cover positive/negative +-- overflow, recovery after it leaves the frame, all-NULL and empty frames. +query +SELECT g, id, try_sum(v) OVER ( + PARTITION BY g ORDER BY id RANGE BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum + +query +SELECT g, id, try_sum(v) OVER ( + PARTITION BY g ORDER BY id ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING) +FROM sliding_integer_sum + +query +SELECT g, id, try_sum(v) OVER ( + PARTITION BY g ORDER BY id ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) +FROM sliding_integer_sum + +-- Legacy sliding sums still run natively, including wrapping overflow. +query +SELECT g, id, + sum(v) OVER (PARTITION BY g ORDER BY id RANGE BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum + +-- Ever-expanding TRY sums retain native overflow-to-NULL behavior. +query +SELECT g, id, try_sum(v) OVER ( + PARTITION BY g ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum + +statement +SET spark.sql.ansi.enabled=true + +query expect_error(ARITHMETIC_OVERFLOW) +SELECT sum(v) OVER (ORDER BY id RANGE BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum WHERE g = 2 + +-- Native expanding sums remain enabled in ANSI and TRY mode. +query +SELECT g, id, + sum(v) OVER (PARTITION BY g ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), + try_sum(v) OVER (PARTITION BY g ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum WHERE id > 1 + +query +SELECT g, id, try_sum(v) OVER ( + PARTITION BY g ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum + +query expect_error(ARITHMETIC_OVERFLOW) +SELECT sum(v) OVER (ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum WHERE g = 1 + +-- Updating before retracting must not raise on the transient union of two frames. +query +SELECT g, id, sum(v) OVER ( + PARTITION BY g ORDER BY id ROWS BETWEEN CURRENT ROW AND CURRENT ROW) +FROM sliding_integer_sum + +statement +CREATE TABLE sliding_integer_sum_peers(id INT, v BIGINT) USING parquet + +statement +INSERT INTO sliding_integer_sum_peers VALUES + (1, 9223372036854775807), (1, 1), (3, -1), (6, NULL), (9, 7) + +-- RANGE peers enter together; a gap removes both overflow-causing values. +query +SELECT id, try_sum(v) OVER (ORDER BY id RANGE BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM sliding_integer_sum_peers + +-- Existing unsupported RANGE upper bounds retain their fallback boundary. +query expect_fallback(Unsupported RANGE frame upper offset) +SELECT id, try_sum(v) OVER (ORDER BY id RANGE BETWEEN 2 PRECEDING AND 1 PRECEDING) +FROM sliding_integer_sum_peers diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 1e82f614ea7..5928aa397c9 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.comet.{CometSortExec, CometWindowExec, CometWindowGr import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.window.{WindowExec => SparkWindowExec} import org.apache.spark.sql.expressions.Window -import org.apache.spark.sql.functions.{count, lead, sum} +import org.apache.spark.sql.functions.{count, expr, lead, sum} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.DecimalType @@ -580,8 +580,12 @@ class CometWindowExecSuite extends CometTestBase { } test("Windows support") { - Seq("true", "false").foreach(aqeEnabled => + for { + aqeEnabled <- Seq("true", "false") + ansiEnabled <- Seq("true", "false") + } { withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansiEnabled, CometConf.COMET_SHUFFLE_ENABLED.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled) { withParquetTable((0 until 10).map(i => (i, 10 - i)), "t1") { // TODO: test nulls @@ -599,17 +603,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) } } } - }) + } + } + } + + for (ansiEnabled <- Seq("true", "false")) { + test(s"sliding integral ROWS sums (ANSI=$ansiEnabled)") { + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansiEnabled, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") { + val values = Seq( + (1, 1, Some(Long.MaxValue)), + (1, 2, Some(1L)), + (1, 3, Some(-1L)), + (1, 4, None), + (1, 5, None), + (2, 1, Some(Long.MinValue)), + (2, 2, Some(-1L)), + (2, 3, Some(1L)), + (2, 4, None), + (2, 5, None), + (3, 1, None), + (3, 2, None), + (3, 3, None)) + withParquetTable(values, "sliding_rows_sum") { + val df = spark.table("sliding_rows_sum").toDF("g", "id", "v") + // DataFrame bounds are literals, so PRECEDING needs no constant folding. + val frame = Window.partitionBy("g").orderBy("id").rowsBetween(-1, Window.currentRow) + checkSparkAnswerAndOperator(df.select($"g", $"id", expr("try_sum(v)").over(frame))) + val widerFrame = + Window.partitionBy("g").orderBy("id").rowsBetween(-2, Window.currentRow) + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") { + checkSparkAnswerAndOperator( + df.select($"g", $"id", expr("try_sum(v)").over(widerFrame))) + } + + // Exercise successful ANSI execution separately from the overflow cases. + val sumInput = if (SQLConf.get.ansiEnabled) df.where($"id" > 1) else df + checkSparkAnswerAndOperator(sumInput.select($"g", $"id", sum("v").over(frame))) + for (dataType <- Seq("tinyint", "smallint", "int")) { + checkSparkAnswerAndOperator(df.select(sum($"id".cast(dataType)).over(frame))) + } + + if (SQLConf.get.ansiEnabled) { + for (group <- Seq(1, 2)) { + val (sparkError, cometError) = + checkSparkAnswerMaybeThrows(df.where($"g" === group).select(sum("v").over(frame))) + assert(sparkError.exists(_.getMessage.contains("ARITHMETIC_OVERFLOW"))) + assert(cometError.exists(_.getMessage.contains("ARITHMETIC_OVERFLOW"))) + } + } + + // Floating-point sums remain native in both modes. + checkSparkAnswerAndOperator( + df.where($"id" > 1) + .select( + sum($"v".cast("double")).over(frame), + expr("try_sum(CAST(v AS DOUBLE))").over(frame))) + } + } + } } test("window: simple COUNT(*) without frame") { diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSlidingSumBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSlidingSumBenchmark.scala new file mode 100644 index 00000000000..f533de2c7de --- /dev/null +++ b/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") { + for ((shape, value) <- Seq( + "cancelling" -> "id % 17 - 8", + "positive" -> "id % 17", + // Partial sums fit in i64, but both sign totals cross the bounds. + "near-limit" -> + "CASE WHEN id % 2 = 0 THEN 4611686018427387904L ELSE -4611686018427387903L END"); + nullPercent <- (if (shape == "cancelling") Seq(0, 50, 100) else Seq(0, 50))) { + withTempPath { dir => + spark + .range(rows) + .selectExpr( + "id", + s"CASE WHEN pmod(id, 100) < $nullPercent THEN NULL ELSE $value END AS v") + .write + .parquet(dir.getCanonicalPath) + withTempTable("sliding_sum_benchmark") { + spark.read + .parquet(dir.getCanonicalPath) + .createOrReplaceTempView("sliding_sum_benchmark") + for ((frameName, bounds) <- selectedFrames) { + def run( + function: String, + ansi: Boolean, + native: Boolean, + inputRows: Long, + verify: Boolean): Seq[Row] = { + // Spark 3.5's withSQLConf returns Unit. + var result = Seq.empty[Row] + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansi.toString, + CometConf.COMET_EXEC_WINDOW_ENABLED.key -> native.toString) { + // Consume the window values with a sink unaffected by ANSI mode. + val filter = if (inputRows < rows) s" WHERE id < $inputRows" else "" + val query = s"SELECT bit_xor(s), count(s) FROM (SELECT $function(v) OVER " + + s"(ORDER BY id ROWS BETWEEN $bounds) AS s " + + s"FROM sliding_sum_benchmark$filter)" + val df = spark.sql(query) + if (verify) { + assert( + df.queryExecution.executedPlan.exists( + _.isInstanceOf[CometWindowExec]) == native, + df.queryExecution.executedPlan.toString) + } + result = df.collect().toSeq + } + result + } + // Spark recomputes suffix frames in O(rows^2). Check a small partition + // against Spark, then compare all native cases on the full partition. + val verificationRows = if (frameName == "suffix") math.min(rows, 1024L) else rows + val expected = run( + "sum", + ansi = true, + native = false, + inputRows = verificationRows, + verify = true) + for ((_, function, ansi, native) <- nativeCases ++ sparkCases.tail) { + assert(run(function, ansi, native, verificationRows, verify = true) == expected) + } + if (verificationRows != rows) { + val legacy = + run("sum", ansi = false, native = true, inputRows = rows, verify = true) + for ((_, function, ansi, native) <- nativeCases.tail) { + assert(run(function, ansi, native, rows, verify = true) == legacy) + } + } + val cases = if (frameName == "suffix") nativeCases else nativeCases ++ sparkCases + val benchmark = new Benchmark( + s"$shape: frame=$frameName, NULL=$nullPercent%", + rows, + output = output) + for ((name, function, ansi, native) <- cases) { + benchmark.addCase(name) { _ => + val _ = run(function, ansi, native, rows, verify = false) + } + } + benchmark.run() + } + } + } + } + } + } +}