tlfua commented on code in PR #7450:
URL: https://github.com/apache/datafusion-comet/pull/7450#discussion_r4460975671
##########
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 @ ..] => {
Review Comment:
Thanks, done in 719c9b1e7. A `Scalar(Int32(NULL))` length now returns NULL
for every row, or a NULL scalar when the string is a scalar too, since
`StringLPad`/`StringRPad` are null-intolerant. This also closes
https://github.com/apache/datafusion-comet/issues/7452. Unit tests cover 2 and
3 arguments for `lpad`, `rpad` and `read_side_padding`, and the SQL files add
`lpad(s, (SELECT max(len) FROM t WHERE len > 100))` with and without a pad.
##########
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))?,
Review Comment:
Agreed, thanks. The scalar is no longer expanded. The padding loop takes the
strings as an iterator, and the scalar arm passes `std::iter::repeat_n(value,
n)` (7c4f4ae0e, following @comphead's suggestion below). The output reservation
no longer counts the scalar per row either.
`wide_scalar_string_is_not_expanded` pads a 1 MiB scalar over 4096 rows to
length 1, which panicked with `offset overflow` before.
https://github.com/apache/datafusion-comet/pull/7316 splits batches at the
planner only for column-free string scalars, so it does not reach a runtime
scalar from `IF`/`CASE WHEN` like this one, which is why the kernel handles it.
A dictionary-typed scalar string is now rejected rather than sent to the
dictionary arm (https://github.com/apache/datafusion-comet/issues/7453).
##########
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,
+ )?)),
+ }
+ }
Review Comment:
Thanks, done in 7c4f4ae0e. The padding loop now takes the strings as an
iterator, an `ascii` flag and a capacity floor. The array arms pass
`string_array.iter()`, `is_ascii()` and `value_data().len()`, so their path is
unchanged from `main`. The scalar arm passes `std::iter::repeat_n(value, n)`,
`value.is_none_or(str::is_ascii)` and `0`, so nothing is expanded and the
output is not sized by the scalar.
One part stays: when every argument is a scalar, the loop runs with `n = 1`
and builds a 1-row array, and `try_from_array` turns it back into a scalar,
since the result has to be a scalar there. The `to_array_of_size` half of the
round trip is gone.
##########
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')
Review Comment:
Good idea, done in 7c4f4ae0e for both files. The queries now run on
`test_lpad`/`test_rpad` with `IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM
t), s)`, so the scalar is `'é'` and reaches the non-ASCII path. With `len` as
the length they also cover a scalar string with an array length. The `CASE
WHEN`, 2-argument, literal length, NULL subquery, mixed batch and NULL length
queries moved over too, and the extra tables are gone.
--
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]
