comphead commented on code in PR #6948:
URL: https://github.com/apache/datafusion-comet/pull/6948#discussion_r4488344603
##########
native/spark-expr/src/comet_scalar_funcs.rs:
##########
@@ -131,7 +131,11 @@ pub fn create_comet_physical_fun_with_eval_mode(
make_comet_scalar_udf!("ceil", spark_ceil, data_type)
}
"spark_sequence" => {
- make_comet_scalar_udf!("spark_sequence", spark_sequence, data_type)
+ // Preserve the exact outer pool, including per-plan tracking and
task-shared policy.
+ Ok(Arc::new(ScalarUDF::new_from_impl(SparkSequence::new(
+ data_type,
+ Arc::clone(&session_state.runtime_env().memory_pool),
+ ))))
}
Review Comment:
Since this branch was cut, #7171 added a separate
`spark_sequence_recomputed` arm here for Spark 4.0.5, 4.1.4 and 4.2.1, built
with `make_comet_scalar_udf!`. A rebase that keeps that arm as is leaves those
versions outside the pool. Could the overflow mode become a field of
`SparkSequence`, so both names go through the reservation, with a test of the
recomputed name under a bounded pool?
##########
native/spark-expr/src/array_funcs/sequence.rs:
##########
@@ -184,39 +267,94 @@ where
None => -1,
}
};
+ let scalar_length = if args
+ .iter()
+ .all(|arg| matches!(arg, ColumnarValue::Scalar(_)))
+ {
+ let s = start.value(0);
+ let e = stop.value(0);
+ Some(sequence_length(s, e, row_step(0, s, e))?)
+ } else {
+ None
+ };
+ let scalar_total = scalar_length
+ .map(|len| batch_length(len as u128 * num_rows as u128))
+ .transpose()?;
+ // Columns admit scratch first, then validate and cache each row length
once before
+ // growing this same reservation for values. Scalars know their complete
size already.
+ let initial = (|| {
+ let bitmap_count = if null_counts.iter().any(|count| *count != 0) {
+ 2
+ } else {
+ 0
+ };
+ let bytes = construction_scratch(num_rows, true, bitmap_count)?;
+ let bytes = bytes
+ .checked_add(value_bytes::<T>(scalar_total.unwrap_or(0))?)
+ .ok_or_else(size_error)?;
+ pool.map(|pool| reserve_construction(pool, bytes))
+ .transpose()
+ })();
+ // Keep error-only typed parsing in the cold helper so it need not borrow
hot input locals.
+ let scratch_error = move |error| {
+ // Only resource exhaustion needs an allocation-free validation pass
to keep Spark
+ // row errors and the batch element ceiling ahead of it. Propagate
other pool errors,
+ // such as an interrupted Spark task, immediately without inspecting
more input rows.
+ if scalar_length.is_none() && matches!(&error,
DataFusionError::ResourcesExhausted(_)) {
+ validate_before_resource_error::<T>(args, num_rows, error)
+ } else {
+ error
+ }
+ };
+ let reservation = initial.map_err(scratch_error)?;
+ // 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, allocating
+ // one accumulator and applying the remaining masks in place.
+ let input_nulls = [
+ start.nulls(num_rows),
+ stop.nulls(num_rows),
+ step.as_ref().and_then(|step| step.nulls(num_rows)),
+ ];
+ let nulls = NullBuffer::union_many(input_nulls.iter().map(Option::as_ref))
+ .filter(|nulls| nulls.null_count() != 0);
+ if nulls
+ .as_ref()
+ .is_some_and(|nulls| nulls.null_count() == num_rows)
+ {
+ // Reuse the accumulator so another all-null bitmap cannot exceed its
capacity allowance.
+ return null_sequence::<T>(child_field, num_rows, nulls);
+ }
// 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 nulls.as_ref().is_some_and(|nulls| nulls.is_null(row)) {
- lengths.push(0);
- continue;
- }
- 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);
- }
- // 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 {
- total_elements: total.to_string(),
- },
- )));
+ let mut lengths = sequence_vec::<usize>(num_rows).map_err(scratch_error)?;
+ let total = if let Some(len) = scalar_length {
+ lengths.resize(num_rows, len);
+ scalar_total.unwrap()
Review Comment:
For all-scalar inputs `scalar_length` is constant, but this still allocates
and reserves a `num_rows` `usize` length cache (64 KiB at 8192 rows). Could the
scalar path pass `cache_lengths = false`? The test that pins "four cached
lengths" would change with it.
##########
native/spark-expr/src/array_funcs/sequence.rs:
##########
@@ -245,6 +383,133 @@ where
)?))
}
+// Registration and reservation have the same lexical lifetime, including
failed try_grow.
+fn reserve_construction(pool: &Arc<dyn MemoryPool>, bytes: usize) ->
Result<MemoryReservation> {
+ let reservation = MemoryConsumer::new("Sequence").register(pool);
+ reservation.try_grow(bytes)?;
+ Ok(reservation)
+}
+
+fn size_error() -> DataFusionError {
+ DataFusionError::ResourcesExhausted("Sequence construction size
overflow".to_string())
+}
+
+fn bitmap_capacity(rows: usize) -> Result<usize> {
+ // Round bitmap bytes to Arrow's 64-byte allocation boundary.
+ Ok(rows.div_ceil(8).checked_add(63).ok_or_else(size_error)? & !63)
+}
+
+fn batch_length(total: u128) -> Result<usize> {
+ if total > i32::MAX as u128 {
+ return Err(DataFusionError::External(Box::new(
+ SparkError::SequenceBatchTooLarge {
+ total_elements: total.to_string(),
+ },
+ )));
+ }
+ Ok(total as usize)
+}
+
+fn value_bytes<T: Integral>(total: usize) -> Result<usize> {
+ total
+ .checked_mul(std::mem::size_of::<T::Native>())
+ .ok_or_else(size_error)
+}
+
+fn construction_scratch(rows: usize, cache_lengths: bool, bitmap_count: usize)
-> Result<usize> {
+ let offsets = rows
+ .checked_add(1)
+ .and_then(|n| n.checked_mul(std::mem::size_of::<i32>()))
+ .ok_or_else(size_error)?;
+ let lengths = if cache_lengths {
+ rows.checked_mul(std::mem::size_of::<usize>())
+ .ok_or_else(size_error)?
+ } else {
+ 0
+ };
+ // Bulk unions use one accumulator; two rounded bitmap sizes allow for
capacity doubling.
+ // The full backing allocation of a borrowed, possibly sliced input bitmap
is not charged.
+ let validity = bitmap_capacity(rows)?
+ .checked_mul(bitmap_count)
+ .ok_or_else(size_error)?;
+ offsets
+ .checked_add(lengths)
+ .and_then(|n| n.checked_add(validity))
+ .ok_or_else(size_error)
+}
+
+#[cold]
+fn validate_before_resource_error<T: Integral>(
+ args: &[ColumnarValue],
+ rows: usize,
+ resource_error: DataFusionError,
+) -> DataFusionError
+where
+ T::Native: Into<i64>,
+{
+ let validation = (|| -> Result<usize> {
+ 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 mut total = 0u128;
+ for row in 0..rows {
+ if start.is_null(row)
+ || stop.is_null(row)
+ || step.as_ref().is_some_and(|step| step.is_null(row))
+ {
+ continue;
+ }
+ let s = start.value(row);
+ let e = stop.value(row);
+ let step = step
+ .as_ref()
+ .map_or_else(|| if s <= e { 1 } else { -1 }, |step|
step.value(row));
+ total += sequence_length(s, e, step)? as u128;
+ }
+ batch_length(total)
+ })();
+ validation.err().unwrap_or(resource_error)
+}
+
+fn sequence_vec<T>(capacity: usize) -> Result<Vec<T>> {
+ let mut values = Vec::new();
+ values.try_reserve_exact(capacity).map_err(|error| {
+ DataFusionError::ResourcesExhausted(format!("Sequence allocation
failed: {error}"))
+ })?;
+ Ok(values)
Review Comment:
On `main` a failed `try_reserve_exact` becomes `SequenceBatchTooLarge`,
whose message tells the user to lower `spark.comet.batchSize`, and its doc in
`native/common/src/error.rs` says it covers allocator refusal. Here allocator
and pool refusals become a plain `ResourcesExhausted` with no hint. Could they
keep the existing variant, or could its doc and message be updated?
--
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]
