mzabaluev opened a new issue, #26581:
URL: https://github.com/apache/datafusion/issues/26581
### Describe the bug
`CaseExpr` (the `NoExpression`, `WithExpression` and
`ExpressionOrExpression` eval methods)
evaluates each `WHEN` branch on a filtered copy of the batch. Per branch, it
calls
`filter_record_batch` once to get the rows the branch takes, and once more
to get the rows left
for the next branch. `CaseBody::project()` (#18329) limits this to the
columns the `CASE`
references. But it collects whole top-level `Column`s, so a `THEN s['a']`
keeps the entire struct
column `s` in the projected batch.
When `s` is wide or nested (lists of structs, maps), every branch copies all
of `s`, twice, to
read one scalar field. The cost becomes proportional to branches × struct
size rather than
branches × field size.
SQL planned through the logical optimizer usually avoids this: leaf
expression pushdown
(`datafusion.optimizer.enable_leaf_expression_pushdown`) moves the
`get_field` calls into a
projection below the `CASE`. Physical plans that are not built from an
optimized DataFusion
logical plan hit it. One example is engines like Comet, that translate
another system's query plans into
DataFusion physical expressions. So does SQL with that option turned off.
### To Reproduce
With the `datafusion` Python package (55.0.0) and `pyarrow`: table `t(id, k,
s)`, where
`k = id % 9` picks the branch and `s` is a struct of eight `bigint` fields
`a`..`h` plus
`items: list<struct<x: bigint, y: string>>` with N items per row. Each
`WHEN` reads a different
scalar field of `s`; none reads `items`.
```python
import time, datafusion, pyarrow as pa, pyarrow.compute as pc
ROWS, ITEMS, BATCH = 4_000_000, 32, 8192
item_t = pa.struct([("x", pa.int64()), ("y", pa.string())])
def make_batch(start, n):
ids = pa.array(range(start, start + n), pa.int64())
fields = [pc.multiply(ids, i + 1) for i in range(8)]
xs = pa.array(range(n * ITEMS), pa.int64())
items = pa.ListArray.from_arrays(
pa.array([i * ITEMS for i in range(n + 1)], pa.int32()),
pa.StructArray.from_arrays([xs, pc.cast(xs, pa.string())],
fields=list(item_t)))
s = pa.StructArray.from_arrays(fields + [items], names=list("abcdefgh")
+ ["items"])
return pa.record_batch([ids, pc.remainder(ids, 9), s], names=["id", "k",
"s"])
ctx = datafusion.SessionContext()
ctx.sql("SET datafusion.optimizer.enable_leaf_expression_pushdown =
false").collect()
ctx.register_record_batches("t", [[make_batch(i, min(BATCH, ROWS - i)) for i
in range(0, ROWS, BATCH)]])
case = "CASE " + " ".join(f"WHEN k = {i} THEN {{{c}}}" for i, c in
enumerate("abcdefgh")) + " ELSE -1 END"
on_struct = "SELECT sum(" + case.format(**{c: f"s['{c}']" for c in
"abcdefgh"}) + ") FROM t"
on_fields = ("SELECT sum(" + case.format(**{c: c for c in "abcdefgh"}) + ")
FROM (SELECT k, "
+ ", ".join(f"s['{c}'] AS {c}" for c in "abcdefgh") + " FROM
t)")
for name, sql in [("CASE over s['a'] .. s['h']", on_struct), ("CASE over
extracted fields", on_fields)]:
times = []
for _ in range(3):
t0 = time.perf_counter(); ctx.sql(sql).collect();
times.append(time.perf_counter() - t0)
print(f"{name}: {min(times):.3f} s")
```
The two queries return the same result. `EXPLAIN` shows `CASE WHEN t.k =
Int64(0) THEN t.s[a] …`
evaluated straight over the scan in the first, and over a `ProjectionExec` of
`get_field(s@1, a) as a, …` in the second.
Best of 3, 4M rows, 14-core Apple M4 Pro:
| `items` per row | `CASE` over `s['a']`..`s['h']` | `CASE` over extracted
fields | Ratio |
| --- | --- | --- | --- |
| 0 | 0.031 s | 0.017 s | 1.8× |
| 8 | 0.069 s | 0.020 s | 3.5× |
| 32 | 0.260 s | 0.022 s | 12× |
| 128 | 1.073 s | 0.024 s | 46× |
With `enable_leaf_expression_pushdown` left at its default (`true`), both
take about 0.02 s.
### Expected behavior
The cost of evaluating `CASE` should depend on the fields the branches read,
not on the size of
the struct those fields live in.
In a real workload this matters a lot. We hit it in a Spark-compatible
engine that builds
DataFusion physical plans directly: a production-derived query projects
about 90 `CASE`
expressions after an `explode`, and about 30 of them read fields of large
nested structs (lists
of structs, maps). A `perf` profile of that stage shows `CASE` evaluation
dominated by
`filter_struct` on those structs. That build also turned off
`CaseBody::project()` because of
#21231, which widens the copy to every column. But even with the projection,
each such `CASE`
still copies its whole struct twice per branch.
### Possible solutions
1. **Project field accesses, not just columns.** In `CaseBody::project()`,
treat a struct field
access chain on a column (`get_field(col, 'a')`, possibly nested) as a
projected leaf: compute
it once on the input batch, put the result in the projected batch, and
rewrite the body to
reference it. The per-branch filters then copy only the fields that are
read. This is the same
rewrite leaf expression pushdown does in the logical plan, applied inside
`CaseExpr` for plans
that don't go through it.
2. **Skip input filtering for cheap branches.** When every `THEN`/`ELSE` is
cheap and infallible
(column references, field accesses, literals), evaluate the branch values
on the full batch
and merge the results by the `WHEN` masks, instead of filtering the
inputs.
3. Separately, filtering nested list/struct arrays through
`MutableArrayData` row ranges is
expensive in itself (`set_bits` and the per-child `build_extend`
closures), which is arrow-rs
territory.
### Additional context
- #18329 added the `CASE` input projection.
- #21231 tracks the projection's handling of custom column expressions.
- #22487 is a similar pattern: a sort key over `get_field` is evaluated
against the materialized
struct rather than the base column.
--
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]
