From 208b97043b98422df071dd1d80423163b7b343c6 Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Mon, 2 Nov 2026 19:23:49 -0800 Subject: [PATCH 1/5] fix: support scalar string argument in native lpad/rpad The string argument of `lpad`/`rpad` can reference columns and still evaluate to a scalar for a whole batch, e.g. `IF(id > 0, (SELECT max(s) FROM t), s)` when every row takes the subquery branch. `read_side_padding` only accepted an array there and failed with `Unsupported arguments`. Expand a scalar string argument to an array, sized to the other array arguments or one row if all arguments are scalars, and reuse the existing kernel. Return a scalar when every argument was a scalar. Closes #7245. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9 --- .../char_varchar_utils/read_side_padding.rs | 111 ++++++++++++++++++ .../expressions/string/string_lpad.sql | 24 ++++ .../expressions/string/string_rpad.sql | 24 ++++ 3 files changed, 159 insertions(+) diff --git a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs index d066076b369..5615d5db949 100644 --- a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs +++ b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs @@ -181,6 +181,29 @@ 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. + // Expand it to an array so the arms above handle it. + [ColumnarValue::Scalar(string), rest @ ..] => { + let num_rows = rest.iter().find_map(|arg| match arg { + ColumnarValue::Array(array) => Some(array.len()), + ColumnarValue::Scalar(_) => None, + }); + let mut expanded = Vec::with_capacity(args.len()); + expanded.push(ColumnarValue::Array( + string.to_array_of_size(num_rows.unwrap_or(1))?, + )); + expanded.extend_from_slice(rest); + let result = spark_read_side_padding2(&expanded, truncate, is_left_pad)?; + match num_rows { + Some(_) => Ok(result), + // Every argument was a scalar, so the result is one too. + None => Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result.to_array(1)?, + 0, + )?)), + } + } other => Err(DataFusionError::Internal(format!( "Unsupported arguments {other:?} for function rpad/lpad/read_side_padding", ))), @@ -585,6 +608,94 @@ 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])]; diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql index 145aa2a42b9..c3deb9d8e4e 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql @@ -59,3 +59,27 @@ SELECT lpad('hi', 5, 'x'), lpad('hello', 3, 'x'), lpad('', 3, 'a'), lpad(NULL, 5 query expect_dispatch(lpad) SELECT lpad('hi', 5, 'xy') FROM test_lpad + +-- The string argument references a column but is a scalar for the whole batch when +-- every row takes the subquery branch (issue #7245). +statement +CREATE TABLE test_lpad_scalar_str(id int, s string) USING parquet + +statement +INSERT INTO test_lpad_scalar_str VALUES (1, 'ab'), (2, 'cd') + +query expect_native(lpad) +SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str), s), 5, '0') FROM test_lpad_scalar_str + +query expect_native(lpad) +SELECT id, lpad(CASE WHEN id > 0 THEN (SELECT max(s) FROM test_lpad_scalar_str) ELSE s END, 5, '0') FROM test_lpad_scalar_str + +query expect_native(lpad) +SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str), s), 1) FROM test_lpad_scalar_str + +query expect_native(lpad) +SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str WHERE id > 100), s), 5, '0') FROM test_lpad_scalar_str + +-- Only some rows take the subquery branch, so the string argument stays an array. +query expect_native(lpad) +SELECT id, lpad(IF(id > 1, (SELECT max(s) FROM test_lpad_scalar_str), s), 5, '0') FROM test_lpad_scalar_str diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql index 6f243867361..7a49280d129 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql @@ -59,3 +59,27 @@ SELECT rpad('hi', 5, 'x'), rpad('hello', 3, 'x'), rpad('', 3, 'a'), rpad(NULL, 5 query expect_dispatch(rpad) SELECT rpad('hi', 5, 'xy') FROM test_rpad + +-- The string argument references a column but is a scalar for the whole batch when +-- every row takes the subquery branch (issue #7245). +statement +CREATE TABLE test_rpad_scalar_str(id int, s string) USING parquet + +statement +INSERT INTO test_rpad_scalar_str VALUES (1, 'ab'), (2, 'cd') + +query expect_native(rpad) +SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str), s), 5, '0') FROM test_rpad_scalar_str + +query expect_native(rpad) +SELECT id, rpad(CASE WHEN id > 0 THEN (SELECT max(s) FROM test_rpad_scalar_str) ELSE s END, 5, '0') FROM test_rpad_scalar_str + +query expect_native(rpad) +SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str), s), 1) FROM test_rpad_scalar_str + +query expect_native(rpad) +SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str WHERE id > 100), s), 5, '0') FROM test_rpad_scalar_str + +-- Only some rows take the subquery branch, so the string argument stays an array. +query expect_native(rpad) +SELECT id, rpad(IF(id > 1, (SELECT max(s) FROM test_rpad_scalar_str), s), 5, '0') FROM test_rpad_scalar_str From e0a5101f9f2f83cc3d89a1f34391161fe8f93617 Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Tue, 3 Nov 2026 12:40:54 -0800 Subject: [PATCH 2/5] fix: handle NULL scalar length and avoid expanding scalar string in native lpad/rpad A NULL scalar length (e.g. from a scalar subquery returning no rows) now yields NULL rows instead of an error. A scalar string argument is read per row rather than expanded with to_array_of_size, which could overflow the 32-bit offsets of a Utf8 array for wide strings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9 --- .../char_varchar_utils/read_side_padding.rs | 197 ++++++++++++++++-- .../expressions/string/string_lpad.sql | 7 + .../expressions/string/string_rpad.sql | 7 + 3 files changed, 188 insertions(+), 23 deletions(-) diff --git a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs index 5615d5db949..44173d47376 100644 --- a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs +++ b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs @@ -18,8 +18,8 @@ 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::{ArrayRef, OffsetSizeTrait}; +use arrow::array::{make_array, new_null_array, Array, AsArray, DictionaryArray}; +use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait}; use arrow::datatypes::DataType; use datafusion::common::{cast::as_generic_string_array, DataFusionError, ScalarValue}; use datafusion::physical_plan::ColumnarValue; @@ -48,6 +48,17 @@ fn spark_read_side_padding2( is_left_pad: bool, ) -> Result { 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::( @@ -183,18 +194,51 @@ 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. - // Expand it to an array so the arms above handle it. - [ColumnarValue::Scalar(string), rest @ ..] => { - let num_rows = rest.iter().find_map(|arg| match arg { + // 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", + ))) + } + }; + let num_rows = match length { ColumnarValue::Array(array) => Some(array.len()), ColumnarValue::Scalar(_) => None, - }); - let mut expanded = Vec::with_capacity(args.len()); - expanded.push(ColumnarValue::Array( - string.to_array_of_size(num_rows.unwrap_or(1))?, - )); - expanded.extend_from_slice(rest); - let result = spark_read_side_padding2(&expanded, truncate, is_left_pad)?; + }; + let len = num_rows.unwrap_or(1); + let result = match string { + ScalarValue::Utf8(value) => pad_strings::( + StringInput::Scalar { + value: value.as_deref(), + len, + }, + truncate, + length.clone(), + pad_string, + is_left_pad, + )?, + ScalarValue::LargeUtf8(value) => pad_strings::( + StringInput::Scalar { + value: value.as_deref(), + len, + }, + truncate, + length.clone(), + pad_string, + is_left_pad, + )?, + other => { + return Err(DataFusionError::Internal(format!( + "Unsupported data type {:?} for function rpad/lpad/read_side_padding", + other.data_type() + ))) + } + }; match num_rows { Some(_) => Ok(result), // Every argument was a scalar, so the result is one too. @@ -217,8 +261,61 @@ fn spark_read_side_padding_internal( pad_string: &str, is_left_pad: bool, ) -> Result { - let string_array = as_generic_string_array::(array)?; + pad_strings( + StringInput::Array(as_generic_string_array::(array)?), + truncate, + pad_type, + pad_string, + is_left_pad, + ) +} + +/// The strings to pad: a string array, or one value shared by `len` rows. +enum StringInput<'a, T: OffsetSizeTrait> { + Array(&'a GenericStringArray), + Scalar { value: Option<&'a str>, len: usize }, +} +impl<'a, T: OffsetSizeTrait> StringInput<'a, T> { + fn len(&self) -> usize { + match self { + Self::Array(array) => array.len(), + Self::Scalar { len, .. } => *len, + } + } + + fn is_ascii(&self) -> bool { + match self { + Self::Array(array) => array.is_ascii(), + Self::Scalar { value, .. } => value.is_none_or(str::is_ascii), + } + } + + /// Bytes of input values, used to size the output when rows pass through untruncated. A + /// scalar counts once: reserving it for every row could exceed the output's offsets. + fn value_bytes(&self) -> usize { + match self { + Self::Array(array) => array.value_data().len(), + Self::Scalar { value, .. } => value.map_or(0, str::len), + } + } + + #[inline] + fn value(&self, i: usize) -> Option<&'a str> { + match self { + Self::Array(array) => array.is_valid(i).then(|| array.value(i)), + Self::Scalar { value, .. } => *value, + } + } +} + +fn pad_strings( + strings: StringInput<'_, T>, + truncate: bool, + pad_type: ColumnarValue, + pad_string: &str, + is_left_pad: bool, +) -> Result { match pad_type { ColumnarValue::Array(array_int) => { let int_pad_array = array_int.as_primitive::(); @@ -236,18 +333,18 @@ fn spark_read_side_padding_internal( max_length = max_length.max(length); } let mut builder = GenericStringBuilder::::with_capacity( - string_array.len(), - data_capacity.max(string_array.value_data().len()), + strings.len(), + data_capacity.max(strings.value_bytes()), ); let padder = Padder { pad: PadPattern::new(pad_string, max_length), - ascii: string_array.is_ascii(), + ascii: strings.is_ascii(), truncate, is_left_pad, }; - for (string, length) in string_array.iter().zip(int_pad_array) { - match (string, length) { + for (i, length) in int_pad_array.iter().enumerate() { + match (strings.value(i), length) { (Some(string), Some(length)) => { if length >= 0 { padder.append(&mut builder, string, length as usize); @@ -266,18 +363,18 @@ fn spark_read_side_padding_internal( let length = 0.max(i32::try_from(const_pad_length)?) as usize; let mut builder = GenericStringBuilder::::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: strings.is_ascii(), truncate, is_left_pad, }; - for string in string_array.iter() { - match string { + for i in 0..strings.len() { + match strings.value(i) { Some(string) => padder.append(&mut builder, string, length), _ => builder.append_null(), } @@ -703,4 +800,58 @@ mod tests { 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::(); + assert_eq!(result.len(), rows); + assert!(result.iter().all(|v| v == Some("x"))); + } + } + + // 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()); + } } diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql index c3deb9d8e4e..17c8f7b030b 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql @@ -83,3 +83,10 @@ SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str WHERE id > 1 -- Only some rows take the subquery branch, so the string argument stays an array. query expect_native(lpad) SELECT id, lpad(IF(id > 1, (SELECT max(s) FROM test_lpad_scalar_str), s), 5, '0') FROM test_lpad_scalar_str + +-- A scalar subquery that returns no rows gives a NULL length (issue #7452). +query expect_native(lpad) +SELECT id, lpad(s, (SELECT max(id) FROM test_lpad_scalar_str WHERE id > 100)) FROM test_lpad_scalar_str + +query expect_native(lpad) +SELECT id, lpad(s, (SELECT max(id) FROM test_lpad_scalar_str WHERE id > 100), '0') FROM test_lpad_scalar_str diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql index 7a49280d129..04d02951b06 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql @@ -83,3 +83,10 @@ SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str WHERE id > 1 -- Only some rows take the subquery branch, so the string argument stays an array. query expect_native(rpad) SELECT id, rpad(IF(id > 1, (SELECT max(s) FROM test_rpad_scalar_str), s), 5, '0') FROM test_rpad_scalar_str + +-- A scalar subquery that returns no rows gives a NULL length (issue #7452). +query expect_native(rpad) +SELECT id, rpad(s, (SELECT max(id) FROM test_rpad_scalar_str WHERE id > 100)) FROM test_rpad_scalar_str + +query expect_native(rpad) +SELECT id, rpad(s, (SELECT max(id) FROM test_rpad_scalar_str WHERE id > 100), '0') FROM test_rpad_scalar_str From f913360c59020fc3cb642934c208e1312a922727 Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Tue, 3 Nov 2026 12:53:24 -0800 Subject: [PATCH 3/5] refactor: pass strings to the padding loop as an iterator; reuse test tables Address review: the padding loop takes an iterator, so the array path is unchanged from main and a scalar string is std::iter::repeat_n. The SQL tests reuse test_lpad/test_rpad, which covers a non-ASCII scalar string and a scalar string with an array length. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9 --- .../char_varchar_utils/read_side_padding.rs | 136 +++++++++--------- .../expressions/string/string_lpad.sql | 24 ++-- .../expressions/string/string_rpad.sql | 24 ++-- 3 files changed, 87 insertions(+), 97 deletions(-) diff --git a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs index 44173d47376..d10f2f9f7af 100644 --- a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs +++ b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs @@ -19,7 +19,7 @@ use arrow::array::builder::GenericStringBuilder; use arrow::array::cast::as_dictionary_array; use arrow::array::types::Int32Type; use arrow::array::{make_array, new_null_array, Array, AsArray, DictionaryArray}; -use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait}; +use arrow::array::{ArrayRef, OffsetSizeTrait}; use arrow::datatypes::DataType; use datafusion::common::{cast::as_generic_string_array, DataFusionError, ScalarValue}; use datafusion::physical_plan::ColumnarValue; @@ -212,26 +212,34 @@ fn spark_read_side_padding2( }; let len = num_rows.unwrap_or(1); let result = match string { - ScalarValue::Utf8(value) => pad_strings::( - StringInput::Scalar { - value: value.as_deref(), - len, - }, - truncate, - length.clone(), - pad_string, - is_left_pad, - )?, - ScalarValue::LargeUtf8(value) => pad_strings::( - StringInput::Scalar { - value: value.as_deref(), - len, - }, - truncate, - length.clone(), - pad_string, - is_left_pad, - )?, + ScalarValue::Utf8(value) | ScalarValue::LargeUtf8(value) => { + let value = value.as_deref(); + // A scalar contributes no capacity floor: reserving its bytes for every + // row could exceed the 32-bit offsets of the output. + let strings = std::iter::repeat_n(value, len); + let ascii = value.is_none_or(str::is_ascii); + if matches!(string, ScalarValue::Utf8(_)) { + pad_strings::( + strings, + ascii, + 0, + truncate, + length.clone(), + pad_string, + is_left_pad, + )? + } else { + pad_strings::( + strings, + ascii, + 0, + truncate, + length.clone(), + pad_string, + is_left_pad, + )? + } + } other => { return Err(DataFusionError::Internal(format!( "Unsupported data type {:?} for function rpad/lpad/read_side_padding", @@ -261,8 +269,11 @@ fn spark_read_side_padding_internal( pad_string: &str, is_left_pad: bool, ) -> Result { - pad_strings( - StringInput::Array(as_generic_string_array::(array)?), + let string_array = as_generic_string_array::(array)?; + pad_strings::( + string_array.iter(), + string_array.is_ascii(), + string_array.value_data().len(), truncate, pad_type, pad_string, @@ -270,47 +281,12 @@ fn spark_read_side_padding_internal( ) } -/// The strings to pad: a string array, or one value shared by `len` rows. -enum StringInput<'a, T: OffsetSizeTrait> { - Array(&'a GenericStringArray), - Scalar { value: Option<&'a str>, len: usize }, -} - -impl<'a, T: OffsetSizeTrait> StringInput<'a, T> { - fn len(&self) -> usize { - match self { - Self::Array(array) => array.len(), - Self::Scalar { len, .. } => *len, - } - } - - fn is_ascii(&self) -> bool { - match self { - Self::Array(array) => array.is_ascii(), - Self::Scalar { value, .. } => value.is_none_or(str::is_ascii), - } - } - - /// Bytes of input values, used to size the output when rows pass through untruncated. A - /// scalar counts once: reserving it for every row could exceed the output's offsets. - fn value_bytes(&self) -> usize { - match self { - Self::Array(array) => array.value_data().len(), - Self::Scalar { value, .. } => value.map_or(0, str::len), - } - } - - #[inline] - fn value(&self, i: usize) -> Option<&'a str> { - match self { - Self::Array(array) => array.is_valid(i).then(|| array.value(i)), - Self::Scalar { value, .. } => *value, - } - } -} - -fn pad_strings( - strings: StringInput<'_, T>, +/// Pads each of `strings`. `ascii` says every string is ASCII, and `min_capacity` is a floor +/// for the output's data buffer, covering rows that pass through untruncated. +fn pad_strings<'a, T: OffsetSizeTrait, I: ExactSizeIterator>>( + strings: I, + ascii: bool, + min_capacity: usize, truncate: bool, pad_type: ColumnarValue, pad_string: &str, @@ -334,17 +310,17 @@ fn pad_strings( } let mut builder = GenericStringBuilder::::with_capacity( strings.len(), - data_capacity.max(strings.value_bytes()), + data_capacity.max(min_capacity), ); let padder = Padder { pad: PadPattern::new(pad_string, max_length), - ascii: strings.is_ascii(), + ascii, truncate, is_left_pad, }; - for (i, length) in int_pad_array.iter().enumerate() { - match (strings.value(i), length) { + for (string, length) in strings.zip(int_pad_array) { + match (string, length) { (Some(string), Some(length)) => { if length >= 0 { padder.append(&mut builder, string, length as usize); @@ -368,13 +344,13 @@ fn pad_strings( ); let padder = Padder { pad: PadPattern::new(pad_string, length), - ascii: strings.is_ascii(), + ascii, truncate, is_left_pad, }; - for i in 0..strings.len() { - match strings.value(i) { + for string in strings { + match string { Some(string) => padder.append(&mut builder, string, length), _ => builder.append_null(), } @@ -840,6 +816,24 @@ mod tests { } } + #[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() { diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql index 17c8f7b030b..d1ffce15507 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql @@ -61,32 +61,30 @@ query expect_dispatch(lpad) SELECT lpad('hi', 5, 'xy') FROM test_lpad -- The string argument references a column but is a scalar for the whole batch when --- every row takes the subquery branch (issue #7245). -statement -CREATE TABLE test_lpad_scalar_str(id int, s string) USING parquet - -statement -INSERT INTO test_lpad_scalar_str VALUES (1, 'ab'), (2, 'cd') +-- every row takes the subquery branch (issue #7245). max(s) is 'é', which is not ASCII, +-- and the predicate holds for every row without being folded. +query expect_native(lpad) +SELECT lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad), s), len, 'xy') FROM test_lpad query expect_native(lpad) -SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str), s), 5, '0') FROM test_lpad_scalar_str +SELECT lpad(CASE WHEN coalesce(len, 0) > -1000 THEN (SELECT max(s) FROM test_lpad) ELSE s END, len, 'xy') FROM test_lpad query expect_native(lpad) -SELECT id, lpad(CASE WHEN id > 0 THEN (SELECT max(s) FROM test_lpad_scalar_str) ELSE s END, 5, '0') FROM test_lpad_scalar_str +SELECT lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad), s), len) FROM test_lpad query expect_native(lpad) -SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str), s), 1) FROM test_lpad_scalar_str +SELECT lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad), s), 3, 'x') FROM test_lpad query expect_native(lpad) -SELECT id, lpad(IF(id > 0, (SELECT max(s) FROM test_lpad_scalar_str WHERE id > 100), s), 5, '0') FROM test_lpad_scalar_str +SELECT lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad WHERE len > 100), s), len, 'x') FROM test_lpad -- Only some rows take the subquery branch, so the string argument stays an array. query expect_native(lpad) -SELECT id, lpad(IF(id > 1, (SELECT max(s) FROM test_lpad_scalar_str), s), 5, '0') FROM test_lpad_scalar_str +SELECT lpad(IF(len > 3, (SELECT max(s) FROM test_lpad), s), len, 'x') FROM test_lpad -- A scalar subquery that returns no rows gives a NULL length (issue #7452). query expect_native(lpad) -SELECT id, lpad(s, (SELECT max(id) FROM test_lpad_scalar_str WHERE id > 100)) FROM test_lpad_scalar_str +SELECT lpad(s, (SELECT max(len) FROM test_lpad WHERE len > 100)) FROM test_lpad query expect_native(lpad) -SELECT id, lpad(s, (SELECT max(id) FROM test_lpad_scalar_str WHERE id > 100), '0') FROM test_lpad_scalar_str +SELECT lpad(s, (SELECT max(len) FROM test_lpad WHERE len > 100), 'x') FROM test_lpad diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql index 04d02951b06..752db6f1025 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql @@ -61,32 +61,30 @@ query expect_dispatch(rpad) SELECT rpad('hi', 5, 'xy') FROM test_rpad -- The string argument references a column but is a scalar for the whole batch when --- every row takes the subquery branch (issue #7245). -statement -CREATE TABLE test_rpad_scalar_str(id int, s string) USING parquet - -statement -INSERT INTO test_rpad_scalar_str VALUES (1, 'ab'), (2, 'cd') +-- every row takes the subquery branch (issue #7245). max(s) is 'é', which is not ASCII, +-- and the predicate holds for every row without being folded. +query expect_native(rpad) +SELECT rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad), s), len, 'xy') FROM test_rpad query expect_native(rpad) -SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str), s), 5, '0') FROM test_rpad_scalar_str +SELECT rpad(CASE WHEN coalesce(len, 0) > -1000 THEN (SELECT max(s) FROM test_rpad) ELSE s END, len, 'xy') FROM test_rpad query expect_native(rpad) -SELECT id, rpad(CASE WHEN id > 0 THEN (SELECT max(s) FROM test_rpad_scalar_str) ELSE s END, 5, '0') FROM test_rpad_scalar_str +SELECT rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad), s), len) FROM test_rpad query expect_native(rpad) -SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str), s), 1) FROM test_rpad_scalar_str +SELECT rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad), s), 3, 'x') FROM test_rpad query expect_native(rpad) -SELECT id, rpad(IF(id > 0, (SELECT max(s) FROM test_rpad_scalar_str WHERE id > 100), s), 5, '0') FROM test_rpad_scalar_str +SELECT rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad WHERE len > 100), s), len, 'x') FROM test_rpad -- Only some rows take the subquery branch, so the string argument stays an array. query expect_native(rpad) -SELECT id, rpad(IF(id > 1, (SELECT max(s) FROM test_rpad_scalar_str), s), 5, '0') FROM test_rpad_scalar_str +SELECT rpad(IF(len > 3, (SELECT max(s) FROM test_rpad), s), len, 'x') FROM test_rpad -- A scalar subquery that returns no rows gives a NULL length (issue #7452). query expect_native(rpad) -SELECT id, rpad(s, (SELECT max(id) FROM test_rpad_scalar_str WHERE id > 100)) FROM test_rpad_scalar_str +SELECT rpad(s, (SELECT max(len) FROM test_rpad WHERE len > 100)) FROM test_rpad query expect_native(rpad) -SELECT id, rpad(s, (SELECT max(id) FROM test_rpad_scalar_str WHERE id > 100), '0') FROM test_rpad_scalar_str +SELECT rpad(s, (SELECT max(len) FROM test_rpad WHERE len > 100), 'x') FROM test_rpad From 6161cfb3d1331b48b0cb0eadca15f3070b68b8e4 Mon Sep 17 00:00:00 2001 From: Tien-Lung Fu <0050211@gmail.com> Date: Thu, 5 Nov 2026 13:22:48 -0800 Subject: [PATCH 4/5] test: cover a scalar string with a NULL length; reword lpad/rpad str reason The serde only rejects a literal str, and the native kernel now handles a str that evaluates to a scalar at runtime, so the unsupported reason says literal rather than scalar. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9 --- spark/src/main/scala/org/apache/comet/serde/strings.scala | 2 +- .../resources/sql-tests/expressions/string/string_lpad.sql | 4 ++++ .../resources/sql-tests/expressions/string/string_rpad.sql | 4 ++++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/strings.scala b/spark/src/main/scala/org/apache/comet/serde/strings.scala index 3ae52735841..bf5218f8f9b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/strings.scala +++ b/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." } diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql index d1ffce15507..369a858eb61 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_lpad.sql @@ -88,3 +88,7 @@ SELECT lpad(s, (SELECT max(len) FROM test_lpad WHERE len > 100)) FROM test_lpad query expect_native(lpad) SELECT lpad(s, (SELECT max(len) FROM test_lpad WHERE len > 100), 'x') FROM test_lpad + +-- A scalar string with a NULL scalar length gives a NULL scalar. +query expect_native(lpad) +SELECT lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad), s), (SELECT max(len) FROM test_lpad WHERE len > 100), 'x') FROM test_lpad diff --git a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql index 752db6f1025..b24cc09ec84 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/string_rpad.sql @@ -88,3 +88,7 @@ SELECT rpad(s, (SELECT max(len) FROM test_rpad WHERE len > 100)) FROM test_rpad query expect_native(rpad) SELECT rpad(s, (SELECT max(len) FROM test_rpad WHERE len > 100), 'x') FROM test_rpad + +-- A scalar string with a NULL scalar length gives a NULL scalar. +query expect_native(rpad) +SELECT rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad), s), (SELECT max(len) FROM test_rpad WHERE len > 100), 'x') FROM test_rpad From 1f54ae08b352daacbf2781cd4fd716371b170559 Mon Sep 17 00:00:00 2001 From: tlfua <0050211@gmail.com> Date: Mon, 9 Nov 2026 15:53:11 -0800 Subject: [PATCH 5/5] refactor: extract a generic helper for the scalar lpad/rpad string Move the scalar string arm into spark_read_side_padding_scalar::, dispatching Utf8 to i32 offsets and LargeUtf8 to i64 offsets in separate arms instead of re-checking the variant. pad_strings now returns ArrayRef, so the all-scalar case converts with ScalarValue::try_from_array(&result, 0). Document that min_capacity only applies when the length is an array. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9 --- .../char_varchar_utils/read_side_padding.rs | 114 +++++++++--------- 1 file changed, 59 insertions(+), 55 deletions(-) diff --git a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs index d10f2f9f7af..3d3d1d5b084 100644 --- a/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs +++ b/native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs @@ -206,54 +206,25 @@ fn spark_read_side_padding2( ))) } }; - let num_rows = match length { - ColumnarValue::Array(array) => Some(array.len()), - ColumnarValue::Scalar(_) => None, - }; - let len = num_rows.unwrap_or(1); - let result = match string { - ScalarValue::Utf8(value) | ScalarValue::LargeUtf8(value) => { - let value = value.as_deref(); - // A scalar contributes no capacity floor: reserving its bytes for every - // row could exceed the 32-bit offsets of the output. - let strings = std::iter::repeat_n(value, len); - let ascii = value.is_none_or(str::is_ascii); - if matches!(string, ScalarValue::Utf8(_)) { - pad_strings::( - strings, - ascii, - 0, - truncate, - length.clone(), - pad_string, - is_left_pad, - )? - } else { - pad_strings::( - strings, - ascii, - 0, - truncate, - length.clone(), - pad_string, - is_left_pad, - )? - } - } - other => { - return Err(DataFusionError::Internal(format!( - "Unsupported data type {:?} for function rpad/lpad/read_side_padding", - other.data_type() - ))) - } - }; - match num_rows { - Some(_) => Ok(result), - // Every argument was a scalar, so the result is one too. - None => Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( - &result.to_array(1)?, - 0, - )?)), + match string { + ScalarValue::Utf8(value) => spark_read_side_padding_scalar::( + value.as_deref(), + truncate, + length, + pad_string, + is_left_pad, + ), + ScalarValue::LargeUtf8(value) => spark_read_side_padding_scalar::( + 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!( @@ -270,7 +241,7 @@ fn spark_read_side_padding_internal( is_left_pad: bool, ) -> Result { let string_array = as_generic_string_array::(array)?; - pad_strings::( + Ok(ColumnarValue::Array(pad_strings::( string_array.iter(), string_array.is_ascii(), string_array.value_data().len(), @@ -278,11 +249,44 @@ fn spark_read_side_padding_internal( pad_type, pad_string, is_left_pad, - ) + )?)) +} + +fn spark_read_side_padding_scalar( + value: Option<&str>, + truncate: bool, + pad_type: &ColumnarValue, + pad_string: &str, + is_left_pad: bool, +) -> Result { + 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::( + 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, and `min_capacity` is a floor -/// for the output's data buffer, covering rows that pass through untruncated. +/// 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>>( strings: I, ascii: bool, @@ -291,7 +295,7 @@ fn pad_strings<'a, T: OffsetSizeTrait, I: ExactSizeIterator Result { +) -> Result { match pad_type { ColumnarValue::Array(array_int) => { let int_pad_array = array_int.as_primitive::(); @@ -333,7 +337,7 @@ fn pad_strings<'a, T: OffsetSizeTrait, I: ExactSizeIterator 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; @@ -355,7 +359,7 @@ fn pad_strings<'a, T: OffsetSizeTrait, I: ExactSizeIterator builder.append_null(), } } - Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + Ok(Arc::new(builder.finish())) } } }