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..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 @@ -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 { 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::( @@ -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::( + 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!( "Unsupported arguments {other:?} for function rpad/lpad/read_side_padding", ))), @@ -195,7 +241,61 @@ fn spark_read_side_padding_internal( is_left_pad: bool, ) -> Result { let string_array = as_generic_string_array::(array)?; + Ok(ColumnarValue::Array(pad_strings::( + 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( + 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. 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, + min_capacity: usize, + 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::(); @@ -213,17 +313,17 @@ 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(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( _ => 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::::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,6 +685,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])]; @@ -592,4 +780,76 @@ 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"))); + } + } + + #[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()); + } } 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 145aa2a42b9..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 @@ -59,3 +59,36 @@ 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). 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 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 lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad), s), len) FROM test_lpad + +query expect_native(lpad) +SELECT lpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_lpad), s), 3, 'x') FROM test_lpad + +query expect_native(lpad) +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 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 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 6f243867361..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 @@ -59,3 +59,36 @@ 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). 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 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 rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad), s), len) FROM test_rpad + +query expect_native(rpad) +SELECT rpad(IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM test_rpad), s), 3, 'x') FROM test_rpad + +query expect_native(rpad) +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 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 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