From 6fc59544a9c20cd307d640b26d827f357654dad4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Fri, 4 Sep 2026 20:36:36 +0800 Subject: [PATCH] fix: return NULL from rpad/lpad when the length column is NULL instead of panicking `spark_read_side_padding_internal` unwrapped every value of the length array, so `rpad(s, len)` / `lpad(s, len)` (and the 3-arg forms with a literal pad) panicked with `called Option::unwrap() on a None value` and failed the task as soon as the length column contained a NULL. Spark's `StringRPad` / `StringLPad` are null-intolerant: a NULL length yields a NULL row. Treat a NULL length like a NULL string and emit a null row, and size the output buffer from the non-null lengths only, since the values under null slots are unspecified. Non-null rows are unchanged. Closes #5672 Co-Authored-By: Claude Fable 5.1 --- .../char_varchar_utils/read_side_padding.rs | 112 +++++++++++++++++- .../expressions/string/string_lpad.sql | 6 +- .../expressions/string/string_rpad.sql | 6 +- .../comet/CometStringExpressionSuite.scala | 21 ++++ 4 files changed, 137 insertions(+), 8 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 565100eeb10..d066076b369 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 @@ -202,11 +202,13 @@ fn spark_read_side_padding_internal( // Every row is padded to its target length, so the sum of the target // lengths sizes the output (exactly, for ASCII input), except when a - // row is longer than its target and passes through untruncated. + // row is longer than its target and passes through untruncated. Null + // lengths produce null rows and are skipped: the values under null + // slots are unspecified and must not size the output. let mut data_capacity = 0usize; let mut max_length = 0usize; - for length in int_pad_array.values() { - let length = (*length).max(0) as usize; + for length in int_pad_array.iter().flatten() { + let length = length.max(0) as usize; data_capacity = data_capacity.saturating_add(length); max_length = max_length.max(length); } @@ -222,15 +224,16 @@ fn spark_read_side_padding_internal( }; for (string, length) in string_array.iter().zip(int_pad_array) { - let length = length.unwrap(); - match string { - Some(string) => { + match (string, length) { + (Some(string), Some(length)) => { if length >= 0 { padder.append(&mut builder, string, length as usize); } else { builder.append_value(""); } } + // Spark's StringRPad/StringLPad are null-intolerant: a null + // string or a null length yields a null row. _ => builder.append_null(), } } @@ -399,6 +402,10 @@ mod tests { ColumnarValue::Scalar(ScalarValue::Utf8(Some(pad.to_string()))) } + fn len_array(lengths: &[Option]) -> ColumnarValue { + ColumnarValue::Array(Arc::new(Int32Array::from(lengths.to_vec())) as ArrayRef) + } + #[test] fn rpad_default_padding() { let args = vec![ @@ -492,4 +499,97 @@ mod tests { vec![Some("a".to_string()), Some(" abc".to_string()), None] ); } + + #[test] + fn rpad_null_length_yields_null_row() { + // Spark's StringRPad is null-intolerant: a NULL length makes only that + // row NULL, the other rows are padded or truncated as usual. + let strings = [Some("abc"), Some("abc"), Some("abcdef"), Some("abc")]; + let lengths = [Some(5), None, Some(2), Some(-1)]; + // 2 args (default pad of ' ') + let args = vec![utf8(&strings), len_array(&lengths)]; + assert_eq!( + result_values(spark_rpad(&args).unwrap()), + vec![ + Some("abc ".to_string()), + None, + Some("ab".to_string()), + Some("".to_string()), + ] + ); + // 3 args + let args = vec![utf8(&strings), len_array(&lengths), pad_scalar("xy")]; + assert_eq!( + result_values(spark_rpad(&args).unwrap()), + vec![ + Some("abcxy".to_string()), + None, + Some("ab".to_string()), + Some("".to_string()), + ] + ); + // read-side padding takes the same path, without truncation + let args = vec![utf8(&strings), len_array(&lengths)]; + assert_eq!( + result_values(spark_read_side_padding(&args).unwrap()), + vec![ + Some("abc ".to_string()), + None, + Some("abcdef".to_string()), + Some("".to_string()), + ] + ); + } + + #[test] + fn lpad_null_length_yields_null_row() { + let strings = [Some("abc"), Some("abc"), Some("abcdef"), Some("abc")]; + let lengths = [Some(5), None, Some(2), Some(-1)]; + // 2 args (default pad of ' ') + let args = vec![utf8(&strings), len_array(&lengths)]; + assert_eq!( + result_values(spark_lpad(&args).unwrap()), + vec![ + Some(" abc".to_string()), + None, + Some("ab".to_string()), + Some("".to_string()), + ] + ); + // 3 args + let args = vec![utf8(&strings), len_array(&lengths), pad_scalar("xy")]; + assert_eq!( + result_values(spark_lpad(&args).unwrap()), + vec![ + Some("xyabc".to_string()), + None, + Some("ab".to_string()), + Some("".to_string()), + ] + ); + } + + #[test] + fn null_string_and_null_length_yields_null_row() { + let strings = [None, None, Some("abc")]; + let lengths = [None, Some(4), None]; + let args = vec![utf8(&strings), len_array(&lengths)]; + assert_eq!( + result_values(spark_rpad(&args).unwrap()), + vec![None, None, None] + ); + let args = vec![utf8(&strings), len_array(&lengths), pad_scalar("x")]; + assert_eq!( + result_values(spark_lpad(&args).unwrap()), + vec![None, None, None] + ); + } + + #[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]); + } } 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 c27d93de621..93867576aac 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 @@ -19,7 +19,7 @@ statement CREATE TABLE test_lpad(s string, len int, pad string) USING parquet statement -INSERT INTO test_lpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x') +INSERT INTO test_lpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x') query expect_fallback(Only scalar values are supported for the `pad` argument) SELECT lpad(s, len, pad) FROM test_lpad @@ -27,6 +27,10 @@ SELECT lpad(s, len, pad) FROM test_lpad query SELECT lpad(s, len) FROM test_lpad +-- column + column + literal +query +SELECT lpad(s, len, 'x') FROM test_lpad + -- column + literal + literal query SELECT lpad(s, 5, '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 4ea06c3b23b..5832f487851 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 @@ -19,7 +19,7 @@ statement CREATE TABLE test_rpad(s string, len int, pad string) USING parquet statement -INSERT INTO test_rpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x') +INSERT INTO test_rpad VALUES ('hi', 5, 'x'), ('hello', 3, 'x'), ('hi', 5, 'xy'), ('', 3, 'a'), (NULL, 5, 'x'), ('hi', 0, 'x'), ('hi', -1, 'x'), ('hi', NULL, 'x'), (NULL, NULL, 'x') query expect_fallback(Only scalar values are supported for the `pad` argument) SELECT rpad(s, len, pad) FROM test_rpad @@ -27,6 +27,10 @@ SELECT rpad(s, len, pad) FROM test_rpad query SELECT rpad(s, len) FROM test_rpad +-- column + column + literal +query +SELECT rpad(s, len, 'x') FROM test_rpad + -- column + literal + literal query SELECT rpad(s, 5, 'x') FROM test_rpad diff --git a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala index 0b5acc70a28..a0fa91b8fe1 100644 --- a/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometStringExpressionSuite.scala @@ -44,6 +44,27 @@ class CometStringExpressionSuite extends CometTestBase { testStringPadding("rpad") } + test("lpad/rpad with NULL length") { + // FuzzDataGenerator never generates NULL integers (#5389), so build the rows explicitly. + // Spark's StringLPad/StringRPad are null-intolerant: a NULL length yields a NULL row. + val data: Seq[(String, Option[Int])] = Seq( + ("abc", Some(5)), + ("abc", None), + (null, None), + (null, Some(5)), + ("abcdef", Some(2)), + ("abc", Some(-1)), + ("abc", Some(0))) ++ edgeCases.flatMap(s => Seq((s, None), (s, Some(4)))) + withParquetTable(data, "tbl") { + for (expr <- Seq("lpad", "rpad")) { + // 2 args (default pad of ' ') + checkSparkAnswerAndOperator(s"SELECT _1, _2, $expr(_1, _2) FROM tbl") + // 3 args with a literal pad + checkSparkAnswerAndOperator(s"SELECT _1, _2, $expr(_1, _2, 'xy') FROM tbl") + } + } + } + test("lpad binary") { testBinaryPadding("lpad") }