Skip to content

feat(devtools): stage_pipeline writes the memo after each step, the selected choice and the enumeration - #1024

Merged
zzylol merged 2 commits into
mainfrom
feat/1003-stage-pipeline-memo
Oct 28, 2026
Merged

zzylol merged 2 commits into
mainfrom
feat/1003-stage-pipeline-memo

Conversation

@zzylol

@zzylol zzylol commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

Closes #1003. Part of #1008.

Why

To review the planner stage by stage (#1008) we need the memo after every step, which candidate the selected plan took for each group, and the candidate DAGs the memo stands for. The asap-stage-pipeline/v1 document had none of these, so neither the DAG viewer (#1005, #1006) nor pricing every enumerated DAG (#1004) could be built.

What

  • stage_pipeline writes an optional memo section (on by default, --no-memo drops it): steps (stage0, pass1, pass2, concrete_states) with each group and its candidates, selected_choice (group → candidate index over the last step), and enumeration (per_step counts up to --memo-enumeration-limit, default 1000; the first --max-enumerated, default 200, choices of one step composed into DAGs). --max-memo-candidates (default 64) caps the candidates listed per group. The format is the shared contract the viewer and pricing PRs build against.
  • The memo is the one production builds: route.rs records it in the new Stage1ToStage3Run::memo_steps; devtools does not re-run any step.
  • Fix in memo::enumeration (feat(optimizer): enumerate a memo's candidate DAGs and compose each choice #1009): enumerate leaves out choices whose candidates read each other in a cycle, and compose returns ComposeError::Cycle for one. Before, a workload with a rollup between two groupings (the scans_projecting_different_labels_are_shared workload) overflowed the stack in compose.
  • New viewer example examples/memo-ratio-of-quantiles.json, planned from memo-examples/promql/ratio-of-quantiles-5m.json (the issue's workload), with an examples.json entry.

How

  1. search_domain snapshots the memo after Pass 1 (with the accuracy-scoped sizes), after local window composition, and after concrete state-group expansion. The Stage 0 and Pass 1 snapshots carry across pane-width rounds with the reused local memo. candidate_connection_search keeps each round's snapshots under its selected plan's namespaced id, and returns the snapshots of the round whose plan won (else the last round's).
  2. Selected choice: build_candidate_inventory_before now also returns each memo candidate's root physical id (memo_candidate_roots, indexed like CandidateMemo::candidates). Terminal-reader factoring maps these ids without removing duplicates. Each selected configured candidate maps through source_candidate to a root id, and from there to its memo candidate. When factoring merged several memo candidates into one root, the one that reads the selected producers' groups wins. selected_choice is null with a reason in three cases: the plan uses two candidates of one group, no memo candidate matches a selected root's producers, or a group needed from the query roots has no selected candidate.
  3. stage_pipeline serializes the snapshots with the existing FlatDag export. It enumerates each step with memo::enumeration::enumerate and composes the listed choices with compose.

Key code interfaces added / modified

crates/plan-selection/src/search/pipeline.rs

/// New type. The memo after each step of the planning round whose plan was
/// selected (or of the last round when none was), as the search built it.
pub struct MemoSteps {
    /// After hash-consing the frontends' roots: groups, no candidates.
    pub stage0: CandidateMemo,                                  // new
    /// After Pass 1 (per-computation candidates, then accuracy-scoped sizes).
    pub pass1: CandidateMemo,                                   // new
    /// After local window composition (pane candidates).
    pub pass2: CandidateMemo,                                   // new
    /// After the Pass 2 memo rules and concrete summary-state groups: the memo the search selects from.
    pub concrete_states: CandidateMemo,                         // new
    /// The selected plan's candidate for each group of concrete_states it uses, or why there is none.
    pub selected_choice: Result<MemoChoice, String>,            // new
}

/// Stays. The Stage 1 → Stage 3 run (other fields unchanged).
pub struct Stage1ToStage3Run<Id> {
    /// The memo after each step and the selected choice; None when the search built no memo.
    pub memo_steps: Option<MemoSteps>,                          // new
}

crates/devtools/src/planning_stages.rs

/// Stays (other fields unchanged).
pub struct PlanningStagesRun {
    /// Copied from Stage1ToStage3Run::memo_steps.
    pub memo_steps: Option<MemoSteps>,                          // new
}

crates/logical-asap-optimization/src/memo/enumeration.rs

/// Stays. A choice that does not compose.
pub enum ComposeError {
    /// A chosen group has no candidate at that index (bad input).
    NoSuchCandidate { group: MemoGroupId, candidate: usize },   // unchanged
    /// A candidate reads a group whose node it does not contain (internal bug in a memo rule).
    ChildNotInCandidate { group: MemoGroupId, candidate: usize, child: MemoGroupId }, // unchanged
    /// The chosen candidates read each other in a cycle (an unsatisfiable choice).
    Cycle { group: MemoGroupId },                               // new
    /// Re-deriving schemas and guarantees failed.
    Recompose(PerComputationCandidateError),                    // unchanged
}
/// Now leaves out choices whose candidates form a cycle.
pub fn enumerate(memo: &CandidateMemo, limit: usize) -> MemoEnumeration; // behavior changed

crates/plan-selection/src/search/candidates/inventory.rs, factoring.rs (crate-private)

/// The inventory, plus each memo candidate's root id indexed like CandidateMemo::candidates.
pub(in crate::search) fn build_candidate_inventory_before(/* unchanged */)
    -> Result<(CandidateInventory, Vec<Vec<PhysicalCandidateId>>), CandidateInventoryError>; // was Result<CandidateInventory, _>
/// Also maps memo_candidate_roots to the factored ids, keeping duplicates.
pub(super) fn factor_terminal_readers(inventory: &mut CandidateInventory,
    memo_candidate_roots: &mut [Vec<PhysicalCandidateId>], deadline: Option<Instant>) -> Result<usize, _>; // new parameter

stage_pipeline flags (new): --no-memo, --max-memo-candidates <n> (64), --memo-enumeration-limit <n> (1000), --max-enumerated <n> (200). Document: the new optional memo section, described in the header of stage_pipeline.rs and in tools/dag-viewer/README.md.

Usage

# The workload: one query every minute, max_error 0.01, failure_probability 0.001, over `data`
# (10^4 series every 5 ms, 2·10^6 rows/s).
stage_pipeline --workload tools/dag-viewer/memo-examples/promql/ratio-of-quantiles-5m.json \
  --deployment-capabilities '{"raw_data_retained":true,"query_time_results_retained":false}' \
  --cost-model aws-fargate-us-east-1 --search candidate-connections --enforce-accuracy-guarantees \
  --out tools/dag-viewer/examples/memo-ratio-of-quantiles.json
# run_search_over → run_stage1_to_stage3_with → candidate_connection_search records MemoSteps for the
# winning pane-width round; stage_pipeline's memo_json writes them. Result (checked):
#   steps: stage0 7 groups / 0 candidates, pass1 7 / 19, pass2 7 / 259, concrete_states 354 / 1738
#   enumeration.per_step: stage0 1, pass1 49 (complete), pass2 and concrete_states ≥ 1000 (incomplete)
#   enumeration.step = "pass1", 49 choices, each composed into a FlatDag
#   selected_choice = {"0":0,"1":0,"2":0,"3":0,"4":0,"5":0,"6":1,"7":0,"8":1,"353":0}

Major algorithm changes

Enumeration of cyclic choices (memo::enumeration)

  • Problem: in the expanded memo of quantile by (host, job) (0.5, data) + quantile by (job) (0.5, data) (+ two range quantiles), group A's candidate reads group B and one of B's candidates reads A. enumerate produced choices taking both, and compose recursed until the stack overflowed.
  • Before: walk tries every candidate of each needed group, and composed memoizes a group only after it finishes. After: when choosing a candidate for a group, walk skips the candidate if one of its children is that group or already reads it through the chosen candidates (reaches). Every edge is checked as its group is chosen, so no cyclic choice is completed, and the work stays bounded by limit. composed keeps a set of the groups it is still composing and returns ComposeError::Cycle when it meets one of them again.
  • Correctness: a cyclic choice stands for no DAG, so leaving it out keeps "choices = candidate DAGs". Tested by choices_whose_candidates_read_each_other_are_left_out.
  • Complexity: O(chosen groups + chosen edges) per candidate tried; at most limit complete choices. Measured: the 4-query rollup workload (scans_projecting_different_labels_are_shared's) enumerates pass1 256, pass2 784 (both complete) and concrete_states up to 1000, with the whole stage_pipeline run taking 30 s (planning included).

Selected choice derivation: new, read-only (see How). It does not change selection.

Before this PR

The document had the Stage 0 DAG, the selected plan's Stage 1/2 DAGs and the search report, but no groups, candidates, rule origins or group → candidate map. compose could overflow the stack on memos with rollups between groupings.

After this PR

Every stage_pipeline document carries the memo after each step, the selected choice, and the enumeration, ready for #1004 (pricing), #1005 and #1006 (viewer). Documents without memo still validate and render.

Evidence

  • Issue's workload in production: the PromQL frontend adds each series' identity as a unique key, so the two quantile_over_time(0.5, data[5m]) share one group. Production builds 7 groups, 19 Pass 1 candidates, 49 Pass 1 choices. The issue's 8 / 26 / 343 were measured without series identity (the expanded_memo test helper), so the test asserts the production numbers.
  • Example document size: 13,818,705 bytes (≈ 12.9 MB of it is the memo: concrete_states has 336 PaneChain groups, each with its full subtree DAG; the document is pretty-printed).

Verification

  • cargo test --release -p asap-devtools --test stage_pipeline: 48/49 pass. The one failure, scans_projecting_different_labels_are_shared (selected cost 0.2100 > 0.1518), fails identically on origin/main (b6d838b). It is pre-existing and unrelated.
    • memo_section_lists_each_step_the_selected_choice_and_the_enumeration: on the issue's workload, the steps are in order, with 7 groups / 19 candidates / 49 complete Pass 1 choices. Every group, candidate and choice DAG is a valid FlatDag, group_of_node names child groups, every selected_choice candidate exists in the last step, the listed step follows the rule, and --no-memo drops the section.
    • composed_selected_choice_is_the_selected_stage1_dag: composing selected_choice gives the selected Stage 1 DAG on the issue's workload and planner-layering Example 1. The comparison applies Stage 2's implement_candidate_operators (exact top-k → Project/Limit/Sort) and recompose_guarantees first, as Stage 1 reconstruction does (Memo: convert memo-group candidates into the vector of all candidate DAGs #1002's acceptance 3).
    • workload_examples_preserve_data_and_valid_materialized_views: now 10 workload examples, including the new one.
  • cargo test --release -p asap-devtools --test viewer_contract: pass.
  • cargo test --release -p asap-integration-tests --test memo_enumeration: 4 pass. New: choices_whose_candidates_read_each_other_are_left_out.
  • cargo test --release -p asap-plan-selection --lib: 379 pass.
  • Viewer: node --test tools/dag-viewer/stages.test.mjs 15 pass. New: the validator accepts the example with and without memo. python3 -m unittest test_viewer test_server OK.
  • cargo fmt --all --check; cargo clippy --release -p asap-plan-selection -p asap-devtools -p asap-integration-tests -p asap-logical-asap-optimization --all-targets -- -D warnings.

Review fixes (second commit): the stricter selected_choice checks above, cycles pruned during enumeration rather than after, aggregate labels naming their measures (Aggregate quantile(q=0.5)), and --no-memo documented as combinable with the cap flags (which it ignores). Re-run after the fixes: memo_enumeration (4), the three memo and example tests of stage_pipeline, asap-logical-asap-optimization --lib enumeration, asap-plan-selection --lib route (18), clippy, fmt and the node tests. Example size after regeneration: 13,818,953 bytes.

Architectural decisions

  • The memo is recorded on the production path (one small MemoSteps struct on the run), not re-run in devtools, so the document shows exactly what the search selected from. It is always recorded: the clones share Rc nodes and cost far less than the MILP.
  • record_round also keeps the latest round's snapshots on the archive, so a search that ends with an error (from any early return) still reports a memo. The reviewer's suggestion to keep them only in memo_rounds would drop the memo on those paths.
  • The selected choice comes from the selected candidate connections through root ids kept across factoring. Rejected alternative: matching the Stage 1 DAG against enumerated choices, which would not scale to the expanded memo.
  • pass2 holds only window composition, because production runs the other Pass 2 rules inside the concrete-state expansion. Each step is the memo exactly as production has it after that call.

Limitations and follow-up

Human review — do not complete with an agent

  • The MVP boundary is correct.
  • New conceptual layers or public interfaces are necessary.
  • The before/after description matches the intended product behavior.
  • Human reviewer:
  • Decision and rationale:

🤖 Generated with Claude Code

…elected choice and the enumeration

The stage-pipeline document gains an optional `memo` section: the memo
the production search built (stage0, pass1, pass2, concrete_states) from
the planning round whose plan was selected, the selected plan's
group -> candidate choice over the last step, and each step's candidate
DAG count with the first choices composed. `--no-memo`,
`--max-memo-candidates`, `--memo-enumeration-limit` and `--max-enumerated`
control it.

The search records the steps in `Stage1ToStage3Run::memo_steps`, and
derives the selected choice from the selected candidate connections
through each memo candidate's root id. `enumerate` leaves out choices
whose candidates read each other in a cycle, and `compose` reports such a
choice instead of overflowing the stack.

Closes #1003. Part of #1008.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
@zzylol zzylol added dag-viewer DAG export and visualization tooling dev-tools Developer tools, diagnostics, and demos enhancement New feature or request labels Oct 28, 2026
Reject a selected choice that misses a needed group or cannot name a memo
candidate, prune cyclic candidates while enumerating (bounded by the
limit), label aggregates by their measures, and loosen the --no-memo usage.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dag-viewer DAG export and visualization tooling dev-tools Developer tools, diagnostics, and demos enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

stage_pipeline: write the memo after each step, the selected choice and the enumeration

1 participant