IgnatiusPang opened a new pull request, #25751:
URL: https://github.com/apache/datafusion/pull/25751
## Rationale for this change
When aggregating empty batches or streaming partitioned state arrays across
distributed executors, several aggregate accumulators in
`datafusion-functions-aggregate` lacked boundary checks before indexing
`states`:
### 1. `FirstValueAccumulator` and `LastValueAccumulator` (`first_last.rs`)
In `TrivialFirstValueAccumulator::merge_batch` and
`TrivialLastValueAccumulator::merge_batch`:
```rust
let flags = states[1].as_boolean();
```
If `states.len() < 2`, direct indexing panics with `index out of bounds: the
len is ... but the index is 1`.
In `FirstValueAccumulator::merge_batch` and
`LastValueAccumulator::merge_batch`:
```rust
let is_set_idx = states.len() - 1;
let flags = states[is_set_idx].as_boolean();
```
If an empty slice `&[]` is received, `states.len() - 1` underflows `usize`
to `usize::MAX`, immediately crashing with an out-of-bounds panic.
### 2. `OrderSensitiveArrayAggAccumulator` and `ArrayAggAccumulator`
(`array_agg.rs`)
In `ensure_sorted_indices`:
```rust
let sorted_len = self.sorted_runs.iter().map(|run| run.len()).sum::<usize>();
let mut unsorted_indices = Vec::with_capacity(self.entries.len() -
sorted_len);
```
If `sorted_len > self.entries.len()`, `self.entries.len() - sorted_len`
underflows, causing an immediate Out-Of-Memory (OOM) abort when attempting to
allocate `usize::MAX` capacity.
Similarly in `ArrayAggAccumulator::retract_batch`:
```rust
let available = front.len() - self.front_offset;
```
and in `ArrayAggAccumulator::evaluate`:
```rust
a.slice(self.front_offset, a.len() - self.front_offset)
```
If `self.front_offset > front.len()`, arithmetic underflow occurs.
### 3. `NthValueAccumulator` (`nth_value.rs`)
In `NthValueAccumulator::merge_batch`:
```rust
if states.is_empty() { return Ok(()); }
let Some(agg_orderings) = states[1].as_list_opt::<i32>() else ...
```
If `states.len() == 1`, `is_empty()` check passes, but accessing `states[1]`
panics out-of-bounds. Replaced with `if states.len() < 2`.
### 4. `take_need` (`first_last/state.rs`)
In `take_need`:
```rust
EmitTo::First(n) => {
let first_n: BooleanBuffer = bool_buf.slice(0, n);
bool_buf_builder.append_buffer(&bool_buf.slice(n, bool_buf.len() - n));
```
If `n > bool_buf.len()`, `bool_buf.slice(0, n)` panics and `bool_buf.len() -
n` underflows. Clamped with `let n = n.min(bool_buf.len());`.
## What changes are included in this PR?
1. **`first_last.rs`**:
- Added early-exit empty checks `if states.is_empty() { return Ok(()); }`
and `if states.len() < 2 { return Ok(()); }`.
- Replaced `states.len() - 1` with `states.len().saturating_sub(1)`.
2. **`array_agg.rs`**:
- Replaced `self.entries.len() - sorted_len` with
`self.entries.len().saturating_sub(sorted_len)`.
- Replaced `front.len() - self.front_offset` with
`front.len().saturating_sub(self.front_offset)`.
- Clamped `front_offset` in `evaluate()` to `a.len()`.
- Replaced `group_rows.len() - 1` with
`group_rows.len().saturating_sub(1)`.
3. **`nth_value.rs`**:
- Replaced `if states.is_empty()` with `if states.len() < 2` before
accessing `states[1]`.
4. **`first_last/state.rs`**:
- Clamped `n` to `bool_buf.len()` in `take_need`.
## Are these changes tested?
Yes, tested across unit and integration suites ensuring zero regression on
standard inputs and panic-free handling on empty/sub-slice boundaries.
## Are there any user-facing changes?
No user-facing SQL API changes. Improves engine robustness against empty
state partitions.
--
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]
