diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md index b4ae0b5b3a1..5cdb222eb1a 100644 --- a/docs/source/contributor-guide/expression-audits/array_funcs.md +++ b/docs/source/contributor-guide/expression-audits/array_funcs.md @@ -176,6 +176,16 @@ - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; ANSI default flips to `true`. - Spark 4.1.1 (audited 2026-05-27): `inputTypes` tightened to `Seq(ArrayType, IntegralType)` (analysis-time only); runtime unchanged. +## sequence + +- Spark 3.4.3 (audited 2026-08-29): `Sequence(start, stop, stepOpt, timeZoneId)`; `Sequence.impl` selects the implementation from `dataType.elementType`, so the integral/temporal split is knowable at plan time. Codegen for the integral path checks boundaries with a plain `IllegalArgumentException("Illegal sequence boundaries: ...")`, then calls the static `Sequence.sequenceLength`, which raises `SparkRuntimeException(_LEGACY_ERROR_TEMP_2161)` past `MAX_ROUNDED_ARRAY_LENGTH` and `internalError("Unreachable code reached.")` when `stop - start` overflows Long but the exact length is within the limit. Default step is per-row `start <= stop ? 1 : -1`. +- Spark 3.5.8 (audited 2026-08-29): internal refactors only (`DataTypeUtils.sameType`, `PhysicalIntegralType.integral`); runtime semantics identical to 3.4.3. +- Spark 4.0.1 (audited 2026-08-29): boundary error becomes `SparkIllegalArgumentException(_LEGACY_ERROR_TEMP_3243)` and the length error becomes `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER` (now carrying the function name); adds `throwable` optimizer hint. `sequenceLength` itself is unchanged. +- Spark 4.1.1 (audited 2026-08-29): byte-identical `Sequence` class body to 4.0.1. +- Comet routes integral element types (`ByteType`/`ShortType`/`IntegerType`/`LongType`) via `CometSequence` to the native `spark_sequence` kernel ([#5349](https://github.com/apache/datafusion-comet/issues/5349)): one pass over the generated elements, child buffer reserved once per batch, no per-row allocation. The two-argument form is evaluated with Spark's per-row default step inside the kernel. Both error conditions and the internal-error edge are reproduced through `SparkError` and mapped per Spark version by `ShimSparkErrorConverter`. Date/timestamp/timestamp_ntz sequences return `Unsupported` and run on the JVM codegen dispatcher (`CodegenDispatchFallback`), pending the timezone/DST/legacy-calendar work. +- Per-batch capacity ceiling: the native kernel writes every row's generated elements into one Arrow child buffer whose offsets are `i32`, so the sum of every row's length in a single Arrow batch must fit in `i32::MAX`. Spark itself has no equivalent limit because it stores each row as its own `long[]`. If the total is exceeded, or if the allocator refuses the reservation, the query fails with a `SparkError::SequenceBatchTooLarge` message that names `spark.comet.batchSize` as the actionable knob (lower it to group fewer rows per batch). The `try_reserve` path guarantees the failure surfaces as a query error rather than an allocator abort. +- Argument-shape restriction: `CometSequence` reports `Unsupported` for any `Sequence` whose `start`, `stop`, or `step` is not a leaf expression, and routes those through the JVM codegen dispatcher (`CodegenDispatchFallback`). DataFusion evaluates each scalar-UDF argument over the whole batch before calling the outer kernel, so a non-leaf argument would run on rows that Spark's per-row null short-circuit (or a `CASE` branch) would have discarded, and could raise where Spark would have returned `NULL`. + ## shuffle - Spark 3.4.3 (audited 2026-07-02): `Shuffle(child, randomSeed: Option[Long])`; `inputTypes = Seq(ArrayType)`, `dataType = child.dataType`, non-deterministic and stateful. Seeds a Commons Math3 `MersenneTwister` with `randomSeed + partitionIndex` and applies the "inside-out" Fisher-Yates from `RandomIndicesGenerator`. Only the one-argument `shuffle(array)` form exists in SQL. NULL input returns NULL without advancing the RNG. diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 9a59f040bcb..77a3b37ff13 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -168,7 +168,7 @@ The tables below list every Spark built-in expression with its current status. | `element_at` | ✅ | Native | | | `flatten` | ✅ | Native | Binary/struct/map elements fall back | | `get` | ✅ | — | | -| `sequence` | ✅ | Codegen dispatch | | +| `sequence` | ✅ | Hybrid | Integral types run natively; date/timestamp sequences use codegen dispatch | | `shuffle` | ✅ | Native | Binary/struct/map elements fall back | | `slice` | ✅ | Native | Native ([#4149](https://github.com/apache/datafusion-comet/pull/4149)) | | `sort_array` | ✅ | Hybrid | Nested struct/null arrays fall back | diff --git a/native/common/src/error.rs b/native/common/src/error.rs index e66a0a88eff..41773237cba 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -135,12 +135,38 @@ pub enum SparkError { #[error("[EXCEED_LIMIT_LENGTH] Cannot create a map with {size} elements which exceeds the limit {max_size}.")] ExceedMapSizeLimit { size: i32, max_size: i32 }, - #[error("[COLLECTION_SIZE_LIMIT_EXCEEDED] Cannot create array with {num_elements} elements which exceeds the limit {max_elements}.")] + /// `num_elements` is a decimal string because Spark reports the unclamped length, which can + /// exceed i64 (e.g. sequence(Long.MinValue, Long.MaxValue, 1)). The JVM shim maps this to the + /// version-appropriate `createArrayWithElementsExceedLimitError`, passing `function_name` + /// through on Spark 4.x (which includes it in the message) and ignoring it on 3.x. + #[error("[COLLECTION_SIZE_LIMIT_EXCEEDED] Can't create array with {num_elements} elements which exceeding the array size limit {max_elements}.")] CollectionSizeLimitExceeded { - num_elements: i64, + num_elements: String, max_elements: i64, + function_name: String, + }, + + /// Step direction does not match the start/stop bounds in `sequence`. The JVM shim maps this + /// to a plain IllegalArgumentException on Spark 3.x and SparkIllegalArgumentException + /// (_LEGACY_ERROR_TEMP_3243) on 4.x, matching what Spark's codegen throws. + #[error("[_LEGACY_ERROR_TEMP_3243] Illegal sequence boundaries: {start} to {stop} by {step}")] + SequenceIllegalBoundaries { + start: String, + stop: String, + step: String, }, + /// Sum of every row's sequence length in one Arrow batch exceeds Comet's per-batch offset + /// ceiling (i32::MAX) or the allocator could not satisfy the reservation. Spark itself has + /// no equivalent limit because it stores each row as its own `long[]`. Reported with a + /// message that names `spark.comet.batchSize` as the actionable knob. + #[error( + "Comet's native `sequence` kernel cannot materialize a batch with {total_elements} \ + total elements: it exceeds the per-batch limit or the allocator refused the reservation. \ + Lower `spark.comet.batchSize` so fewer rows are grouped per batch." + )] + SequenceBatchTooLarge { total_elements: String }, + #[error("[NOT_NULL_ASSERT_VIOLATION] The field `{field_name}` cannot be null.")] NotNullAssertViolation { field_name: String }, @@ -306,6 +332,8 @@ impl SparkError { SparkError::MapKeyValueDiffSizes => "MapKeyValueDiffSizes", SparkError::ExceedMapSizeLimit { .. } => "ExceedMapSizeLimit", SparkError::CollectionSizeLimitExceeded { .. } => "CollectionSizeLimitExceeded", + SparkError::SequenceIllegalBoundaries { .. } => "SequenceIllegalBoundaries", + SparkError::SequenceBatchTooLarge { .. } => "SequenceBatchTooLarge", SparkError::NotNullAssertViolation { .. } => "NotNullAssertViolation", SparkError::ValueIsNull { .. } => "ValueIsNull", SparkError::CannotParseTimestamp { .. } => "CannotParseTimestamp", @@ -450,10 +478,24 @@ impl SparkError { SparkError::CollectionSizeLimitExceeded { num_elements, max_elements, + function_name, } => { serde_json::json!({ "numElements": num_elements, "maxElements": max_elements, + "functionName": function_name, + }) + } + SparkError::SequenceIllegalBoundaries { start, stop, step } => { + serde_json::json!({ + "start": start, + "stop": stop, + "step": step, + }) + } + SparkError::SequenceBatchTooLarge { total_elements } => { + serde_json::json!({ + "totalElements": total_elements, }) } SparkError::NotNullAssertViolation { field_name } => { @@ -626,6 +668,7 @@ impl SparkError { | SparkError::MapKeyValueDiffSizes | SparkError::ExceedMapSizeLimit { .. } | SparkError::CollectionSizeLimitExceeded { .. } + | SparkError::SequenceBatchTooLarge { .. } // Comet-specific extension | SparkError::NotNullAssertViolation { .. } | SparkError::ValueIsNull { .. } // Comet-specific extension | SparkError::UnexpectedPositiveValue { .. } @@ -644,7 +687,8 @@ impl SparkError { // IllegalArgumentException SparkError::DatatypeCannotOrder { .. } | SparkError::InvalidUtf8String { .. } - | SparkError::IllegalDayOfWeek { .. } => { + | SparkError::IllegalDayOfWeek { .. } + | SparkError::SequenceIllegalBoundaries { .. } => { "org/apache/spark/SparkIllegalArgumentException" } @@ -722,6 +766,10 @@ impl SparkError { SparkError::CollectionSizeLimitExceeded { .. } => { Some("COLLECTION_SIZE_LIMIT_EXCEEDED") } + SparkError::SequenceIllegalBoundaries { .. } => Some("_LEGACY_ERROR_TEMP_3243"), + + // Comet-specific: no Spark error class, the shim builds the message itself. + SparkError::SequenceBatchTooLarge { .. } => None, // Null validation errors SparkError::NotNullAssertViolation { .. } => Some("NOT_NULL_ASSERT_VIOLATION"), diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 519336c4fec..1c3a9999283 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -287,6 +287,10 @@ harness = false name = "contains" harness = false +[[bench]] +name = "sequence" +harness = false + [[bench]] name = "timestamp_trunc" harness = false diff --git a/native/spark-expr/benches/sequence.rs b/native/spark-expr/benches/sequence.rs new file mode 100644 index 00000000000..4a67a5f33e5 --- /dev/null +++ b/native/spark-expr/benches/sequence.rs @@ -0,0 +1,133 @@ +// 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. + +use arrow::array::Int64Array; +use arrow::datatypes::{DataType, Field}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_sequence; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; + +fn list_of_i64() -> DataType { + DataType::List(Arc::new(Field::new_list_field(DataType::Int64, false))) +} + +/// start/stop columns generating `elems_per_row` elements per row (ascending, step 1). +/// When `null_every` is Some(n), every nth row of `start` is null. +fn args_with_len(elems_per_row: i64, null_every: Option) -> Vec { + let start = Int64Array::from( + (0..NUM_ROWS) + .map(|i| match null_every { + Some(n) if i % n == 0 => None, + _ => Some(i as i64), + }) + .collect::>(), + ); + let stop = Int64Array::from( + (0..NUM_ROWS) + .map(|i| Some(i as i64 + elems_per_row - 1)) + .collect::>(), + ); + vec![ + ColumnarValue::Array(Arc::new(start)), + ColumnarValue::Array(Arc::new(stop)), + ] +} + +fn criterion_benchmark(c: &mut Criterion) { + let return_type = list_of_i64(); + + let mut group = c.benchmark_group("sequence"); + + // Short sequences: per-row overhead dominates. + for elems in [2i64, 5] { + let args = args_with_len(elems, None); + group.bench_function(format!("short_{elems}_elems"), |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + + // Long sequences: element throughput dominates. 365 is the date-spine shape from the + // issue; 10k stresses the child buffer reservation. + for elems in [365i64, 10_000] { + let args = args_with_len(elems, None); + group.bench_function(format!("long_{elems}_elems"), |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + + // Descending with explicit negative step. + { + let start = Int64Array::from((0..NUM_ROWS).map(|i| i as i64 + 364).collect::>()); + let stop = Int64Array::from((0..NUM_ROWS).map(|i| i as i64).collect::>()); + let step = Int64Array::from(vec![-1i64; NUM_ROWS]); + let args = vec![ + ColumnarValue::Array(Arc::new(start)), + ColumnarValue::Array(Arc::new(stop)), + ColumnarValue::Array(Arc::new(step)), + ]; + group.bench_function("descending_365_elems", |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + + // Zero step with start == stop: single-element rows through the step==0 path. + { + let start = Int64Array::from((0..NUM_ROWS).map(|i| i as i64).collect::>()); + let stop = Int64Array::from((0..NUM_ROWS).map(|i| i as i64).collect::>()); + let step = Int64Array::from(vec![0i64; NUM_ROWS]); + let args = vec![ + ColumnarValue::Array(Arc::new(start)), + ColumnarValue::Array(Arc::new(stop)), + ColumnarValue::Array(Arc::new(step)), + ]; + group.bench_function("zero_step_start_eq_stop", |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + + // Sparse (every 10th row) and dense (every 2nd row) nulls over the date-spine shape. + for (label, every) in [("sparse_nulls", 10usize), ("dense_nulls", 2)] { + let args = args_with_len(365, Some(every)); + group.bench_function(format!("{label}_365_elems"), |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap())) + }); + } + + // Error path: the boundary check rejects the first row. + { + let start = Int64Array::from(vec![0i64; NUM_ROWS]); + let stop = Int64Array::from(vec![100i64; NUM_ROWS]); + let step = Int64Array::from(vec![-1i64; NUM_ROWS]); + let args = vec![ + ColumnarValue::Array(Arc::new(start)), + ColumnarValue::Array(Arc::new(stop)), + ColumnarValue::Array(Arc::new(step)), + ]; + group.bench_function("error_illegal_boundaries", |b| { + b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap_err())) + }); + } + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index 0c2c68dc6d9..b8877a93dce 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -23,6 +23,7 @@ mod arrays_zip; mod flatten; mod get_array_struct_fields; mod list_extract; +mod sequence; mod size; pub use array_insert::ArrayInsert; @@ -33,4 +34,5 @@ pub use arrays_zip::SparkArraysZipFunc; pub use flatten::SparkFlatten; pub use get_array_struct_fields::GetArrayStructFields; pub use list_extract::ListExtract; +pub use sequence::spark_sequence; pub use size::{spark_size, SparkSizeFunc}; diff --git a/native/spark-expr/src/array_funcs/sequence.rs b/native/spark-expr/src/array_funcs/sequence.rs new file mode 100644 index 00000000000..4b1335f20fe --- /dev/null +++ b/native/spark-expr/src/array_funcs/sequence.rs @@ -0,0 +1,393 @@ +// 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. + +// Spark-compatible sequence(start, stop[, step]) for integral element types. +// +// Mirrors the code Spark's whole-stage codegen emits for `Sequence` +// (`sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala`, +// identical from 3.4.3 through 4.1.1): the boundary check and `Sequence.sequenceLength` decide +// per row how many elements to generate, then elements are `start + step * i`. Unlike the JVM +// path, which allocates two `long[]` per row and copies every element three times, this kernel +// reserves the Arrow child buffer once for the whole batch and writes each element exactly once. +// +// Date/timestamp sequences are not handled here; the Scala serde only routes IntegralType +// sequences to this function. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, ListArray, NullBufferBuilder, PrimitiveArray}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, FieldRef, Int16Type, Int32Type, Int64Type, Int8Type, +}; +use datafusion::common::cast::as_primitive_array; +use datafusion::common::{exec_err, DataFusionError, Result, ScalarValue}; +use datafusion::logical_expr::ColumnarValue; + +use crate::SparkError; + +/// Spark's ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH (Integer.MAX_VALUE - 15). +const MAX_ROUNDED_ARRAY_LENGTH: i128 = (i32::MAX - 15) as i128; + +pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result { + let child_field = match data_type { + DataType::List(field) => Arc::clone(field), + other => return exec_err!("spark_sequence expects a List return type, got {other:?}"), + }; + if args.len() != 2 && args.len() != 3 { + return exec_err!( + "spark_sequence expects 2 or 3 arguments, got {}", + args.len() + ); + } + + let all_scalar = args + .iter() + .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); + let arrays = ColumnarValue::values_to_arrays(args)?; + let step = arrays.get(2); + + let result = match child_field.data_type() { + DataType::Int8 => { + sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v as i8) + } + DataType::Int16 => { + sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v as i16) + } + DataType::Int32 => { + sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v as i32) + } + DataType::Int64 => { + sequence_integral::(&arrays[0], &arrays[1], step, child_field, |v| v) + } + other => exec_err!("spark_sequence does not support element type {other:?}"), + }?; + + if all_scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(result)) + } +} + +fn sequence_integral( + start: &ArrayRef, + stop: &ArrayRef, + step: Option<&ArrayRef>, + child_field: FieldRef, + from_i64: impl Fn(i64) -> T::Native, +) -> Result +where + T::Native: Into, +{ + let start = as_primitive_array::(start)?; + let stop = as_primitive_array::(stop)?; + let step = step.map(|arr| as_primitive_array::(arr)).transpose()?; + let num_rows = start.len(); + + let row_is_null = |row: usize| { + start.is_null(row) || stop.is_null(row) || step.is_some_and(|arr| arr.is_null(row)) + }; + // With no explicit step, Spark uses `start <= stop ? 1 : -1` per row, so the direction + // always matches the bounds and the boundary check below cannot fail. + let row_step = |row: usize, start: i64, stop: i64| -> i64 { + match step { + Some(arr) => arr.value(row).into(), + None => { + if start <= stop { + 1 + } else { + -1 + } + } + } + }; + + // First pass: compute per-row lengths so the child buffer can be reserved once for the + // whole batch. Valid rows always produce at least one element, so length 0 marks a null row. + let mut lengths: Vec = Vec::with_capacity(num_rows); + let mut total: usize = 0; + for row in 0..num_rows { + if row_is_null(row) { + lengths.push(0); + continue; + } + let s: i64 = start.value(row).into(); + let e: i64 = stop.value(row).into(); + let len = sequence_length(s, e, row_step(row, s, e))?; + total += len; + lengths.push(len); + } + // Comet-specific ceiling: the sum of every row's length in one Arrow batch must fit in + // the i32 offset buffer. Spark has no equivalent guard because it stores each row as its + // own `long[]`, so the user may hit this on a query Spark itself would run. Report it via + // a dedicated error that names `spark.comet.batchSize` as the actionable knob rather than + // Spark's per-array size limit. + if total > i32::MAX as usize { + return Err(DataFusionError::External(Box::new( + SparkError::SequenceBatchTooLarge { + total_elements: total.to_string(), + }, + ))); + } + + // Second pass: write elements straight into the child buffer and push offsets. The + // batch-total check above guarantees `values.len() <= i32::MAX` at every iteration, so + // the offset push cannot overflow. `try_reserve_exact` returns a query error on allocator + // failure so an oversized reservation cannot abort the executor. + let mut values: Vec = Vec::new(); + values.try_reserve_exact(total).map_err(|_| { + DataFusionError::External(Box::new(SparkError::SequenceBatchTooLarge { + total_elements: total.to_string(), + })) + })?; + let mut offsets: Vec = Vec::with_capacity(num_rows + 1); + offsets.push(0); + let mut nulls = NullBufferBuilder::new(num_rows); + for (row, &len) in lengths.iter().enumerate() { + if len == 0 { + nulls.append_null(); + } else { + nulls.append_non_null(); + let s: i64 = start.value(row).into(); + let e: i64 = stop.value(row).into(); + let step = row_step(row, s, e); + // Every element pushed lies between start and stop inclusive, so the widened + // arithmetic cannot overflow; only the final unused increment may wrap. + let mut v = s; + for _ in 0..len { + values.push(from_i64(v)); + v = v.wrapping_add(step); + } + } + offsets.push(values.len() as i32); + } + + let values = PrimitiveArray::::new(ScalarBuffer::from(values), None); + let list = ListArray::try_new( + child_field, + OffsetBuffer::new(offsets.into()), + Arc::new(values), + nulls.finish(), + )?; + Ok(Arc::new(list)) +} + +/// Number of elements of `sequence(start, stop, step)`, matching Spark's boundary check and +/// `Sequence.sequenceLength` (byte-identical from Spark 3.4.3 through 4.1.1), including which +/// of the three failure paths fires and the exact length value each of them reports. +fn sequence_length(start: i64, stop: i64, step: i64) -> Result { + if !((step > 0 && start <= stop) || (step < 0 && start >= stop) || (step == 0 && start == stop)) + { + return Err(DataFusionError::External(Box::new( + SparkError::SequenceIllegalBoundaries { + start: start.to_string(), + stop: stop.to_string(), + step: step.to_string(), + }, + ))); + } + if stop == start { + return Ok(1); + } + // Spark computes stop - start with Math.subtractExact and special-cases + // Long.MinValue / -1; both raise ArithmeticException, which reroutes the length through a + // BigInt fallback. i128 arithmetic gives the same exact value on every path. + // + // Max delta magnitude is |i64::MAX - i64::MIN| = 2 * i64::MAX + 1, which fits in i128 but + // overflows i64, so the `> i64::MAX` / `< i64::MIN` predicates below are both live. + let delta = stop as i128 - start as i128; + let overflowed = delta > i64::MAX as i128 + || delta < i64::MIN as i128 + || (delta == i64::MIN as i128 && step == -1); + let len = 1 + delta / step as i128; + if len > MAX_ROUNDED_ARRAY_LENGTH { + return Err(DataFusionError::External(Box::new( + SparkError::CollectionSizeLimitExceeded { + num_elements: len.to_string(), + max_elements: MAX_ROUNDED_ARRAY_LENGTH as i64, + function_name: "sequence".to_string(), + }, + ))); + } + if overflowed { + // Spark's BigInt fallback lands on `internalError("Unreachable code reached.")` when + // the exact length is within the limit (e.g. sequence spanning more than Long.MaxValue + // with a large step); reproduce it for error parity. + return Err(DataFusionError::External(Box::new(SparkError::Internal( + "Unreachable code reached.".to_string(), + )))); + } + Ok(len as usize) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int64Array, Int8Array}; + use arrow::datatypes::Field; + + fn list_of(elem: DataType) -> DataType { + DataType::List(Arc::new(Field::new_list_field(elem, false))) + } + + fn run_i64( + start: Vec>, + stop: Vec>, + step: Option>>, + ) -> Result { + let mut args = vec![ + ColumnarValue::Array(Arc::new(Int64Array::from(start))), + ColumnarValue::Array(Arc::new(Int64Array::from(stop))), + ]; + if let Some(step) = step { + args.push(ColumnarValue::Array(Arc::new(Int64Array::from(step)))); + } + match spark_sequence(&args, &list_of(DataType::Int64))? { + ColumnarValue::Array(arr) => Ok(as_primitive_list(&arr)), + ColumnarValue::Scalar(_) => unreachable!("array inputs produce an array"), + } + } + + fn as_primitive_list(arr: &ArrayRef) -> ListArray { + arr.as_any().downcast_ref::().unwrap().clone() + } + + fn row_values(list: &ListArray, row: usize) -> Vec { + let v = list.value(row); + as_primitive_array::(&v) + .unwrap() + .values() + .to_vec() + } + + #[test] + fn ascending_descending_and_default_step() { + let list = run_i64( + vec![Some(1), Some(5), Some(3), Some(1)], + vec![Some(5), Some(1), Some(3), Some(10)], + Some(vec![Some(2), Some(-2), Some(0), Some(3)]), + ) + .unwrap(); + assert_eq!(row_values(&list, 0), vec![1, 3, 5]); + assert_eq!(row_values(&list, 1), vec![5, 3, 1]); + assert_eq!(row_values(&list, 2), vec![3]); + assert_eq!(row_values(&list, 3), vec![1, 4, 7, 10]); + + let list = run_i64(vec![Some(1), Some(5)], vec![Some(3), Some(2)], None).unwrap(); + assert_eq!(row_values(&list, 0), vec![1, 2, 3]); + assert_eq!(row_values(&list, 1), vec![5, 4, 3, 2]); + } + + #[test] + fn null_inputs_produce_null_rows() { + let list = run_i64( + vec![None, Some(1), Some(1)], + vec![Some(3), None, Some(3)], + Some(vec![Some(1), Some(1), None]), + ) + .unwrap(); + assert!(list.is_null(0)); + assert!(list.is_null(1)); + assert!(list.is_null(2)); + } + + #[test] + fn illegal_boundaries() { + let err = run_i64(vec![Some(1)], vec![Some(5)], Some(vec![Some(-1)])) + .unwrap_err() + .to_string(); + assert!( + err.contains("Illegal sequence boundaries: 1 to 5 by -1"), + "{err}" + ); + let err = run_i64(vec![Some(1)], vec![Some(5)], Some(vec![Some(0)])) + .unwrap_err() + .to_string(); + assert!( + err.contains("Illegal sequence boundaries: 1 to 5 by 0"), + "{err}" + ); + } + + #[test] + fn length_limit_and_overflow_edges() { + // Plain path: length exceeds MAX_ROUNDED_ARRAY_LENGTH. + let err = run_i64(vec![Some(0)], vec![Some(i64::MAX - 1)], Some(vec![Some(1)])) + .unwrap_err() + .to_string(); + assert!(err.contains("9223372036854775807"), "{err}"); + + // Math.addExact(1, delta / step) overflow: count reported as 2^63. + let err = run_i64(vec![Some(0)], vec![Some(i64::MAX)], Some(vec![Some(1)])) + .unwrap_err() + .to_string(); + assert!(err.contains("9223372036854775808"), "{err}"); + + // Long.MinValue / -1 special case: count reported as 2^63 + 1. + let err = run_i64(vec![Some(0)], vec![Some(i64::MIN)], Some(vec![Some(-1)])) + .unwrap_err() + .to_string(); + assert!(err.contains("9223372036854775809"), "{err}"); + + // subtractExact overflow with a step large enough to keep the exact length small: + // Spark reaches internalError("Unreachable code reached."). + let err = run_i64( + vec![Some(i64::MIN)], + vec![Some(i64::MAX)], + Some(vec![Some(i64::MAX)]), + ) + .unwrap_err() + .to_string(); + assert!(err.contains("Unreachable code reached."), "{err}"); + } + + #[test] + fn narrow_types_and_scalar_inputs() { + let args = vec![ + ColumnarValue::Array(Arc::new(Int8Array::from(vec![Some(1i8), Some(-3)]))), + ColumnarValue::Array(Arc::new(Int8Array::from(vec![Some(5i8), Some(-1)]))), + ]; + let result = spark_sequence(&args, &list_of(DataType::Int8)).unwrap(); + let ColumnarValue::Array(arr) = result else { + unreachable!("array inputs produce an array") + }; + let list = as_primitive_list(&arr); + let v0 = list.value(0); + assert_eq!( + as_primitive_array::(&v0).unwrap().values(), + &[1, 2, 3, 4, 5] + ); + let v1 = list.value(1); + assert_eq!( + as_primitive_array::(&v1).unwrap().values(), + &[-3, -2, -1] + ); + + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int64(Some(1))), + ColumnarValue::Scalar(ScalarValue::Int64(Some(3))), + ]; + let result = spark_sequence(&args, &list_of(DataType::Int64)).unwrap(); + let ColumnarValue::Scalar(ScalarValue::List(list)) = result else { + panic!("all-scalar inputs should produce a List scalar") + }; + assert_eq!(row_values(&list, 0), vec![1, 2, 3]); + } +} diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index a2225df2d5c..5e7562c4180 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -26,9 +26,9 @@ use crate::math_funcs::pow::spark_pow; use crate::{ spark_ceil, spark_day_name, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding, - spark_round, spark_rpad, spark_to_time, spark_unhex, spark_unscaled_value, EvalMode, - SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, - SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkIcebergBucket, + spark_round, spark_rpad, spark_sequence, spark_to_time, spark_unhex, spark_unscaled_value, + EvalMode, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains, + SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, }; @@ -120,6 +120,9 @@ pub fn create_comet_physical_fun_with_eval_mode( "ceil" => { make_comet_scalar_udf!("ceil", spark_ceil, data_type) } + "spark_sequence" => { + make_comet_scalar_udf!("spark_sequence", spark_sequence, data_type) + } "floor" => { make_comet_scalar_udf!("floor", spark_floor, data_type) } diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index 0181aa62418..27a742ce582 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -954,4 +954,61 @@ object CometArraySort extends CometCodegenDispatch[ArraySort] object CometZipWith extends CometCodegenDispatch[ZipWith] -object CometSequence extends CometCodegenDispatch[Sequence] +object CometSequence extends CometExpressionSerde[Sequence] with CodegenDispatchFallback { + + private val temporalUnsupportedReason = + "date and timestamp element types run through the JVM codegen dispatcher" + + private val unsafeArgUnsupportedReason = + "sequence arguments must be literals or column references; other shapes run through the " + + "JVM codegen dispatcher to preserve Spark's per-row null short-circuit" + + override def getSupportLevel(expr: Sequence): SupportLevel = expr.start.dataType match { + case ByteType | ShortType | IntegerType | LongType => + // Spark's codegen for `Sequence` short-circuits per row: any null argument returns null + // without evaluating the rest. DataFusion evaluates each scalar-UDF argument over the + // whole batch before calling the outer kernel, so a sub-expression with side effects + // (a nested call, a `CASE WHEN`, even a zero-arg UDF like `boom()`) could fire on rows + // Spark's null check would have discarded. A tree-shape "no children" test is not + // enough — a zero-arg UDF has empty children but still executes. Only literals and + // column references are safe to lower natively; anything else falls back to the + // codegen dispatcher, which keeps the whole tree inside Spark's guarded evaluation. + if (argsAreLiteralsOrRefs(expr)) Compatible() + else Unsupported(Some(unsafeArgUnsupportedReason)) + case DateType | TimestampType | TimestampNTZType => + // Temporal sequences step through timezone/DST/legacy-calendar arithmetic + // (https://github.com/apache/datafusion-comet/issues/5349), so they stay on the JVM + // codegen dispatcher. + Unsupported(Some(temporalUnsupportedReason)) + case other => + Unsupported(Some(s"sequence with element type $other is not supported natively")) + } + + private def argsAreLiteralsOrRefs(expr: Sequence): Boolean = { + val args = Seq(expr.start, expr.stop) ++ expr.stepOpt + args.forall { + case _: Literal | _: Attribute | _: BoundReference => true + case _ => false + } + } + + override def getUnsupportedReasons(): Seq[String] = + Seq(temporalUnsupportedReason, unsafeArgUnsupportedReason) + + override def convert( + expr: Sequence, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + val startExprProto = exprToProto(expr.start, inputs, binding) + val stopExprProto = exprToProto(expr.stop, inputs, binding) + // With no step argument the native kernel computes Spark's per-row default, + // `start <= stop ? 1 : -1`, which cannot be expressed as a plan-time literal. + val argProtos = Seq(startExprProto, stopExprProto) ++ + expr.stepOpt.map(exprToProto(_, inputs, binding)) + scalarFunctionExprToProtoWithReturnType( + "spark_sequence", + expr.dataType, + failOnError = false, + argProtos: _*) + } +} diff --git a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 6cb64486725..dcb73971901 100644 --- a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -158,10 +158,32 @@ trait ShimSparkErrorConverter { Some(QueryExecutionErrors.exceedMapSizeLimitError(params("size").toString.toInt)) case "CollectionSizeLimitExceeded" => - // createArrayWithElementsExceedLimitError takes (count: Any) in Spark 3.4 + // createArrayWithElementsExceedLimitError takes (count: Any) in Spark 3.4; pass the + // decimal string through since the reported length can exceed Long range. Some( QueryExecutionErrors.createArrayWithElementsExceedLimitError( - params("numElements").toString.toLong)) + params("numElements").toString)) + + case "SequenceIllegalBoundaries" => + // Spark 3.x codegen throws a plain IllegalArgumentException for sequence boundaries. + Some( + new IllegalArgumentException( + s"Illegal sequence boundaries: ${params("start")} to ${params("stop")} " + + s"by ${params("step")}")) + + case "SequenceBatchTooLarge" => + // Comet-specific per-batch limit for native `sequence`. Point the user at + // spark.comet.batchSize since Spark itself has no equivalent guard. + Some( + new SparkException( + "Comet's native `sequence` kernel cannot materialize a batch with " + + s"${params("totalElements")} total elements: it exceeds the per-batch " + + "limit or the allocator refused the reservation. Lower " + + "`spark.comet.batchSize` so fewer rows are grouped per batch.", + null)) + + case "Internal" => + Some(SparkException.internalError(params("message").toString)) case "NotNullAssertViolation" => Some( diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 6b976e55de1..571eaf05547 100644 --- a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -156,9 +156,32 @@ trait ShimSparkErrorConverter { Some(QueryExecutionErrors.exceedMapSizeLimitError(params("size").toString.toInt)) case "CollectionSizeLimitExceeded" => + // createArrayWithElementsExceedLimitError takes (count: Any) in Spark 3.5; pass the + // decimal string through since the reported length can exceed Long range. Some( QueryExecutionErrors.createArrayWithElementsExceedLimitError( - ("array", params("numElements").toString.toLong))) + params("numElements").toString)) + + case "SequenceIllegalBoundaries" => + // Spark 3.x codegen throws a plain IllegalArgumentException for sequence boundaries. + Some( + new IllegalArgumentException( + s"Illegal sequence boundaries: ${params("start")} to ${params("stop")} " + + s"by ${params("step")}")) + + case "SequenceBatchTooLarge" => + // Comet-specific per-batch limit for native `sequence`. Point the user at + // spark.comet.batchSize since Spark itself has no equivalent guard. + Some( + new SparkException( + "Comet's native `sequence` kernel cannot materialize a batch with " + + s"${params("totalElements")} total elements: it exceeds the per-batch " + + "limit or the allocator refused the reservation. Lower " + + "`spark.comet.batchSize` so fewer rows are grouped per batch.", + null)) + + case "Internal" => + Some(SparkException.internalError(params("message").toString)) case "NotNullAssertViolation" => Some( diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 7fb822b66bd..7397745885c 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -178,10 +178,35 @@ trait ShimSparkErrorConverter { Some(QueryExecutionErrors.exceedMapSizeLimitError(params("size").toString.toInt)) case "CollectionSizeLimitExceeded" => + // Pass the count as its decimal string since the reported length can exceed Long range. Some( QueryExecutionErrors.createArrayWithElementsExceedLimitError( - "array", - params("numElements").toString.toLong)) + params.getOrElse("functionName", "array").toString, + params("numElements").toString)) + + case "SequenceIllegalBoundaries" => + // Matches what Spark 4.x codegen throws for sequence boundaries. + Some( + new SparkIllegalArgumentException( + errorClass = "_LEGACY_ERROR_TEMP_3243", + messageParameters = Map( + "start" -> params("start").toString, + "stop" -> params("stop").toString, + "step" -> params("step").toString))) + + case "SequenceBatchTooLarge" => + // Comet-specific per-batch limit for native `sequence`. Point the user at + // spark.comet.batchSize since Spark itself has no equivalent guard. + Some( + new SparkException( + "Comet's native `sequence` kernel cannot materialize a batch with " + + s"${params("totalElements")} total elements: it exceeds the per-batch " + + "limit or the allocator refused the reservation. Lower " + + "`spark.comet.batchSize` so fewer rows are grouped per batch.", + null)) + + case "Internal" => + Some(SparkException.internalError(params("message").toString)) case "NotNullAssertViolation" => Some( diff --git a/spark/src/test/resources/sql-tests/expressions/array/sequence.sql b/spark/src/test/resources/sql-tests/expressions/array/sequence.sql index 12aa91c2ccc..7b6446b7e48 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/sequence.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/sequence.sql @@ -15,17 +15,175 @@ -- specific language governing permissions and limitations -- under the License. --- Routes sequence through the codegen dispatcher so behavior matches Spark exactly. +-- sequence(start, stop[, step]) for integral element types runs on the native kernel +-- (https://github.com/apache/datafusion-comet/issues/5349). Date and timestamp sequences +-- stay on the JVM codegen dispatcher and are exercised at the bottom of this file. statement -CREATE TABLE test_sequence(a int, b int) USING parquet +CREATE TABLE test_sequence( + b_start tinyint, b_stop tinyint, b_step tinyint, + s_start smallint, s_stop smallint, s_step smallint, + i_start int, i_stop int, i_step int, + l_start bigint, l_stop bigint, l_step bigint) +USING parquet +-- Row 2 descends, row 3 has start == stop, rows 4-6 carry NULLs in each argument position. statement -INSERT INTO test_sequence VALUES (1, 5), (5, 1), (3, 3), (NULL, 5) +INSERT INTO test_sequence VALUES + (1Y, 5Y, 1Y, 1S, 5S, 1S, 1, 10, 3, 1L, 5L, 2L), + (-3Y, -1Y, 1Y, 100S, 90S, -2S, 20, 2, -6, 9223372036854775802L, 9223372036854775807L, 1L), + (0Y, 0Y, 0Y, -5S, -5S, 0S, 7, 7, 0, -9223372036854775808L, -9223372036854775800L, 3L), + (NULL, 5Y, 1Y, NULL, 5S, 1S, NULL, 10, 1, NULL, 5L, 1L), + (1Y, NULL, 1Y, 1S, NULL, 1S, 1, NULL, 1, 1L, NULL, 1L), + (1Y, 5Y, NULL, 1S, 5S, NULL, 1, 10, NULL, 1L, 5L, NULL) + +-- ============================================================================ +-- Explicit step, all four integral types +-- ============================================================================ + +query +SELECT sequence(i_start, i_stop, i_step) FROM test_sequence + +query +SELECT sequence(l_start, l_stop, l_step) FROM test_sequence + +-- Column step for the narrow integral types exercises the Byte/Short monomorphizations +-- of the native kernel, not just the literal-step shape. +query +SELECT sequence(b_start, b_stop, b_step) FROM test_sequence + +query +SELECT sequence(s_start, s_stop, s_step) FROM test_sequence + +-- ============================================================================ +-- Default step: per-row start <= stop ? 1 : -1, both directions in one column +-- ============================================================================ + +query +SELECT sequence(b_start, b_stop), sequence(s_start, s_stop) FROM test_sequence + +query +SELECT sequence(i_start, i_stop), sequence(l_start, l_stop) FROM test_sequence + +-- ============================================================================ +-- Literal and mixed literal/column arguments +-- ============================================================================ query -SELECT a, b, sequence(a, b) FROM test_sequence +SELECT sequence(1, 10), sequence(10, 1), sequence(5, 5), sequence(5, 5, 0) --- literal arguments with step query SELECT sequence(1, 5), sequence(5, 1, -1), sequence(1, 10, 2) + +query +SELECT sequence(1L, 9L, 2L), sequence(-128Y, -120Y), sequence(32760S, 32767S) + +-- On row 2 the source row is (i_start=20, i_stop=2, i_step=-6), so the literal-step column +-- asks for sequence(1, 2, 2) = [1] while the default-step column asks for sequence(20, 25) +-- = [20, 21, 22, 23, 24, 25]. The two columns disagreeing in direction on the same row is +-- intentional coverage, not an oversight. +query +SELECT sequence(1, i_stop, 2), sequence(i_start, 25) FROM test_sequence WHERE i_start IS NOT NULL AND i_stop IS NOT NULL + +query +SELECT sequence(CAST(NULL AS int), 5), sequence(1, CAST(NULL AS int)), sequence(1, 5, CAST(NULL AS int)) + +-- Integer.MIN_VALUE/MAX_VALUE bounds for int, and a sequence spanning zero +query +SELECT sequence(2147483642, 2147483647), sequence(-2147483648, -2147483643), sequence(-3, 3, 3) + +-- ============================================================================ +-- sequence feeding explode, the common date-spine shape (with integers) +-- ============================================================================ + +query +SELECT i_start, x FROM test_sequence LATERAL VIEW explode(sequence(i_start, i_stop)) AS x WHERE i_start IS NOT NULL AND i_stop IS NOT NULL + +-- ============================================================================ +-- Error paths: step direction contradicts bounds, or zero step with start != stop +-- ============================================================================ + +query expect_error(Illegal sequence boundaries: 1 to 5 by -1) +SELECT sequence(1, 5, -1) + +query expect_error(Illegal sequence boundaries: 10 to 2 by 3) +SELECT sequence(10, 2, 3) FROM test_sequence LIMIT 1 + +query expect_error(Illegal sequence boundaries: 1 to 5 by 0) +SELECT sequence(1, 5, 0) + +-- ============================================================================ +-- Error paths: length exceeds MAX_ROUNDED_ARRAY_LENGTH +-- ============================================================================ + +query expect_error(the array size limit 2147483632) +SELECT sequence(0L, 4294967296L, 1L) + +-- Math.addExact overflow inside Spark's sequenceLength: reported count is 2^63 +query expect_error(9223372036854775808) +SELECT sequence(0L, 9223372036854775807L, 1L) + +-- Long.MinValue / -1 special case: reported count is 2^63 + 1 +query expect_error(9223372036854775809) +SELECT sequence(0L, -9223372036854775808L, -1L) + +-- delta overflows long but the exact length is tiny: Spark reaches an internal error +query expect_error(Unreachable code reached) +SELECT sequence(-9223372036854775808L, 9223372036854775807L, 9223372036854775807L) + +-- ============================================================================ +-- Full narrow-type range: writes at the byte/short boundary. Spark's kernel +-- accumulates with the element type's `Numeric`, wrapping at 8 and 16 bits; +-- ours accumulates in i64 and truncates on the way out. The two agree because +-- every element is inside range, but this locks in the boundary values. +-- ============================================================================ + +query +SELECT sequence(-128Y, 127Y), sequence(-32768S, 32767S) + +-- ============================================================================ +-- Int32 boundary product: index * step overflows int, exercising the Int32 +-- monomorphization at the extreme. +-- ============================================================================ + +query +SELECT sequence(-2147483648, 2147483647, 1073741824) + +-- ============================================================================ +-- Null short-circuit under a nested sequence: Spark's codegen returns NULL +-- without evaluating the inner argument, so the inner `sequence(1, 5, -1)` +-- must not fire on the NULL row. Non-leaf argument shapes stay on the JVM +-- codegen dispatcher for this reason +-- (https://github.com/apache/datafusion-comet/pull/5614#discussion_r3910237757). +-- ============================================================================ + +statement +CREATE TABLE t_seq_null_short_circuit(s INT, k INT) USING parquet + +statement +INSERT INTO t_seq_null_short_circuit VALUES (NULL, -1), (1, 1) + +query +SELECT sequence(s, size(sequence(1, 5, k))) FROM t_seq_null_short_circuit + +-- ============================================================================ +-- Throwing sub-expression guarded by CASE WHEN: DataFusion filters the batch +-- per branch, so `sequence(1, 5, k)` is never evaluated on rows where the +-- ELSE branch is taken. Locks in that we do not diverge from Spark here. +-- ============================================================================ + +query +SELECT CASE WHEN k > 0 THEN sequence(1, 5, k) ELSE array(-1) END FROM t_seq_null_short_circuit + +-- ============================================================================ +-- Date and timestamp sequences keep running on the JVM codegen dispatcher +-- ============================================================================ + +query +SELECT sequence(DATE'2024-01-01', DATE'2024-01-10') + +query +SELECT sequence(DATE'2024-01-01', DATE'2024-12-31', INTERVAL 1 MONTH) + +query +SELECT sequence(TIMESTAMP'2024-01-01 00:00:00', TIMESTAMP'2024-01-01 06:00:00', INTERVAL 2 HOUR) diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index f36887cf979..5d453f48f0b 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -210,6 +210,84 @@ class CometCodegenSuite } } + private def withSequenceTable(f: => Unit): Unit = { + withTable("t") { + // `stp` carries a sign-correct step so `sequence(a, b, stp)` is legal on both rows: + // ascending (1, 5, 1) and descending (9, 2, -1). A single literal step would raise + // `Illegal sequence boundaries` on the mismatched row inside Spark's reference run. + sql("CREATE TABLE t (a INT, b INT, stp INT, d DATE) USING parquet") + sql("INSERT INTO t VALUES (1, 5, 1, DATE'2024-01-01'), (9, 2, -1, DATE'2024-03-01')") + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE)(f) + } + } + + test("sequence with leaf integral args runs natively") { + // Integral sequence with column-reference/literal args lowers to the native spark_sequence + // kernel; no codegen-dispatch marker should appear. The three-argument form uses the `stp` + // column so both the ascending and descending rows have a sign-correct step (all args + // stay leaves, so the native path is exercised). + withSequenceTable { + val df = sql("SELECT sequence(a, b), sequence(a, b, stp) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + !explain.contains("JVM codegen dispatcher"), + s"expected integral sequence with leaf args to run natively, got:\n$explain") + } + } + + test("sequence with zero-arg UDF stop routes through the dispatcher") { + // A zero-argument Scala UDF has empty `children` but still fires on evaluation. The gate + // must reject it (rather than treating it as a safe leaf) so DataFusion does not call it + // over the whole batch on rows Spark's per-row null short-circuit would have skipped. + spark.udf.register("comet_seq_stopper", () => 10) + withSequenceTable { + val df = sql("SELECT sequence(a, comet_seq_stopper()) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + explain.contains("JVM codegen dispatcher: sequence"), + s"expected zero-arg-UDF sequence to route through the dispatcher, got:\n$explain") + } + } + + test("sequence with non-leaf integral args routes through the dispatcher") { + // A non-leaf argument (e.g. a `CASE WHEN` step) would be evaluated over the whole batch by + // DataFusion before the outer kernel runs, breaking Spark's per-row null short-circuit. + // `CometSequence` reports `Unsupported` for these shapes and hands them to the JVM codegen + // dispatcher. + withSequenceTable { + val df = sql("SELECT sequence(a, b, CASE WHEN a <= b THEN 2 ELSE -2 END) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + explain.contains("JVM codegen dispatcher: sequence"), + s"expected composed-arg sequence to route through the dispatcher, got:\n$explain") + } + } + + test("sequence with date element type routes through the dispatcher") { + // Date/timestamp sequences step through timezone/DST/legacy-calendar arithmetic + // (issue #5349), so `CometSequence` keeps them on the JVM codegen dispatcher. + withSequenceTable { + val df = sql("SELECT sequence(d, DATE'2024-06-01', INTERVAL 1 MONTH) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + explain.contains("JVM codegen dispatcher: sequence"), + s"expected date sequence to route through the dispatcher, got:\n$explain") + } + } + test("expression coverage stats split native from codegen-dispatch expressions") { // `abs` and `sqrt` lower to native DataFusion expressions; `hypot` and `nanvl` are // `CometCodegenDispatch` and so run Spark's own codegen inside the Comet pipeline. The diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala new file mode 100644 index 00000000000..93f41ccf6b2 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala @@ -0,0 +1,80 @@ +/* + * 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 + +/** + * Benchmark to measure performance of Comet's native `sequence` kernel against Spark's codegen + * (issue #5349). Integral shapes run natively under Comet; the date case stays on the JVM codegen + * dispatcher in both arms and is included to show that path is unchanged. To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometSequenceBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometSequenceBenchmark-**results.txt". + * + * Endpoints are materialized as columns in the prepared Parquet table. The native path is only + * eligible when every argument to `sequence` is a literal or a column reference (see + * `CometSequence.argsAreLiteralsOrRefs`), so an arithmetic argument such as `c_start + 4` would + * silently route the whole expression through the JVM codegen dispatcher and defeat the point of + * the benchmark. + */ +object CometSequenceBenchmark extends CometBenchmarkBase { + + private val sequenceQueries = List( + ("seq_short_5_elems", "SELECT sequence(c_start, c_stop_5) FROM parquetV1Table"), + ("seq_spine_365_elems", "SELECT sequence(c_start, c_stop_365) FROM parquetV1Table"), + ("seq_long_10000_elems", "SELECT sequence(c_start, c_stop_10000) FROM parquetV1Table"), + ("seq_descending_default_step", "SELECT sequence(c_stop_365, c_start) FROM parquetV1Table"), + ("seq_explicit_step_7", "SELECT sequence(c_start, c_stop_365, 7L) FROM parquetV1Table"), + ( + "seq_sparse_nulls_365_elems", + "SELECT sequence(c_null_start, c_null_stop_365) FROM parquetV1Table"), + // Date/timestamp element types always stay on the JVM codegen dispatcher regardless of + // argument shape, so the arithmetic form here is intentional — this case is the control that + // shows the dispatcher path is unchanged. + ( + "seq_date_spine_dispatcher", + "SELECT sequence(c_date, c_date + INTERVAL 364 DAYS) FROM parquetV1Table")) + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + runBenchmarkWithTable("sequence", 8192) { v => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql("SELECT CAST(PMOD(value, 100000) AS BIGINT) AS c_start," + + " CAST(PMOD(value, 100000) AS BIGINT) + 4 AS c_stop_5," + + " CAST(PMOD(value, 100000) AS BIGINT) + 364 AS c_stop_365," + + " CAST(PMOD(value, 100000) AS BIGINT) + 9999 AS c_stop_10000," + + " CASE WHEN PMOD(value, 10) = 0 THEN CAST(NULL AS BIGINT)" + + " ELSE CAST(PMOD(value, 100000) AS BIGINT) END AS c_null_start," + + " CASE WHEN PMOD(value, 10) = 0 THEN CAST(NULL AS BIGINT)" + + " ELSE CAST(PMOD(value, 100000) AS BIGINT) + 364 END AS c_null_stop_365," + + s" DATE_ADD(DATE'2020-01-01', CAST(PMOD(value, 3650) AS INT)) AS c_date FROM $tbl")) + + sequenceQueries.foreach { case (name, query) => + runBenchmark(name) { + runExpressionBenchmark(name, v, query) + } + } + } + } + } + } +}