diff --git a/datafusion/functions/src/unicode/common.rs b/datafusion/functions/src/unicode/common.rs index f42b871a45e81..859374cdba3bb 100644 --- a/datafusion/functions/src/unicode/common.rs +++ b/datafusion/functions/src/unicode/common.rs @@ -44,7 +44,13 @@ pub(crate) fn pad_data_capacity(length_array: &Int64Array) -> usize { length_array .iter() .flatten() - .fold(0, |acc, len| acc.saturating_add(len as usize)) + // 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) } /// A trait for `left` and `right` byte slicing operations @@ -229,3 +235,27 @@ fn general_left_right_view( Ok(Arc::new(array) as ArrayRef) } } + +#[cfg(test)] +mod tests { + use arrow::array::Int64Array; + + use super::pad_data_capacity; + + #[test] + fn test_pad_data_capacity() { + let capacity = + |lengths: Vec>| 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 + ); + } +} diff --git a/datafusion/sqllogictest/test_files/string/string_literal.slt b/datafusion/sqllogictest/test_files/string/string_literal.slt index 690c6492a463e..aa6b9d0023b73 100644 --- a/datafusion/sqllogictest/test_files/string/string_literal.slt +++ b/datafusion/sqllogictest/test_files/string/string_literal.slt @@ -397,6 +397,17 @@ true true true query error lpad requested length 2147483648 too large SELECT lpad('hi', column1) FROM (VALUES (2147483648), (1)) AS t(column1); +# a negative length column is treated as 0 on the array path +query T +SELECT lpad('hi', column1, 'x') FROM (VALUES (-1), (3)) AS t(column1); +---- +(empty) +xhi + +# a length column far above i32::MAX is an error on the array path +query error lpad requested length 9223372036854775807 too large +SELECT lpad('hi', column1) FROM (VALUES (9223372036854775807), (1)) AS t(column1); + # a large scalar target length skips the scalar fast path query I SELECT character_length(lpad('x', 16385, 'a')); @@ -817,6 +828,17 @@ true true true query error rpad requested length 2147483648 too large SELECT rpad('hi', column1) FROM (VALUES (2147483648), (1)) AS t(column1); +# a negative length column is treated as 0 on the array path +query T +SELECT rpad('hi', column1, 'x') FROM (VALUES (-1), (3)) AS t(column1); +---- +(empty) +hix + +# a length column far above i32::MAX is an error on the array path +query error rpad requested length 9223372036854775807 too large +SELECT rpad('hi', column1) FROM (VALUES (9223372036854775807), (1)) AS t(column1); + # a large scalar target length skips the scalar fast path query I SELECT character_length(rpad('x', 16385, 'a'));