feat(cubestore): top-k merge strategies via CUBESTORE_TOPK_STRATEGY

[waralexrom]

Summary

The distributed value-ordered top-k path (GROUP BY x ORDER BY agg(...) LIMIT k)
ran through a per-row streaming merge on the router that was ~25× slower than
ClickHouse on the same data. This adds two opt-in router-merge strategies behind
a single config flag CUBESTORE_TOPK_STRATEGY, leaving the current behavior as
the default.

Changes

• CUBESTORE_TOPK_STRATEGY config flag (ConfigObj, lenient parsing, plumbed
  to both router and worker planners) with three values:
  • streaming (default) — the original streaming NRA merge, unchanged.
  • vectorized — same streaming NRA merge (keeps early termination, so router
    memory stays bounded), but the per-row machinery (Mutex<Vec<Group>>,
    per-row 1-element accumulator updates, ScalarValue keys) is rewritten over
    an arrow RowConverter + typed ScalarValue arithmetic.
  • full_merge — ClickHouse-style plan: workers send their groups unsorted, the
    router re-aggregates with a single vectorized hash aggregate and takes the
    top-k with a fetch-limited sort. Drops early termination (materializes every
    distinct group on the router).
• HLL merge and unsupported key/agg types fall back to the default strategy.
• vectorized keeps MIN/MAX exact (skips NaN like arrow); float SUM is reordered
  (documented), so it is not bit-identical to streaming.

Testing

• Unit tests: V1↔V2 equivalence across the existing top-k cases + direct typed
  arithmetic tests (NaN min/max, decimal scale, sum/null).
• New SQL test sql::tests::topk_full_merge: real distributed execution over a
  multi-partition UNION, asserting Sum/Min/Max + ASC/DESC + HAVING match the
  default merge.

[claude]

Claude finished @waralexrom's task in 2m 24s —— View job

---

Code review

Verdict: nice work — 1 medium, 2 low. Click to expand.

• Read changed files (config, execute.rs, plan.rs, sql/mod.rs)
• Check for correctness, perf, edge-case issues
• Post review

Overall

Clean, well-scoped change. Three points worth raising:

• The opt-in/default split is the right call: the V1 path stays byte-identical, and both V2 strategies are gated by a single env var with lenient parsing and a sane fallback.
• The test setup — running every existing top-k case through both V1 and V2 inside run_topk_as_batch (with a dedicated typed scalar_arithmetic test for the float/decimal cases) — is the most valuable part of this PR. It makes V2 hard to silently break.
• The topk_v2_supported gate is conservative in the right direction: HLL merge, dictionary/nested keys, and non-numeric value types all fall back to V1.

Findings

Medium — Integer SUM silently wraps on overflow (scalar_add_assign, execute.rs:1387-1394)

The macro uses wrapping_add for all integer widths (and wrapping_add again for Decimal128/Decimal256 on execute.rs:1408,1414):

*$d = Some($d.unwrap_or(0).wrapping_add($s.unwrap_or(0)));

V1 routes integer SUM through DataFusion's sum accumulator, which returns an error on overflow. V2 silently wraps to a negative value, which is observable when the top-k boundary group's true sum exceeds i64::MAX (and changes which group lands in position k). The PR description documents float reordering but not integer overflow behavior.

The doc comment on scalar_combine says "MIN/MAX, by contrast, are exact" — but it's silent on integer SUM, which arguably should be exact too. Two options:

1. Use checked_add and surface an ArrowError::ComputeError-style DataFusionError (matches DF behavior).
2. Keep wrapping_add but extend the doc comment to call this out explicitly and add a topk_v2_supported gate that documents it as expected behavior.

Fix this →

Low — process_batch_rows allocates an OwnedRow per row for the probe (execute.rs process_batch_rows)

let probe = self.scratch_rows.row(r).owned();
let (gid, is_new) = match self.group_index.get(&probe) {

Every row pays an OwnedRow allocation, even on a hit (i.e. the common case after the first batch). With std::collections::HashMap and Hash/Eq on Row<'_>/OwnedRow, you can avoid this with RawEntry/hashbrown::HashMap::raw_entry to look up by a borrowed Row. Given this is the hot path the perf story is built around, it's worth either doing that or hashbrown's Equivalent-based lookup; otherwise the "vectorized" label is a bit aspirational on heavy-hit workloads. Not a correctness bug.

Low — plan_topk_full_merge carries an unused agg_fun parameter

agg_fun is only read via vec![None; agg_fun.len()], which is equivalent to vec![None; lower_node.aggregate_expr.len()]. The parameter can be dropped at the call site. Tiny but visible from the signature.

Things that look right, just flagging for awareness

• update_frontier reading the last row of every (incoming, possibly non-empty) batch — correct given streams are monotone in the order-by aggregate and clamp_to_neutral_bound handles non-order-by columns the same way V1 does.
• The two-pass round in update (refresh all frontiers, then process all rows) is documented for why it matters; nice to see the rationale inline.
• plan_topk_full_merge correctly inserts CoalescePartitionsExec before the Single-mode re-aggregate — the comment explains why (no EnforceDistribution in Cube), which will save the next person a debugging session.
• HAVING in full-merge correctly reuses lower_node.schema for create_physical_expr, matching V1.
• env_topk_strategy warning-on-unknown (and not panicking) is consistent with the prior group_by_limit fix on this branch.

· branch: `cubestore-topk-optimization`

[claude]

Claude finished @waralexrom's task in 1m 14s —— View job

---

Code review (follow-up on b4b7964)

Verdict: prior medium and one low addressed; 1 low remains, 1 minor new note. Click to expand.

• Re-read execute.rs diff for the refactor commit
• Verify the integer-SUM concern from the prior review
• Check whether plan_topk_full_merge still carries the unused agg_fun
• Note one new observation about the hash table choice

What b4b7964 changed

TopKStateV2.group_index went from HashMap<OwnedRow, usize> to HashMap<u64, SmallVec<[usize; 1]>>, with the probe borrowed from scratch_rows and equality verified against key_store.row(gid). This kills the per-row OwnedRow allocation that was flagged previously — good.

The scalar_add_assign doc comment now explicitly calls out that integer/decimal SUM uses wrapping_add to match DataFusion's SumAccumulator (add_wrapping). That's a reasonable resolution of the prior medium — V2 and the default streaming path now have a documented, consistent wrap-on-overflow contract rather than V2 silently diverging.

Still applies

Low — plan_topk_full_merge still takes an unused agg_fun (rust/cubestore/cubestore/src/queryplanner/topk/plan.rs:766)

Only used at line 839 as vec![None; agg_fun.len()], equivalent to vec![None; lower_node.aggregate_expr.len()]. The parameter can be dropped at the call site.

Fix this →

One note on the new hash-keyed map

group_index: HashMap<u64, SmallVec<[usize; 1]>> with DefaultHasher (std SipHash) hashes the row once into a u64, then HashMap's RandomState re-hashes that u64 on every lookup/insert. Functionally fine — collisions across distinct keys are ~2⁻⁶⁴ and the SmallVec collision chain handles them. Two cheap improvements if you want to squeeze more out of the hot path:

• Use BuildHasherDefault<NoOpHasher> (identity hasher) so the second hashing pass is a no-op.
• Or switch to hashbrown::HashTable::find / entry with an explicit hash, which avoids both the double-hash and the SmallVec wrapper. This is what arrow-row/DataFusion use internally for the same pattern.

Not blocking — the change as-is is already a clear win over the previous OwnedRow allocation per row.

Looks correct

• The SmallVec<[usize; 1]> collision chain plus key_store.row(g) == probe equality check is safe — group identity remains tied to the actual Row bytes, not the hash.
• Doc comment on scalar_add_assign accurately reflects DataFusion's SumAccumulator overflow semantics; V1↔V2 stay consistent for integer/decimal SUM.

· branch: `cubestore-topk-optimization`

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 58.74%. Comparing base (6697321) to head (b4b7964).

Additional details and impacted files

@@                       Coverage Diff                       @@
##           cubestore-query-performance   #11152      +/-   ##
===============================================================
- Coverage                        58.80%   58.74%   -0.07%     
===============================================================
  Files                              217      217              
  Lines                            17511    17495      -16     
  Branches                          3563     3558       -5     
===============================================================
- Hits                             10298    10278      -20     
- Misses                            6695     6699       +4     
  Partials                           518      518              

Flag	Coverage Δ
cube-backend	58.74% <ø> (-0.07%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[claude]

Claude finished @waralexrom's task in 1m 43s —— View job

---

Code review (re-run on 343a6d4)

Verdict: prior concerns resolved or were incorrect; no new blockers. 1 minor optional note.

• Diff vs origin/master across all 7 changed files
• Re-verified concerns from the two prior reviews
• Spot-checked plan_topk_full_merge, TopKStateV2::process_batch_rows, scalar_add_assign

Resolution of prior findings

• Medium — Integer/decimal SUM overflow (scalar_add_assign). Resolved as "documented intentional behavior" in the doc comment at execute.rs:1495: V2 uses wrapping_add to match DataFusion's SumAccumulator::add_wrapping. V1↔V2 are now consistent (both wrap), and the divergence note belongs to the API contract rather than a silent footgun.
• Low — per-row OwnedRow allocation. Resolved in 343a6d4: group_index: HashMap<u64, SmallVec<[usize; 1]>> with the probe borrowed from scratch_rows and equality checked against key_store.row(gid) (execute.rs:1254-1281). Hot path is now allocation-free on hits.
• Low — "plan_topk_full_merge carries an unused agg_fun". ❌ Incorrect in the prior review — agg_fun is also used at plan.rs:860 (&agg_fun[c.agg_index].0 in make_sort_expr), so the parameter is necessary. Apologies for the noise.

Still applies (minor, optional)

Low — DefaultHasher + HashMap<u64, ...> double-hashes (execute.rs:1256-1260)

The probe row is hashed once with std::collections::hash_map::DefaultHasher (SipHash); then HashMap's RandomState re-hashes that u64 on every lookup/insert. Functionally fine — collisions across distinct keys are ~2⁻⁶⁴ and the SmallVec chain plus the key_store.row(g) == probe equality check handle them. Two cheap improvements if you want to squeeze more out of the hot path:

• Use BuildHasherDefault<NoOpHasher> (identity hasher) so the second pass is a no-op.
• Or switch to hashbrown::HashTable::find / entry with an explicit hash, avoiding both the double-hash and the SmallVec wrapper. This is what arrow-row / DataFusion use internally for the same pattern.

Not blocking.

Looks correct

• update_frontier refresh-then-process two-pass over nodes (execute.rs:1211-1224) — correctly avoids the under-estimation race on groups seen only on an already-processed node; the inline rationale will save someone a debugging session later.
• nodes_seen for a new gid is initialized with finished_nodes.clone() rather than vec![false; ..] — correct: a node that's already drained cannot contribute to the group, so it shouldn't be considered "pending" for early-termination math.
• topk_v2_supported correctly falls back to V1 for HLL Merge, dictionary/nested keys, and non-numeric value types.
• plan_topk_full_merge inserts CoalescePartitionsExec ahead of the Single-mode re-aggregate (plan.rs:834-842) — necessary because EnforceDistribution is disabled in CubeStore. Comment captures the "why".
• HAVING in full-merge correctly creates the physical predicate against lower_node.schema (plan.rs:846), matching V1 semantics.
• env_topk_strategy warns and falls back on unknown values rather than panicking — consistent with the prior group_by_limit env-var handling on this branch.

Test coverage (V1↔V2 equivalence across the existing top-k cases, typed scalar_arithmetic for NaN min/max + decimal scale + sum/null, and the new sql::tests::topk_full_merge distributed end-to-end test) is the load-bearing piece here. Nice work.

· branch: cubestore-topk-optimization ·