compute: let hinted TopK stages emit the rows they keep

[frankmcsherry]

A basic TopK stage emits negations of the rows it drops, and the stage's result is its input plus those negations. That keeps the reduce's state proportional to the dropped rows, which is small when groups are barely larger than the limit and large otherwise. For DISTINCT ON over groups of eight rows with an EXPECTED GROUP SIZE hint, the final stage's arrangement held 8.75M records for a 1.25M-row result, and the consolidate that cancels the negations processed 18.75M records and was the largest operator in the dataflow. Every hierarchical stage after the first is in the same position: it sees at most sixteen times the limit per key and drops most of it.

This change lets a stage emit the rows it keeps instead, chosen per stage:

• The renderer computes the rows a stage keeps per key (literal limit plus offset) and reads the hint through the plan's bucket list, which is empty exactly when the hint was at most sixteen, the fan-in of one stage. In that case the lone final stage emits kept rows when twice its kept count fits in sixteen; otherwise, and whenever there are bucket stages, stages emit negated dropped rows as before. A kept-row stage whose keys hold one row each would copy its whole input into its output arrangement, so the choice needs the hint's evidence. BasicTopKPlan is part of the durably stored LIR schema and is unchanged.
• A kept-row final stage comes out of a reduce already consolidated, so the final consolidate is skipped.
• The monotonic TopK path keeps the dropped-rows form, which its retraction loop relies on.

The behaviour is behind enable_compute_topk_retained_stages, default off in production and on in the CI configuration, and registered with the parallel workload's flag flipper.

Hydration of SELECT DISTINCT ON (g) g, v FROM ... OPTIONS (DISTINCT ON INPUT GROUP SIZE = 8) ORDER BY g, v DESC indexed by g, 10M rows in 1.25M groups, one worker, median of three:

	before	after
hydration wall	8.10s	3.89s
TopK final consolidate	2109ms	absent
Reduced TopK input	1105ms	640ms
reduce output arrangement	8.75M records	1.25M records

Without a hint the plan is unchanged (25.0s before and after; that case is dominated by seven bucket stages that cannot shrink eight-row groups, a separate question about the default bucket schedule). Measured on a local build with the other hydration PRs in flight applied.

test/sqllogictest/topk.slt passes with the flag on; CI runs the whole suite with it on.

Keying the input for the first stage once

The initial map hashed each row and packed (hash, group), and the first stage's map immediately replaced the hash by hash % modulus and repacked. With a hint and no bucket stages the final stage's modulus is 1, so the hash was computed only to be replaced by zero. The initial map now keys rows for the first stage directly, and a stage rekeys only when its input is keyed for an earlier stage, as the min/max hierarchy already does. This is independent of the flag and removes one FlatMap operator (441ms on the measured shape) from every basic TopK dataflow; hinted topk goes from 4.73s to 4.21s.

Byte slicing in the initial map

When the group key is a column prefix, the initial map now forms (bucket, group) by pushing the bucket and appending the row's leading bytes, instead of unpacking the row and packing the group datums; other group keys keep the datum path. The map goes from 532ms to 240ms on the measured shape; hinted topk from 4.21s to 3.89s.

This branch carries RowRef::split_at_datum from #38665 as its own commit so it builds on main alone; whichever PR lands first brings the helper.

[def-]

QA LLM Review

1. HIGH -- Basic TopK with OFFSET and no LIMIT applies the offset per hash bucket, returning almost nothing

src/compute/src/render/top_k.rs:350

When limit is None, the bucket loop is skipped, so the final stage is also the first and is rendered with modulus = None ((!first).then_some(1) at line 472), leaving the input keyed as the initial map left it. But the initial map keyed rows by hash % buckets[0], and buckets is non-empty for every plan without a group-size hint. The final reduce therefore groups by (hash_bucket, group_key) and skips offset rows per hash bucket instead of per group, so an OFFSET without LIMIT silently drops nearly all rows.

Details

Repro, from test/sqllogictest/order_by.slt:721 (9 rows in fizz):

CREATE VIEW fizzoffsetview AS SELECT a, b FROM fizz OFFSET 6 ROWS;
SELECT count(b), count(a) FROM fizzoffsetview;  -- expected 3 3, now 0 0

Every row lands in its own bucket of buckets[0] = 268435456, so each key sees offset = 6 > 1 row and emits nothing. SELECT val1 FROM (SELECT val1, val2 FROM baz ORDER BY val2 DESC OFFSET 7 ROWS) goes from two rows to zero the same way; where a row has multiplicity 2 the group loses exactly one copy, so the corruption is partial rather than total. TopKPlan::create_from sends every offset > 0 query down the Basic path regardless of monotonicity, so this is any view or subquery with OFFSET and no LIMIT. Not gated by enable_compute_topk_retained_stages; this is the input-keying commit alone. buildkite/test/slt-1 on this branch fails with 18 output mismatches in order_by.slt, all of them OFFSET-without-LIMIT cases.

Fix: pick the initial key from the same condition that decides whether bucket stages are rendered.

-        let first_modulus = buckets.first().copied().unwrap_or(1);
+        // Bucket stages are only rendered when there is a limit; without one the final stage
+        // is the first, and its modulus is 1.
+        let first_modulus = match &limit {
+            Some(_) => buckets.first().copied().unwrap_or(1),
+            None => 1,
+        };

2. MEDIUM -- The kept-rows heuristic compares against 16 rather than the hint, so a small hint with a large limit copies the whole input into the output arrangement

src/compute/src/render/top_k.rs:403

kept.saturating_mul(2) <= 16 substitutes 16 for the group-size hint, but an empty bucket list only proves the hint is at most 16, never that it is that large. Whenever the user's hint is below 2 * (limit + offset), the stage picks Kept in exactly the situation the surrounding comment warns against, and the reduce's output arrangement grows to a copy of the input instead of staying near empty.

Details

The stage's own rationale is "kept rows cost less state when a key's group holds at least twice what the stage keeps", i.e. it wants 2 * kept <= group_size, and the renderer checks 2 * kept <= 16 instead. With OPTIONS (LIMIT INPUT GROUP SIZE = 8) and LIMIT 8, kept = 8 passes the guard, every key has must_shrink == false, and the !must_shrink arm at line 816 pushes every source row into target. State for that TopK goes from Arranged TopK input plus a near-empty reduce output to two full copies of the input, and the skipped TopK final consolidate does not offset that. Milder ratios have the same shape: hint 8 with LIMIT 5 retains 5/8 of the input where the negated form retained 3/8.

The renderer cannot make this call from the bucket list alone: an empty list is a ceiling on the hint, and the decision needs a floor on the group size. 2 * kept <= 16 is sound only for the single hint value 16. Either restrict Kept to cases where a copy of the input is acceptable regardless of the hint, or plumb the hint's value to the renderer, since no threshold on kept alone bounds the worst case (a hint of 1 makes every key's group unshrinkable for any limit).