sunchao commented on code in PR #6948:
URL: https://github.com/apache/datafusion-comet/pull/6948#discussion_r4315003915
##########
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:
[P2] Preserve downstream spilling before propagating this refusal. When
`sequence` feeds a native sort, the sorter can retain enough earlier batches
that this values reservation fails. `SortExec` polls its input before reaching
its own reservation/spill logic, so propagating the error here aborts a query
that could release memory by spilling. 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. The head fails on
the third batch with `ResourcesExhausted` for `Sequence`, before any spill.
Each sequence batch fits comfortably in the pool by itself. Please coordinate
admission with reclamation or protect construction headroom before downstream
buffering consumes it, and add a spilling integration regression test.
Evidence: Ran a disposable Rust integration test using locked DataFusion
55.2.0 / Arrow 59.3.0: `cargo test -p datafusion-comet-spark-expr --locked
--test review_sequence_admission review_sequence_before_spill -- --nocapture`.
It compiled `sequence.rs` from exact base
e555dba6f5c8c62ceca8c9c6e311a51b8fbfeb70 alongside the head factory. Fixture:
one MemTable partition containing 64 batches of 64 rows; for batch n and row i,
BIGINT id=4095-64*n-i, a=n%2, b=a+63. Query: `SELECT id, spark_sequence(a, b)
AS vals FROM t ORDER BY vals, id`. Configuration: target_partitions=1,
batch_size=64, sort_spill_reservation_bytes=65536, GreedyMemoryPool=220000.
Base: Ok(4096), spill_count=47. Head: only 128 input rows consumed,
spill_count=0, refusal for 32.0 KB with 772 bytes of sequence scratch reserved
and 19.1 KB available. At pool limit 480000, base again returns 4096 rows with
15 spills while head fails before its first spill. Assertions passed for both
comparisons. Reproducer saved at `/tmp/
comet-6948-review_sequence_admission.rs`; output at
`/tmp/comet-6948-repro-final.log`. The disposable checkout test was removed.
--
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]
