diff --git a/docs/source/contributor-guide/roadmap.md b/docs/source/contributor-guide/roadmap.md index b8ae8a310f..d738095b41 100644 --- a/docs/source/contributor-guide/roadmap.md +++ b/docs/source/contributor-guide/roadmap.md @@ -28,7 +28,7 @@ Native window execution runs by default (`spark.comet.exec.window.enabled`). The `dense_rank`, `row_number`, `percent_rank`, `cume_dist`, `ntile`), value functions (`lag`, `lead`, `nth_value`, `first_value`, `last_value`), and the `count`, `min`, `max`, `sum`, and `avg` aggregates are accelerated. Remaining work is to close the gaps that still fall back to Spark: statistical aggregates (`stddev`, `variance`, -`corr`, `covar`) and `collect_list` / `collect_set` as window functions ([#4766]), `GROUPS` frames ([#4836]), `RANGE` frames with explicit date or +`corr`, `covar`) and `collect_list` / `collect_set` as window functions ([#4766]), `GROUPS` frames ([#4836]), `RANGE` frames with explicit decimal offsets ([#4834]), `first_value` / `last_value` on `RANGE` frames with a literal offset ([#4835]), non-literal `lag` / `lead` default values ([#4268]), and `WindowGroupLimitExec` ([#4837]). See the [window function compatibility guide](../user-guide/latest/compatibility/operators.md) for the complete list of diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index eec21d283b..ff71f27a3a 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -45,7 +45,7 @@ incorrect result. When any single window expression in a `WindowExec` falls back support as the batch aggregates, so these fall back in both contexts. - `sum` or `avg` on `DECIMAL` with a sliding (non ever-expanding) frame, because the sliding path would wrap on overflow instead of returning Spark's `NULL`. -- `RANGE` frame with an explicit offset when the `ORDER BY` column is `DATE` or `DECIMAL` +- `RANGE` frame with an explicit offset when the `ORDER BY` column is `DECIMAL` ([#4834](https://github.com/apache/datafusion-comet/issues/4834)). - `first_value` / `last_value` on a `RANGE` frame with a literal offset ([#4835](https://github.com/apache/datafusion-comet/issues/4835)). diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index a23a7b1f67..cadb1143f4 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -38,7 +38,9 @@ use crate::execution::{ }; use crate::jvm_bridge::{jni_call, JVMClasses}; use arrow::compute::CastOptions; -use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit, DECIMAL128_MAX_PRECISION}; +use arrow::datatypes::{ + DataType, Field, FieldRef, IntervalDayTime, Schema, TimeUnit, DECIMAL128_MAX_PRECISION, +}; use arrow::ffi_stream::FFI_ArrowArrayStream; use datafusion::functions_aggregate::bit_and_or_xor::{bit_and_udaf, bit_or_udaf, bit_xor_udaf}; use datafusion::functions_aggregate::count::count_udaf; @@ -2751,6 +2753,20 @@ impl PhysicalPlanner { WindowFrameType::Range => WindowFrameUnits::Range, }; + // For RANGE frames the offset scalar must be arithmetic-compatible with + // the ORDER BY column, not just numeric. arrow-arith rejects e.g. + // `Date32 + Int32` — it requires an Interval RHS. Capture the ORDER BY + // datatype once so both bound branches can coerce accordingly. + let range_order_by_type = if matches!(units, WindowFrameUnits::Range) { + sort_exprs + .first() + .map(|s| s.expr.data_type(&input_schema)) + .transpose() + .map_err(|e| GeneralError(format!("Failed to resolve ORDER BY type: {e}")))? + } else { + None + }; + let lower_bound: WindowFrameBound = match spark_window_frame .lower_bound .as_ref() @@ -2786,10 +2802,14 @@ impl PhysicalPlanner { } } WindowFrameUnits::Range => { - let scalar = match offset.range_offset.as_ref() { + let raw = match offset.range_offset.as_ref() { Some(lit) => numeric_literal_to_scalar(lit)?, None => ScalarValue::Int64(Some(offset.offset.abs())), }; + let scalar = coerce_range_offset_for_order_by( + raw, + range_order_by_type.as_ref(), + )?; WindowFrameBound::Preceding(scalar) } WindowFrameUnits::Groups => { @@ -2845,10 +2865,12 @@ impl PhysicalPlanner { } } WindowFrameUnits::Range => { - let scalar = match offset.range_offset.as_ref() { + let raw = match offset.range_offset.as_ref() { Some(lit) => numeric_literal_to_scalar(lit)?, None => ScalarValue::Int64(Some(offset.offset)), }; + let scalar = + coerce_range_offset_for_order_by(raw, range_order_by_type.as_ref())?; WindowFrameBound::Following(scalar) } WindowFrameUnits::Groups => { @@ -3394,6 +3416,37 @@ fn numeric_literal_to_scalar( Ok(scalar) } +/// arrow-arith only provides `Date32 + Interval` kernels, not `Date32 + Int32`. +/// Spark ships `n PRECEDING/FOLLOWING` as an IntegerType literal, so a DATE +/// ORDER BY frame would blow up at execution with +/// "Invalid date arithmetic operation: Date32 + Int32". Wrap integer offsets +/// as IntervalDayTime (days component) so the frame comparator can compute +/// `current_date +/- offset` natively. +fn coerce_range_offset_for_order_by( + scalar: ScalarValue, + order_by_type: Option<&DataType>, +) -> Result { + let Some(dt) = order_by_type else { + return Ok(scalar); + }; + match (dt, &scalar) { + (DataType::Date32, ScalarValue::Int32(Some(v))) => Ok(ScalarValue::IntervalDayTime(Some( + IntervalDayTime::new(*v, 0), + ))), + (DataType::Date32, ScalarValue::Int64(Some(v))) => { + let days = i32::try_from(*v).map_err(|_| { + GeneralError(format!( + "DATE RANGE frame offset {v} does not fit in i32 days" + )) + })?; + Ok(ScalarValue::IntervalDayTime(Some(IntervalDayTime::new( + days, 0, + )))) + } + _ => Ok(scalar), + } +} + /// A physical join filter rewritter which rewrites the column indices in the expression /// to use the new column indices. See `rewrite_physical_expr`. struct JoinFilterRewriter<'a> { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala index 500f9996b9..4c21eca90b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{DataType, DateType, Decimal, DecimalType, DoubleType, LongType, NumericType} +import org.apache.spark.sql.types.{DataType, Decimal, DecimalType, DoubleType, LongType, NumericType} import com.google.common.base.Objects @@ -364,26 +364,19 @@ object CometWindowExec extends CometOperatorSerde[WindowExec] { } // Comet's native window planner ships RANGE frame offsets as - // ScalarValue::Int64, but a couple of ORDER BY types don't tolerate that: - // - DATE: arrow-arith requires an Interval RHS for Date32 arithmetic, - // so execution fails with - // Invalid date arithmetic operation: Date32 + Int32 - // - DECIMAL: Spark decimal arithmetic widens precision on +/-, so the - // computed boundary (e.g. Decimal(11,0)) doesn't match the current - // value's precision (e.g. Decimal(10,0)) and the comparator fails - // with "Uncomparable values". - // Fall back to Spark for those shapes. UNBOUNDED / CURRENT ROW bounds - // don't evaluate the offending arithmetic and stay native. + // ScalarValue::Int64, which the DECIMAL comparator can't handle: Spark + // decimal arithmetic widens precision on +/-, so the computed boundary + // (e.g. Decimal(11,0)) doesn't match the current value's precision + // (e.g. Decimal(10,0)) and the comparator fails with "Uncomparable + // values". Fall back for DECIMAL until the native side normalizes + // precision. DATE is now handled natively by wrapping the offset as + // IntervalDayTime in the planner. UNBOUNDED / CURRENT ROW bounds don't + // evaluate the offending arithmetic and stay native regardless. f match { case SpecifiedWindowFrame(RangeFrame, lb, ub) if !lb.isInstanceOf[SpecialFrameBoundary] || !ub.isInstanceOf[SpecialFrameBoundary] => windowExpr.windowSpec.orderSpec.headOption.map(_.dataType) match { - case Some(DateType) => - withFallbackReason( - windowExpr, - "RANGE frame with explicit offset on DATE ORDER BY is not supported") - return None case Some(_: DecimalType) => withFallbackReason( windowExpr, 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 0e6fcf8946..3b0f6e62d7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -1274,4 +1274,153 @@ class CometWindowExecSuite extends CometTestBase { checkSparkAnswerAndOperator(df) } } + + // --------------------------------------------------------------------------- + // Tests for issue #4834: RANGE frame with explicit PRECEDING/FOLLOWING offset + // on DATE ORDER BY. Today Comet falls back to Spark; these tests are pinned + // with `checkSparkAnswerAndOperator` (assertCometNative = true) so they fail + // until the native side handles Date32 + Interval arithmetic and the Scala + // fallback in CometWindowExec is removed. + // --------------------------------------------------------------------------- + + private def writeDateWindowFixture(dir: java.io.File): Unit = { + (0 until 30) + .map { i => + val day = (i % 28) + 1 + (i % 3, java.sql.Date.valueOf(f"2024-01-$day%02d"), i * 2) + } + .toDF("a", "d", "c") + .repartition(3) + .write + .mode("overwrite") + .parquet(dir.toString) + } + + test("window: RANGE BETWEEN N PRECEDING AND N FOLLOWING with DATE ORDER BY") { + withTempDir { dir => + writeDateWindowFixture(dir) + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN 2 PRECEDING AND 2 FOLLOWING) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } + + test("window: RANGE BETWEEN N PRECEDING AND CURRENT ROW with DATE ORDER BY") { + withTempDir { dir => + writeDateWindowFixture(dir) + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } + + test("window: RANGE BETWEEN CURRENT ROW AND N FOLLOWING with DATE ORDER BY") { + withTempDir { dir => + writeDateWindowFixture(dir) + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN CURRENT ROW AND 3 FOLLOWING) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } + + test("window: RANGE BETWEEN UNBOUNDED PRECEDING AND N FOLLOWING with DATE ORDER BY") { + withTempDir { dir => + writeDateWindowFixture(dir) + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN UNBOUNDED PRECEDING AND 2 FOLLOWING) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } + + test("window: RANGE BETWEEN with DATE ORDER BY and duplicate dates") { + // RANGE (not ROWS) must aggregate all peer rows that share the same date. + // With duplicates we can tell the two apart: ROWS would only include the + // literal N rows, RANGE includes every row whose date is within +/- N days. + withTempDir { dir => + Seq( + (1, java.sql.Date.valueOf("2024-01-01"), 10), + (1, java.sql.Date.valueOf("2024-01-01"), 20), + (1, java.sql.Date.valueOf("2024-01-02"), 30), + (1, java.sql.Date.valueOf("2024-01-02"), 40), + (1, java.sql.Date.valueOf("2024-01-05"), 50), + (1, java.sql.Date.valueOf("2024-01-05"), 60)) + .toDF("a", "d", "c") + .write + .mode("overwrite") + .parquet(dir.toString) + + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } + + test("window: RANGE BETWEEN with DATE ORDER BY and NULL dates") { + // Null dates land at the head under ASC NULLS FIRST (Spark's default). + // Verifies the native frame comparator handles Date32 nulls the same way + // Spark does when computing the boundary. + withTempDir { dir => + Seq( + (1, Option.empty[java.sql.Date], 10), + (1, Option(java.sql.Date.valueOf("2024-01-01")), 20), + (1, Option(java.sql.Date.valueOf("2024-01-02")), 30), + (1, Option.empty[java.sql.Date], 40), + (1, Option(java.sql.Date.valueOf("2024-01-05")), 50)) + .toDF("a", "d", "c") + .write + .mode("overwrite") + .parquet(dir.toString) + + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } + + test("window: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW with DATE ORDER BY") { + // Regression baseline: no explicit offset, so the current fallback in + // CometWindowExec does NOT trigger and this already runs natively. Keep + // this passing throughout the fix. + withTempDir { dir => + writeDateWindowFixture(dir) + spark.read.parquet(dir.toString).createOrReplaceTempView("window_test") + val df = sql(""" + SELECT a, d, c, + SUM(c) OVER (PARTITION BY a ORDER BY d + RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) as sum_c + FROM window_test + """) + checkSparkAnswerAndOperator(df) + } + } }