comphead commented on code in PR #7450:
URL: https://github.com/apache/datafusion-comet/pull/7450#discussion_r4460105015
##########
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:
Could these queries reuse `test_lpad` instead of a second table? It already
has NULLs, empty strings, non-ASCII values and NULL, zero and negative lengths,
so `(SELECT max(s) FROM test_lpad)` gives a multi-byte scalar (`é`), which
reaches the `ascii == false` path that `'ab'` and `'cd'` cannot. A predicate
that holds for every row and is not folded, such as `coalesce(len, 0) > -1000`,
would keep the branch scalar for the whole batch. With `len` as the length
argument, that also covers a scalar string with an array length, which today
only the Rust tests check. I have not run it. The same applies to
`string_rpad.sql`.
##########
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:
Adding to the earlier comment about the expansion:
`spark_read_side_padding_internal` sizes its builder with
`data_capacity.max(string_array.value_data().len())`, so after the expansion a
1 MiB scalar over 64 rows also reserves at least 64 MiB of output, even when
`rpad(..., 1)` produces 64 bytes.
One way to avoid both is for the loops in `spark_read_side_padding_internal`
to take the strings as an iterator plus the `ascii` flag and a capacity floor.
The array arms would pass `string_array.iter()`, `string_array.is_ascii()` and
`value_data().len()`, and this arm
`std::iter::repeat(string.as_deref()).take(n)`, `string.is_ascii()` and `0`. No
expanded array is built, and the all-scalar case becomes the same loop with `n
= 1`, so the `to_array(1)` and `try_from_array` round trip at the end goes
away. I have not tried it.
--
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]
