mbutrovich commented on code in PR #26958:
URL: https://github.com/apache/datafusion/pull/26958#discussion_r4485137830
##########
datafusion/functions/src/unicode/lpad.rs:
##########
@@ -345,4 +347,49 @@ mod tests {
Ok(())
}
+
+ fn array_args(lengths: Vec<i64>) -> Vec<ColumnarValue> {
+ vec![
+ ColumnarValue::Array(Arc::new(StringArray::from(vec!["hi";
lengths.len()]))),
+ ColumnarValue::Array(Arc::new(Int64Array::from(lengths))),
+ ]
+ }
+
+ #[test]
+ fn test_array_length_out_of_range() -> Result<()> {
+ test_function!(
+ LPadFunc::new(),
+ array_args(vec![-1]),
+ Ok(Some("")),
+ &str,
+ Utf8,
+ StringArray
+ );
+ test_function!(
+ LPadFunc::new(),
+ array_args(vec![1 << 62]),
+ exec_err!(
+ "lpad requested length {} too large, maximum allowed length is
{}",
+ 1i64 << 62,
+ i32::MAX
+ ),
+ &str,
+ Utf8,
+ StringArray
+ );
+ test_function!(
+ LPadFunc::new(),
+ array_args(vec![i64::MAX, i64::MAX]),
+ exec_err!(
+ "lpad requested length {} too large, maximum allowed length is
{}",
+ i64::MAX,
+ i32::MAX
+ ),
+ &str,
+ Utf8,
+ StringArray
+ );
+
+ Ok(())
+ }
Review Comment:
Should we replace this test and the matching one in `rpad.rs` with a unit
test of `pad_data_capacity`?
The new sqllogictest cases already cover the user-visible fix. Each of the
four fails on `main` with `capacity overflow`, for both `lpad` and `rpad`.
These unit tests check the same results again. The `1 << 62`, `i64::MAX`, and
`[i64::MAX, i64::MAX]` cases also all take the same branch of the filter.
What no test in the PR checks is either half of the fix. If I remove
`.min(i32::MAX as usize)`, every test in the PR still passes. The same happens
if I remove the `.filter(...)`, and then a single `-1` row reserves about 2 GiB
up front. A direct test checks the filter's bounds and the cap without
allocating anything:
```rust
#[cfg(test)]
mod tests {
use arrow::array::Int64Array;
use super::pad_data_capacity;
#[test]
fn test_pad_data_capacity() {
let capacity =
|lengths: Vec<Option<i64>>|
pad_data_capacity(&Int64Array::from(lengths));
assert_eq!(capacity(vec![]), 0);
// Negative lengths produce "" and lengths above i32::MAX are an
error
assert_eq!(capacity(vec![Some(3), None, Some(-1), Some(i64::MIN)]),
3);
assert_eq!(
capacity(vec![Some(3), Some(i32::MAX as i64 + 1),
Some(i64::MAX)]),
3
);
assert_eq!(
capacity(vec![Some(i32::MAX as i64), Some(1)]),
i32::MAX as usize
);
}
}
```
With that in `common.rs`, the import changes in `lpad.rs` and `rpad.rs` can
go back to what they were.
##########
datafusion/functions/src/unicode/common.rs:
##########
@@ -44,7 +44,10 @@ pub(crate) fn pad_data_capacity(length_array: &Int64Array)
-> usize {
length_array
.iter()
.flatten()
- .fold(0, |acc, len| acc.saturating_add(len as usize))
+ // Out-of-range lengths produce "", an error or null, so add no bytes
+ .filter(|len| (0..=i32::MAX as i64).contains(len))
+ .fold(0usize, |acc, len| acc.saturating_add(len as usize))
+ .min(i32::MAX as usize)
Review Comment:
Could this comment say what each out-of-range length becomes? As I read
[`pad_impl`](https://github.com/apache/datafusion/blob/668468808f6dbb2d206e854251956912fc7e5302/datafusion/functions/src/unicode/pad.rs#L397-L410),
a negative length appends `""` and a length above `i32::MAX` returns the
error. No out-of-range length produces a null. A null comes only from a null
string, length, or fill.
Could the `.min` also get a short comment? After the filter every length is
in range, but many of them can still add up to more than one buffer should
reserve up front, and that isn't obvious from the code.
```suggestion
// Negative lengths produce "" and lengths above i32::MAX are an
error,
// so neither adds bytes
.filter(|len| (0..=i32::MAX as i64).contains(len))
.fold(0usize, |acc, len| acc.saturating_add(len as usize))
// Many in-range lengths can still add up to more than one buffer
// should reserve up front
.min(i32::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]
