Skip to content

perf: evaluate elementwise scalar functions once per distinct dictionary value - #24588

Open
radmirnovii wants to merge 7 commits into
apache:mainfrom
radmirnovii:dict-peel-elementwise
Open

perf: evaluate elementwise scalar functions once per distinct dictionary value#24588
radmirnovii wants to merge 7 commits into
apache:mainfrom
radmirnovii:dict-peel-elementwise

Conversation

@radmirnovii

Copy link
Copy Markdown

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:

// 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 FFI: FFI_ScalarUDF silently drops producer overrides of defaulted trait methods #22330.
  • Multiple dictionary arguments, run-end encoding, and extension-typed fields
    take the unpeeled path.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NysGeXTG5opiJKBApsAe5H

@radmirnovii

Copy link
Copy Markdown
Author

Per-function numbers behind "a flat ~0.4–0.8 µs per cold batch": the five
functions with hand-written arms, each by its own arm and by the mechanism,
512 distinct values of 29 characters over 8192 rows. All four columns come
from one run pinned to one core, with a counter in that run confirming the
kernel saw 512 values on every cold pass — and that the arm's cold and warm
columns agree, as they must for a path that cannot remember anything.

per batch own arm, cold mechanism, cold own arm, warm mechanism, warm
initcap 49.5 µs 50.3 µs 49.5 µs 0.72 µs
reverse 16.7 µs 17.6 µs 17.2 µs 0.73 µs
lower 2.26 µs 2.84 µs 2.04 µs 0.73 µs
ascii 1.68 µs 2.32 µs 1.56 µs 0.73 µs
character_length 1.71 µs 2.31 µs 1.75 µs 0.78 µs

An arm is nearly free where it exists — it already holds the dictionary, so
peeling it is a downcast and a rewrap, about 0.2–0.4 µs over the kernel alone.
The mechanism reaches the same values through a memo lookup and a tier
decision: a flat 0.6–0.9 µs per batch whatever the function, invisible for
initcap/reverse (2–5%), a quarter to a third of the call for the cheap
three — which is why those are not adopters here. What no arm can do is the
warm column: a remembered dictionary costs ~0.75 µs per batch (a read lock, a
pointer comparison and an Arc clone) because nothing is computed.

One measurement note, since it cost this table a rewrite: criterion calls the
benchmark routine afresh per sample, so a batch cursor declared inside it
restarts every sample and stops walking the batches in a clean cycle — the
memo then serves a third of the "cold" passes. The benchmark keeps its cursor
outside the routine and counts the values handed to the kernel in the same
run, so cold means cold.

@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels Aug 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

functions Changes to functions implementation logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant