andygrove commented on code in PR #5614:
URL: https://github.com/apache/datafusion-comet/pull/5614#discussion_r3906787356
##########
spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala:
##########
@@ -158,10 +158,21 @@ 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 "Internal" =>
Review Comment:
The `case "Internal"` arm changes behavior well beyond `sequence`. There are
around twenty `SparkError::Internal` producers today in `temporal.rs`,
`numeric.rs`, `conversion_funcs/string.rs` and `rlike.rs`, and all of them
previously fell through to the `None` branch in `SparkErrorConverter`, which
renders as `new SparkException(msgParams.mkString(", "))`, so users saw
`(message,<text>)`. After this they all become `[INTERNAL_ERROR] <text>`.
That is a clear improvement and I am not asking you to revert it. Could you
call it out in the PR description though? It is a user-visible message change
for a set of expressions that have nothing to do with `sequence`, and right now
the "How are these changes tested?" section does not mention it. It would also
be worth a pass over the Spark SQL suite diffs to confirm nothing was matching
on the old shape.
The same arm is added to the 3.5 and 4.x shims, so this applies to all three.
##########
spark/src/test/resources/sql-tests/expressions/array/sequence.sql:
##########
@@ -15,17 +15,131 @@
-- 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 a, b, sequence(a, b) FROM test_sequence
+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 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)
Review Comment:
The fixture is thorough on the error paths. There are three shapes I think
are worth pinning down that it does not reach today.
The first is full narrow-type range, `sequence(-128Y, 127Y)` and
`sequence(-32768S, 32767S)`. That is the one place where Spark's `arr(i) =
start + step * num.fromInt(i)` genuinely wraps at 8 and 16 bits while your
kernel accumulates in i64 and truncates on the way out. The two agree because
the true value is always in range, but it is the case I would most want a
regression test on, and `sequence(-128Y, -120Y)` does not get there.
The second is a step whose product with the index overflows int, something
like `sequence(-2147483648, 2147483647, 1073741824)`, which exercises the
`Int32` monomorphization at the boundary.
The third is `sequence` under a `CASE WHEN` where the throwing branch is not
taken, for example `SELECT CASE WHEN step > 0 THEN sequence(1, 5, step) ELSE
array(-1) END FROM t` with rows carrying negative and zero steps. DataFusion
filters the batch before evaluating each `then` branch so this works today, but
it is the one construct where an eagerly evaluated throwing expression would
diverge from Spark, and it would be cheap insurance.
I ran all three locally against this branch and they pass on both 4.1 and
3.5, so this is about locking the behavior in rather than chasing a suspected
bug.
##########
native/spark-expr/src/array_funcs/sequence.rs:
##########
@@ -0,0 +1,388 @@
+// 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<ColumnarValue> {
+ 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::<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)
+ }
+ 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<T: ArrowPrimitiveType>(
+ start: &ArrayRef,
+ stop: &ArrayRef,
+ step: Option<&ArrayRef>,
+ 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 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<usize> = 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[]`. Report it with the same size-limit error class Spark
would raise for a
+ // single overlong row so the user sees a familiar message.
+ if total > i32::MAX as usize {
+ return Err(DataFusionError::External(Box::new(
+ SparkError::CollectionSizeLimitExceeded {
+ num_elements: total.to_string(),
+ max_elements: i32::MAX as i64,
+ function_name: "sequence".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.
+ let mut values: Vec<T::Native> = Vec::with_capacity(total);
Review Comment:
Nice work on the error-parity side of this, the `sequenceLength` port
matches Spark on all three failure paths and I checked it against 3.4.3 through
4.1.1.
The thing I keep coming back to is `Vec::with_capacity(total)`. `sequence`
is the first expression we have made native where the output size is unbounded
relative to the input size. Every other `with_capacity` in `array_funcs/` is
sized by `row_count` or `args.len()`, but here `total` is the sum of every
row's generated length, so a single batch can ask for up to `i32::MAX`
elements, which is 16 GiB for `bigint`. Your own benchmark shows the shape:
`seq_long_10000_elems` materializes 8192 x 10000 x 8 bytes, so 655 MB in one
allocation, where Spark holds one row's `long[]` at a time. That is also the
one row in your table with no speedup, which makes me wonder whether the
large-per-row case is paying for itself at all.
Two things I would like to see. Could the allocation go through
`try_reserve` so an oversized batch surfaces as a query error rather than an
allocator abort that takes the executor down? And could the batch ceiling be
documented somewhere the user can find it?
On the ceiling specifically, the message a user gets today is misleading.
`sequence(0, 262143)` over a full 8192-row batch lands on exactly `2147483648`
total elements and trips the check, even though every individual array is well
inside Spark's limit and Spark itself would run the query. The shim ignores the
`max_elements` you pass, so the message reads "Can't create array with
2147483648 elements which exceeding the array size limit 2147483632", which
points the user at a per-array limit that they have not actually exceeded and
gives them nothing actionable. The real fix on their side is to lower
`spark.comet.batchSize`. Given that, is `Compatible()` the right support level,
or should this at least get a compatible note and a line in the audit entry?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
