You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/24 15:36:49 UTC
(datafusion) branch gh-readonly-queue/main/pr-22305-94c58d086d1fc7f4ea6fb0a5f6d82805ef2b7a98 created (now 7ad8e2ce50)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22305-94c58d086d1fc7f4ea6fb0a5f6d82805ef2b7a98
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 7ad8e2ce50 fix array_repeat capacity overflow on constant scalar with large count (#22305)
This branch includes the following new commits:
new 7ad8e2ce50 fix array_repeat capacity overflow on constant scalar with large count (#22305)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix array_repeat capacity overflow on constant scalar with large count (#22305)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22305-94c58d086d1fc7f4ea6fb0a5f6d82805ef2b7a98
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7ad8e2ce501e58f87bdd039cd252c2441bc22d6b
Author: Zhen Chen <cz...@163.com>
AuthorDate: Sun May 24 23:36:30 2026 +0800
fix array_repeat capacity overflow on constant scalar with large count (#22305)
## Which issue does this PR close?
- Closes #22228 .
## Rationale for this change
`array_repeat` still panics for oversized repeat counts in the
constant-scalar path. The simplest reproducer is:
```sql
SELECT array_repeat(1, 9223372036854775807)
```
Unlike the previously reported `array_repeat` overflow cases, this path
does not sum counts across rows and does not multiply nested list
lengths, but it still reaches an unchecked `Vec` preallocation and
panics with `capacity overflow`.
This change makes `array_repeat` reject oversized output lengths up
front and return a normal execution error instead of panicking.
## What changes are included in this PR?
This PR adds explicit bounds checks in repeat.rs so `array_repeat`
validates requested output sizes before allocating buffers.
The main changes are:
- Move repeat-length accumulation into shared checked helpers.
- Reject oversized output lengths with:
`array_repeat: requested length exceeds maximum array size`
- Guard both scalar and list repeat paths so they fail consistently
before hitting unchecked allocation or arithmetic overflow.
- Reuse precomputed outer offsets for the list path instead of
rebuilding them from unchecked lengths.
## Are these changes tested?
Yes.
This PR adds a regression test in repeat.rs covering the constant-scalar
reproducer with `i64::MAX` as the repeat count and verifies that
`array_repeat` returns an execution error rather than panicking.
Validated with:
```bash
cargo test -p datafusion-functions-nested scalar_count_exceeding_max_array_size_returns_error --lib
```
## Are there any user-facing changes?
Yes.
Previously, oversized `array_repeat` calls could panic the process.
After this change, they return a regular execution error:
```text
array_repeat: requested length exceeds maximum array size
```
---
datafusion/functions-nested/src/repeat.rs | 143 ++++++++++++++++-----
.../sqllogictest/test_files/array/array_repeat.slt | 3 +
2 files changed, 111 insertions(+), 35 deletions(-)
diff --git a/datafusion/functions-nested/src/repeat.rs b/datafusion/functions-nested/src/repeat.rs
index ceec748a6e..825530923f 100644
--- a/datafusion/functions-nested/src/repeat.rs
+++ b/datafusion/functions-nested/src/repeat.rs
@@ -38,8 +38,12 @@ use datafusion_expr::{
};
use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
use datafusion_macros::user_doc;
+use std::mem::size_of;
use std::sync::Arc;
+const ARRAY_REPEAT_LENGTH_EXCEEDED: &str =
+ "array_repeat: requested length exceeds maximum array size";
+
make_udf_expr_and_func!(
ArrayRepeat,
array_repeat,
@@ -175,28 +179,12 @@ fn general_repeat<O: OffsetSizeTrait>(
array: &ArrayRef,
count_array: &Int64Array,
) -> Result<ArrayRef> {
- let total_repeated_values: usize = (0..count_array.len())
- .map(|i| get_count_with_validity(count_array, i))
- .sum();
+ let (offsets, total_repeated_values) = build_repeat_offsets::<O>(count_array)?;
let mut take_indices = Vec::with_capacity(total_repeated_values);
- let mut offsets = Vec::with_capacity(count_array.len() + 1);
- offsets.push(O::zero());
- let mut running_offset = 0usize;
for idx in 0..count_array.len() {
let count = get_count_with_validity(count_array, idx);
- running_offset = running_offset.checked_add(count).ok_or_else(|| {
- DataFusionError::Execution(
- "array_repeat: running_offset overflowed usize".to_string(),
- )
- })?;
- let offset = O::from_usize(running_offset).ok_or_else(|| {
- DataFusionError::Execution(format!(
- "array_repeat: offset {running_offset} exceeds the maximum value for offset type"
- ))
- })?;
- offsets.push(offset);
take_indices.extend(std::iter::repeat_n(idx as u64, count));
}
@@ -231,23 +219,23 @@ fn general_list_repeat<O: OffsetSizeTrait>(
count_array: &Int64Array,
) -> Result<ArrayRef> {
let list_offsets = list_array.value_offsets();
+ let (outer_offsets, outer_total) = build_repeat_offsets::<O>(count_array)?;
// calculate capacities for pre-allocation
- let mut outer_total = 0usize;
let mut inner_total = 0usize;
for i in 0..count_array.len() {
let count = get_count_with_validity(count_array, i);
- if count > 0 {
- outer_total += count;
- if list_array.is_valid(i) {
- let len = list_offsets[i + 1].to_usize().unwrap()
- - list_offsets[i].to_usize().unwrap();
- inner_total += len * count;
- }
+ if count > 0 && list_array.is_valid(i) {
+ let len = list_offsets[i + 1].to_usize().unwrap()
+ - list_offsets[i].to_usize().unwrap();
+ inner_total =
+ checked_repeat_len_add(inner_total, checked_repeat_len_mul(len, count)?)?;
+ ensure_array_repeat_output_len::<O>(inner_total)?;
}
}
// Build inner structures
+ ensure_vec_capacity::<O>(checked_repeat_len_add(outer_total, 1)?)?;
let mut inner_offsets = Vec::with_capacity(outer_total + 1);
let mut take_indices = Vec::with_capacity(inner_total);
let mut inner_nulls = BooleanBufferBuilder::new(outer_total);
@@ -262,11 +250,8 @@ fn general_list_repeat<O: OffsetSizeTrait>(
let row_len = end - start;
for _ in 0..count {
- inner_running = inner_running.checked_add(row_len).ok_or_else(|| {
- DataFusionError::Execution(
- "array_repeat: inner offset overflowed usize".to_string(),
- )
- })?;
+ inner_running = checked_repeat_len_add(inner_running, row_len)?;
+ ensure_array_repeat_output_len::<O>(inner_running)?;
let offset = O::from_usize(inner_running).ok_or_else(|| {
DataFusionError::Execution(format!(
"array_repeat: offset {inner_running} exceeds the maximum value for offset type"
@@ -299,16 +284,85 @@ fn general_list_repeat<O: OffsetSizeTrait>(
list_array.data_type().to_owned(),
true,
)),
- OffsetBuffer::<O>::from_lengths(
- count_array
- .iter()
- .map(|c| c.map(|v| if v > 0 { v as usize } else { 0 }).unwrap_or(0)),
- ),
+ OffsetBuffer::new(outer_offsets.into()),
Arc::new(inner_list),
count_array.nulls().cloned(),
)?))
}
+fn build_repeat_offsets<O: OffsetSizeTrait>(
+ count_array: &Int64Array,
+) -> Result<(Vec<O>, usize)> {
+ let mut offsets = Vec::with_capacity(count_array.len() + 1);
+ offsets.push(O::zero());
+ let mut running_offset = 0usize;
+
+ for idx in 0..count_array.len() {
+ let count = get_count_with_validity(count_array, idx);
+ running_offset = checked_repeat_len_add(running_offset, count)?;
+ ensure_array_repeat_output_len::<O>(running_offset)?;
+ let offset = O::from_usize(running_offset).ok_or_else(|| {
+ DataFusionError::Execution(format!(
+ "array_repeat: offset {running_offset} exceeds the maximum value for offset type"
+ ))
+ })?;
+ offsets.push(offset);
+ }
+
+ Ok((offsets, running_offset))
+}
+
+fn checked_repeat_len_add(lhs: usize, rhs: usize) -> Result<usize> {
+ lhs.checked_add(rhs).ok_or_else(|| {
+ DataFusionError::Execution(ARRAY_REPEAT_LENGTH_EXCEEDED.to_string())
+ })
+}
+
+fn checked_repeat_len_mul(lhs: usize, rhs: usize) -> Result<usize> {
+ lhs.checked_mul(rhs).ok_or_else(|| {
+ DataFusionError::Execution(ARRAY_REPEAT_LENGTH_EXCEEDED.to_string())
+ })
+}
+
+fn ensure_array_repeat_output_len<O: OffsetSizeTrait>(len: usize) -> Result<()> {
+ if len > max_array_repeat_output_len::<O>() {
+ return Err(DataFusionError::Execution(
+ ARRAY_REPEAT_LENGTH_EXCEEDED.to_string(),
+ ));
+ }
+
+ Ok(())
+}
+
+fn ensure_vec_capacity<T>(len: usize) -> Result<()> {
+ if len > max_vec_elements::<T>() {
+ return Err(DataFusionError::Execution(
+ ARRAY_REPEAT_LENGTH_EXCEEDED.to_string(),
+ ));
+ }
+
+ Ok(())
+}
+
+fn max_array_repeat_output_len<O: OffsetSizeTrait>() -> usize {
+ max_offset_elements::<O>().min(max_vec_elements::<u64>())
+}
+
+fn max_offset_elements<O: OffsetSizeTrait>() -> usize {
+ if size_of::<O>() == size_of::<i32>() {
+ i32::MAX as usize
+ } else {
+ i64::MAX as usize
+ }
+}
+
+fn max_vec_elements<T>() -> usize {
+ let element_size = size_of::<T>();
+ (isize::MAX as usize)
+ .checked_div(element_size)
+ .unwrap_or(usize::MAX)
+}
+
/// Helper function to get count from count_array at given index
/// Return 0 for null values or non-positive count.
#[inline]
@@ -320,3 +374,22 @@ fn get_count_with_validity(count_array: &Int64Array, idx: usize) -> usize {
if c > 0 { c as usize } else { 0 }
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::array_repeat_inner;
+ use arrow::array::{ArrayRef, Int64Array};
+ use std::sync::Arc;
+
+ #[test]
+ fn scalar_count_exceeding_max_array_size_returns_error() {
+ let element: ArrayRef = Arc::new(Int64Array::from(vec![1]));
+ let count: ArrayRef = Arc::new(Int64Array::from(vec![i64::MAX]));
+
+ let err = array_repeat_inner(&[element, count]).unwrap_err();
+ assert_eq!(
+ err.to_string(),
+ "Execution error: array_repeat: requested length exceeds maximum array size"
+ );
+ }
+}
diff --git a/datafusion/sqllogictest/test_files/array/array_repeat.slt b/datafusion/sqllogictest/test_files/array/array_repeat.slt
index 8052f09cb3..9f17c449c8 100644
--- a/datafusion/sqllogictest/test_files/array/array_repeat.slt
+++ b/datafusion/sqllogictest/test_files/array/array_repeat.slt
@@ -43,6 +43,9 @@ select
----
[[1], [1], [1], [1], [1]] [[1.1, 2.2, 3.3], [1.1, 2.2, 3.3], [1.1, 2.2, 3.3]] [[NULL, NULL], [NULL, NULL], [NULL, NULL]] [[[1, 2], [3, 4]], [[1, 2], [3, 4]]]
+query error DataFusion error: Execution error: array_repeat: requested length exceeds maximum array size
+select array_repeat(1, 9223372036854775807);
+
query ????
select
array_repeat(arrow_cast([1], 'LargeList(Int64)'), 5),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org