SP-9: AST-aware hierarchical stage assignment — QueryExpr stage split + recursive PromQL serialiser

[zzylol]

Closes #90.

What this PR actually does

SP-9 introduces per-stage sub-plans produced by walking the full query algebra AST, replacing the flat SP-3 assignment that put everything on the Agent stage.

---

New types in types.rs

• StageResourceBudgets — per-stage memory/CPU caps (agent_memory_bytes, agent_cpu_micros_per_sample, backend_memory_bytes, precompute_memory_bytes)
• AgentSubPlan, BackendSubPlan, PrecomputeSubPlan, DbSubPlan — per-stage result structs
• StagedPlan — wraps the four sub-plans + a deferral_log: Vec<String> for observability
• CollectionPlan::staged_plan: Option<StagedPlan> — Some when a query_string was supplied, None otherwise (SP-3 flat assignment remains the fallback)

---

planner/stage_split.rs (new, ~980 lines)

split_expr_by_stage(&QueryExpr, &StageResourceBudgets) → StagedPlan

The function takes a QueryExpr (from the general algebra, PR #108) and performs a bottom-up walk, assigning each node to a stage:

QueryExpr node	Default stage
Source, Filter, Window, sketch Agg	Agent OTel Collector
Partition, Merge, Dedup, Agg { Exact(Sum|Count|Min|Max) }	Backend OTel Collector
HistogramQuantile, PromQLSubquery, BinaryOp, Sort+Limit (TopK)	Precompute Engine
Agg { Exact(Avg) }	DB-side (not mergeable across agents)

Budget deferral chain — when a sketch op's estimated memory exceeds the agent cap it is demoted to Backend; if it also exceeds the backend cap it moves to Precompute. Every demotion is logged in StagedPlan::deferral_log.

expr_to_promql(&QueryExpr) → String — recursive descent serialiser covering:

• histogram_quantile(φ, rate(…[w])), quantile_over_time(φ, …[w])
• expr[range:step] subqueries
• Vector binary ops with on(…)/ignoring(…)/group_left/group_right
• Aggregate GROUP BY → op by (keys) (…)
• Sort+Limit → topk(n, …)

---

split_expr_by_stage takes &QueryExpr, not &SketchExpr

The initial commit used SketchExpr directly. A follow-up commit (Option B) ported everything to QueryExpr:

• parse_query_sketch still returns SketchExpr (parsers unchanged)
• handle_plan converts via QueryExpr::from_sketch_expr() before calling split_expr_by_stage
• All node assignments and expr_to_promql operate purely on QueryExpr

---

config/backend.rs

generate_backend_config_staged() added: emits a dedup processor before the merge processor when BackendSubPlan::has_dedup is true.

config/precompute.rs

build_precompute_jobs() uses staged_plan.precompute.query_expr (the upper sub-tree serialised as PromQL) instead of the hardcoded quantile_over_time(0.99, …) template when the staged plan is present and active.

main.rs

handle_plan now runs:

parse_query_sketch(qs)
  → QueryExpr::from_sketch_expr()
  → QueryOptimizer::optimize()
  → split_expr_by_stage(&budgets)
  → plan.staged_plan = Some(StagedPlan)

SP-3 flat assignment is unchanged and remains the fallback when query_string is absent.

---

Test coverage — 16 new tests in planner/stage_split.rs (303 total)

Test	What it covers
ddsketch_agg_goes_to_agent	sketch Agg → Agent stage
hll_stays_at_agent_by_default	HLL Agg → Agent stage
aggregate_sum_goes_to_backend	Exact(Sum) → Backend
aggregate_without_group_by_avg_goes_to_db	Exact(Avg) → DB (non-mergeable)
partition_group_by_goes_to_backend	Partition node → Backend
topk_goes_to_precompute	TopK / Sort+Limit → Precompute
histogram_quantile_activates_precompute	HistogramQuantile → Precompute
binary_op_activates_precompute	BinaryOp → Precompute
ddsketch_deferred_to_backend_when_agent_budget_exceeded	budget deferral Agent→Backend
ddsketch_deferred_to_precompute_when_both_budgets_exceeded	budget deferral Agent→Backend→Precompute
label_matchers_and_tree	ScalarExpr AND-tree → AgentSubPlan::label_filters
promql_ddsketch_with_filter_and_window	expr_to_promql for DDSketch + Window
promql_histogram_quantile	expr_to_promql for histogram_quantile
promql_subquery_renders_range	expr_to_promql for subquery [range:step]
promql_binary_op_renders_operator	expr_to_promql for vector binary ops
promql_topk_wraps_inner	expr_to_promql for Sort+Limit → topk

---

Backward compatibility

• Analyzer::analyze() is unchanged; all existing callers compile without modification.
• CollectionPlan::staged_plan defaults to None — no change to existing plan constructors.
• SP-3 flat assignment is always the fallback when query_string is absent.

🤖 Generated with Claude Code