viirya opened a new pull request, #6480:
URL: https://github.com/apache/datafusion-comet/pull/6480
## Which issue does this PR close?
Closes #6472.
## Rationale for this change
On Spark 4.x with the codegen dispatcher enabled (the default), a cast to a
collated string type runs through the dispatcher, but `instr`,
`substring_index`, `trim`/`ltrim`/`rtrim` with a trim string, `greatest` and
`least` above it still run natively. Their native kernels search, trim or order
strings by raw bytes, so they ignore the collation and return wrong results.
For example, under `UTF8_LCASE`, `instr('a', 'A')` returns 0 in Comet and 1 in
Spark, and `least('a', 'A')` returns `'A'` in Comet and `'a'` in Spark.
## What changes are included in this PR?
- `instr`, `substring_index`, `trim`, `ltrim`, `rtrim`, `greatest` and
`least` get their own serdes. They report `Incompatible` when an input has a
non-`UTF8_BINARY` collation (checked at any nesting level, so struct inputs of
`greatest`/`least` are covered too). All of them mix in
`CodegenDispatchFallback`, so by default these cases run through the JVM
codegen dispatcher (Spark's own implementation) and stay in the Comet pipeline.
With the dispatcher disabled they fall back to Spark. `allowIncompatible=true`
still opts into the native kernel. This follows the approach used for the array
functions in #6471.
- For `trim`/`ltrim`/`rtrim`, only the form with an explicit trim string is
gated. Without a trim string they remove only spaces, which does not depend on
the collation, so `trim(x)` on collated input keeps running natively. I checked
this against Spark for `UTF8_LCASE`, `UNICODE_CI`, `UTF8_BINARY_RTRIM` and
`UNICODE`.
- Non-collated input is converted the same way as before.
The gate covers every non-`UTF8_BINARY` collation, not only the ones listed
in the issue. Running the native kernels on collated input also gives wrong
answers in these cases: `trim(BOTH ...)` and `btrim` under `UTF8_BINARY_RTRIM`,
and `greatest`/`least` under `UNICODE`.
## How are these changes tested?
New SQL file tests under
`spark/src/test/resources/sql-tests/expressions/string/`, all with
`MinSparkVersion: 4.0`:
- `string_search_trim_collation.sql`: `UTF8_LCASE`, `UTF8_BINARY_RTRIM` and
`UNICODE_CI` input for each function, asserted with `expect_dispatch`, plus
struct inputs for `greatest`/`least`. It also checks that `UTF8_BINARY` input
and `trim` without a trim string stay native (`expect_native`).
- `string_search_trim_collation_disabled.sql`: with the dispatcher disabled,
collated input falls back to Spark (`expect_fallback`).
- `string_search_trim_collation_allow_incompatible.sql`:
`allowIncompatible=true` runs the native kernels (`expect_native`), on data
where the bytewise and collation-aware answers agree.
Without the source change, `string_search_trim_collation.sql` fails on
result mismatches and `string_search_trim_collation_disabled.sql` fails because
the expressions run natively instead of falling back. The new tests and the
existing `CometSqlFileTestSuite`, `CometStringExpressionSuite`,
`CometExpressionSuite`, `GenerateDocsSuite`, `QueryPlanSerdeSuite`,
`CometCodegenSuite` and `CometCodegenSourceSuite` pass on Spark 4.1 and 4.0.
This pull request and its description were written by Isaac.
--
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]
