repr: decode rows through a predicted per-column class

[antiguru]

Decoding a Row into [Datum] costs about twice what it needs to, and the cause is the shape of the dispatch. read_datum matches over roughly thirty tag classes per datum, which makes its body too large for a caller to inline, so every datum pays a function call returning a 48-byte Datum through a hidden pointer. Selecting the arm from the column's class rather than from its tag shrinks the match to nine arms, which does inline.

DatumVec learns the class of each column from the rows it decodes and keeps it as a field, so borrow_with improves without any caller changing.

Why a prediction rather than the column types

The class is never trusted. Every arm checks that the tag really belongs to the predicted class and returns None otherwise, so Prediction::decode falls back to read_datum for that datum and learns the column from its tag. A wrong prediction costs throughput and never correctness.

That is what makes learning viable, and learning is what makes the change small. Passing types down instead would need them to exist at the call sites, and they do not: BuildDesc is an id and a plan, and src/compute-types/src/plan.rs carries no column types at all. ReprRelationType is available at the dataflow boundary, on index_exports and index_imports.on_type, but not per plan node, so threading types through the physical plan would be a protocol change. Measuring the checked form against a form that is simply told the types put the check at 0 to 1 percent on integer columns and 8 to 10 percent on mixed ones, so the protocol change would buy very little.

Measurements

cargo bench -p mz-repr --bench predict, throughput in Melem/s, batches of 1024 rows pinned to one core. general is borrow_with_general, which decodes every datum through read_datum; predicted is borrow_with.

shape	general	predicted	ratio
4 x Int64	61.7	100.0	1.62x
8 x Int64	60.2	109.6	1.82x
32 x Int64	46.2	60.2	1.30x
Int64/String/Float64	89.3	134.9	1.51x
the same, 12 columns	92.0	164.9	1.79x
8 x Numeric, no fast arm	63.0	62.2	0.99x

The last row is the shape with nothing to gain, and it is there to show the fallback does not cost anything; it sits at parity. Run to run variance on this machine is about three percent, so these come from a longer measurement window after a first run showed a twenty percent swing on the integer shapes. The table predates the two tag-invariant commits below. Re-running the benchmark after them held every ratio inside that same three percent, so the const generic they introduce is neutral.

The decoder in isolation reaches 2.0x to 2.5x. The rest is borrow_with itself: the per-row mem::take and the repurpose_allocation in DatumVecBorrow::drop cost about 23 percent at arity 8, independently measured, and that share grows as the decode gets cheaper. Worth a separate change.

Three results that shaped the design and are worth recording, since each looked promising and was not:

• Payload width is irrelevant. A probe enum padded to 24, 32 and 40 bytes measures the same at every width, and Datum is 48 bytes because of Range, not Numeric. Shrinking Datum was the original hypothesis and it buys nothing.
• Batching buys nothing. Decoding a batch of rows in one call measures the same as decoding them one at a time.
• Column-major traversal buys nothing. Hoisting the dispatch entirely out of the inner loop is a wash on mixed columns and worse on integers, because it costs a cursor per row and re-reads the batch once per column. The dispatch does not need hoisting, it needs to be small.

Details worth a reviewer's attention

Five things are load-bearing, and each was measured rather than assumed.

Unknown and Other are separate classes and both are needed. Unknown always misses, so a column learns on its first datum; with only one variant the first row would learn nothing. Other never misses, so a column of a type no fast arm covers settles instead of paying a classification on every row.

A Null tag is accepted by every class, since any column can be nullable, but classify returns None for it. A column whose first value is null therefore stays unclassified and tries again on the next row, rather than being mispredicted by the null.

A relation every column of which has settled on Other uses the general decoder directly, guarded by a check that the row's arity has not changed. Predicting such a relation cannot pay for itself, and the per-datum bookkeeping measured 9 percent slower than not predicting at all. The condition is "every column has settled on Other" and deliberately not "no column has a fast arm": the latter also holds while columns are still unclassified, which would stop the learning loop from running again and leave a relation whose fast columns happen to be null in its first row decoding generally for the life of the decoder.

Whether every column has settled is rescanned only when a class changed or truncation dropped one, which keeps the scan off the settled path. Rescanning on every row measured 2 to 3 percent slower on the integer shapes.

There is no arm for Numeric, even though decimals are common. Such an arm can only defer to the general coefficient decode, so it puts a dispatch and a tag compare in front of identical work; it measured 24 percent slower than not covering the type. Removing it also made every other shape faster, since the match got smaller.

decode also truncates the class vector to the row's arity, so a change in arity self-corrects.

borrow_with_general is retained only so the benchmark can measure against the path this replaces. It should go once the change is settled.

Tag invariants

The tag families are addressed by arithmetic rather than by name, and this change adds a third
consumer of that. push_datum writes first + n for a body of n bytes, Tag::actual_int_length
inverts it by spelling the length out per variant, and the predicted decoder derives it from the
tag's distance from its family's first tag. Nothing tied the three together except the declaration
order of the enum and a comment asking the reader to preserve it. Discriminants are not durable, so
this is not a compatibility question. It is only about the encoder and the decoder agreeing on how
many bytes a tag claims.

Two commits make that checkable rather than documented.

varint takes the family width and the width of the array it decodes into as one const generic, so
the caller's from_le_bytes fixes both. It replaces a hand-passed width: u8, which is the actual
footgun: a value one too large claims the first tag of the next family and then reads a body longer
than the integer it is decoding into. That can no longer be spelled.

assert_consecutive! pins all thirteen families at the enum definition, ten variable-length integer
families and three length-prefixed ones, and reports the family it broke by name. These guard
push_datum's arithmetic as well, which predates this change.

Each guard was verified by mutation rather than assumed to work:

mutation	caught by
widen the Int32 family to five bytes	test: tag NonNegativeInt64_0 claimed by more than one family
classify Numeric as Bool in expected_class	test: left: Some(Other) right: Some(Bool)
drop Numeric from expected_class, standing in for a new variant	compile error, non-exhaustive match
swap NonNegativeInt64_16 and NonNegativeInt64_24	compile error, tags are not consecutive: NonNegativeInt64_0, ...
insert a variant into the middle of the String family	compile error, both guards fire

The third row is the one that matters most. classify ends in a catch-all, so before this a Tag
variant added to a family it does not know would silently become Other and decode generally
forever.

Generating the tag table from a declaration was considered and rejected for now. Roughly thirty-five
of the hundred tags have bespoke decode bodies, so a generator would cover the mechanical families
and leave the rest hand-written, and macro-built variant names would stop being greppable while
read_datum's arms still name them. The narrower version worth doing later is to generate the
variable-length integer block alone, which would delete actual_int_length and make contiguity hold
by construction instead of by assertion. That earns its keep once the function has a fourth consumer.

Tests

Every datum decodes correctly whether or not the prediction is right, so tests that assert only on decoded datums cannot tell a working fast path from an inert one. Prediction therefore carries #[cfg(test)] counters of hits, misses and rows that skipped prediction, and the tests assert on them.

src/repr/src/row/predict.rs adds:

• matches_general_decoder, a proptest that decodes generated rows through the learned sequence and through every fixed class including deliberately wrong ones, asserting equality with read_datum throughout.
• prediction_settles, which asserts that after the first row every datum is a hit and no row falls back wholesale.
• first_row_teaching_nothing_still_learns and uncovered_column_does_not_mask_a_learnable_one, covering a first row that classifies nothing, either because the only column is null or because the one column with a fast arm is.
• null_teaches_nothing and arity_change_settles.
• classify_matches_every_tag, which states the whole tag-to-class mapping in a second, exhaustive match and asserts classify agrees for every tag. Tag derives TryFromPrimitive, so scanning the byte range enumerates the enum without a generator having to produce a datum of each type. The exhaustiveness is the point: adding a Tag variant fails to compile here.
• varint_len_agrees_with_tag, which asserts for every tag that the length the decoder computes by arithmetic equals the one Tag::actual_int_length gives by name, and that no tag is claimed by two families.

src/repr/benches/predict.rs is new and covers integer, mixed and numeric-only shapes, the last being the no-gain case.

Not in this change

No feature flag. mz_repr has no dyncfg access, so gating needs either an atomic static set at startup or gating at the call sites. Correctness does not depend on a flag, because the per-datum check makes a wrong prediction a throughput question only, but this should probably not go to production ungated. Happy to add whichever form reviewers prefer.

The share of clusterd on-CPU time spent decoding rows has not been measured, so what this is worth end to end is unknown.

Part of: CPU-223

🤖 Generated with Claude Code

[def-]

QA LLM Review

1. MEDIUM -- A first row that classifies no column disables the prediction for the life of the DatumVec

src/repr/src/row/predict.rs:276

If the first row a DatumVec decodes classifies no column into a fast arm, because every fast-armed column is null in it, the prediction never learns and that DatumVec decodes generally forever, so the change buys nothing for it. any_fast is only ever set inside the learning loop, and the loop is skipped whenever any_fast is false and classes is non-empty, so the general path it falls into can never turn the flag on. This contradicts the field comment at datum_vec.rs:28 that the prediction "settles after the first row".

Details

The trigger is one null in the first row of a narrow relation. Two cases I ran against this branch:

// arity 1, first row NULL, then ten Int64 rows
// -> classes [Unknown], any_fast=false, general path forever
// arity 2 (numeric, int8), first row (3, NULL), then ten (3, 5) rows
// -> classes [Other, Unknown], any_fast=false, general path forever

Recovery only happens if the row arity changes, which for a fixed schema it never does. Output stays correct, and the general path is not slower than the code it replaced, so this is lost speedup rather than a regression: at int8 the benchmark measures 84.6 vs 172.6 Melem/s, so an affected operator keeps the 1x number.

The condition to gate on is "every column has settled on Other", not "no column has a fast arm". Replacing the flag makes both cases above learn on the second row, and re-running the benchmark shows no throughput change on any shape (all deltas within noise, including numeric8 and the 1-fast-of-8 shape):

-    any_fast: bool,
+    /// Whether every column has settled on `Other`.
+    all_other: bool,
@@
-        if !self.any_fast && !self.classes.is_empty() {
+        if self.all_other {
@@
             if out.len() - before != self.classes.len() {
                 self.classes.clear();
+                self.all_other = false;
             }
             return;
@@
                     if let Some(learned) = classify(data[at]) {
                         classes[idx] = learned;
-                        self.any_fast |= learned != DatumClass::Other;
                     }
@@
         classes.truncate(idx);
+        self.all_other = !classes.is_empty() && classes.iter().all(|c| *c == DatumClass::Other);

2. LOW -- Every test still passes if the fast path never fires

src/repr/src/row/predict.rs:351

None of the four tests can tell a prediction hit from a miss, so the whole suite is green even when the optimization is completely inert. matches_general_decoder asserts only on decoded datums, which are identical either way by design. prediction_settles asserts classes is unchanged across rows, but a miss re-runs classify and writes back the same class, so the assertion holds under a total fast-path failure too. null_teaches_nothing and arity_change_settles likewise only inspect classes.

Details

Concretely: finding 1 leaves the decoder in exactly that inert state and no test notices. The same would be true if a tag constant or a varint_len bound were mistyped so an arm always returned None. A #[cfg(test)] hit counter on Prediction, asserted non-zero in prediction_settles, closes this.

3. LOW -- varint_len's doc gives a value that would over-match the next tag group

src/repr/src/row/predict.rs:93

"width is the group's widest encoding, three for the 16-bit groups and so on" mixes up the byte width with the tag count. The call sites pass 8 for the 64-bit groups and 4 for the 32-bit groups, i.e. the widest encoding in bytes, so the 16-bit groups would need 2, not 3. NonNegativeInt16_16 is immediately followed by NonNegativeInt32_0, so an Int16 arm written from this recipe would accept an Int32 tag, return len == 3, and panic in extend::<2> on raw[..3].

[antiguru]

All three findings are correct. Fixed in 2839358.

1. First row that classifies no column

Reproduced before fixing. The mechanism is that decode_one checks the Null tag before matching on the class, so a null datum returns Some and never reaches the branch that learns. With arity one and a null first row, classes stays [Unknown], so the guard !any_fast && !classes.is_empty() sends every later row to the general path, which cannot learn.

Both of your cases now exist as tests, and both failed on the previous commit with exactly the states you predicted:

first_row_teaching_nothing_still_learns          left: [Unknown]        right: [Int64]
uncovered_column_does_not_mask_a_learnable_one   left: [Other, Unknown] right: [Other, Int64]

Took the all_other predicate as suggested. The reason the old condition was wrong is worth stating in the code, so the field comment now says the condition is "every column has settled on Other" and not "no column has a fast arm", since the latter also holds while columns are unclassified.

One correction to the claim that it costs nothing. Rescanning on every row measured 2 to 3 percent slower on the integer shapes, largest where throughput is highest, which is consistent with one extra compare per datum:

shape	before	unconditional rescan	rescan only on change
4 x Int64	101.0	97.7	100.0
8 x Int64	110.0	107.5	109.6
32 x Int64	59.5	59.1	60.2

all_other can only turn on when a class changed or truncation dropped the last column with a fast arm, and it cannot turn off in the loop, since a decoder for which it is already on never reaches the loop. So the scan is guarded on those two conditions, which keeps it off the settled path.

Worth flagging on measurement: my first run of this comparison showed a twenty percent swing on the integer shapes, including on general, which is untouched code. That was machine noise. Run to run variance here is about three percent, so the numbers above come from a longer measurement window, and I would not read anything into a delta smaller than that.

2. Every test passes if the fast path never fires

Correct, and finding 1 is the proof: the decoder was inert for those two shapes and nothing went red. prediction_settles in particular holds under total failure because a miss re-learns the same class.

Added #[cfg(test)] counters on Prediction for hits, misses and rows that skipped prediction. prediction_settles now asserts that after the first row the miss count is zero, the general-row count is zero, and the hit count equals arity times rows, so an inert fast path fails it. The two new tests from finding 1 assert on the same counters rather than only on the classes.

3. varint_len doc

Correct. The call sites pass the widest encoding in bytes, and there is no Int16 arm today, so this is purely a trap for whoever adds one. The doc now names the unit and the consequence:

width is the widest encoding the group has in BYTES, so two for a 16-bit group, four for a 32-bit one and eight for a 64-bit one. Passing the number of tags in the group instead would accept the first tag of the next group, whose length then overruns the array extend fills.

🤖 Generated with Claude Code