radmirnovii opened a new pull request, #24588:
URL: https://github.com/apache/datafusion/pull/24588
## Which issue does this PR close?
Related to #20935, which asks for exactly this: *"Ideally this would be a
reusable pattern (perhaps a helper or wrapper) that individual string
functions
can opt into, rather than duplicating the logic in every function."*
Related to #19458; its first half — delivering the dictionary to the
function —
merged as #22905. This is the other half: using what is delivered, once for
every function.
Note: this branch includes the benchmark commit from #24586 until it merges.
## Rationale for this change
A scalar function over a dictionary-encoded column runs once per row,
although
the batch holds no more distinct values than the dictionary does. Eleven
functions already work around this with hand-written handling in their
kernels:
```rust
// datafusion/functions/src/unicode/reverse.rs
DataType::Dictionary(_, _) => {
let dictionary = args[0].as_any_dictionary();
let converted = reverse(&[Arc::clone(dictionary.values())])?;
Ok(dictionary.with_values(converted))
}
```
That arm sits in nine functions across eight files, with a
`ScalarValue::Dictionary` twin in `bit_length` and `octet_length`. Every new
function writes it again, and a function without one — `encode`, the hashes,
the regex functions — cannot have the optimization at all. This does it once,
in `ScalarFunctionExpr`, for any function that declares the property.
## What changes are included in this PR?
`ScalarUDFImpl::evaluates_elementwise` (default `false`) declares that each
output row depends only on the corresponding input row. Over a dictionary
argument the physical layer then unwraps the call and re-maps the result
through the keys, choosing the cheapest sound tier per batch:
- **remembered** — these values were evaluated before; nothing is computed;
- **as they are** — values passed unchanged when no key is null or `f` is
strict;
- **compacted** — only the values the batch references, null keys redirected
to
one appended NULL slot (correct even where `f(NULL)` is not NULL);
- **expanded** — flat returns only: the materialized column, exactly what
coercion produces today. A dictionary return declines instead and the
function sees the column as it arrived.
Results are reused across batches because a batch carries its own keys but
the
dictionary of its whole column chunk — the Parquet reader hands every batch
of
a chunk the same value buffers. Results are keyed on the memory the values
occupy and hold it alive, bounded to eight dictionaries and 4 MiB per
expression (a Parquet dictionary page is at most 1 MiB by default). Hits take
only a read lock, so the partitions sharing an expression do not serialize on
each other; a first sighting records a hash, and the second, which proves the
dictionary repeats, buys evaluating all of it.
Three functions opt in: **`encode`** (had no dictionary handling and could
not
have; gains encoding preservation plus the declaration, kernel untouched),
**`reverse`** and **`initcap`** (arms since #23930; gain only the
declaration).
The arms stay: the physical layer still hands the dictionary over wherever it
declines — extension metadata, a key type too narrow for the NULL slot.
## Are these changes tested?
26 unit tests in `scalar_function.rs` cover each tier and its boundaries —
null keys under strict and non-strict functions, garbage under null keys, the
NULL slot overflowing a narrow key type, the profitability bound, two
dictionary arguments, dictionary scalars, extension metadata, errors from
referenced vs unreferenced values, memoization across batches but not across
dictionaries or past its byte budget, and concurrent hits over one shared
expression. `functions.slt` adds the narrow-key decline for `reverse` end to
end, and existing dictionary coverage for `reverse`/`initcap` now runs
through
the generic path unchanged. For `encode`, `expr.slt` pins results and the
plan — the cast is to `Dictionary(Int32, BinaryView)`, not away from the
encoding.
## Are there any user-facing changes?
A new trait method with a default; no existing implementation changes. All
three functions return exactly what they did — only how often they compute it
changes.
## Benchmarks
The benchmark lands separately in #24586, measuring today's paths — the
hand-written arm (whose cold and warm batches cost the same, an arm cannot
reuse anything) and the cast every unpreserved function pays. This PR extends
it with `encode`'s dictionary-typed groups, which only become expressible
here. 8192-row batches, medians, pinned to one core. `cold`: a dictionary per
batch; `warm`: one shared across batches, as a Parquet column chunk delivers
them. *no preservation*: the dictionary cast away, one call per row — what
`encode` did before this change (`reverse`'s arm already costs about the cold
column).
| encode, per batch | cold | warm | no preservation |
|---|---:|---:|---:|
| 8 distinct values | 68 µs | 44 µs | 585 µs |
| 256 distinct values | 77 µs | 44 µs | 546 µs – 1.10 ms |
| 512 distinct values | 108 µs | 43 µs | 589 µs |
| 8192 distinct, all different | 0.87–1.09 ms | 44 µs | 584 µs |
| reverse, per batch | cold | warm | no preservation |
|---|---:|---:|---:|
| 8 distinct values | 1.75 µs | 0.69 µs | 288 µs |
| 256 distinct values | 10.0 µs | 0.70 µs | 291 µs |
| 512 distinct values | 17.1 µs | 0.68 µs | 285 µs |
| 8192 distinct, all different | 242 µs | 0.68 µs | 287 µs |
Two cells above are ranges because they are bimodal across repeated runs of
the same binary, and the bimodality is a property of the heap, not of either
path: `encode`'s expand tier at full cardinality measured 0.87–1.09 ms across
run contexts, and the *cast path itself* flipped the same way at 256 distinct
(546 µs alone, 1.10 ms inside the full suite) while doing byte-identical work
to its stable neighbours. Both paths allocate ~450 KB per batch there; the
mechanism's own work in that cell is microseconds (an aborted compaction scan
and a hash), and a column with no repeated values should not be
dictionary-encoded in the first place. Against the hand-written arms the
mechanism costs a flat ~0.6–0.9 µs per cold batch and repays it on the first
repeated dictionary; per-function numbers in the first comment.
## What this does not do
- The pass/compact threshold assumes a value costs ~7 ns; cheap functions can
lose above the batch size (`ascii` up to 4.6x in that band). A per-value
cost
model does not exist yet, so the remaining nine arms should be measured
before opting in, not converted in bulk.
- The eleven functions keep their hand-written handling; removing it is a
separate change.
- FFI is untouched: foreign UDFs inherit the default and lose only the
optimization. Tracked by #22330.
- Multiple dictionary arguments, run-end encoding, and extension-typed fields
take the unpeeled path.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_01NysGeXTG5opiJKBApsAe5H
--
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]
