[zzylol]

Why

#901: the SQL per-group top-k idiom

SELECT city, host, c FROM (SELECT city, host, COUNT(*) AS c,
  ROW_NUMBER() OVER (PARTITION BY city ORDER BY COUNT(*) DESC) AS rn
  FROM metrics WHERE ts > now() - INTERVAL '5' MINUTE GROUP BY city, host) WHERE rn <= 5

was planned as a count, then SQLWindowFunc, then Filter rn <= 5. Three problems followed: no heavy-hitter sketch (CMS+heap, CountSketch+heap) was offered for it; no exact plan compiled (the executor has no window function), so there was no exact answer to compare against; and Stage 3's compile check passed a plan whose count sketch was built at ingestion time while it rejected the same plan with the sketch built at query time. The selected plan (P42-m3) then failed to install: the executor cannot build stored output 2 ... unsupported weighted frequency family.

What

• The SQL frontend's output for the idiom is canonicalized into the operator PromQL topk by (city) lowers to: Aggregate(Reduce(by=[city]), [TopK{5}]) over the COUNT(*) GROUP BY city, host aggregate. Pass 1 now offers the same choices for it as for the global ORDER BY c DESC LIMIT 5 top-k (exact, CMS+heap, CountSketch+heap, Cms, UnivMon).
• A per-group ranking the heavy-hitter TopK does not cover (e.g. ORDER BY AVG(x), a tie-breaking second key, a SELECT list without the partition column) becomes the generic per-partition Limit{k}{Sort}, which the executor runs. So every such query has an exact plan that compiles.
• SQL top-k answers now carry the group keys that name each ranked item (host) after [keys.., topk_k]. This holds in exact plans and from heap sketches, as PromQL top-k answers carry series labels. Before, a SQL top-k answered only the counts.
• The Planner's executor compile check rejects a plan whose stored-output build the executor cannot compile. Before, it accepted such a plan. Its verdict no longer depends on the materialization.
• Heavy-hitter promotion is applied only where no consumer reads the node's columns by position: a query's root, or a branch of a Concat that nothing reads by position either. SELECT host FROM (SELECT host, COUNT(*) AS c … ORDER BY c DESC LIMIT 5) used to return the counts under the name host. It now stays an exact Limit{Sort}. This applies to PromQL too: a ranking that lowers to Limit{Sort} and sits under a reader, as in count(topk(5, sum by (job) (m))), is no longer promoted. That keeps it from losing its columns, at the cost of its heap-sketch choices. PromQL top-ks that the frontend lowers straight to TopK (over counts or rates) are unaffected.

Ties. ROW_NUMBER numbers rows that tie on the ordering value in no particular order. rn <= k therefore keeps an arbitrary subset of the rows tied at the k-th place, and so do PromQL topk, the exact Limit{Sort} and a heap sketch. Those may each keep different tied rows. The tests use data whose counts are distinct at the k-th place, so they are deterministic.

How

• asap_types::ir::canonicalize: a new local rule, try_rewrite_row_number_top_k, matches Project{P}{Filter{rn <= k | rn < k+1 | rn = 1}{[Project{Q}]{SQLWindowFunc{RowNumber, p, o}{X}}}}, where P keeps only columns of X and drops rn. It rewrites the chain to Limit{k,p}{Sort{o,p}{Project{P}{X}}} when P keeps p and o, and to Project{P}{Limit{k,p}{Sort{o,p}{X}}} otherwise. The existing heavy-hitter rule then promotes the first form, and now maps the partition through the projection.
• canonicalize computes once the set of nodes that a non-Concat consumer reads by position (read_by_position). Heavy-hitter promotion is skipped for those nodes.
• Stage 2 implement_topk (exact) and Stage 1 ranked_as_topk (heap sketch) keep the item columns for SQL as well as for PromQL.
• PlannerExecutorCompileCheck::compiles maps InstallablePlanError::StoredOutputNotExecutable to Err("stored output N: …"), the same form check_compiles uses. It used to treat that error as "not installable for another reason" and return Ok.

Key code interfaces added / modified

crates/types/src/ir/canonicalize.rs

/// Rewrite the DAG under `root` into its canonical form (bottom-up). Signature unchanged;
/// now also rewrites SQL's ROW_NUMBER per-partition top-k, and promotes a heavy-hitter
/// only where no consumer reads its columns by position.
pub fn canonicalize(root: Rc<OperatorNode>) -> Result<Rc<OperatorNode>, SchemaDerivationError>; // behavior changed

crates/planner/src/installable.rs

/// The Planner's executor as Stage 3's compile check. Signature unchanged; `Err` now also
/// when the executor cannot build one of the plan's stored outputs ("stored output N: …").
impl ExecutorCompileCheck for PlannerExecutorCompileCheck {
    fn compiles(&self, roots: &[QueryRoot]) -> Result<(), String>; // behavior changed
}

Output format: a SQL top-k root outputs [partition keys.., topk_k, item group keys..]. It used to output [partition keys.., topk_k]. Examples: ["city", "topk_5", "host"] for the per-city query, ["topk_5", "host"] for the global one.

Usage (end to end, the issue's workload; see crates/planner/tests/sql_per_group_topk.rs):

// The workload: the issue's Q1 (global top 5 hosts) and Q2 (top 5 hosts per city), every
// 10 s, over `metrics(ts, city, host, value)` ingesting 666.67 rows/s.
let workload = workload(r#"{"sql": "datafusion_sql"}"#, &[TOP_5_HOSTS, TOP_5_HOSTS_PER_CITY], EXACT);
// The SQL frontend lowers Q2 to Project/Filter(rn <= 5)/Project/SQLWindowFunc/Aggregate;
// canonicalize (run by the frontend) turns it into Aggregate(Reduce([city]), [TopK{5}]).
// Stages 1-3 plan it; Stage 3's PlannerExecutorCompileCheck compiles each candidate,
// stored-output builds included.
let plan: InstallablePlan = plan_installable(UserInput::new(
    &workload, FrontendInput::Sql { catalog: &catalog }, PlanningModels::builtin())).await?;
// The selected exact plan compiles; running Q2 over the test rows answers
// (city, topk_5, host) rows: nyc h1 9, h2 8, h3 7, h4 6, h5 5; sf s1 4, s2 2.
plan.check_compiles()?;

stage_pipeline on the issue's workload with --exhaustive (the issue's command) selects P7-m3: Q1 exact (Count acc) · Q2 exact (Count acc) · Count acc for Q1 in 30 × 10s panes, Count acc for Q2 in 30 × 10s panes · ingestion time: Count acc ×30 panes, Count acc ×30 panes. It reports no stored_outputs_error.

Major algorithm changes

Canonicalization of SQL's ROW_NUMBER per-partition top-k

• Problem: in Q2 above, the window function and the rn filter were left in place. The plan therefore had no TopK target for Pass 1 to offer heap sketches for, and no exact plan the executor could run: all 22 all-exact plans were rejected with value operation has no native implementation.
• Before: canonicalize only rewrote Limit{Sort{[Project]Aggregate([Count|Sum])}} into Aggregate([TopK]). A ROW_NUMBER rewrite had existed and was removed because dropping the visible rn column broke outer projections.
• After (try_rewrite_row_number_top_k, then try_promote_additive_top_ranking, both in apply_local_rules):
  1. Match Project{P} over Filter{Column(rn) <= k} (or < k+1, or = 1, with k ≥ 1). Optionally pass through the derived table's Project{Q}. Under that is SQLWindowFunc{RowNumber, args: [], partition p, order o ≠ []} over X, where rn is the window's appended column.
  2. Map every item of P through Q to a column of X. Give up if P holds rn or a computed item.
  3. If P keeps every column of p and o, emit Limit{k,p'}{Sort{o',p'}{Project{P}{X}}}, with p' and o' remapped. This node has P's schema. Otherwise, emit Project{P}{Limit{k,p}{Sort{o,p}{X}}}.
  4. The same loop then tries heavy-hitter promotion on the new node. Promotion runs only when columns_may_change, meaning the original node is not read by position (read_by_position). It now maps the partition through the projection.
• Correctness: the rewrite keeps exactly the rows with rn <= k for every partition, up to which tied rows are kept, which ROW_NUMBER leaves unspecified anyway. The generic form keeps P's schema exactly. The promoted form is used only at a root or in a Concat branch, as before for the global top-k. Tests:
  • the canonicalize unit tests rownumber_*, a_visible_row_number_is_left_alone, filter_on_a_non_rownumber_column_is_left_alone and does_not_promote_a_ranking_read_by_position, plus idempotence checks;
  • cross_language::sql_rownumber_count_topk_is_promqls_topk_by;
  • the end-to-end tests in sql_per_group_topk (exact answers, and SQL equal to PromQL on the same data).
• Complexity: one extra O(nodes + edges) pass per query root to compute read_by_position. The rule itself is O(|P| + |p| + |o|) per Project. On the issue's workload, the exhaustive search now enumerates 702 logical and 1980 physical candidates, up from 234 and 660, because Q2 has 9 Pass 1 choices instead of 3. stage_pipeline --exhaustive takes 41.6 s in a debug build, up from 17.0 s. The default search takes 2.6 s.
• Assumptions and limits:
  • Only ROW_NUMBER is handled. RANK/DENSE_RANK would keep ties.
  • Only a literal rn <= k, rn < k or rn = 1 works, with every SELECT item a plain column. k >= rn, BETWEEN and a conjunction with another predicate keep the window.
  • A visible rn keeps the window, which has no executor implementation, so no plan of such a query compiles.
  • Under a positional consumer the top-k stays exact (Limit{Sort}) and is not offered heap sketches.

Executor compile check covers stored-output builds

• Problem: PlannerExecutorCompileCheck::compiles builds the installable plan first. That step compiles each stored output's build and fails with StoredOutputNotExecutable when the executor refuses one. The check treated that failure as "not installable for another reason" and returned Ok. A Cms that the executor cannot build (unsupported weighted frequency family) therefore passed when it was kept at ingestion time (P42-m3, selected) and failed when it was built in the query run (P42, P42-m1).
• Before / after: Err(_) => Ok(()) becomes Err(StoredOutputNotExecutable { node_id, reason }) => Err("stored output {node_id}: {reason}"). Other installability errors (e.g. a scalar root) still pass to plan_installable, which reports them.
• Correctness: Stage 3 now selects only plans whose every stored output and query run compile. Test: sql_per_group_topk::a_plan_whose_stored_output_does_not_build_is_rejected fails before this change, when plan_installable returns Installable(StoredOutputNotExecutable), and passes after, when every candidate is rejected as ExecutorCannotCompile.
• Complexity: unchanged. The build was already compiled; only its result is now used.
• Assumptions and limits: a plan the executor cannot build is rejected, never repaired. There is no execution fallback.
• Exposed elsewhere: stage_pipeline's mixed-language test plans a SQL COUNT(*) over a table whose ts time index is nullable. The executor builds no count over a nullable time column (summary time column must be a timestamp). On main, the query-time plan was rejected, but the ingestion-time plan was selected even though it could not be installed. Now both are rejected. That test's table now declares ts non-null, which is what a time index is. The executor limitation itself is unchanged.

Before this PR

On the issue's command (--exhaustive): 234 logical and 660 physical candidates. Q2's Pass 1 choices were exact (Count acc), Cms and UnivMon. All 22 all-exact plans were rejected with executor_cannot_compile. The selected P42-m3 (Q1 Cms · Q2 Cms, Cms at ingestion time) could not be installed (stored_outputs_error). SQL top-k answers had no item column.

After this PR

On the same command: 702 logical and 1980 physical candidates. Q2's choices match Q1's: exact (Count acc), CMS+heap (Count acc), CMS+heap + Cms, CMS+heap + UnivMon, CountSketch+heap (Count acc), CountSketch+heap + Cms, CountSketch+heap + UnivMon, Cms, UnivMon. All 22 all-exact plans compile; 20 lost on cost, 1 lost on a cost tie to the enumeration order, and one was selected. The Cms plans are rejected as executor_cannot_compile whether built at ingestion time (stored output 2: … unsupported weighted frequency family) or at query time. The selected plan is P7-m3, both queries exact from ingestion-time count accumulators in 30 × 10 s panes, and it installs. The sketch plans are still rejected (count_min_weights_not_proven_non_negative, compile errors) because of #886.

#886 (sketch candidate placement for SQL top-k by COUNT(*)) is not fixed here. Q2 now goes through the same top-k candidate generation as Q1, so #886's misplaced candidates (a per-host Cms, a heap over already exact counts) now apply to Q2 too, and the #886 fix will cover both. The compile-check fix means the per-host Cms/UnivMon candidates #886 asks about are now rejected at every materialization, never selected through an ingestion-time build.

Verification

Run locally on the rebased head:

• canonicalize unit tests (asap-types): ROW_NUMBER shapes, the rn bounds (= 1, <= 0, < 1, = 2), positional readers, Concat branches, idempotence.
• asap-devtools --test cross_language (including a nested PromQL ranking that stays Limit{Sort}) and --test stage_pipeline.
• cargo nextest run -p asap-planner --test sql_per_group_topk: 6 passed:
  • exact per-city top 5;
  • SQL equals PromQL on the same data;
  • the global top 5 answers its hosts;
  • the generic forms (a tie-breaking key, SELECT host, c) run exactly;
  • the approximate plan installs;
  • stored-output compile consistency.
• cargo nextest run -p asap-integration-tests --test per_computation_candidates_sql --test stage2_physical_asap_optimization_sql --test sql_to_post_asap --test sql_to_physical: 23 passed.
• asap-frontend-sql --test sql_lowering, asap-planner --test runtime_install --test e2e_plan, asap-physical-asap-optimization, asap-logical-asap-optimization, integration promql_to_post_asap and operator_sharing: all passed.
• An independent review found no blocking issues. This PR addresses its findings: rn = 1, Concat branches seen through a reader that reads by position, the PromQL behavior change stated and tested, and end-to-end tests of the generic forms.
• cargo fmt --check, plus cargo clippy --all-targets --all-features --locked -- -D warnings on the touched crates: clean.
• End to end: stage_pipeline on the issue's workload, numbers above.