fix: avoid lpad/rpad panic on out-of-range length arrays

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<F: LeftRightSlicer>(
         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<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
+        );
+    }
+}

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'));