You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/11 10:58:39 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6833-52a3a336a79c8736b5824274f26d4d9d1d1666aa created (now 50d73b37e4)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6833-52a3a336a79c8736b5824274f26d4d9d1d1666aa
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 50d73b37e4 perf: integral sequence optimisation (#6833)
This branch includes the following new commits:
new 50d73b37e4 perf: integral sequence optimisation (#6833)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: perf: integral sequence optimisation (#6833)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6833-52a3a336a79c8736b5824274f26d4d9d1d1666aa
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 50d73b37e461f7aa2ced66287939bf3537d160b7
Author: nat-openai <na...@openai.com>
AuthorDate: Sun Oct 11 10:58:19 2026 +0000
perf: integral sequence optimisation (#6833)
* perf: optimise integral sequence generation
Reduce the per-row and per-element cost of integral sequence(start, stop,
step) while preserving the existing Spark-compatible results and errors.
- Borrow typed scalar arguments instead of materialising full input columns.
Keep all-scalar results as List scalars and preserve constant-folding and
caller expansion behaviour.
- Calculate ordinary sequence lengths with checked i64 arithmetic, including
unit-step fast paths. Retain exact wide arithmetic for overflow so Spark's
boundary, collection-size and internal-error cases report the same errors.
- Reserve the child values buffer once and initialise fixed slices using
independent index arithmetic. This removes per-element capacity checks and
lets LLVM vectorise stores. Wrapping intermediates preserve values near
integer bounds; publish the vector only after every value is initialised.
- Union and reuse Arrow validity buffers once per batch instead of checking
each argument and rebuilding validity for every row. Return all-null lists
directly after input type and length validation, preserving null suppression
of invalid underlying values and supporting sliced input bitmaps.
- Retain per-row lengths between sizing and generation so short sequences do
not repeat length arithmetic. This avoids the short-row regressions found
during measurement without introducing separate generation kernels.
Preserve validation order: every row is checked before the aggregate Arrow
List offset ceiling, so a later Spark row error still precedes that ceiling.
Keep existing scalar result shape, null placement, allocator-failure handling
and i32 offset limits. This change introduces no sequence byte budget, buffer
ownership policy, cancellation polling or JNI callbacks.
Extend behavioural tests across all scalar/array argument combinations and
integral widths, extreme values and steps, sliced null bitmaps, empty and
mismatched arrays, length overflow and row-versus-batch error precedence.
Extend the Criterion benchmark to 32 shapes covering one-element and long
sequences, default and explicit steps, mixed directions, scalar expansion,
null densities and errors. Record the performance audit: long column-input
sequences improve by approximately 5x on Apple M3 Max, with no reproducible
regression in the measured shapes. Repeated all-scalar batch expansion is
effectively unchanged.
* test: restore integral sequence benchmark coverage
Keep the optimisation split's benchmark coverage aligned with the original
sequence proposal while retaining the existing execution and memory policy.
- Add the JVM benchmark queries and matching input columns for sequences of
2, 1,000 and 50,000 elements. These fill the short and long column-input
gaps without changing Spark session setup or introducing byte-budget
configuration or cancellation-related settings.
- Add 2-element and 10,000-element all-scalar cases to the native benchmark,
retaining the one-element case. Exercise both one-row and 8,192-row output
expansion through spark_sequence(...).into_array(rows), increasing native
coverage from 32 to 36 shapes while measuring equivalent caller work.
- Clarify that the recorded approximately 5x speedup compares column-valued
inputs against the previous native kernel, rather than against Spark's
implementation or all-scalar batch expansion.
The production kernel is unchanged. Validate the Scala benchmark with the
reactor test-compile and Spotless checks, compile the native release benchmark,
and check Rust and Markdown formatting. The restored cases are ready for
benchmark execution; this commit makes no new performance claims for them.
---
.../expression-audits/array_funcs.md | 1 +
native/spark-expr/benches/sequence.rs | 72 ++-
native/spark-expr/src/array_funcs/sequence.rs | 498 ++++++++++++++++-----
.../sql/benchmark/CometSequenceBenchmark.scala | 6 +
4 files changed, 461 insertions(+), 116 deletions(-)
diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md
index 50ac45b852..9c918cbd87 100644
--- a/docs/source/contributor-guide/expression-audits/array_funcs.md
+++ b/docs/source/contributor-guide/expression-audits/array_funcs.md
@@ -180,6 +180,7 @@
## sequence
+- Performance (tuned 2026-10-09): checked i64 length arithmetic, borrowed scalar inputs, bulk validity union, and vectorisable child-buffer stores; approximately 5× faster than the previous native kernel for column-valued inputs producing 365-element and 10,000-element sequences on arm64, with short/null-heavy shapes also improved. Benchmark: `benches/sequence.rs`.
- 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 `internal [...]
- 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.
diff --git a/native/spark-expr/benches/sequence.rs b/native/spark-expr/benches/sequence.rs
index 4a67a5f33e..200500a0a7 100644
--- a/native/spark-expr/benches/sequence.rs
+++ b/native/spark-expr/benches/sequence.rs
@@ -18,6 +18,7 @@
use arrow::array::Int64Array;
use arrow::datatypes::{DataType, Field};
use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::ScalarValue;
use datafusion::physical_plan::ColumnarValue;
use datafusion_comet_spark_expr::spark_sequence;
use std::hint::black_box;
@@ -57,7 +58,7 @@ fn criterion_benchmark(c: &mut Criterion) {
let mut group = c.benchmark_group("sequence");
// Short sequences: per-row overhead dominates.
- for elems in [2i64, 5] {
+ for elems in [1i64, 2, 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()))
@@ -111,6 +112,33 @@ fn criterion_benchmark(c: &mut Criterion) {
});
}
+ // Explicit steps on one-element rows, and non-unit steps in both directions.
+ for (label, elems) in [
+ ("explicit_step_one_element", 1i64),
+ ("mixed_non_unit_steps", 17),
+ ] {
+ let start = Int64Array::from(vec![0i64; NUM_ROWS]);
+ let step = Int64Array::from(
+ (0..NUM_ROWS)
+ .map(|i| if i % 2 == 0 { 3i64 } else { -3 })
+ .collect::<Vec<_>>(),
+ );
+ let stop = Int64Array::from(
+ step.values()
+ .iter()
+ .map(|step| step * (elems - 1))
+ .collect::<Vec<_>>(),
+ );
+ let args = vec![
+ ColumnarValue::Array(Arc::new(start)),
+ ColumnarValue::Array(Arc::new(stop)),
+ ColumnarValue::Array(Arc::new(step)),
+ ];
+ group.bench_function(label, |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]);
@@ -126,6 +154,48 @@ fn criterion_benchmark(c: &mut Criterion) {
});
}
+ // Scalar steps avoid the cost of materialising an input column.
+ for elems in [1i64, 2, 5, 365, 10_000] {
+ let mut args = args_with_len(elems, None);
+ args.push(ColumnarValue::Scalar(ScalarValue::Int64(Some(1))));
+ group.bench_function(format!("scalar_step_{elems}_elems"), |b| {
+ b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap()))
+ });
+ }
+
+ // Include the caller's expansion so scalar and array results do equivalent work.
+ for rows in [1, NUM_ROWS] {
+ for elems in [1i64, 2, 5, 365, 10_000] {
+ let args = vec![
+ ColumnarValue::Scalar(ScalarValue::Int64(Some(0))),
+ ColumnarValue::Scalar(ScalarValue::Int64(Some(elems - 1))),
+ ];
+ group.bench_function(format!("all_scalar_{rows}_rows_{elems}_elems"), |b| {
+ b.iter(|| {
+ black_box(
+ spark_sequence(&args, &return_type)
+ .unwrap()
+ .into_array(rows)
+ .unwrap(),
+ )
+ })
+ });
+ }
+ }
+
+ for (label, every) in [
+ ("sparse_nulls", 10usize),
+ ("dense_nulls", 2),
+ ("all_nulls", 1),
+ ] {
+ for elems in [1i64, 2, 5] {
+ let args = args_with_len(elems, Some(every));
+ group.bench_function(format!("{label}_{elems}_elems"), |b| {
+ b.iter(|| black_box(spark_sequence(&args, &return_type).unwrap()))
+ });
+ }
+ }
+
group.finish();
}
diff --git a/native/spark-expr/src/array_funcs/sequence.rs b/native/spark-expr/src/array_funcs/sequence.rs
index 4b1335f20f..17d519c7ce 100644
--- a/native/spark-expr/src/array_funcs/sequence.rs
+++ b/native/spark-expr/src/array_funcs/sequence.rs
@@ -29,19 +29,19 @@
use std::sync::Arc;
-use arrow::array::{Array, ArrayRef, ListArray, NullBufferBuilder, PrimitiveArray};
-use arrow::buffer::{OffsetBuffer, ScalarBuffer};
+use arrow::array::{Array, ArrayRef, ListArray, PrimitiveArray};
+use arrow::buffer::{NullBuffer, 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::common::{exec_err, internal_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;
+const MAX_ROUNDED_ARRAY_LENGTH: i64 = (i32::MAX - 15) as i64;
pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result<ColumnarValue> {
let child_field = match data_type {
@@ -58,22 +58,29 @@ pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result<Co
let all_scalar = args
.iter()
.all(|arg| matches!(arg, ColumnarValue::Scalar(_)));
- let arrays = ColumnarValue::values_to_arrays(args)?;
- let step = arrays.get(2);
+ let num_rows = args
+ .iter()
+ .find_map(|arg| match arg {
+ ColumnarValue::Array(array) => Some(array.len()),
+ ColumnarValue::Scalar(_) => None,
+ })
+ .unwrap_or(1);
+ for arg in args {
+ if let ColumnarValue::Array(array) = arg {
+ if array.len() != num_rows {
+ return internal_err!(
+ "Arguments has mixed length. Expected length: {num_rows}, found length: {}",
+ array.len()
+ );
+ }
+ }
+ }
let result = match child_field.data_type() {
- DataType::Int8 => {
- sequence_integral::<Int8Type>(&arrays[0], &arrays[1], step, child_field, |v| v as i8)
- }
- DataType::Int16 => {
- sequence_integral::<Int16Type>(&arrays[0], &arrays[1], step, child_field, |v| v as i16)
- }
- DataType::Int32 => {
- sequence_integral::<Int32Type>(&arrays[0], &arrays[1], step, child_field, |v| v as i32)
- }
- DataType::Int64 => {
- sequence_integral::<Int64Type>(&arrays[0], &arrays[1], step, child_field, |v| v)
- }
+ DataType::Int8 => sequence_integral::<Int8Type>(args, num_rows, child_field),
+ DataType::Int16 => sequence_integral::<Int16Type>(args, num_rows, child_field),
+ DataType::Int32 => sequence_integral::<Int32Type>(args, num_rows, child_field),
+ DataType::Int64 => sequence_integral::<Int64Type>(args, num_rows, child_field),
other => exec_err!("spark_sequence does not support element type {other:?}"),
}?;
@@ -86,59 +93,114 @@ pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result<Co
}
}
-fn sequence_integral<T: ArrowPrimitiveType>(
- start: &ArrayRef,
- stop: &ArrayRef,
- step: Option<&ArrayRef>,
+trait Integral: ArrowPrimitiveType {
+ fn scalar(value: &ScalarValue) -> Result<Option<Self::Native>>;
+ fn from_i64(value: i64) -> Self::Native;
+}
+macro_rules! integral {
+ ($ty:ty, $variant:ident, $native:ty) => {
+ impl Integral for $ty {
+ fn scalar(value: &ScalarValue) -> Result<Option<Self::Native>> {
+ match value {
+ ScalarValue::$variant(value) => Ok(*value),
+ _ => exec_err!(
+ "spark_sequence expects {} arguments, got {}",
+ Self::DATA_TYPE,
+ value.data_type()
+ ),
+ }
+ }
+ fn from_i64(value: i64) -> Self::Native {
+ value as $native
+ }
+ }
+ };
+}
+integral!(Int8Type, Int8, i8);
+integral!(Int16Type, Int16, i16);
+integral!(Int32Type, Int32, i32);
+integral!(Int64Type, Int64, i64);
+
+enum Input<'a, T: Integral> {
+ Scalar(Option<T::Native>),
+ Array(&'a PrimitiveArray<T>),
+}
+impl<'a, T: Integral> Input<'a, T>
+where
+ T::Native: Into<i64>,
+{
+ fn new(value: &'a ColumnarValue) -> Result<Self> {
+ match value {
+ ColumnarValue::Scalar(value) => Ok(Self::Scalar(T::scalar(value)?)),
+ ColumnarValue::Array(array) => Ok(Self::Array(as_primitive_array::<T>(array)?)),
+ }
+ }
+ #[inline(always)]
+ fn value(&self, row: usize) -> i64 {
+ match self {
+ Self::Scalar(value) => value.unwrap().into(),
+ Self::Array(array) => array.values()[row].into(),
+ }
+ }
+ fn nulls(&self, len: usize) -> Option<NullBuffer> {
+ match self {
+ Self::Scalar(None) => Some(NullBuffer::new_null(len)),
+ Self::Scalar(Some(_)) => None,
+ Self::Array(array) => array.nulls().cloned(),
+ }
+ }
+}
+
+fn sequence_integral<T: Integral>(
+ args: &[ColumnarValue],
+ num_rows: usize,
child_field: FieldRef,
- from_i64: impl Fn(i64) -> T::Native,
) -> Result<ArrayRef>
where
T::Native: Into<i64>,
{
- let start = as_primitive_array::<T>(start)?;
- let stop = as_primitive_array::<T>(stop)?;
- let step = step.map(|arr| as_primitive_array::<T>(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 start = Input::<T>::new(&args[0])?;
+ let stop = Input::<T>::new(&args[1])?;
+ let step = args.get(2).map(Input::<T>::new).transpose()?;
+ // Reuse validity once per batch instead of checking three inputs and rebuilding a
+ // bitmap row by row. Arrow unions sliced bitmaps with bulk word operations.
+ let nulls = NullBuffer::union(
+ start.nulls(num_rows).as_ref(),
+ stop.nulls(num_rows).as_ref(),
+ );
+ let step_nulls = step.as_ref().and_then(|step| step.nulls(num_rows));
+ let nulls = NullBuffer::union(nulls.as_ref(), step_nulls.as_ref())
+ .filter(|nulls| nulls.null_count() != 0);
+ if nulls
+ .as_ref()
+ .is_some_and(|nulls| nulls.null_count() == num_rows)
+ {
+ return Ok(Arc::new(ListArray::new_null(child_field, num_rows)));
+ }
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
- }
- }
+ match &step {
+ Some(step) => step.value(row),
+ None if start <= stop => 1,
+ None => -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<usize> = Vec::with_capacity(num_rows);
- let mut total: usize = 0;
+ // Validate every row before checking the batch total, preserving the first Spark error.
+ // Retain lengths so short sequences do not repeat per-row arithmetic during generation.
+ let mut lengths = Vec::with_capacity(num_rows);
+ let mut total = 0usize;
for row in 0..num_rows {
- if row_is_null(row) {
+ if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row)) {
lengths.push(0);
continue;
}
- let s: i64 = start.value(row).into();
- let e: i64 = stop.value(row).into();
+ let s = start.value(row);
+ let e = stop.value(row);
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.
+ // Arrow's List offsets must fit i32; Spark's per-row array limit is checked above.
if total > i32::MAX as usize {
return Err(DataFusionError::External(Box::new(
SparkError::SequenceBatchTooLarge {
@@ -147,94 +209,99 @@ where
)));
}
- // 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<T::Native> = Vec::new();
values.try_reserve_exact(total).map_err(|_| {
DataFusionError::External(Box::new(SparkError::SequenceBatchTooLarge {
total_elements: total.to_string(),
}))
})?;
- let mut offsets: Vec<i32> = Vec::with_capacity(num_rows + 1);
+ let slots = &mut values.spare_capacity_mut()[..total];
+ let mut offsets = Vec::with_capacity(num_rows + 1);
offsets.push(0);
- let mut nulls = NullBufferBuilder::new(num_rows);
+ let mut written = 0;
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();
+ if len != 0 {
+ let s = start.value(row);
+ let e = stop.value(row);
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);
+ // Fixed bounds and independent arithmetic let LLVM vectorise the stores without
+ // a capacity check for each element. Wrapping intermediates preserve i64 extremes.
+ for (i, slot) in slots[written..written + len].iter_mut().enumerate() {
+ slot.write(T::from_i64(s.wrapping_add(step.wrapping_mul(i as i64))));
}
+ written += len;
}
- offsets.push(values.len() as i32);
+ offsets.push(written as i32);
}
-
+ // SAFETY: the sizing pass reserved `total` slots and the loop initialised each exactly
+ // once. No fallible operation exposes the partially initialised vector.
+ unsafe { values.set_len(total) };
let values = PrimitiveArray::<T>::new(ScalarBuffer::from(values), None);
- let list = ListArray::try_new(
+ Ok(Arc::new(ListArray::try_new(
child_field,
OffsetBuffer::new(offsets.into()),
Arc::new(values),
- nulls.finish(),
- )?;
- Ok(Arc::new(list))
+ nulls,
+ )?))
}
-/// 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.
+/// Match released Spark 3.4--4.1 Sequence.sequenceLength, including its wide-arithmetic
+/// fallback error when subtraction overflows even though the resulting length would fit.
+#[inline(always)]
fn sequence_length(start: i64, stop: i64, step: i64) -> Result<usize> {
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(),
- },
- )));
+ return Err(boundary_error(start, stop, step));
}
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 let Some(len) = stop
+ .checked_sub(start)
+ .and_then(|delta| match step {
+ 1 => Some(delta),
+ -1 => delta.checked_neg(),
+ _ => delta.checked_div(step),
+ })
+ .and_then(|quotient| quotient.checked_add(1))
+ {
+ if len <= MAX_ROUNDED_ARRAY_LENGTH {
+ return Ok(len as usize);
+ }
+ return Err(length_error(len as i128));
}
- 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(),
- ))));
+ sequence_length_overflow(start, stop, step)
+}
+
+#[cold]
+fn boundary_error(start: i64, stop: i64, step: i64) -> DataFusionError {
+ DataFusionError::External(Box::new(SparkError::SequenceIllegalBoundaries {
+ start: start.to_string(),
+ stop: stop.to_string(),
+ step: step.to_string(),
+ }))
+}
+
+#[cold]
+fn sequence_length_overflow(start: i64, stop: i64, step: i64) -> Result<usize> {
+ // Only the rare overflow/error path needs wide division. There is no separate add-overflow
+ // branch: its exact length necessarily exceeds Spark's maximum and is rejected here.
+ let len = 1 + (stop as i128 - start as i128) / step as i128;
+ if len > MAX_ROUNDED_ARRAY_LENGTH as i128 {
+ return Err(length_error(len));
}
- Ok(len as usize)
+ Err(DataFusionError::External(Box::new(SparkError::Internal(
+ "Unreachable code reached.".to_string(),
+ ))))
+}
+
+#[cold]
+fn length_error(len: i128) -> DataFusionError {
+ DataFusionError::External(Box::new(SparkError::CollectionSizeLimitExceeded {
+ num_elements: len.to_string(),
+ max_elements: MAX_ROUNDED_ARRAY_LENGTH,
+ function_name: "sequence".to_string(),
+ }))
}
#[cfg(test)]
@@ -390,4 +457,205 @@ mod tests {
};
assert_eq!(row_values(&list, 0), vec![1, 2, 3]);
}
+ #[test]
+ fn scalar_and_array_combinations_preserve_each_integral_type() {
+ macro_rules! verify {
+ ($ty:ty, $variant:ident, $native:ty) => {{
+ for (start, stop, step) in [
+ (<$native>::MIN, <$native>::MIN + 3, 1),
+ (<$native>::MAX, <$native>::MAX - 3, -1),
+ (0, 0, 0),
+ (-3, 4, 2),
+ (<$native>::MIN, -1, <$native>::MAX),
+ (0, <$native>::MIN, <$native>::MIN),
+ ] {
+ for mask in 0..8 {
+ let args = [start, stop, step]
+ .into_iter()
+ .enumerate()
+ .map(|(index, value)| {
+ let scalar = ScalarValue::$variant(Some(value));
+ if mask & (1 << index) == 0 {
+ ColumnarValue::Scalar(scalar)
+ } else {
+ ColumnarValue::Array(scalar.to_array_of_size(3).unwrap())
+ }
+ })
+ .collect::<Vec<_>>();
+ let result = spark_sequence(&args, &list_of(<$ty>::DATA_TYPE)).unwrap();
+ assert_eq!(matches!(result, ColumnarValue::Scalar(_)), mask == 0);
+ let array = result.into_array(3).unwrap();
+ let list = as_primitive_list(&array);
+ let expected_len = if start == stop {
+ 1
+ } else {
+ (1 + (stop as i128 - start as i128) / step as i128) as usize
+ };
+ for row in 0..3 {
+ let values = list.value(row);
+ let values = as_primitive_array::<$ty>(&values).unwrap();
+ assert_eq!(values.len(), expected_len);
+ for (i, actual) in values.values().iter().enumerate() {
+ assert_eq!(
+ *actual as i128,
+ start as i128 + step as i128 * i as i128
+ );
+ }
+ }
+ }
+ }
+ }};
+ }
+ verify!(Int8Type, Int8, i8);
+ verify!(Int16Type, Int16, i16);
+ verify!(Int32Type, Int32, i32);
+ verify!(Int64Type, Int64, i64);
+ }
+
+ #[test]
+ fn sliced_null_inputs_skip_invalid_underlying_values() {
+ use arrow::buffer::NullBuffer;
+ // Slice at a non-byte-aligned offset. The null start hides invalid boundaries;
+ // the null stop hides a length overflow; the null step hides zero-step boundaries.
+ let arrays = [
+ Int64Array::new(
+ vec![99, 5, i64::MIN, 1, 7, 99].into(),
+ Some(NullBuffer::from(vec![true, false, true, true, true, true])),
+ ),
+ Int64Array::new(
+ vec![99, 1, i64::MAX, 5, 7, 99].into(),
+ Some(NullBuffer::from(vec![true, true, false, true, true, true])),
+ ),
+ Int64Array::new(
+ vec![99, 1, 1, 0, 0, 99].into(),
+ Some(NullBuffer::from(vec![true, true, true, false, true, true])),
+ ),
+ ];
+ let args = arrays
+ .into_iter()
+ .map(|array| ColumnarValue::Array(Arc::new(array.slice(1, 4))))
+ .collect::<Vec<_>>();
+ let result = spark_sequence(&args, &list_of(DataType::Int64))
+ .unwrap()
+ .into_array(4)
+ .unwrap();
+ let list = as_primitive_list(&result);
+ assert_eq!(
+ list.nulls().unwrap(),
+ &NullBuffer::from(vec![false, false, false, true])
+ );
+ assert_eq!(list.value_offsets(), &[0, 0, 0, 0, 1]);
+ assert_eq!(row_values(&list, 3), vec![7]);
+ assert_eq!(list.values().null_count(), 0);
+
+ for index in 0..3 {
+ let mut args = vec![
+ ColumnarValue::Scalar(ScalarValue::Int64(Some(1))),
+ ColumnarValue::Scalar(ScalarValue::Int64(Some(5))),
+ ColumnarValue::Scalar(ScalarValue::Int64(Some(-1))),
+ ];
+ args[index] = ColumnarValue::Scalar(ScalarValue::Int64(None));
+ let result = spark_sequence(&args, &list_of(DataType::Int64)).unwrap();
+ assert!(
+ matches!(result, ColumnarValue::Scalar(ScalarValue::List(ref list)) if list.is_null(0))
+ );
+ }
+ }
+
+ #[test]
+ fn empty_and_mismatched_arrays() {
+ let empty = ColumnarValue::Array(Arc::new(Int64Array::from(Vec::<i64>::new())));
+ let scalar = ColumnarValue::Scalar(ScalarValue::Int64(Some(1)));
+ let result = spark_sequence(&[empty.clone(), scalar.clone()], &list_of(DataType::Int64))
+ .unwrap()
+ .into_array(0)
+ .unwrap();
+ let list = as_primitive_list(&result);
+ assert!(list.is_empty());
+ assert_eq!(list.value_offsets(), &[0]);
+ assert!(list.values().is_empty());
+ let one = ColumnarValue::Array(Arc::new(Int64Array::from(vec![1])));
+ assert!(spark_sequence(&[empty, scalar, one], &list_of(DataType::Int64)).is_err());
+ }
+
+ #[test]
+ fn row_errors_precede_aggregate_limit() {
+ let error = run_i64(
+ vec![Some(0), Some(0), Some(5)],
+ vec![Some(1_073_741_823), Some(1_073_741_823), Some(1)],
+ Some(vec![Some(1), Some(1), Some(1)]),
+ )
+ .unwrap_err()
+ .to_string();
+ assert!(
+ error.contains("Illegal sequence boundaries: 5 to 1 by 1"),
+ "{error}"
+ );
+ let error = run_i64(
+ vec![Some(0), Some(0)],
+ vec![Some(1_073_741_823), Some(1_073_741_823)],
+ None,
+ )
+ .unwrap_err()
+ .to_string();
+ assert!(
+ error.contains("2147483648") && error.contains("spark.comet.batchSize"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn checked_length_matches_exact_released_spark_arithmetic() {
+ let values = [
+ i64::MIN,
+ i64::MIN + 1,
+ -(i32::MAX as i64),
+ -1,
+ 0,
+ 1,
+ i32::MAX as i64,
+ i64::MAX - 1,
+ i64::MAX,
+ ];
+ for start in values {
+ for stop in values {
+ for step in values {
+ let result = sequence_length(start, stop, step);
+ let legal = (step > 0 && start <= stop)
+ || (step < 0 && start >= stop)
+ || (step == 0 && start == stop);
+ if !legal {
+ assert!(result
+ .unwrap_err()
+ .to_string()
+ .contains("Illegal sequence boundaries"));
+ } else if start == stop {
+ assert_eq!(result.unwrap(), 1);
+ } else {
+ let delta = stop as i128 - start as i128;
+ let len = 1 + delta / step as i128;
+ if len > MAX_ROUNDED_ARRAY_LENGTH as i128 {
+ let error = result.unwrap_err().to_string();
+ assert!(
+ error.contains(&len.to_string()),
+ "{start}, {stop}, {step}: {error}"
+ );
+ } else if delta < i64::MIN as i128 || delta > i64::MAX as i128 {
+ assert!(result
+ .unwrap_err()
+ .to_string()
+ .contains("Unreachable code reached."));
+ } else {
+ assert_eq!(result.unwrap(), len as usize);
+ }
+ }
+ }
+ }
+ }
+ assert_eq!(
+ sequence_length(0, MAX_ROUNDED_ARRAY_LENGTH - 1, 1).unwrap(),
+ MAX_ROUNDED_ARRAY_LENGTH as usize
+ );
+ assert!(sequence_length(0, MAX_ROUNDED_ARRAY_LENGTH, 1).is_err());
+ }
}
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
index 1b559ee6dd..0674890f1a 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometSequenceBenchmark.scala
@@ -37,9 +37,12 @@ package org.apache.spark.sql.benchmark
object CometSequenceBenchmark extends CometBenchmarkBase {
private val sequenceQueries = List(
+ ("seq_short_2_elems", "SELECT sequence(c_start, c_stop_2) FROM parquetV1Table"),
("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_1000_elems", "SELECT sequence(c_start, c_stop_1000) FROM parquetV1Table"),
("seq_long_10000_elems", "SELECT sequence(c_start, c_stop_10000) FROM parquetV1Table"),
+ ("seq_long_50000_elems", "SELECT sequence(c_start, c_stop_50000) 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"),
(
@@ -59,9 +62,12 @@ object CometSequenceBenchmark extends CometBenchmarkBase {
prepareTable(
dir,
spark.sql("SELECT CAST(PMOD(value, 100000) AS BIGINT) AS c_start," +
+ " CAST(PMOD(value, 100000) AS BIGINT) + 1 AS c_stop_2," +
" 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) + 999 AS c_stop_1000," +
" CAST(PMOD(value, 100000) AS BIGINT) + 9999 AS c_stop_10000," +
+ " CAST(PMOD(value, 100000) AS BIGINT) + 49999 AS c_stop_50000," +
" 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)" +
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org