Repository navigation
Conversation
The string argument of `lpad`/`rpad` can reference columns and still evaluate to a scalar for a whole batch, e.g. `IF(id > 0, (SELECT max(s) FROM t), s)` when every row takes the subquery branch. `read_side_padding` only accepted an array there and failed with `Unsupported arguments`. Expand a scalar string argument to an array, sized to the other array arguments or one row if all arguments are scalars, and reuse the existing kernel. Return a scalar when every argument was a scalar. Closes apache#7245. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9
| // 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 @ ..] => { |
There was a problem hiding this comment.
Could we also handle a NULL scalar length here? A length subquery such as (SELECT max(id) FROM t WHERE id > 100) can return Scalar(Int32(NULL)), causing both lpad and rpad to fail with Unsupported arguments instead of returning NULL.
This is an existing kernel gap that the new scalar-string path still reaches. Please add NULL propagation and a regression test for this case.
There was a problem hiding this comment.
Thanks, done in 719c9b1. 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 #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.
| }); | ||
| let mut expanded = Vec::with_capacity(args.len()); | ||
| expanded.push(ColumnarValue::Array( | ||
| string.to_array_of_size(num_rows.unwrap_or(1))?, |
There was a problem hiding this comment.
Would it be possible to read the scalar string directly when building each output row? I tested a 1 MiB scalar with 64 array lengths of 1: it expands to 64 MiB of input for only 64 bytes of output. At 4,096 rows, both lpad and rpad panic with offset overflow.
Avoiding the intermediate expansion could help with both memory usage and the overflow tracked in #7247.
There was a problem hiding this 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) (7c4f4ae, 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.
#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 (#7453).
comphead
left a comment
There was a problem hiding this comment.
Thanks for splitting this out of #5889. CaseWhenExpr::evaluate_eagerly returns the chosen branch's ColumnarValue as it is when every row picks one branch, so the IF and CASE WHEN fixtures do reach the new arm. The comments below are about the expansion and the fixture coverage.
| // 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, | ||
| )?)), | ||
| } | ||
| } |
There was a problem hiding this 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.
There was a problem hiding this comment.
Thanks, done in 7c4f4ae. 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.
| -- 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') |
There was a problem hiding this 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.
There was a problem hiding this comment.
Good idea, done in 7c4f4ae 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.
…ative lpad/rpad A NULL scalar length (e.g. from a scalar subquery returning no rows) now yields NULL rows instead of an error. A scalar string argument is read per row rather than expanded with to_array_of_size, which could overflow the 32-bit offsets of a Utf8 array for wide strings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9
… tables Address review: the padding loop takes an iterator, so the array path is unchanged from main and a scalar string is std::iter::repeat_n. The SQL tests reuse test_lpad/test_rpad, which covers a non-ASCII scalar string and a scalar string with an array length. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9
coderfender
left a comment
There was a problem hiding this comment.
Thanks for this, and for following up on the earlier threads so quickly. I checked the new arms against Spark 4.1.1. StringLPad and StringRPad are null intolerant and UTF8String.lpad/rpad count code points, and the scalar arm reuses Padder unchanged, so the results line up. It's also nice that the scalar is no longer expanded or reserved per row, and the array arms do the same work as before.
Two things I'd like to see before merge.
-
spark/src/main/scala/org/apache/comet/serde/strings.scala:564is outside the diff, but this PR makes it stale.PadReasons.literalStrReasonsays "Scalar values are not supported for thestrargument." The check behind it only rejects aLiteral, and the native kernel now handles astrthat evaluates to a scalar at runtime. Since this reason ends up on the generated compatibility page, could we reword it to something like "Literal values are not supported for thestrargument"? Separately, now that the kernel handlesScalar(str)with both a scalar and an array length, is theLiteralcheck still needed? If lifting it is out of scope here, could you open a follow-up issue and link it? -
In
string_lpad.sqlandstring_rpad.sql(around line 90), the NULL length arm has its own branch for a scalar string (read_side_padding.rs:58-60), and only the Rust testnull_scalar_length_yields_nullreaches it. Would you add one query that combines the two? For example: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
That checks the scalar in, scalar out NULL result end to end against Spark.
Both bugs also ship in branch-1.2 (and branch-1.1), and default configs reach them, so this looks like a backport candidate. The new queries go into fixtures that every Spark profile runs, so it would be good to run the other profiles before it goes in the queue.
…reason The serde only rejects a literal str, and the native kernel now handles a str that evaluates to a scalar at runtime, so the unsupported reason says literal rather than scalar. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9
|
Thanks for the careful check against Spark 4.1.1. Both done in 0ff089d.
On the other Spark profiles: the fork can't see |
0ff089d to
6161cfb
Compare
|
@coderfender Thanks again for the careful review. Both points are addressed (details in my reply above), and the follow-up for lifting the Note: I only rewrote the commit author email on this branch, with no code changes, so |
| ))) | ||
| } | ||
| }; | ||
| match num_rows { |
There was a problem hiding this comment.
DataFusion's ScalarFunctionExpr::evaluate already turns a one-row array into a scalar when every argument is a scalar. A local probe with this block removed still got a scalar for multi-row batches. Could this just return result and drop num_rows?
There was a problem hiding this 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::Scalarfor a scalar string, so this keeps the scalar arm consistent with it. - The Rust unit tests call
spark_read_side_padding2directly 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_positionandsort_array.
Following coderfender's suggestion, 1f54ae08b 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.
| // 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(_)) { |
There was a problem hiding this comment.
nit: Could Utf8 and LargeUtf8 be handled in separate arms calling a small pad_scalar::<T> helper? That drops the matches! re-check and the duplicated seven-argument call.
There was a problem hiding this comment.
Thanks, done in 1f54ae08b. 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.
|
|
||
| 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." |
There was a problem hiding this comment.
Separately, lpad/rpad don't go through withNullShortCircuit, so in ANSI mode lpad(s, CAST(t AS INT)) evaluates the cast on rows where s is NULL, which Spark skips. Worth a follow-up issue?
There was a problem hiding this comment.
Good catch, thanks. I opened #7583 to track it, since it is a separate change from the scalar str fix here. I'll follow up with a fix there.
coderfender
left a comment
There was a problem hiding this comment.
Thanks for the quick turnaround. Both points from my last round look good at this head. literalStrReason reads right now, #7525 tracks lifting the Literal check, and the new query covers the scalar string with a NULL length end to end. Two small things I'd still like to see before merge.
-
native/spark-expr/src/static_invoke/char_varchar_utils/read_side_padding.rs:214-241. The scalar arm matchesUtf8 | LargeUtf8together and then checksmatches!(string, ScalarValue::Utf8(_))again, and the twopad_stringscalls take the same seven arguments apart from the offset type. Could the body move into a generic helper next tospark_read_side_padding_internal, something likespark_read_side_padding_scalar::<T>(value: Option<&str>, length, truncate, pad_string, is_left_pad)? The arm would then dispatch the way the array arms above do, withScalarValue::Utf8(value) => spark_read_side_padding_scalar::<i32>(...)andScalarValue::LargeUtf8(value) => spark_read_side_padding_scalar::<i64>(...). It also removes a small trap. Theelsepicksi64for anything that isn'tUtf8, so a variant added to the pattern later would quietly get aLargeUtf8builder. While you're there,pad_stringscould returnArrayRef, since both of its branches build an array. The tail here would then beScalarValue::try_from_array(&result, 0)without theto_array(1)step, the same shape asspark_array_positionandsort_array. -
read_side_padding.rs:284. The doc saysmin_capacityis a floor for the output's data buffer, but only the branch with an array length reads it. The branch with a scalar length sizes the builder fromstrings.len().saturating_mul(length)and ignores it, asmaindid. Could the doc say it only applies when the length is an array? Otherwise a reader may think the0passed from the scalar arm matters for a literal length, or change that branch to honor it, which would change the reservation on the commonlpad(col, 5, 'x')path.
|
@tlfua I added |
Move the scalar string arm into spark_read_side_padding_scalar::<T>, dispatching Utf8 to i32 offsets and LargeUtf8 to i64 offsets in separate arms instead of re-checking the variant. pad_strings now returns ArrayRef, so the all-scalar case converts with ScalarValue::try_from_array(&result, 0). Document that min_capacity only applies when the length is an array. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 285ed31f-1f7b-4322-85ae-0d80fd06dac9
|
@coderfender Thanks for the review and for adding
The fork CI run for this commit passed: https://github.com/tlfua/datafusion-comet/actions/runs/42724088108 Since this PR comes from a fork, the upstream workflows for the new commit are waiting for approval: https://github.com/apache/datafusion-comet/actions/runs/42724097999. Could a committer please approve them when convenient, so the |
|
done thank you @tlfua |
coderfender
left a comment
There was a problem hiding this comment.
Thanks for the refactor. The scalar arm now dispatches to spark_read_side_padding_scalar::<T> the same way the array arms do, pad_strings returns an ArrayRef, and the min_capacity doc matches what the code does. Results are unchanged. LGTM once the run-all-spark-profiles run is green.
|
Thanks for the approval and for merging A note on the CI run for the merge commit: Spark 4.2, JDK 17 [exec] failed, and This looks unrelated to this PR, which only touches lpad/rpad. It appeared after the merge with |
Which issue does this PR close?
Closes #7245.
Closes #7452.
Related to #5889: as discussed in this comment, the padding fix is split out of that PR so it can land before the struct-typed scalar subquery support.
Rationale for this change
The string argument of
lpad/rpadcan reference columns and still evaluate to a scalar for a whole batch. For example,IF(id > 0, (SELECT max(s) FROM t), s)returns a scalar when every row in the batch takes the subquery branch. The serde only rejects a literal string, so this shape reaches the nativeread_side_paddingkernel. Every arm there expects the string to be an array, so the query fails with:@sunchao raised this in the #5889 review and suggested handling a runtime scalar string in the kernel.
What changes are included in this PR?
spark_read_side_padding2gets new arms for a scalar string argument,Utf8andLargeUtf8, which call a genericspark_read_side_padding_scalar::<T>helper. The padding loop now takes the strings as an iterator, anasciiflag and a capacity floor. The array arms passstring_array.iter(), so their path is the same as onmain. The scalar arm passesstd::iter::repeat_n(value, n)with no capacity floor, so the scalar is never expanded into an array. Expanding it would needlen(s) * rowsbytes of input and reserve as much output, which overflows the 32-bit offsets of aUtf8array for a wide string (for example 1 MiB × 4096 rows). Truncation,len <= 0and empty pad behave as before.Unsupported argumentserror (Native lpad and rpad fail when the length argument is a NULL scalar, such as a scalar subquery that returns no rows #7452).StringLPad/StringRPad.The serde only changes the
strunsupported reason, from "Scalar values" to "Literal values", since it rejects only a literalstrand the kernel now handles astrthat is a scalar at runtime. Literal strings and non-literal padding still use codegen dispatch, as documented inexpressions.md. Lifting the literalstrrestriction is tracked in #7525.How are these changes tested?
New unit tests in
read_side_padding.rs:lpad,rpad,rpadtruncation, andread_side_paddingLargeUtf8scalar stringlpad,rpadandread_side_paddingoffset overflowbefore the second commit)These failed on
mainwith theUnsupported argumentserror above, and pass with this change.SQL file tests in
expressions/string/string_lpad.sqlandstring_rpad.sql, usingexpect_native(lpad|rpad)on the existingtest_lpad/test_rpadtables.IF(coalesce(len, 0) > -1000, (SELECT max(s) FROM t), s)takes the subquery branch for every row, so the string is a scalar for the whole batch.max(s)is'é', so the non-ASCII path is covered.IFandCASE WHEN, withlenas the length (scalar string, array length)Locally (macOS, Rust 1.99.1):
cargo test -p datafusion-comet-spark-expr --lib read_side_padding: 20 passedcargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings: cleancargo fmt --all -- --check: cleanI did not run
CometSqlFileTestSuitelocally. The fork CI on the default Spark 4.1 profile passed, includingstring_lpad.sqlandstring_rpad.sqlin the[expressions]job: https://github.com/tlfua/datafusion-comet/actions/runs/42112117022. I did not run the other Spark profiles.AI disclosure: the change, the tests and this description were written with an AI coding assistant, and reviewed by me.