andygrove commented on code in PR #6819: URL: https://github.com/apache/datafusion-comet/pull/6819#discussion_r4233959866
########## docs/source/contributor-guide/memory_management.md: ########## @@ -370,6 +370,32 @@ unchanged. `spill_replay.rs` describes how the wrapper recognizes these requests on in DataFusion. [Issue #6583](https://github.com/apache/datafusion-comet/issues/6583) tracks removing it. +## Integral sequence output allowance Review Comment: This section adds a budget, but the `Who allocates what` table above still lists the operator pool as the only thing bounding native allocations. Could it get a row for the sequence allowance, in whatever form it ends up? ########## native/spark-expr/src/array_funcs/sequence.rs: ########## @@ -54,187 +200,389 @@ pub fn spark_sequence(args: &[ColumnarValue], data_type: &DataType) -> Result<Co args.len() ); } + for arg in args { + if let ColumnarValue::Array(array) = arg { + if array.len() != rows { + return exec_err!( + "spark_sequence argument has {} rows, expected {rows}", + array.len() + ); + } + } + } + check_cancelled(check)?; + let result = match child_field.data_type() { + DataType::Int8 => sequence_integral::<Int8Type>(args, rows, child_field, pool, check), + DataType::Int16 => sequence_integral::<Int16Type>(args, rows, child_field, pool, check), + DataType::Int32 => sequence_integral::<Int32Type>(args, rows, child_field, pool, check), + DataType::Int64 => sequence_integral::<Int64Type>(args, rows, child_field, pool, check), + other => exec_err!("spark_sequence does not support element type {other:?}"), + }?; + Ok(ColumnarValue::Array(result)) +} - let all_scalar = args - .iter() - .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); - let arrays = ColumnarValue::values_to_arrays(args)?; - let step = arrays.get(2); +fn check_cancelled(check: Option<&CancellationCheck>) -> Result<()> { + if let Some(check) = check { + check()?; + } + Ok(()) +} - let result = match child_field.data_type() { - DataType::Int8 => { - sequence_integral::<Int8Type>(&arrays[0], &arrays[1], step, child_field, |v| v as i8) +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 + } } - DataType::Int16 => { - sequence_integral::<Int16Type>(&arrays[0], &arrays[1], step, child_field, |v| v as i16) + }; +} +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)?)), } - DataType::Int32 => { - sequence_integral::<Int32Type>(&arrays[0], &arrays[1], step, child_field, |v| v as i32) + } + #[inline(always)] + fn is_null(&self, row: usize) -> bool { + match self { + Self::Scalar(value) => value.is_none(), + Self::Array(array) => array.nulls().is_some_and(|nulls| nulls.is_null(row)), } - DataType::Int64 => { - sequence_integral::<Int64Type>(&arrays[0], &arrays[1], step, child_field, |v| v) + } + #[inline(always)] + fn value(&self, row: usize) -> i64 { + match self { + Self::Scalar(value) => value.unwrap().into(), + Self::Array(array) => array.values()[row].into(), } - other => exec_err!("spark_sequence does not support element type {other:?}"), - }?; + } + fn has_nulls(&self) -> bool { + match self { + Self::Scalar(value) => value.is_none(), + Self::Array(array) => array.null_count() != 0, + } + } +} - if all_scalar { - Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( - &result, 0, - )?)) - } else { - Ok(ColumnarValue::Array(result)) +#[derive(Clone, Copy, Default)] +struct Row { + start: i64, + step: i64, + len: usize, +} + +// Keep the borrowed input access in the typed row loops. An out-of-line tuple return adds +// a call and stack writes for each row in both sizing and generation. +#[inline(always)] +fn row_bounds<T: Integral>( + start: &Input<'_, T>, + stop: &Input<'_, T>, + step: Option<&Input<'_, T>>, + has_nulls: bool, + row: usize, +) -> Option<(i64, i64, i64)> +where + T::Native: Into<i64>, +{ + if has_nulls + && (start.is_null(row) || stop.is_null(row) || step.is_some_and(|s| s.is_null(row))) + { + return None; } + let start = start.value(row); + let stop = stop.value(row); + let step = match step { + Some(step) => step.value(row), + None if start <= stop => 1, + None => -1, + }; + Some((start, stop, step)) } -fn sequence_integral<T: ArrowPrimitiveType>( - start: &ArrayRef, - stop: &ArrayRef, - step: Option<&ArrayRef>, +fn sequence_integral<T: Integral>( + args: &[ColumnarValue], + rows: usize, child_field: FieldRef, - from_i64: impl Fn(i64) -> T::Native, + pool: &Arc<SequenceMemoryPool>, + check: Option<&CancellationCheck>, ) -> 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)) + let start = Input::<T>::new(&args[0])?; + let stop = Input::<T>::new(&args[1])?; + let step = args.get(2).map(Input::<T>::new).transpose()?; + let has_nulls = + start.has_nulls() || stop.has_nulls() || step.as_ref().is_some_and(Input::has_nulls); + let scalar = args + .iter() + .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); + let repeated = if scalar && rows != 0 { + Some( + match row_bounds(&start, &stop, step.as_ref(), has_nulls, 0) { + Some((start, stop, step)) => Row { + start, + step, + len: sequence_length(start, stop, step)?, + }, + None => Row::default(), + }, + ) + } else { + None }; - // 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 - } + + // Constant-size scratch: validate every row in input order before checking either aggregate + // limit. Recompute the cheap lengths while generating; no row-length Vec or output allocation + // is made during sizing. u128 keeps the exact batch count even for huge scalar requests. + let (total, nullable) = if let Some(row) = repeated { + ((row.len as u128) * (rows as u128), row.len == 0) + } else { + let mut total = 0u128; + let mut nullable = false; + for row in 0..rows { + if row != 0 && row % CHECK_ROWS == 0 { + check_cancelled(check)?; + } + match row_bounds(&start, &stop, step.as_ref(), has_nulls, row) { + Some((start, stop, step)) => total += sequence_length(start, stop, step)? as u128, + None => nullable = true, } } + (total, nullable) }; - - // 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 { + if total > i32::MAX as u128 { 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); + let total = total as usize; + let value_bytes = total + .checked_mul(std::mem::size_of::<T::Native>()) + .ok_or_else(size_error)?; + let offset_count = rows.checked_add(1).ok_or_else(size_error)?; + let offset_bytes = offset_count + .checked_mul(std::mem::size_of::<i32>()) + .ok_or_else(size_error)?; + let null_bytes = if nullable { rows.div_ceil(8) } else { 0 }; + let value_layout = buffer_layout(value_bytes)?; + let offset_layout = buffer_layout(offset_bytes)?; + let null_layout = buffer_layout(null_bytes)?; + let required = value_layout + .size() + .checked_add(offset_layout.size()) + .and_then(|n| n.checked_add(null_layout.size())) + .ok_or_else(size_error)?; + check_cancelled(check)?; + let mut reservation = pool.reserve(required)?; + let mut values = SequenceBuffer::new(value_layout, &mut reservation)?; + let mut offsets = SequenceBuffer::new(offset_layout, &mut reservation)?; + let mut validity = if nullable { + Some(SequenceBuffer::new(null_layout, &mut reservation)?) + } else { + None + }; + let validity_slots = validity + .as_mut() + .map(|buffer| buffer.slots::<u8>(null_bytes)); + // The bitmap is small relative to its offsets, but bound cancellation latency here too. + if let Some(slots) = validity_slots { + for chunk in slots.chunks_mut(CHECK_VALUES) { + check_cancelled(check)?; + for slot in chunk { + slot.write(0); } } - offsets.push(values.len() as i32); } - - let values = PrimitiveArray::<T>::new(ScalarBuffer::from(values), None); - let list = ListArray::try_new( + let validity_slots = validity + .as_mut() + .map(|buffer| buffer.slots::<u8>(null_bytes)); + let value_slots = values.slots::<T::Native>(total); + let offset_slots = offsets.slots::<i32>(offset_count); + offset_slots[0].write(0); + let mut written = 0; + let mut until_check = CHECK_VALUES; + let mut validity_slots = validity_slots; + for row in 0..rows { Review Comment: Rows where `start == stop` are 20 to 29 percent slower than main in interleaved criterion runs, whatever the step is. `zero_step_start_eq_stop` goes from 45.5 to 57.2 µs, and a one-element shape with the default step goes from 37.9 to 48.9 µs. Two- and five-element rows are faster. `dense_nulls_2_elems`, `dense_nulls_5_elems` and `sparse_nulls_5_elems` are also 6 to 8 percent slower, and `all_scalar_8192_rows_365_elems` about 5 percent. `optimizing_expressions.md` asks for no meaningful regression on any shape. Could you add a one-element shape to `benches/sequence.rs`, look at that case, and share the criterion comparison against main? ########## native/spark-expr/src/array_funcs/sequence.rs: ########## @@ -15,35 +15,181 @@ // 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. +//! Spark-compatible integral sequence with whole-invocation admission. +use std::fmt::{Debug, Formatter}; +use std::hash::{Hash, Hasher}; 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::{BooleanBuffer, 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::logical_expr::ColumnarValue; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use datafusion_comet_jni_bridge::{check_exception, errors::CometError, JVMClasses}; +use jni::objects::{Global, JObject}; +use jni::signature::{Primitive, ReturnType}; +use super::sequence_memory::{ + buffer_layout, default_sequence_pool, size_error, SequenceBuffer, SequenceMemoryPool, +}; 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; +const CHECK_ROWS: usize = 4096; +const CHECK_VALUES: usize = 65_536; +type CancellationCheck = Arc<dyn Fn() -> Result<()> + Send + Sync>; + +/// The result must remain an array of number_rows rows. An immutable UDF could be folded to a +/// list scalar whose subsequent broadcast allocates an unadmitted child buffer. Volatile here +/// describes admission/cancellation, even though the generated values themselves are deterministic. +pub struct SparkSequence { + data_type: DataType, + signature: Signature, + pool: Arc<SequenceMemoryPool>, + check: Option<CancellationCheck>, +} + +impl SparkSequence { + pub fn new(data_type: DataType, pool: Arc<SequenceMemoryPool>) -> Self { + Self { + data_type, + signature: Signature::variadic_any(Volatility::Volatile), + pool, + check: None, + } + } + + /// Capture Spark's task lifecycle signal, already obtained on the executor task thread by + /// createPlan. The method lookup is cached once per expression; execution never consults + /// TaskContext.get(), which is unset on Tokio workers. + pub fn with_task_context( Review Comment: This would be the only native code that polls `TaskContext.isInterrupted`. If cancellation latency matters for long native batches, it matters for sorts, joins and scans too, so it seems better solved once in the execution loop than inside one kernel. At the default allowance a single call takes tens of milliseconds at most, so the polling buys little here. It adds at least six JNI calls per batch plus one per 65,536 generated values, and the criterion numbers don't include them because the benchmark leaves `check` unset. Could we drop it from this PR and open an issue for a general mechanism if it turns out to be needed? -- 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]
