| // 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[]`, 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<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); | |
| 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::<T>::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<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(), | |
| }, | |
| ))); | |
| } | |
| 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<Option<i64>>, | |
| stop: Vec<Option<i64>>, | |
| step: Option<Vec<Option<i64>>>, | |
| ) -> Result<ListArray> { | |
| 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::<ListArray>().unwrap().clone() | |
| } | |
| fn row_values(list: &ListArray, row: usize) -> Vec<i64> { | |
| let v = list.value(row); | |
| as_primitive_array::<Int64Type>(&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::<Int8Type>(&v0).unwrap().values(), | |
| &[1, 2, 3, 4, 5] | |
| ); | |
| let v1 = list.value(1); | |
| assert_eq!( | |
| as_primitive_array::<Int8Type>(&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]); | |
| } | |
| } |