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

crates/devtools/src/planning_stages.rs

@@ -7,7 +7,7 @@ use asap_logical_asap_optimization::per_computation_candidates::PerComputationCa
 pub use asap_logical_asap_optimization::workload_candidates::Pass2Candidate;
 use asap_physical_asap_optimization::PhysicalWorkloadPlan;
 use asap_plan_selection::pass::{run_stage1_to_stage3_with, Stage1ToStage3Search};
-pub use asap_plan_selection::pass::{LogicalCandidate, PrunedCandidate};
+pub use asap_plan_selection::pass::{LogicalCandidate, MemoSteps, PrunedCandidate};
 use asap_plan_selection::PlanningModels;
 use asap_plan_selection::{Selection, SelectionError};
 use asap_types::ir::schema::schema_support::with_promql_series_identity;
@@ -29,6 +29,8 @@ pub struct PlanningStagesRun {
     pub stage2_physical_asap_optimization: Vec<PhysicalWorkloadPlan>,
     pub stage2_pruned: Vec<PrunedCandidate>,
     pub stage3_plan_selection: Result<Selection, SelectionError>,
+    /// The memo after each step and the selected plan's choice over it.
+    pub memo_steps: Option<MemoSteps>,
 }
 
 /// Lower a PromQL workload, each series' identity as a column (the row
@@ -205,6 +207,7 @@ pub fn run_search_over(
         },
         stage2_pruned: run.stage2_pruned,
         stage3_plan_selection: run.stage3_plan_selection,
+        memo_steps: run.memo_steps,
     })
 }
 

crates/devtools/tests/stage_pipeline.rs

@@ -1362,7 +1362,7 @@ fn workload_examples_preserve_data_and_valid_materialized_views() {
         .iter()
         .filter(|e| e.get("workload").is_some())
         .collect();
-    assert_eq!(workload_examples.len(), 9);
+    assert_eq!(workload_examples.len(), 10);
     // Each production search has a wall-clock budget. Bound fixture
     // concurrency so this acceptance test does not starve its own solvers.
     for batch in workload_examples.chunks(2) {
@@ -2206,3 +2206,169 @@ fn a_classic_histogram_quantile_gets_a_selected_plan() {
         assert!(selection["selected"].is_string(), "{query}: {selection}");
     }
 }
+
+/// The issue's example workload (#1003): one ratio of per-series quantiles,
+/// every minute, ε = 0.01, δ = 0.001.
+const MEMO_RATIO_OF_QUANTILES: &str =
+    "../../tools/dag-viewer/memo-examples/promql/ratio-of-quantiles-5m.json";
+
+/// The `memo` section: the memo after each step, the selected choice and
+/// the enumeration, on the issue's example.
+#[test]
+fn memo_section_lists_each_step_the_selected_choice_and_the_enumeration() {
+    let out = std::env::temp_dir().join(format!(
+        "stage_pipeline_{}_memo_ratio.json",
+        std::process::id()
+    ));
+    let status = Command::new(env!("CARGO_BIN_EXE_stage_pipeline"))
+        .args(["--workload", MEMO_RATIO_OF_QUANTILES, "--out"])
+        .arg(&out)
+        .status()
+        .unwrap();
+    assert!(status.success());
+    let document: Value = serde_json::from_str(&std::fs::read_to_string(&out).unwrap()).unwrap();
+    std::fs::remove_file(out).unwrap();
+    let memo = &document["memo"];
+    let steps = memo["steps"].as_array().unwrap();
+    let names: Vec<_> = steps.iter().map(|s| s["step"].as_str().unwrap()).collect();
+    assert_eq!(names, ["stage0", "pass1", "pass2", "concrete_states"]);
+    let candidates = |step: &Value| -> u64 {
+        step["groups"]
+            .as_array()
+            .unwrap()
+            .iter()
+            .map(|g| g["candidate_count"].as_u64().unwrap())
+            .sum()
+    };
+    // The PromQL frontend gives each series its identity as a column, so the
+    // two `quantile_over_time(0.5, data[5m])` have a unique key and share one
+    // group: 7 groups (scan, range, two quantile groups, ratio, sum, avg),
+    // 7 + 7 + 5 Pass 1 candidates, 7 × 7 candidate DAGs.
+    assert_eq!(steps[0]["groups"].as_array().unwrap().len(), 7);
+    assert_eq!(candidates(&steps[0]), 0);
+    assert_eq!(steps[1]["groups"].as_array().unwrap().len(), 7);
+    assert_eq!(candidates(&steps[1]), 19);
+    let per_step = &memo["enumeration"]["per_step"];
+    assert_eq!(
+        per_step["stage0"],
+        serde_json::json!({"count": 1, "complete": true})
+    );
+    assert_eq!(
+        per_step["pass1"],
+        serde_json::json!({"count": 49, "complete": true})
+    );
+    // The listed step is the last one enumerated completely within 200
+    // choices (Stage 0 has no candidates to choose), else the last step.
+    let listed = ["pass1", "pass2", "concrete_states"]
+        .into_iter()
+        .rfind(|step| {
+            per_step[step]["complete"] == true && per_step[step]["count"].as_u64().unwrap() <= 200
+        })
+        .unwrap_or("concrete_states");
+    assert_eq!(memo["enumeration"]["step"], listed);
+    if listed != "concrete_states" {
+        assert_eq!(
+            memo["enumeration"]["choices"].as_array().unwrap().len() as u64,
+            per_step[listed]["count"].as_u64().unwrap()
+        );
+    }
+    // Every group's DAG and every listed candidate's DAG is a valid FlatDag,
+    // and each boundary names one of the candidate's child groups.
+    for step in steps {
+        for group in step["groups"].as_array().unwrap() {
+            validated_dag(&group["dag"], 1);
+            for candidate in group["candidates"].as_array().unwrap() {
+                validated_dag(&candidate["dag"], 1);
+                let children = candidate["children"].as_array().unwrap();
+                for child in candidate["group_of_node"].as_object().unwrap().values() {
+                    assert!(children.contains(child));
+                }
+            }
+        }
+    }
+    // Every candidate of the selected choice exists in the last step.
+    let last = steps.last().unwrap()["groups"].as_array().unwrap();
+    let choice = memo["selected_choice"].as_object().unwrap();
+    assert!(!choice.is_empty());
+    for (group, index) in choice {
+        let group = &last[group.parse::<usize>().unwrap()];
+        assert!(index.as_u64().unwrap() < group["candidate_count"].as_u64().unwrap());
+    }
+    let enumeration = &memo["enumeration"];
+    let choices = enumeration["choices"].as_array().unwrap();
+    assert!(!choices.is_empty() && choices.len() <= 200);
+    for choice in choices {
+        validated_dag(&choice["dag"], 1);
+    }
+    // `--no-memo` leaves the section out.
+    let out = std::env::temp_dir().join(format!(
+        "stage_pipeline_{}_no_memo_ratio.json",
+        std::process::id()
+    ));
+    let status = Command::new(env!("CARGO_BIN_EXE_stage_pipeline"))
+        .args(["--workload", MEMO_RATIO_OF_QUANTILES, "--no-memo", "--out"])
+        .arg(&out)
+        .status()
+        .unwrap();
+    assert!(status.success());
+    let document: Value = serde_json::from_str(&std::fs::read_to_string(&out).unwrap()).unwrap();
+    std::fs::remove_file(out).unwrap();
+    assert!(document.get("memo").is_none());
+}
+
+/// Composing the selected choice over the last step's memo gives the
+/// selected plan's Stage 1 DAG, on the issue's example and #509 Example 1.
+/// Stage 1's DAG is rebuilt from the candidates' operator implementations
+/// (`implement_candidate_operators`: an exact top-k becomes
+/// `Project(Limit(Sort))`), so the composed query DAGs are compared after
+/// the same implementation.
+#[test]
+fn composed_selected_choice_is_the_selected_stage1_dag() {
+    use asap_devtools::planning_stages::run_default_search_over;
+    use asap_logical_asap_optimization::memo::enumeration::compose;
+    use asap_logical_asap_optimization::per_computation_candidates::recompose_guarantees;
+    use asap_types::ir::flat::flatten;
+    let ratio: PlanningWorkload =
+        serde_json::from_str(&std::fs::read_to_string(MEMO_RATIO_OF_QUANTILES).unwrap()).unwrap();
+    for workload in [ratio, example1()] {
+        let roots = lower_promql_roots(&workload).unwrap();
+        let entries: Vec<_> = workload.query_workload.entries().collect();
+        let data = workload.data_workload.clone().unwrap_or_default();
+        let run = run_default_search_over(
+            roots,
+            &entries,
+            &data,
+            asap_plan_selection::PlanningModels::builtin(),
+        )
+        .unwrap();
+        let steps = run.memo_steps.as_ref().unwrap();
+        let choice = steps.selected_choice.as_ref().unwrap();
+        let memo = &steps.concrete_states;
+        let composed = compose(memo, choice).unwrap();
+        let composed: Vec<_> = (0..entries.len())
+            .map(|query| {
+                let (_, node) = composed
+                    .iter()
+                    .find(|(group, _)| memo.groups[group.0 as usize].roots_of.contains(&query))
+                    .unwrap();
+                QueryRoot::Operator(
+                    asap_physical_asap_optimization::implement_candidate_operators(node, &[])
+                        .unwrap(),
+                )
+            })
+            .collect::<Vec<_>>();
+        // As reconstruction does, over the implemented operators.
+        let composed = recompose_guarantees(&composed).unwrap();
+        let [selected] = &run.stage1_logical_asap_optimization[..] else {
+            panic!("the default search keeps the selected Stage 1 candidate")
+        };
+        let selected = selected.roots.as_ref().unwrap();
+        assert_eq!(composed.len(), selected.len());
+        for (composed, selected) in composed.into_iter().zip(selected) {
+            assert_eq!(
+                serde_json::to_value(flatten(&[composed]).0).unwrap(),
+                serde_json::to_value(flatten(std::slice::from_ref(selected)).0).unwrap()
+            );
+        }
+    }
+}

crates/integration-tests/src/lib.rs

@@ -648,6 +648,7 @@ pub mod planning {
             stage3_selected: Some(selected),
             stage1_selected: Some(selected_logical),
             stage3_plan_selection: Ok(selection),
+            memo_steps: None,
         }
     }
 }

crates/integration-tests/tests/memo_enumeration.rs

@@ -193,3 +193,53 @@ fn expanded_memo_enumeration_stops_at_the_limit_and_every_candidate_composes() {
         }
     }
 }
+
+/// Candidates may read each other: a coarse quantile's rollup reads the finer
+/// grouping, whose candidates may read the coarse one. A choice taking both
+/// composes no DAG: `compose` says so, and `enumerate` leaves it out.
+#[test]
+fn choices_whose_candidates_read_each_other_are_left_out() {
+    let queries = [
+        "quantile_over_time(0.5, data[1m])",
+        "quantile_over_time(0.9, data[5m])",
+        "quantile by (host, job) (0.5, data)",
+        "quantile by (job) (0.5, data)",
+    ];
+    let entry = QueryWorkloadEntry {
+        query: Default::default(),
+        requirements: QueryRequirements {
+            accuracy: approximate(),
+            max_latency_ms: None,
+        },
+        predictability: Default::default(),
+        recurrence: QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval(
+            10_000,
+        ))),
+        range_evaluation: None,
+    };
+    let memo = expanded_memo(roots(&queries), &vec![entry; queries.len()]);
+    let mutual = (0..memo.groups.len())
+        .flat_map(|group| (0..memo.candidates[group].len()).map(move |index| (group, index)))
+        .find_map(|(group, index)| {
+            memo.candidates[group][index]
+                .children
+                .iter()
+                .find_map(|&child| {
+                    memo.candidates[child.0 as usize]
+                        .iter()
+                        .position(|c| c.children.contains(&MemoGroupId(group as u32)))
+                        .map(|back| (child, back))
+                })
+                .map(|(child, back)| [(MemoGroupId(group as u32), index), (child, back)])
+        })
+        .expect("a pair of groups whose candidates read each other");
+    assert!(matches!(
+        compose(&memo, &mutual.into()),
+        Err(asap_logical_asap_optimization::memo::enumeration::ComposeError::Cycle { .. })
+    ));
+    let enumeration = enumerate(&memo, 2_000);
+    assert!(!enumeration.choices.is_empty());
+    for choice in &enumeration.choices {
+        compose(&memo, choice).unwrap();
+    }
+}

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

@@ -11,10 +11,12 @@
 //! that differ only in what recomposing guarantees re-derives compose equal
 //! DAGs. Composing checks only that a choice composes; input contracts,
 //! accuracy at the roots, Stage 2 legality and Stage 3 cost stay with the
-//! validator and the later stages. The memo must be acyclic, as every memo
-//! step builds it.
+//! validator and the later stages. The memo's groups are acyclic, as every
+//! memo step builds them, but candidates may read each other (a rollup of
+//! a finer grouping whose candidate reads the coarser one): a choice taking
+//! such candidates together composes no DAG and is not enumerated.
 
-use std::collections::{BTreeMap, HashMap};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::rc::Rc;
 
 use asap_types::ir::{OperatorNode, QueryRoot};
@@ -52,13 +54,18 @@ pub enum ComposeError {
         candidate: usize,
         child: MemoGroupId,
     },
+    /// The chosen candidates read each other in a cycle (e.g. a coarse
+    /// group's rollup reads a finer group whose chosen candidate reads it).
+    #[error("group {group:?} reads itself through the chosen candidates")]
+    Cycle { group: MemoGroupId },
     /// Re-deriving the composed nodes' schemas and guarantees failed.
     #[error("recomposing the composed DAG: {0}")]
     Recompose(#[from] PerComputationCandidateError),
 }
 
 /// Every choice of one candidate per needed group, starting from the query
-/// roots, stopping after `limit` choices.
+/// roots, stopping after `limit` choices. A choice whose candidates read
+/// each other in a cycle composes no DAG and is left out.
 pub fn enumerate(memo: &CandidateMemo, limit: usize) -> MemoEnumeration {
     fn walk(
         memo: &CandidateMemo,
@@ -81,6 +88,15 @@ pub fn enumerate(memo: &CandidateMemo, limit: usize) -> MemoEnumeration {
             }
         };
         for (index, candidate) in memo.candidates[group.0 as usize].iter().enumerate() {
+            // A candidate reading a group that already reads this one closes
+            // a cycle; every edge is checked as its group is chosen.
+            if candidate
+                .children
+                .iter()
+                .any(|&child| reaches(memo, choice, child, group))
+            {
+                continue;
+            }
             choice.insert(group, index);
             let mut next = pending.clone();
             next.extend(candidate.children.iter().rev());
@@ -113,10 +129,14 @@ pub fn compose(
         group: MemoGroupId,
         choice: &MemoChoice,
         done: &mut HashMap<MemoGroupId, Rc<OperatorNode>>,
+        composing: &mut HashSet<MemoGroupId>,
     ) -> Result<Rc<OperatorNode>, ComposeError> {
         if let Some(found) = done.get(&group) {
             return Ok(Rc::clone(found));
         }
+        if !composing.insert(group) {
+            return Err(ComposeError::Cycle { group });
+        }
         let index = choice.get(&group).copied().unwrap_or(0);
         let candidate =
             memo.candidates[group.0 as usize]
@@ -125,7 +145,7 @@ pub fn compose(
                     group,
                     candidate: index,
                 })?;
-        let own_nodes: std::collections::HashSet<_> = OperatorNode::reachable(&candidate.node)
+        let own_nodes: HashSet<_> = OperatorNode::reachable(&candidate.node)
             .iter()
             .map(Rc::as_ptr)
             .collect();
@@ -139,18 +159,22 @@ pub fn compose(
                     child,
                 });
             }
-            replacements.insert(Rc::as_ptr(old), composed(memo, child, choice, done)?);
+            replacements.insert(
+                Rc::as_ptr(old),
+                composed(memo, child, choice, done, composing)?,
+            );
         }
         let node = replace(&candidate.node, &replacements);
         done.insert(group, Rc::clone(&node));
         Ok(node)
     }
     let mut done = HashMap::new();
+    let mut composing = HashSet::new();
     let (groups, roots): (Vec<_>, Vec<_>) = root_groups(memo)
         .map(|group| {
             Ok((
                 group,
-                QueryRoot::Operator(composed(memo, group, choice, &mut done)?),
+                QueryRoot::Operator(composed(memo, group, choice, &mut done, &mut composing)?),
             ))
         })
         .collect::<Result<Vec<_>, ComposeError>>()?
@@ -167,6 +191,24 @@ pub fn compose(
         .collect())
 }
 
+/// Whether `from` is `to` or reads it through the candidates `choice` names.
+fn reaches(memo: &CandidateMemo, choice: &MemoChoice, from: MemoGroupId, to: MemoGroupId) -> bool {
+    let mut pending = vec![from];
+    let mut seen = HashSet::new();
+    while let Some(group) = pending.pop() {
+        if group == to {
+            return true;
+        }
+        if !seen.insert(group) {
+            continue;
+        }
+        if let Some(&index) = choice.get(&group) {
+            pending.extend(&memo.candidates[group.0 as usize][index].children);
+        }
+    }
+    false
+}
+
 fn root_groups(memo: &CandidateMemo) -> impl DoubleEndedIterator<Item = MemoGroupId> + '_ {
     memo.groups
         .iter()