0lai0 commented on code in PR #5614: URL: https://github.com/apache/datafusion-comet/pull/5614#discussion_r3913994254
########## 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: Thanks @andygrove for review. I switched to `try_reserve_exact` and added SequenceBatchTooLarge pointing at spark.comet.batchSize. Documented the per-batch ceiling in `array_funcs.md` and a new Limitations section. Leaf-arg integral sequences stay Compatible(). -- 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]
