Skip to content

Joint planning search: stream Stage 1 to Stage 3 and reject permanently invalid query candidates (methods 2 + incumbent) - #792

Merged
zzylol merged 1 commit into
mainfrom
feat/joint-search-1-streaming
Oct 21, 2026
Merged

zzylol merged 1 commit into
mainfrom
feat/joint-search-1-streaming

Conversation

@zzylol

@zzylol zzylol commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Why

Stage 1 to Stage 3 build every whole-workload candidate, price each one, and keep all of them in Vecs in Stage1ToStage3Run. A five-query prefix of ASAPQuery's random mixed workload (seed-862, workload-01) takes 146 s and 26 GB, and six queries do not finish. This PR implements the first steps of docs/design_docs/proposals/joint-planning-search.md:

  • the streaming incumbent from the search loop;
  • method 2, rejecting permanently invalid branches.

The selected plan does not change.

Part of #673.

What

  • Streaming selection. Stage 3 checks and prices each physical candidate as Stage 2 builds it. Only the best candidate so far is kept; under cost model 3's pareto-front, the Pareto front so far is kept instead. The run and the selection keep their full candidate lists only when the caller asks for every candidate (developer tools, tests).
  • Method 2: early rejection of permanently invalid query candidates. Each query's own candidates (one Pass 1 choice for every target it reads) are planned alone before composition. Every Stage 2 materialization of each of these is checked, both for the composed candidate and for its single-query Pass 2 candidates.
    • A whole-workload choice is skipped, together with its Pass 2 candidates, only when it picks a query candidate that fails a check no remaining choice can repair:
      • the executor has no kernel for a per-series aggregate it reads;
      • a Count-Min reads weights that are not proven non-negative;
      • the output has no guarantee, or a bound that stays unknown at every size, and the deployment's accuracy model and evidence judge only the guarantee.
    • Failures that resizing could repair are deferred, as are failures in a different family or a different materialization. Examples are a known bound that misses its target, or a failure-probability resize that fails. Every completed plan still goes through the full Stage 3 check.
  • Exhaustive mode. Stage1ToStage3Search { exhaustive: true, .. } composes every choice, as before. The equivalence tests compare the default search against it.
  • Selection report. Selection::search and SelectedPlanReport::search give:
    • the number of candidates Stage 3 checked;
    • for each query, its own candidates checked and the ones discarded, each with its cause and reason.
  • When every candidate of some query is discarded, NoValidCandidate names that query (q<n>) with its first discard.
  • What the selection keeps (for readers of the report). With the default search, Selection.rejected keeps only the first rejection of each cause, plus the losing members of the Pareto front. Selection.costs and Selection.accepted keep only the selected candidate, or the Pareto front under pareto-front. Only keep_every_candidate (developer tools, select_plan) lists every candidate as before.
  • Measured cost table (feat: price every operator from the measured cost table, with calibration and predicted usage #789). When the deployment's measured table prices some operator, the formulas are calibrated by one factor fitted over every candidate that passes the checks before pricing. The default search fits that factor over the same unpruned set, so prices match the exhaustive search's. That costs one extra pass over every candidate (Stage 1, Stage 2 and the checks, no ranking), so there is no speedup in that case; memory still streams. Without a measured table, the factor is 1 and no pass runs.

How

  • select_plan is now PlanSelector::new + offer per candidate + finish. The checks, pricing and tie order are those of the old loop.
    • Price, latency and memory rankings replace the incumbent only when the new candidate is strictly better. The earlier candidate still wins ties within tolerance, as before.
    • pareto-front drops an offered candidate that a front member dominates, and drops front members the new candidate dominates. It then selects with the old choose, over the front in candidate order.
    • select_plan keeps every candidate, so its output is unchanged.
  • QueryCandidates::check checks each query's own candidates:
    1. It enumerates the digits that decide the query's plan: the targets it reads, plus any target whose fused computation absorbs one of them.
    2. It plans the query alone through stage1_candidates_of_choice (Pass 1 and Pass 2), then optimize_physical_asap, then checked_candidate.
    3. A candidate is discarded only if every alternative fails with CheckFailure::size_independent.
    4. invalid_prefix(choice) returns the shortest prefix of digits that decides an invalid query. stage1_logical_asap_optimization_skipping then jumps past every choice that shares that prefix, so the product is not walked one choice at a time. The order of the remaining candidates is unchanged, so ties resolve as before.
  • Accuracy precondition. AccuracyModel::may_accept_regardless_of_guarantee and AccuracyEvidenceProvider::may_provide_evidence both default to true. With either at true, no accuracy failure counts as permanent, because a deployment model or its evidence could accept a larger summary of the same plan, and sharing can give it one. Only the built-in model and NoAccuracyEvidence return false.
  • run_stage1_to_stage3 (and the facade's Stage1ToStage3Pass) use the default search. Tests and helpers that inspect every candidate now call run_stage1_to_stage3_with(.., Stage1ToStage3Search::exhaustive()), which is the old behavior.

Key code interfaces added / modified

crates/plan-selection/src/pass/stage1_to_stage3.rs

/// How `run_stage1_to_stage3_with` searches, and what it keeps.          // new
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Stage1ToStage3Search {
    /// Keep every Stage 1 and Stage 2 candidate and every Stage 3 rejection in the run (developer tools); default false: only the best so far.
    pub keep_every_candidate: bool,                                  // new
    /// Compose every whole-workload choice without checking each query's own candidates first: the exhaustive search tests compare against; default false.
    pub exhaustive: bool,                                            // new
}
impl Stage1ToStage3Search {
    /// The exhaustive search keeping every candidate: the planner before this PR.
    pub fn exhaustive() -> Self;                                     // new
}

/// Stage 1 to Stage 3 over the frontends' roots with the default search (each query's candidates checked first, only the best kept).
pub fn run_stage1_to_stage3<Id: Clone + Debug>(roots: Vec<(Id, QueryRoot)>, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>) -> Result<Stage1ToStage3Run<Id>, PerComputationCandidateError>; // was: exhaustive, kept every candidate
/// `run_stage1_to_stage3` with the search `search`.
pub fn run_stage1_to_stage3_with<Id: Clone + Debug>(roots: Vec<(Id, QueryRoot)>, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, search: Stage1ToStage3Search) -> Result<Stage1ToStage3Run<Id>, PerComputationCandidateError>; // new
/// After Pass 1, with the search `search`.
pub fn run_stage1_to_stage3_from_pass1<Id: Clone + Debug>(per_computation: PerComputationCandidates<Id>, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, search: Stage1ToStage3Search) -> Stage1ToStage3Run<Id>; // was: no `search`

/// The output of Stage 1 to Stage 3 for one workload. The type stays.
pub struct Stage1ToStage3Run<Id> {
    /// Stage 1, Pass 1: the per-computation candidates.
    pub stage1_pass1_per_computation: PerComputationCandidates<Id>, // unchanged
    /// Stage 1, Pass 1: the size of the Cartesian product of the targets' candidates.
    pub combinations: usize,                                         // unchanged
    /// Stage 1: every logical candidate; empty unless every candidate is kept.
    pub stage1_logical_asap_optimization: Vec<LogicalCandidate>,     // was: always filled
    /// Stage 2: every physical candidate; empty unless every candidate is kept.
    pub stage2_physical_asap_optimization: Vec<PhysicalCandidate>,   // was: always filled
    /// Stage 2: the logical candidates it could not implement; empty unless every candidate is kept.
    pub stage2_pruned: Vec<PrunedCandidate>,                         // was: always filled
    /// Stage 3: plan selection.
    pub stage3_plan_selection: Result<Selection, SelectionError>,    // unchanged
    /// Stage 3: the selected physical candidate as Stage 2 built it (`selected()` reads it).
    pub stage3_selected: Option<PhysicalCandidate>,                  // new
    /// Stage 3: the selected candidate's logical roots with confirmed weights.
    pub stage3_selected_logical: Option<Vec<QueryRoot>>,             // unchanged
}

crates/plan-selection/src/lib.rs

/// Stage 3 over candidates offered one at a time; keeps only the best so far (or the Pareto front so far) unless `keep_every_candidate`.
pub struct PlanSelector<'a> { /* private */ }                        // new
impl<'a> PlanSelector<'a> {
    /// A selector for `entries` over `data` under `models`; bad settings or capabilities are an error here (bad input).
    pub fn new(entries: &'a [QueryWorkloadEntry], data: &'a DataWorkload, models: PlanningModels<'a>, keep_every_candidate: bool) -> Result<Self, SelectionError>; // new
    /// Check and price the next candidate in candidate order; keep it if it can still be selected.
    pub fn offer(&mut self, candidate: &PhysicalCandidate) -> Result<(), SelectionError>; // new
    /// Whether a candidate of Stage 1 candidate `from_logical` is kept.
    pub fn keeps_from_logical(&self, from_logical: &str) -> bool;    // new
    /// The selection, and the selected candidate as Stage 2 built it.
    pub fn finish(self) -> Result<(Selection, PhysicalCandidate), SelectionError>; // new
    /// Whether the measured cost table prices some operator, so the formulas are calibrated over every candidate first.
    pub fn needs_calibration(&self) -> bool;                         // new
    /// A candidate's part of the calibration fit; empty when it fails the checks before pricing.
    pub fn calibration_fit(&self, candidate: &PhysicalCandidate) -> Result<CalibrationFit, SelectionError>; // new
    /// Calibrate with the sum of every candidate's fit, before the first `offer`.
    pub fn set_calibration(&mut self, fit: CalibrationFit);         // new
}
/// The measured and formula CPU time of the nodes the measured table prices. The type stays; it becomes public.
pub struct CalibrationFit { /* private */ }                          // was: private
impl CalibrationFit {
    /// Add another candidate's CPU time.
    pub fn add(&mut self, other: CalibrationFit);                    // was: private
}
/// Stage 3 over every candidate at once, keeping every candidate (a `PlanSelector` with `keep_every_candidate`).
pub fn select_plan(cands: &[PhysicalCandidate], entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>) -> Result<Selection, SelectionError>; // unchanged signature

/// Stage 3's result. The type stays.
pub struct Selection {
    /// The selected candidate's id.
    pub selected: String,                                            // unchanged
    /// The plan to run.
    pub plan: Vec<QueryRoot>,                                        // unchanged
    /// Every valid candidate's cost when every candidate is kept; otherwise the selected candidate's and the front's.
    pub costs: BTreeMap<String, CandidateCost>,                      // was: always every valid candidate
    /// The accuracy accepted per query, for the candidates in `costs`.
    pub accepted: BTreeMap<String, Vec<AcceptedAccuracy>>,           // was: always every valid candidate
    /// Every rejection when every candidate is kept; otherwise the first rejection of each cause, and the front's losers.
    pub rejected: Vec<Rejection>,                                    // was: always every rejection
    /// Under `pareto-front`, the front in candidate order.
    pub front: Vec<String>,                                          // unchanged
    /// The selected plan's pane retention.
    pub pane_retention_ms: BTreeMap<PhysicalASAPNodeId, u64>,        // unchanged
    /// What the search checked and discarded.
    pub search: SearchReport,                                        // new
}

/// What the search behind a selection checked and discarded, for the selection report.
#[derive(Serialize, Deserialize)]
pub struct SearchReport {                                            // new
    /// The physical candidates Stage 3 checked, valid or not.
    pub candidates_checked: u64,                                     // new
    /// Per query in workload order; empty for an exhaustive search or `select_plan`.
    pub queries: Vec<QuerySearchReport>,                             // new
}
/// One query's own candidates, checked before composition.
pub struct QuerySearchReport {                                       // new
    /// How many of its own candidates were checked.
    pub candidates_checked: u64,                                     // new
    /// The ones that failed a check no remaining choice can repair.
    pub discarded: Vec<DiscardedQueryCandidate>,                     // new
}
/// One discarded query candidate.
pub struct DiscardedQueryCandidate {                                 // new
    /// (target index, candidate index) for each target the query reads.
    pub choice: Vec<(usize, usize)>,                                 // new
    /// The check it failed (an executor or accuracy cause; reflects the plan, not bad input).
    pub cause: RejectionCause,                                       // new
    /// The failure, with its numbers.
    pub reason: String,                                              // new
}

crates/plan-selection/src/accuracy.rs

pub trait AccuracyModel {
    fn satisfies(&self, guarantee: &AccuracyGuarantee, target: &AccuracyTarget) -> bool; // unchanged
    fn accepts_output_regardless_of_guarantee(&self, output: &QueryRoot, target: &AccuracyTarget) -> bool; // unchanged
    /// Whether `accepts_output_regardless_of_guarantee` may accept some output; default true (no accuracy failure is treated as permanent). Built-in: false.
    fn may_accept_regardless_of_guarantee(&self) -> bool;            // new, default true
}
pub trait AccuracyEvidenceProvider {
    fn operand_domain(&self, operand: &OperatorNode) -> Option<OperandDomain>; // unchanged
    fn output_guarantee(&self, output: &QueryRoot) -> Option<AccuracyGuarantee>; // unchanged
    /// Whether this provider may return any evidence; default true. `NoAccuracyEvidence`: false.
    fn may_provide_evidence(&self) -> bool;                          // new, default true
}

crates/logical-asap-optimization/src/workload_candidates.rs

/// Stage 1 enumeration without the choices `invalid_prefix` rejects (Some(n): every choice agreeing on the first n digits is invalid), nor their Pass 2 candidates.
pub fn stage1_logical_asap_optimization_skipping<'a, Id: Clone>(per_computation: &'a PerComputationCandidates<Id>, entries: &[QueryWorkloadEntry], invalid_prefix: impl FnMut(&[usize]) -> Option<usize> + 'a) -> impl Iterator<Item = Stage1Candidate> + 'a; // new
/// The Stage 1 candidates of one Pass 1 choice: the composed candidate, then its Pass 2 candidates.
pub fn stage1_candidates_of_choice<Id: Clone>(per_computation: &PerComputationCandidates<Id>, entries: &[QueryWorkloadEntry], choice: &[usize]) -> Vec<Stage1Candidate>; // new

crates/planner/src/installable.rs

/// Stage 3's record of the selected plan, for reporting. The type stays.
pub struct SelectedPlanReport {
    /// The selected candidate's id.
    pub selected_candidate_id: String,                               // unchanged
    /// Its price per second.
    pub cost_per_second: Option<f64>,                                // unchanged
    /// The resources Stage 3 predicts it uses.
    pub predicted_usage: Option<PredictedUsage>,                     // unchanged
    /// What the search checked and discarded; `#[serde(default)]`.
    pub search: SearchReport,                                        // new
}

Usage

The user plans five queries of seed-862 workload-01 with plan_installable, each repeating at its interval with a 5% error target, under cost model 1 (Fargate prices):

// The user's workload: q01-sum-fine, q02-sum-coarse, q03-quantile,
// q04-quantile, q05-quantile; built by the example converter.
let workload: PlanningWorkload = /* five RepeatingEntry, max_error 0.05 */;
// plan_installable lowers PromQL, then Stage1ToStage3Pass calls
// run_stage1_to_stage3 (default search): QueryCandidates::check plans each
// query alone and discards q03..q05's per-series summary candidates of
// `avg by (...) (quantile_over_time(...))` (no guarantee, or an unknown
// relative bound); Stage 1 skips every choice picking them; PlanSelector
// prices the rest one at a time and keeps only the cheapest.
let plan = plan_installable(UserInput::new(&workload, FrontendInput::Promql { histograms: None }, models)).await?;
// What the user gets back: the same installable plan as before this PR,
// and the selection report.
let report = plan.selection.unwrap();
assert_eq!(report.cost_per_second, Some(3.7117181201488165e-6));
assert_eq!(report.search.candidates_checked, 264);         // exhaustive: 126,984
let discarded: Vec<_> = report.search.queries.iter()
    .map(|q| (q.candidates_checked, q.discarded.len())).collect();
assert_eq!(discarded, [(2, 0), (2, 0), (3, 2), (3, 2), (3, 2)]);
// e.g. q03's discard: cause NoAccuracyGuarantee, "the output has no guarantee,
// and the deployment has no evidence for it, so it cannot meet target ...".

(The numbers come from running the local converter against this commit. The converter isn't part of this PR.)

Before this PR

Prefix of seed-862 workload-01 Time Peak memory
3 0.72 s 0.15 GB
4 6.5 s 1.58 GB
5 146 s 26.2 GB
6 killed at 300 s 58.6 GB at the kill
10 (9 PromQL + 1 SQL) not run (6 does not finish)

Each query's invalid summary candidates multiply the product, and every candidate is held until the end.

After this PR

Same machine, release build of the same converter, plan_installable end to end. Measured on the rebased head (after #787 and #789); the 9-query row is the 10-query prefix without its SQL query:

Prefix of seed-862 workload-01 Time Peak memory Candidates checked by Stage 3
3 0.11 s 24 MB 132
4 0.23 s 24 MB 132
5 0.25 s 25 MB 264
6 1.60 s 30 MB 2,056
9 (the PromQL queries) 128 s 217 MB 131,088
  • Plans match. For prefixes 3, 4 and 5, the installable plan JSON is byte-identical to main's, apart from the selection report: same DAG, same stored outputs, and the same cost (3.7117181201488165e-6 for 5). Only the candidate ids differ (for 5, P61-m8 vs P4813-m8), because skipped choices are no longer numbered.

  • Streaming alone, without method 2, against the default search. The exhaustive search with only the best kept (Stage1ToStage3Search { exhaustive: true, keep_every_candidate: false }) selects the same plan and physical DAG, at the same cost, as the default search:

    • 5 queries: exhaustive checks 126,984 candidates in 79 s at 81 MB; default checks 264 in 0.17 s;
    • 6 queries: exhaustive checks 3,042,072 candidates in 2,132 s at 220 MB; default checks 2,056 in 1.4 s.
  • The 9-query PromQL prefix still takes 128 s. The remaining Cartesian product is the work for methods 5, 6, 3 and 4 (following PRs).

  • Kubernetes mixin recording rules (kubernetes_mixin_rules.txt, rules that plan alone, 5% target, 1,000 series per metric):

    • the first 5 rules plan jointly in 65 s at 2.3 GB;
    • 10 rules do not finish in 600 s;
    • 20 rules run out of memory.

    These rules have many targets per query, so the product is large, and this PR does not reduce it. They were not measured on main.

Verification

  • cargo test --workspace --locked: passes on the rebased head.
  • New crates/integration-tests/tests/joint_search_equivalence.rs. Each check compares the default search with Stage1ToStage3Search::exhaustive(), the pre-PR path: every candidate priced, then selected at once by select_plan. It asserts the same plan (flattened), the same cost, the same physical DAG and materialization, and the same Pareto front (or both NoValidCandidate).
    • planner_layering_examples_select_the_exhaustive_plan: Examples 1 and 2 and Pattern B under cost model 1, cost model 2, and cost model 3 with each objective (64 GB, which rejects Pattern B's ingestion-time panes), plus Pattern A as given under cost model 1.

    • random_small_workloads_select_the_exhaustive_plan: 60 seeded random workloads of 2 to 4 PromQL queries, mixing sum, avg-of-quantile, quantile, topk, rate, max, count, increase, stddev_over_time and topk(k, sum by (...) (rate(c_total[w]))), with random targets. Each runs under one cost model in turn, with 1 GB cluster memory so that limits reject some candidates. The test also asserts that every early-rejection cause occurs (no executor kernel, Count-Min weights not proven non-negative, no guarantee, unknown bound).

    • a_calibrated_cost_model_selects_the_exhaustive_plan: random workloads under a measured cost table, so the calibration is fitted over every candidate.

    • Ignored in debug and run here in release:

      • pattern_a_selects_the_exhaustive_plan_under_every_cost_model: Pattern A as given and monthly, under every cost model, plus Pattern B; 439 s.
      • random_small_workloads_select_the_exhaustive_plan_under_every_cost_model: all 60 seeds under all 5 cost models, 300 comparisons; 14 s.

      Both pass. Pattern A at rest is left out because its exhaustive search takes hours.

    • Method 2 behavior:

      • a_query_candidate_without_a_guarantee_is_discarded_before_composition;
      • a_known_bound_that_misses_is_not_discarded_before_composition (a resize could repair it);
      • a_model_that_may_accept_any_output_discards_nothing (a custom accuracy model);
      • a_query_with_no_valid_candidate_is_named_in_the_failure.
  • Unit tests.
    • streaming_selection_selects_what_select_plan_selects: same selection, cost and front under every cost model, keeping at most the front.
    • skipping_enumeration_drops_only_the_invalid_prefixes.
  • Existing tests. Those that inspect every candidate now use the exhaustive mode and pass unchanged. deployment_accuracy_model now reads its candidate lists from an exhaustive run instead of the facade. e2e_plan checks search.candidates_checked.
  • Prefix measurements (tables above): /usr/bin/time around the converter. Plans were compared as JSON.
  • Not run: clippy, fmt check, GitHub CI.

Limitations and follow-up

  • No component split, lower bound, dominance, local materialization or budget yet. Those are methods 5, 3, 4, 6 and 7, and the budget, in the following PRs.
  • devtools stage_pipeline still enumerates exhaustively through run_stage2_and_stage3. It keeps every candidate, so its output is unchanged.

@zzylol zzylol added enhancement New feature or request optimizer Workload optimization and plan selection search Candidate-plan search space and algorithms labels Oct 21, 2026
…y invalid query candidates

Implements methods 2 and the streaming incumbent of
docs/design_docs/proposals/joint-planning-search.md.

- Stage 3 becomes a streaming selector (PlanSelector): each physical
  candidate is checked and priced as Stage 2 builds it, and only the best
  so far (cost model 3 pareto-front: the front so far) is kept.
  select_plan keeps every candidate, as before.
- Each query's own candidates are checked planned alone before
  composition; a whole-workload choice that picks one failing a check no
  remaining choice can repair (no executor kernel, Count-Min weights not
  proven non-negative, no guarantee or an unknown bound, with accuracy
  models that judge only the guarantee) is skipped with its Pass 2
  candidates.
- Stage1ToStage3Search selects the exhaustive search and keeping every
  candidate for tests and developer tools; the selection reports what the
  search checked and discarded per query.
- Equivalence tests against the exhaustive search on the planner-layering
  examples and seeded random workloads under every cost model.
@zzylol
zzylol force-pushed the feat/joint-search-1-streaming branch from 216d60f to d757aa1 Compare October 21, 2026 01:42
@zzylol
zzylol merged commit 463e3f5 into main Oct 21, 2026
1 of 4 checks passed
zzylol added a commit that referenced this pull request Oct 21, 2026
…'s numbering

#795's tests were written against the exhaustive search's candidate
ids (P5-m1); with #792's default search the same plan is P3-m1, and a
failed selection lists the first rejection of each cause.
zzylol added a commit that referenced this pull request Oct 21, 2026
…'s numbering

#795's tests were written against the exhaustive search's candidate
ids (P5-m1); with #792's default search the same plan is P3-m1, and a
failed selection lists the first rejection of each cause.
@zzylol
zzylol deleted the feat/joint-search-1-streaming branch November 3, 2026 02:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request optimizer Workload optimization and plan selection search Candidate-plan search space and algorithms

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant