fix: return an error instead of aborting in repeat for huge LargeUtf8 results

datafusion/functions/src/string/repeat.rs

@@ -22,6 +22,12 @@ 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. 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;
+
 #[user_doc(
     doc_section(label = "String Functions"),
     description = "Returns a string with an input string repeated a specified number.",
@@ -125,7 +131,7 @@ impl ScalarUDFImpl for RepeatFunc {
                         compute_repeat(s, count, i32::MAX as usize)?,
                     )),
                     ScalarValue::LargeUtf8(Some(s)) => ScalarValue::LargeUtf8(Some(
-                        compute_repeat(s, count, i64::MAX as usize)?,
+                        compute_repeat(s, count, LARGE_UTF8_MAX_SIZE)?,
                     )),
                     _ => {
                         return internal_err!(
@@ -225,7 +231,16 @@ fn repeat(string_array: &ArrayRef, count_array: &ArrayRef) -> Result<ArrayRef> {
         LargeUtf8 => {
             let string_arr = string_array.as_string::<i64>();
             let (total_capacity, max_item_capacity) =
-                calculate_capacities(&string_arr, number_array, i64::MAX as usize)?;
+                calculate_capacities(&string_arr, number_array, LARGE_UTF8_MAX_SIZE)?;
+            // The builder allocates the whole output up front, so the cap must
+            // also bound the total, not only each row.
+            if total_capacity > LARGE_UTF8_MAX_SIZE {
+                return exec_err!(
+                    "string size overflow on repeat, max size is {}, but got {}",
+                    LARGE_UTF8_MAX_SIZE,
+                    total_capacity
+                );
+            }
             let builder = GenericStringArrayBuilder::<i64>::with_capacity(
                 string_array.len(),
                 total_capacity,
@@ -454,6 +469,21 @@ mod tests {
             Utf8,
             StringArray
         );
+        test_function!(
+            RepeatFunc::new(),
+            vec![
+                ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("a")))),
+                ColumnarValue::Scalar(ScalarValue::Int64(Some(1 << 60))),
+            ],
+            exec_err!(
+                "string size overflow on repeat, max size is {}, but got {}",
+                u32::MAX,
+                1usize << 60
+            ),
+            &str,
+            LargeUtf8,
+            LargeStringArray
+        );
 
         Ok(())
     }
@@ -519,6 +549,27 @@ mod tests {
         );
     }
 
+    #[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()
+            .strip_backtrace();
+        assert_eq!(
+            err,
+            format!(
+                "Execution error: string size overflow on repeat, max size is {}, but got {}",
+                u32::MAX,
+                n * u32::MAX as usize
+            )
+        );
+    }
+
     #[test]
     fn test_repeat_sliced_large_string_with_null_offset() {
         let (strings, counts) =