fix array_repeat scalar path overflows total repeated-value count#22274
fix array_repeat scalar path overflows total repeated-value count#22274xiedeyantu wants to merge 3 commits into
Conversation
kosiew
left a comment
There was a problem hiding this comment.
@xiedeyantu
Thanks for tightening up the overflow handling here. I found two remaining paths where we can still hit unchecked allocation or Arrow offset construction before returning the intended execution error.
| "array_repeat: total repeated values overflowed usize", | ||
| )?; | ||
|
|
||
| let mut take_indices = Vec::with_capacity(total_repeated_values); |
There was a problem hiding this comment.
The scalar/non-list path still allocates take_indices before confirming that the result fits in the returned List offset type.
For example, on the default List<i32> path, a count like 2147483648 does not overflow usize, so checked_sum_counts succeeds. After that, Vec::with_capacity(total_repeated_values) can try to reserve around 16 GiB before the later O::from_usize(running_offset) error is reached.
Could we validate O::from_usize(total_repeated_values) first, and return the existing offset error before doing any capacity reservation or extension?
There was a problem hiding this comment.
The list-input path still has a couple of unchecked capacity and offset construction spots.
outer_total + 1 can overflow when the checked sum is exactly usize::MAX, and the outer offsets are later rebuilt with OffsetBuffer::<O>::from_lengths(...), which can panic on usize or offset overflow in Arrow.
Since this function is aiming to make capacity and offset calculations overflow-safe, could we avoid the unchecked + 1 and avoid from_lengths for untrusted counts? It looks like we can reuse the already checked totals and build or validate the outer offsets with checked_add plus O::from_usize before allocating.
@kosiew Thanks for the review! I’ll wait until #22305 is merged before making adjustments to this PR, as I’m concerned about potential merge conflicts. |
| DataFusionError::Execution( | ||
| "array_repeat: outer total overflowed usize".to_string(), | ||
| ) |
There was a problem hiding this comment.
| DataFusionError::Execution( | |
| "array_repeat: outer total overflowed usize".to_string(), | |
| ) | |
| exec_datafusion_err!("array_repeat: running_offset overflowed usize") |
Same for others
| } | ||
| } | ||
|
|
||
| fn checked_sum_counts( |
There was a problem hiding this comment.
Can probably just inline this since its used only once
Which issue does this PR close?
Rationale for this change
The scalar path of array_repeat can overflow while computing the total number of repeated values. Previously, the implementation used unchecked accumulation for these totals, which could lead to overflow behavior that was not explicitly handled. This change makes the capacity and offset calculations overflow-safe and returns a clear execution error when the total output size exceeds usize.
What changes are included in this PR?
Are these changes tested?
Are there any user-facing changes?