mbutrovich commented on code in PR #26957:
URL: https://github.com/apache/datafusion/pull/26957#discussion_r4485254564
##########
datafusion/functions/src/string/repeat.rs:
##########
@@ -519,6 +552,35 @@ mod tests {
);
}
+ #[test]
+ fn test_repeat_large_string_array_exceeds_max_size() {
+ let strings: ArrayRef =
Arc::new(LargeStringArray::from(vec![Some("a")]));
+ let counts: ArrayRef = Arc::new(Int64Array::from(vec![Some(1 << 60)]));
+
+ let err = super::repeat(&strings, &counts).unwrap_err().to_string();
+ assert!(
+ err.contains("string size overflow on repeat"),
+ "unexpected error: {err}"
+ );
+ }
+
+ #[test]
+ fn test_repeat_large_string_array_total_exceeds_max_size() {
+ // Each row is under the cap, but together they exceed it.
+ let strings: ArrayRef =
+ Arc::new(LargeStringArray::from(vec![Some("a"), Some("a")]));
+ let counts: ArrayRef = Arc::new(Int64Array::from(vec![
+ Some(3_000_000_000),
+ Some(3_000_000_000),
+ ]));
+
+ let err = super::repeat(&strings, &counts).unwrap_err().to_string();
+ assert!(
+ err.contains("string size overflow on repeat"),
+ "unexpected error: {err}"
+ );
+ }
Review Comment:
Could we fold these two tests into one? With the total check in place,
`test_repeat_large_string_array_exceeds_max_size` still passes if the array
path's per-row cap goes back to `i64::MAX`, because a single row of `1 << 60`
bytes also exceeds the total. So
`test_repeat_large_string_array_total_exceeds_max_size` already covers what the
single-row test checks.
I'd also suggest different inputs for the total test. If the total check
regresses, two rows of 3 GB pass the per-row check, so the builder allocates
about 6 GB
([here](https://github.com/apache/datafusion/blob/1b181b6901dd003e6380d650e4328afa584a18de/datafusion/functions/src/string/repeat.rs#L247-L250))
and the loop fills it before the assertion runs. With 2^17 rows at exactly the
cap, the total is about 512 TiB, so a regression panics at once on the refused
allocation, as in the issue. Running this version with the total check removed
fails in 0.02s with `failed to allocate memory for layout Layout { size:
562949953290240, .. }`. It also checks that a row of exactly `u32::MAX` bytes
is accepted. Asserting the full message checks that the total check is the one
that fired. If the per-row check becomes `>=`, the test fails with `but got
4294967295`.
```suggestion
#[test]
fn test_repeat_large_string_array_total_exceeds_max_size() {
// Each row is at the cap, but the total is larger than any
allocation
// can satisfy, so a missing total check fails at once.
let n = 1 << 17;
let strings: ArrayRef =
Arc::new(LargeStringArray::from(vec![Some("a"); n]));
let counts: ArrayRef = Arc::new(Int64Array::from(vec![Some(u32::MAX
as i64); n]));
let err = super::repeat(&strings, &counts).unwrap_err().to_string();
assert_eq!(
err,
format!(
"Execution error: string size overflow on repeat, max size
is {}, but got {}",
u32::MAX,
n * u32::MAX as usize
)
);
}
```
##########
datafusion/functions/src/string/repeat.rs:
##########
@@ -22,6 +22,15 @@ use datafusion_expr_common::signature::{Coercion,
TypeSignatureClass};
use datafusion_expr_common::sort_properties::ExprProperties;
use datafusion_macros::user_doc;
+/// Largest result, in bytes, that `repeat` produces for `LargeUtf8`, both for
a
+/// single value and for the total of an array.
+///
+/// `LargeUtf8` offsets would allow up to `i64::MAX` bytes, but the result is
+/// allocated up front, and an allocation larger than the address space aborts
+/// the process instead of returning an error. This cap matches DuckDB's string
+/// size limit and is about twice the `Utf8` limit of `i32::MAX`.
+const LARGE_UTF8_MAX_SIZE: usize = u32::MAX as usize;
Review Comment:
Could we trim this comment to what holds for the code as it is now? The
`i64::MAX` offset limit and the comparisons with DuckDB and `Utf8` explain how
the value was picked, and the PR description already covers that. The reason
the cap exists is the up-front allocation.
```suggestion
/// Largest result, in bytes, that `repeat` produces for `LargeUtf8`, both
for a
/// single value and for the total of an array. The result is allocated up
front,
/// and an allocation the allocator refuses aborts or panics instead of
returning
/// an error.
const LARGE_UTF8_MAX_SIZE: usize = u32::MAX as usize;
```
--
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]
