tlfua commented on code in PR #7450:
URL: https://github.com/apache/datafusion-comet/pull/7450#discussion_r4519895414
##########
native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs:
##########
@@ -181,6 +192,70 @@ 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",
+ )))
+ }
+ };
+ 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(_)) {
Review Comment:
Thanks, done in
[1f54ae08b](https://github.com/apache/datafusion-comet/pull/7450/commits/1f54ae08b352daacbf2781cd4fd716371b170559).
`Utf8` and `LargeUtf8` now have separate arms that call a generic
`spark_read_side_padding_scalar::<T>` helper, so the `matches!` re-check and
the duplicated call are gone, and any other scalar type returns an error. I
named it `spark_read_side_padding_scalar` rather than `pad_scalar` because the
tests already have a `pad_scalar` helper, and the longer name matches
`spark_read_side_padding_internal`.
##########
native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs:
##########
@@ -181,6 +192,70 @@ 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",
+ )))
+ }
+ };
+ 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::<i32, _>(
+ strings,
+ ascii,
+ 0,
+ truncate,
+ length.clone(),
+ pad_string,
+ is_left_pad,
+ )?
+ } else {
+ pad_strings::<i64, _>(
+ 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 {
Review Comment:
Thanks for probing this. You're right that `ScalarFunctionExpr::evaluate`
converts a one-row array to a scalar when every argument is a scalar, so the
query result would be the same without it. I'd still like to keep the
conversion in the kernel, for three reasons:
- The NULL-length arm a few lines above already returns a
`ColumnarValue::Scalar` for a scalar string, so this keeps the scalar arm
consistent with it.
- The Rust unit tests call `spark_read_side_padding2` directly and expect a
scalar back when every argument is a scalar. Without the conversion they would
see a one-row array.
- It's the same pattern other kernels in this crate use, such as
`spark_array_position` and `sort_array`.
Following coderfender's suggestion,
[1f54ae08b](https://github.com/apache/datafusion-comet/pull/7450/commits/1f54ae08b352daacbf2781cd4fd716371b170559)
also simplifies the tail: `pad_strings` now returns an `ArrayRef`, so the
all-scalar case is just `ScalarValue::try_from_array(&result, 0)` with no
`to_array(1)` step. If you'd still prefer to drop it and rely on DataFusion,
I'm happy to change it.
##########
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."
Review Comment:
Good catch, thanks. I opened
https://github.com/apache/datafusion-comet/issues/7583 to track it, since it is
a separate change from the scalar `str` fix here. I'll follow up with a fix
there.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
