fix: support scalar string argument in native lpad/rpad

native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs

@@ -18,7 +18,7 @@
 use arrow::array::builder::GenericStringBuilder;
 use arrow::array::cast::as_dictionary_array;
 use arrow::array::types::Int32Type;
-use arrow::array::{make_array, Array, AsArray, DictionaryArray};
+use arrow::array::{make_array, new_null_array, Array, AsArray, DictionaryArray};
 use arrow::array::{ArrayRef, OffsetSizeTrait};
 use arrow::datatypes::DataType;
 use datafusion::common::{cast::as_generic_string_array, DataFusionError, ScalarValue};
@@ -48,6 +48,17 @@ fn spark_read_side_padding2(
     is_left_pad: bool,
 ) -> Result<ColumnarValue, DataFusionError> {
     match args {
+        // Spark's StringRPad/StringLPad are null-intolerant, so a NULL length makes every
+        // row NULL. A scalar subquery that returns no rows produces such a length.
+        [string, ColumnarValue::Scalar(ScalarValue::Int32(None)), ..] => match string {
+            ColumnarValue::Array(array) => Ok(ColumnarValue::Array(new_null_array(
+                array.data_type(),
+                array.len(),
+            ))),
+            ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar(ScalarValue::try_from(
+                scalar.data_type(),
+            )?)),
+        },
         [ColumnarValue::Array(array), ColumnarValue::Scalar(ScalarValue::Int32(Some(length)))] => {
             match array.data_type() {
                 DataType::Utf8 => spark_read_side_padding_internal::<i32>(
@@ -181,6 +192,41 @@ fn spark_read_side_padding2(
                 ))),
             }
         }
+        // The string argument can be a scalar even when it references columns, e.g.
+        // `IF(p, (SELECT max(s) FROM t), s)` when every row takes the subquery branch.
+        // Every row reads the one value rather than an expanded copy, which for a wide string
+        // could exceed the 32-bit offsets of the input array.
+        [ColumnarValue::Scalar(string), length, rest @ ..] => {
+            let pad_string = match rest {
+                [] => SPACE,
+                [ColumnarValue::Scalar(ScalarValue::Utf8(Some(pad)))] => pad.as_str(),
+                _ => {
+                    return Err(DataFusionError::Internal(format!(
+                        "Unsupported arguments {args:?} for function rpad/lpad/read_side_padding",
+                    )))
+                }
+            };
+            match string {
+                ScalarValue::Utf8(value) => spark_read_side_padding_scalar::<i32>(
+                    value.as_deref(),
+                    truncate,
+                    length,
+                    pad_string,
+                    is_left_pad,
+                ),
+                ScalarValue::LargeUtf8(value) => spark_read_side_padding_scalar::<i64>(
+                    value.as_deref(),
+                    truncate,
+                    length,
+                    pad_string,
+                    is_left_pad,
+                ),
+                other => Err(DataFusionError::Internal(format!(
+                    "Unsupported data type {:?} for function rpad/lpad/read_side_padding",
+                    other.data_type()
+                ))),
+            }
+        }
         other => Err(DataFusionError::Internal(format!(
             "Unsupported arguments {other:?} for function rpad/lpad/read_side_padding",
         ))),
@@ -195,7 +241,61 @@ fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
     is_left_pad: bool,
 ) -> Result<ColumnarValue, DataFusionError> {
     let string_array = as_generic_string_array::<T>(array)?;
+    Ok(ColumnarValue::Array(pad_strings::<T, _>(
+        string_array.iter(),
+        string_array.is_ascii(),
+        string_array.value_data().len(),
+        truncate,
+        pad_type,
+        pad_string,
+        is_left_pad,
+    )?))
+}
 
+fn spark_read_side_padding_scalar<T: OffsetSizeTrait>(
+    value: Option<&str>,
+    truncate: bool,
+    pad_type: &ColumnarValue,
+    pad_string: &str,
+    is_left_pad: bool,
+) -> Result<ColumnarValue, DataFusionError> {
+    let num_rows = match pad_type {
+        ColumnarValue::Array(array) => Some(array.len()),
+        ColumnarValue::Scalar(_) => None,
+    };
+    // A scalar contributes no capacity floor: reserving its bytes for every row could
+    // exceed the 32-bit offsets of the output.
+    let result = pad_strings::<T, _>(
+        std::iter::repeat_n(value, num_rows.unwrap_or(1)),
+        value.is_none_or(str::is_ascii),
+        0,
+        truncate,
+        pad_type.clone(),
+        pad_string,
+        is_left_pad,
+    )?;
+    match num_rows {
+        Some(_) => Ok(ColumnarValue::Array(result)),
+        // Every argument was a scalar, so the result is one too.
+        None => Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
+            &result, 0,
+        )?)),
+    }
+}
+
+/// Pads each of `strings`. `ascii` says every string is ASCII. When `pad_type` is an array,
+/// `min_capacity` is a floor for the output's data buffer, covering rows that pass through
+/// untruncated. When `pad_type` is a scalar, the buffer is sized from the length alone and
+/// `min_capacity` is ignored.
+fn pad_strings<'a, T: OffsetSizeTrait, I: ExactSizeIterator<Item = Option<&'a str>>>(
+    strings: I,
+    ascii: bool,
+    min_capacity: usize,
+    truncate: bool,
+    pad_type: ColumnarValue,
+    pad_string: &str,
+    is_left_pad: bool,
+) -> Result<ArrayRef, DataFusionError> {
     match pad_type {
         ColumnarValue::Array(array_int) => {
             let int_pad_array = array_int.as_primitive::<Int32Type>();
@@ -213,17 +313,17 @@ fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
                 max_length = max_length.max(length);
             }
             let mut builder = GenericStringBuilder::<T>::with_capacity(
-                string_array.len(),
-                data_capacity.max(string_array.value_data().len()),
+                strings.len(),
+                data_capacity.max(min_capacity),
             );
             let padder = Padder {
                 pad: PadPattern::new(pad_string, max_length),
-                ascii: string_array.is_ascii(),
+                ascii,
                 truncate,
                 is_left_pad,
             };
 
-            for (string, length) in string_array.iter().zip(int_pad_array) {
+            for (string, length) in strings.zip(int_pad_array) {
                 match (string, length) {
                     (Some(string), Some(length)) => {
                         if length >= 0 {
@@ -237,29 +337,29 @@ fn spark_read_side_padding_internal<T: OffsetSizeTrait>(
                     _ => builder.append_null(),
                 }
             }
-            Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+            Ok(Arc::new(builder.finish()))
         }
         ColumnarValue::Scalar(const_pad_length) => {
             let length = 0.max(i32::try_from(const_pad_length)?) as usize;
 
             let mut builder = GenericStringBuilder::<T>::with_capacity(
-                string_array.len(),
-                string_array.len().saturating_mul(length),
+                strings.len(),
+                strings.len().saturating_mul(length),
             );
             let padder = Padder {
                 pad: PadPattern::new(pad_string, length),
-                ascii: string_array.is_ascii(),
+                ascii,
                 truncate,
                 is_left_pad,
             };
 
-            for string in string_array.iter() {
+            for string in strings {
                 match string {
                     Some(string) => padder.append(&mut builder, string, length),
                     _ => builder.append_null(),
                 }
             }
-            Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+            Ok(Arc::new(builder.finish()))
         }
     }
 }
@@ -585,11 +685,171 @@ mod tests {
         );
     }
 
+    fn utf8_scalar(value: Option<&str>) -> ColumnarValue {
+        ColumnarValue::Scalar(ScalarValue::Utf8(value.map(str::to_string)))
+    }
+
+    fn scalar_result(col: ColumnarValue) -> ScalarValue {
+        match col {
+            ColumnarValue::Scalar(value) => value,
+            ColumnarValue::Array(array) => panic!("expected a scalar, got {array:?}"),
+        }
+    }
+
+    // A conditional such as `IF(p, (SELECT max(s) FROM t), s)` evaluates to a
+    // scalar when every row in the batch takes the subquery branch (#7245).
+    #[test]
+    fn scalar_string_with_scalar_length() {
+        let args = vec![utf8_scalar(Some("cd")), len_scalar(5), pad_scalar("0")];
+        assert_eq!(
+            scalar_result(spark_lpad(&args).unwrap()),
+            ScalarValue::Utf8(Some("000cd".to_string()))
+        );
+        assert_eq!(
+            scalar_result(spark_rpad(&args).unwrap()),
+            ScalarValue::Utf8(Some("cd000".to_string()))
+        );
+        let args = vec![utf8_scalar(Some("abcdef")), len_scalar(3)];
+        assert_eq!(
+            scalar_result(spark_rpad(&args).unwrap()),
+            ScalarValue::Utf8(Some("abc".to_string()))
+        );
+        assert_eq!(
+            scalar_result(spark_read_side_padding(&args).unwrap()),
+            ScalarValue::Utf8(Some("abcdef".to_string()))
+        );
+    }
+
+    #[test]
+    fn scalar_string_with_array_length() {
+        let lengths = [Some(3), None, Some(1), Some(-1)];
+        let args = vec![utf8_scalar(Some("ab")), len_array(&lengths)];
+        assert_eq!(
+            result_values(spark_rpad(&args).unwrap()),
+            vec![
+                Some("ab ".to_string()),
+                None,
+                Some("a".to_string()),
+                Some("".to_string()),
+            ]
+        );
+        let args = vec![
+            utf8_scalar(Some("ab")),
+            len_array(&lengths),
+            pad_scalar("x"),
+        ];
+        assert_eq!(
+            result_values(spark_lpad(&args).unwrap()),
+            vec![
+                Some("xab".to_string()),
+                None,
+                Some("a".to_string()),
+                Some("".to_string()),
+            ]
+        );
+    }
+
+    #[test]
+    fn null_scalar_string_yields_null() {
+        let args = vec![utf8_scalar(None), len_scalar(5), pad_scalar("0")];
+        assert_eq!(
+            scalar_result(spark_lpad(&args).unwrap()),
+            ScalarValue::Utf8(None)
+        );
+        let args = vec![utf8_scalar(None), len_array(&[Some(2), None])];
+        assert_eq!(result_values(spark_rpad(&args).unwrap()), vec![None, None]);
+    }
+
+    #[test]
+    fn large_utf8_scalar_string() {
+        let args = vec![
+            ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("x".to_string()))),
+            len_scalar(3),
+            pad_scalar("ab"),
+        ];
+        assert_eq!(
+            scalar_result(spark_lpad(&args).unwrap()),
+            ScalarValue::LargeUtf8(Some("abx".to_string()))
+        );
+    }
+
     #[test]
     fn all_null_lengths_yield_all_null_rows() {
         let args = vec![utf8(&[Some("abc"), Some("")]), len_array(&[None, None])];
         assert_eq!(result_values(spark_rpad(&args).unwrap()), vec![None, None]);
         let args = vec![utf8(&[Some("abc"), Some("")]), len_array(&[None, None])];
         assert_eq!(result_values(spark_lpad(&args).unwrap()), vec![None, None]);
     }
+
+    // A scalar subquery that returns no rows yields a NULL scalar length (#7452).
+    #[test]
+    fn null_scalar_length_yields_null() {
+        let null_length = ColumnarValue::Scalar(ScalarValue::Int32(None));
+        let strings = utf8(&[Some("abc"), None]);
+        for args in [
+            vec![strings.clone(), null_length.clone()],
+            vec![strings, null_length.clone(), pad_scalar("x")],
+        ] {
+            assert_eq!(result_values(spark_rpad(&args).unwrap()), vec![None, None]);
+            assert_eq!(result_values(spark_lpad(&args).unwrap()), vec![None, None]);
+            assert_eq!(
+                result_values(spark_read_side_padding(&args).unwrap()),
+                vec![None, None]
+            );
+        }
+        let args = vec![utf8_scalar(Some("abc")), null_length, pad_scalar("x")];
+        assert_eq!(
+            scalar_result(spark_lpad(&args).unwrap()),
+            ScalarValue::Utf8(None)
+        );
+    }
+
+    // Expanding a 1 MiB scalar to 4,096 rows would need 4 GiB of input, past the 32-bit
+    // offsets of a Utf8 array, for 4 KiB of output.
+    #[test]
+    fn wide_scalar_string_is_not_expanded() {
+        let wide = "x".repeat(1 << 20);
+        let rows = 4096;
+        let lengths = len_array(&vec![Some(1); rows]);
+        for pad in [spark_lpad, spark_rpad] {
+            let args = vec![utf8_scalar(Some(&wide)), lengths.clone(), pad_scalar("0")];
+            let result = pad(&args).unwrap().to_array(rows).unwrap();
+            let result = result.as_string::<i32>();
+            assert_eq!(result.len(), rows);
+            assert!(result.iter().all(|v| v == Some("x")));
+        }
+    }
+
+    #[test]
+    fn non_ascii_scalar_string_with_array_length() {
+        let args = vec![
+            utf8_scalar(Some("é")),
+            len_array(&[Some(3), Some(1), Some(0), None]),
+            pad_scalar("中"),
+        ];
+        assert_eq!(
+            result_values(spark_lpad(&args).unwrap()),
+            vec![
+                Some("中中é".to_string()),
+                Some("é".to_string()),
+                Some("".to_string()),
+                None
+            ]
+        );
+    }
+
+    // A dictionary scalar is rejected rather than routed to the dictionary arm.
+    #[test]
+    fn dictionary_scalar_string_is_unsupported() {
+        let string = ScalarValue::Dictionary(
+            Box::new(DataType::Int32),
+            Box::new(ScalarValue::Utf8(Some("ab".to_string()))),
+        );
+        let args = vec![
+            ColumnarValue::Scalar(string),
+            len_array(&[Some(4)]),
+            pad_scalar("x"),
+        ];
+        assert!(spark_rpad(&args).is_err());
+    }
 }

spark/src/main/scala/org/apache/comet/serde/strings.scala

@@ -561,7 +561,7 @@ object CometRLike
 }
 
 private object PadReasons {
-  val literalStrReason = "Scalar values are not supported for the `str` argument."
+  val literalStrReason = "Literal values are not supported for the `str` argument."
   val nonLiteralPadReason = "Only scalar values are supported for the `pad` argument."
 }