viirya opened a new pull request, #6471:
URL: https://github.com/apache/datafusion-comet/pull/6471
## Which issue does this PR close?
Closes #6470.
## 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 `array_contains`,
`arrays_overlap`, `array_distinct` and `array_union` above it still run
natively. Their native kernels compare strings by raw bytes, so they ignore the
collation and return wrong results. For example, under `UTF8_LCASE`,
`array_contains(array('a'), 'A')` returns `false` in Comet and `true` in Spark.
Under `UTF8_BINARY_RTRIM`, `'x '` and `'x'` are treated as different.
## What changes are included in this PR?
- `CometArrayContains`, `CometArraysOverlap`, `CometArrayUnion` and a new
`CometArrayDistinct` serde report `Incompatible` when the array element type
contains a non-`UTF8_BINARY` collated string, at any nesting level. All four
mix in `CodegenDispatchFallback`, so by default these cases run through the JVM
codegen dispatcher (Spark's own comparison) and stay in the Comet pipeline.
With the dispatcher disabled they fall back to Spark. `allowIncompatible=true`
still opts into the native kernel. This is the same approach `array_intersect`,
`array_join` and `array_max`/`array_min` already use.
- `ArrayDistinct` was mapped to a plain
`CometScalarFunction("array_distinct")`. It now has its own serde. Non-collated
input is converted the same way as before.
- `map_contains_key` is rewritten to `array_contains(map_keys(...), key)`,
so a collated `map_contains_key` is now dispatched too. The existing test that
checks the nested-map literal guard now sets
`ArrayContains.allowIncompatible=true` to keep reaching that guard, like its
floating-point counterpart does, and also checks the default-config answer.
I checked the other expressions that compare, hash, dedupe or search
strings, using dispatched collated casts (`UTF8_LCASE`, `UTF8_BINARY_RTRIM`,
`UNICODE_CI`). These already match Spark: `=`, `IN`/`InSet`, `array_except`,
`array_intersect`, `sort_array`, `array_max`, `contains`, `startswith`,
`endswith`, `locate`, `find_in_set`, `replace`, `split`, `translate`, `like`,
`upper`, `hash`, `xxhash64` and map key lookup. `array_position` and
`array_remove` fall back to Spark and are correct.
The following string expressions still run natively on collated input and
return wrong results. They are not array functions, so they are left for
separate work:
- `instr` and `substring_index` (`UTF8_LCASE`, `UNICODE_CI`)
- `trim`/`ltrim`/`rtrim` with an explicit trim string (all three collations)
- `greatest` and `least` (`UTF8_LCASE`, `UNICODE_CI`; `least` also for
`UTF8_BINARY_RTRIM`)
## How are these changes tested?
New SQL file tests under
`spark/src/test/resources/sql-tests/expressions/array/`, all with
`MinSparkVersion: 4.0`:
- `array_element_equality_collation.sql`: `UTF8_LCASE` and
`UTF8_BINARY_RTRIM` input for each of the four functions, plus collated strings
inside struct elements. Each case asserts `expect_dispatch(...)`. A
non-collated control asserts `expect_native(...)`.
- `array_element_equality_collation_disabled.sql`: with the dispatcher
disabled, each function falls back with the dispatcher-disabled reason.
- `array_element_equality_collation_allow_incompatible.sql`:
`allowIncompatible=true` keeps the native kernels.
Without the fix, `array_element_equality_collation.sql` fails with a result
mismatch and `array_element_equality_collation_disabled.sql` fails because only
the cast is reported as the fallback. With the fix, both pass.
`CometSqlFileTestSuite` (array files), `CometArrayExpressionSuite`,
`CometMapExpressionSuite` and `GenerateDocsSuite` pass on Spark 4.1. The new
SQL files and `CometMapExpressionSuite` also pass on Spark 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]
