nat-openai commented on code in PR #6948:
URL: https://github.com/apache/datafusion-comet/pull/6948#discussion_r4316878470
##########
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()
+ } else {
+ let mut total = 0u128;
+ 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 as u128;
+ lengths.push(len);
+ }
+ batch_length(total)?
+ };
+ if scalar_length.is_none() {
+ if let Some(reservation) = &reservation {
+ reservation.try_grow(value_bytes::<T>(total)?)?;
Review Comment:
> With 64 batches of 64 rows, 64-element BIGINT sequences, a 220,000-byte
pool, and 64 KiB of sort spill reservation, the exact base kernel returns all
4,096 rows with 47 spills.
This does feel a little bit like a contrived example. On the main branch,
the `sequence` memory reservations are untracked, and with such small
sequences, there's probably ample system headroom for them to succeed.
Combining that with a very small memory pool (for tracked memory) ensures that
we can demonstrate spilling in the sort operator while hitting the limit when
we try to reserve space in `sequence`.
That's not to say I'm discounting this point in general though, and I'm
considering those proposals.
> Please coordinate admission with reclamation
IIUC, this would mean adding pretty novel cross-native-operator memory
coordination? If so, I image we need to think about the design of that pretty
carefully.
> or protect construction headroom before downstream buffering consumes it
Codex points me toward [`reserve_memory_for_merge` in sort.rs as
](https://github.com/apache/datafusion/blob/55.2.0/datafusion/physical-plan/src/sorts/sort.rs#L780-L795)
as a precedent for this sort of thing. I'm trying to understand how this
buffer could be sized such that it provides enough space for any computable
input without starving other operators... guarantees seem hard to come by and
it might remain possible to come up with adversarial counter examples (large
sequences; small pools) that break any assumptions.
--
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]
