david-mollitor-db opened a new pull request, #58841: URL: https://github.com/apache/spark/pull/58841
### What changes were proposed in this pull request? A `CASE WHEN` whose branches are all `key = literal THEN constant` on a single key is really a lookup table — e.g. mapping status/country/category codes to labels (often generated by BI tools or dbt, with tens to hundreds of arms). Today `CaseWhen` whole-stage codegen emits an if/else-if chain (`CaseWhen.multiBranchesCodegen`); for a string key that is up to N `UTF8String` comparisons per row, growing linearly with the branch count. This adds a **codegen-only** recognizer in `CaseWhen.doGenCode`: when the CASE is a lookup on a binary-collation string key, it builds a constant key→value table once on the driver and emits a boxing-free, constant-size inline hash probe (a power-of-two open-addressed `int[]` bucket table plus parallel key/value `ArrayData`) instead of the O(N) chain. The generated code is O(1) in size and O(1) per row. Non-lookup CASEs fall back to the existing chain unchanged. The interpreted `eval` path and the `CaseWhen` node itself are unchanged, so there is **no plan-shape / golden-file impact** — this is purely a code-generation substitution. **Scope.** The probe is used only when every branch condition is `EqualTo(key, literal)` (either argument order) on one deterministic key, every branch value is foldable, the key is a binary-collation `StringType`, and there are at least 10 distinct non-null keys (first-match-wins on duplicates; null literal keys dropped). All other cases keep the chain. Float/double (NaN and `-0.0`/`+0.0` equality) and non-binary-collation strings (collation-aware equality) are excluded. **String keys only — by measurement.** A benchmark (10M rows) shows the probe beats the chain for string keys — about **1.1× / 1.6× / 2.6×** at 10 / 50 / 100 branches on matching data (the probe stays flat while the chain grows linearly), with no regression on all-miss data. Integral/temporal keys were measured to *not* benefit — their comparisons are cheap and well branch-predicted, so the chain ties or wins and even regresses on the common all-miss case — so those keys keep the chain. **Reuse.** `GetMapValue` already has this machinery for foldable map-column lookups (`PrebuiltHashExecutor`, SPARK-55959). This change extracts its core (bucket build, the mirrored driver/codegen hash, and the probe loop) into a shared `PrebuiltHashProbe` used by both `GetMapValue` and `CaseWhen`, so there is a single boxing-free probe implementation rather than two copies. **Correctness.** Null key → ELSE (matching `EqualTo` semantics); a constant `NULL` branch value → `NULL` on a hit; duplicate literal keys keep the first branch's value; a missing key with no ELSE → `NULL`; `EqualNullSafe` and any non-equality / mixed branch bail to the chain. **Config.** A new internal kill-switch `spark.sql.optimizer.caseWhenLookup.enabled` (default true) forces the chain when set to false. ### Why are the changes needed? Lookup-shaped `CASE WHEN` expressions (code→label maps) are common in ETL/BI SQL and can have hundreds of arms. For string keys the current O(N) chain does up to N string comparisons per row, which dominates the per-row cost. The hash probe makes membership O(1) per row and keeps the generated method O(1) in size. ### Does this PR introduce _any_ user-facing change? No. Query results are unchanged; this is an internal whole-stage code-generation optimization, gated by an internal config. ### How was this patch tested? - New unit tests in `ConditionalExpressionSuite`: a probe-vs-chain correctness matrix over string keys (hits at different bucket positions, miss→ELSE, miss with no ELSE→NULL, null key→ELSE/NULL, constant-`NULL` branch value, duplicate keys first-wins, swapped `EqualTo` argument order); a test asserting the hash probe is generated under codegen (and the chain when the config is disabled); and bail tests for ineligible key types (integer/long/float/non-binary-collation string) and shapes (non-foldable value, non-equality branch, non-deterministic key, below threshold). `checkEvaluation` exercises interpreted, `CODEGEN_ONLY`, no-codegen, unsafe-projection, and optimizer paths. - Existing `ComplexTypeSuite` and `CollectionExpressionsSuite` (covering `GetMapValue` and `ElementAt`) validate that the `PrebuiltHashProbe` extraction is behavior-preserving. - Added `CaseWhenBenchmark`; committable numbers to be regenerated on GitHub Actions. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac 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]
