refactor(stage1): Pass 1 adds a memo group for every intermediate result (#1203 PR 2a-1)

crates/devtools/src/bin/asap_planner.rs

@@ -967,7 +967,7 @@ impl MemoNodes {
 /// children.
 fn memo_groups_json(table: &mut MemoNodes, memo: &Memo, previous: Option<&Memo>) -> Vec<Value> {
     let ids = |groups: &[MemoGroupId]| groups.iter().map(|g| g.0).collect::<Vec<_>>();
-    memo.groups
+    memo.groups()
         .iter()
         .zip(&memo.candidates)
         .enumerate()
@@ -996,7 +996,7 @@ fn memo_groups_json(table: &mut MemoNodes, memo: &Memo, previous: Option<&Memo>)
                                 .iter()
                                 .find(|child| {
                                     Rc::ptr_eq(
-                                        &memo.groups[child.0 as usize].logical_node,
+                                        &memo.groups()[child.0 as usize].logical_node,
                                         original,
                                     )
                                 })

crates/devtools/src/memo_pricing.rs

@@ -367,7 +367,7 @@ fn group_of(
     node: &Rc<OperatorNode>,
     query: usize,
 ) -> Result<MemoGroupId, String> {
-    let node_of = |g: &&MemoGroupId| &memo.groups[g.0 as usize].logical_node;
+    let node_of = |g: &&MemoGroupId| &memo.groups()[g.0 as usize].logical_node;
     groups
         .iter()
         .find(|g| Rc::ptr_eq(node_of(g), node))

crates/devtools/src/planning_stages.rs

@@ -19,7 +19,7 @@ use asap_types::ir::schema::schema_support::with_promql_series_identity;
 use asap_types::ir::schema::{DataType, Field, Schema};
 use asap_types::ir::QueryRoot;
 use asap_types::physical::PhysicalWorkloadPlan;
-use asap_types::stages::{AssignmentRealizer, DagOfLogicalMemoGroups, Memo};
+use asap_types::stages::{AssignmentRealizer, Memo};
 use asap_types::workload::bounds::with_confirmed_weights_at;
 use asap_types::workload::{DataWorkload, PlanningWorkload, QueryWorkloadEntry};
 use serde_json::Value;
@@ -305,15 +305,6 @@ pub fn run_search_over(
     })
 }
 
-/// Stage 0's memo groups, `dag`, without candidates.
-fn stage0_memo_groups(dag: DagOfLogicalMemoGroups) -> Memo {
-    Memo {
-        candidates: vec![Vec::new(); dag.groups.len()],
-        groups: dag.groups,
-        query_roots: dag.query_roots,
-    }
-}
-
 /// The stages over `roots` up to and including `until`, as the facade runs
 /// them. A stage's failure is reported in the run, as the facade does.
 // ponytail: copies run_planning's stage order, which has no stop point;
@@ -329,7 +320,7 @@ pub fn run_stages_until(
     let stage0_dag = generate_dag_of_logical_memo_groups(&roots);
     let mut stages = PlanningStagesRun {
         run_until: until,
-        stage0_memo_groups: Some(stage0_memo_groups(stage0_dag.clone())),
+        stage0_memo_groups: Some(Memo::new(stage0_dag.clone())),
         pass2_local_window_composition: None,
         stage0_query_language_frontends: roots.clone(),
         run: PlanningRun {

crates/devtools/src/planning_stages_text.rs

@@ -242,7 +242,7 @@ fn write_memo(out: &mut String, memo: &Memo, previous: Option<&Memo>, with_candi
     // The group lines share one set of column widths, with each group's
     // candidates printed under its line.
     let group_rows: Vec<Vec<String>> = memo
-        .groups
+        .groups()
         .iter()
         .enumerate()
         .map(|(id, group)| {
@@ -525,9 +525,12 @@ fn root_label(node: &Rc<OperatorNode>) -> String {
 fn chain_label(memo: &Memo, root: &Rc<OperatorNode>, input_groups: &[MemoGroupId]) -> String {
     let (flat, originals) = flatten(&[QueryRoot::Operator(root.clone())]);
     let is_input = |id: usize| {
-        input_groups
-            .iter()
-            .any(|group| Rc::ptr_eq(&memo.groups[group.0 as usize].logical_node, &originals[id]))
+        input_groups.iter().any(|group| {
+            Rc::ptr_eq(
+                &memo.groups()[group.0 as usize].logical_node,
+                &originals[id],
+            )
+        })
     };
     let mut labels = Vec::new();
     let mut pending = vec![flat.nodes.len() - 1];
@@ -611,8 +614,9 @@ mod tests {
             has_line(&["IngestionTime  MatAtIngestion(Memory)", "<- {C1}"]),
             "{text}"
         );
+        // G4 and G5 are groups Pass 1 adds beneath the query's aggregates.
         assert!(
-            text.contains("selected choice: G0:c0 G1:c0 G2:c0 G3:c0\n"),
+            text.contains("selected choice: G0:c0 G1:c0 G2:c0 G3:c0 G4:c0 G5:c0\n"),
             "{text}"
         );
         assert!(text.contains("accepted accuracy:\n    q1: "), "{text}");

crates/devtools/tests/asap_planner.rs

@@ -1524,7 +1524,7 @@ fn widened_scans_keep_memo_groups_and_stage2_candidates_acyclic() {
             .flat_map(|candidate| candidate.input_groups.iter().map(|g| g.0 as usize))
             .collect()
     };
-    for group in 0..memo.groups.len() {
+    for group in 0..memo.groups().len() {
         assert!(
             !reaches(&group_reads, group, group),
             "memo group {group} reads itself"
@@ -2067,11 +2067,14 @@ fn memo_section_lists_each_step_and_the_selected_choice() {
     // group: 7 groups (scan, range, two quantile groups, ratio, sum, avg),
     // 5 + 5 + 5 Pass 1 candidates (each quantile group: exact, DDSketch and
     // a KLL sized for the whole result for the whole budget, DDSketch and a
-    // KLL for half of it).
+    // KLL for half of it). Each summary candidate is one operator over the
+    // groups Pass 1 adds beneath it: the 4 builds both quantiles share,
+    // each quantile's 4 estimates, and the `avg`'s 6 (two states, their
+    // finalizes, a `Project` and the division), one candidate each.
     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]), 15);
+    assert_eq!(steps[1]["groups"].as_array().unwrap().len(), 25);
+    assert_eq!(candidates(&steps[1]), 33);
     // The node table is a valid FlatDag (children before parents) holding
     // each node once; every group's and candidate's `node` names a node of
     // it, and each boundary names a node and one of the candidate's child

crates/devtools/tests/memo_pricing.rs

@@ -8,7 +8,7 @@ use asap_devtools::planning_stages::{lower_promql_roots, run_stages_until, RunUn
 use asap_logical_asap_optimization::enumerate;
 use asap_logical_asap_optimization::generate_dag_of_logical_memo_groups;
 use asap_logical_asap_optimization::test_support::generate_per_computation_candidates;
-use asap_logical_asap_optimization::test_support::try_populate_pass1;
+use asap_logical_asap_optimization::test_support::try_add_pass1_candidates;
 use asap_logical_asap_optimization::DeploymentCapabilities;
 use asap_planner::PlanningModels;
 use asap_types::ir::properties::AccuracyTarget;
@@ -85,7 +85,7 @@ fn planned(workload: &PlanningWorkload, series_identity: bool) -> Planned {
         roots,
         entries: workload.query_workload.entries().collect(),
         data: workload.data_workload.clone().unwrap(),
-        memo: try_populate_pass1(initial, &inventory).unwrap(),
+        memo: try_add_pass1_candidates(initial, &inventory, &[], &Default::default()).unwrap(),
     }
 }
 

crates/integration-tests/src/lib.rs

@@ -343,14 +343,15 @@ pub mod planning {
     ) -> asap_types::stages::Memo {
         use asap_logical_asap_optimization::test_support::generate_per_computation_candidates;
         use asap_logical_asap_optimization::test_support::{
-            add_local_window_composition, try_populate_pass1,
+            add_local_window_composition, try_add_pass1_candidates,
         };
         let initial = asap_logical_asap_optimization::generate_dag_of_logical_memo_groups(&roots);
         let inventory = generate_per_computation_candidates(
             &asap_logical_asap_optimization::generate_dag_of_logical_memo_groups(&roots),
         )
         .unwrap();
-        let mut memo = try_populate_pass1(initial, &inventory).unwrap();
+        let mut memo =
+            try_add_pass1_candidates(initial, &inventory, &[], &Default::default()).unwrap();
         add_local_window_composition(&mut memo, entries);
         asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups(memo)
             .unwrap()
@@ -451,7 +452,7 @@ pub mod planning {
             .candidates
             .iter()
             .enumerate()
-            .filter(|(group, _)| memo.groups[*group].reached_by_queries() == [q])
+            .filter(|(group, _)| memo.groups()[*group].reached_by_queries() == [q])
             .flat_map(|(group, candidates)| {
                 candidates
                     .iter()
@@ -463,9 +464,9 @@ pub mod planning {
             .collect();
         // Group indices don't give the order: Pass 2 appends groups at the
         // end. Depth from the root does.
-        let mut depth = vec![usize::MAX; memo.groups.len()];
+        let mut depth = vec![usize::MAX; memo.groups().len()];
         let mut pending: std::collections::VecDeque<_> = memo
-            .groups
+            .groups()
             .iter()
             .enumerate()
             .filter(|(group, _)| {
@@ -481,7 +482,7 @@ pub mod planning {
             }
             depth[group] = d;
             pending.extend(
-                memo.groups[group]
+                memo.groups()[group]
                     .child_groups
                     .iter()
                     .map(|child| (child.0 as usize, d + 1)),
@@ -534,19 +535,32 @@ pub mod planning {
                 )
             })
             .collect();
-        for (state, group) in memo.groups.iter().enumerate() {
+        for (state, group) in memo.groups().iter().enumerate() {
             if !is_build(&group.logical_node) || group.reached_by_queries().len() < queries {
                 continue;
             }
+            // Each query's consumer of the state, and its own summary candidate
+            // above, which reads the consumer's group (Pass 1 adds the consumer
+            // as a group beneath the query's aggregate).
             let reading: Option<std::collections::BTreeMap<usize, usize>> = (0..queries)
                 .map(|q| {
-                    own_choices(memo, q, |c| {
+                    let reader = own_choices(memo, q, |c| {
                         c.input_groups.iter().any(|child| child.0 as usize == state)
                     })
                     .into_iter()
-                    .next()
+                    .next()?;
+                    let top = own_builds[q].iter().copied().find(|&(group, index)| {
+                        asap_logical_asap_optimization::test_support::groups_read_through_added_groups(
+                            memo,
+                            &memo.candidates[group][index].input_groups,
+                        )
+                        .added_groups
+                        .contains(&asap_types::stages::MemoGroupId(reader.0 as u32))
+                    })?;
+                    Some([top, reader])
                 })
-                .collect();
+                .collect::<Option<Vec<_>>>()
+                .map(|pairs| pairs.into_iter().flatten().collect());
             if let Some(choice) = reading {
                 plans.push((
                     format!("shared{state}"),

crates/integration-tests/tests/finer_aggregate_reuse.rs

@@ -5,7 +5,7 @@ use asap_devtools::planning_stages::run_search_over;
 use asap_executor::values::Value;
 use asap_logical_asap_optimization::generate_dag_of_logical_memo_groups;
 use asap_logical_asap_optimization::test_support::{
-    generate_per_computation_candidates, try_populate_pass1,
+    generate_per_computation_candidates, try_add_pass1_candidates,
 };
 use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode};
 use asap_types::stages::{CandidateForMemoGroup, Memo, OriginOfCandidateForMemoGroup};
@@ -53,29 +53,36 @@ fn expanded(queries: &[String]) -> Memo {
         &asap_logical_asap_optimization::generate_dag_of_logical_memo_groups(&roots),
     )
     .unwrap();
-    let memo = try_populate_pass1(initial, &inventory).unwrap();
+    let memo = try_add_pass1_candidates(initial, &inventory, &[], &Default::default()).unwrap();
     asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups(memo)
         .unwrap()
 }
 
 /// The aggregate group of query `query` alone and its rollup candidates.
 fn rollups(memo: &Memo, query: usize) -> (usize, Vec<&CandidateForMemoGroup>) {
-    let owner = memo
-        .groups
-        .iter()
-        .position(|group| {
-            group.reached_by_queries() == [query]
-                && matches!(
-                    group.logical_node.non_asap(),
-                    Some(NonASAPOp::Aggregate { .. })
-                )
+    let rollups_of = |group: usize| -> Vec<&CandidateForMemoGroup> {
+        memo.candidates[group]
+            .iter()
+            .filter(|c| c.origin == OriginOfCandidateForMemoGroup::Rollup)
+            .collect()
+    };
+    // The query's own aggregate group, or, for a rollup of summary states,
+    // the group of the readout Pass 1 added beneath it, which the rule
+    // rewrites.
+    let own = |group: &asap_types::stages::MemoGroup| group.reached_by_queries() == [query];
+    let owner = (0..memo.groups().len())
+        .find(|&group| own(&memo.groups()[group]) && !rollups_of(group).is_empty())
+        .or_else(|| {
+            memo.groups().iter().position(|group| {
+                own(group)
+                    && matches!(
+                        group.logical_node.non_asap(),
+                        Some(NonASAPOp::Aggregate { .. })
+                    )
+            })
         })
         .expect("the query's own aggregate group");
-    let candidates = memo.candidates[owner]
-        .iter()
-        .filter(|c| c.origin == OriginOfCandidateForMemoGroup::Rollup)
-        .collect();
-    (owner, candidates)
+    (owner, rollups_of(owner))
 }
 
 /// Whether `candidate` reads another group's result through a `Filter`
@@ -125,7 +132,7 @@ fn assert_same_rows(memo: &Memo, owner: usize, candidate: &CandidateForMemoGroup
     };
     for rows in row_sets() {
         let before = physical_common::execute_raw_rows_at(
-            &memo.groups[owner].logical_node,
+            &memo.groups()[owner].logical_node,
             rows.clone(),
             1000,
         );

crates/integration-tests/tests/having_shared_aggregate.rs

@@ -5,7 +5,7 @@ use asap_devtools::planner_fuzz::{read_seed_file, FuzzSettings, FuzzWorkload, Pl
 use asap_executor::values::Value;
 use asap_logical_asap_optimization::generate_dag_of_logical_memo_groups;
 use asap_logical_asap_optimization::test_support::{
-    generate_per_computation_candidates, try_populate_pass1,
+    generate_per_computation_candidates, try_add_pass1_candidates,
 };
 use asap_types::ir::NonASAPOp;
 use asap_types::stages::Memo;
@@ -51,7 +51,7 @@ fn expanded(queries: &[String]) -> Memo {
         .collect();
     let dag = generate_dag_of_logical_memo_groups(&roots);
     let inventory = generate_per_computation_candidates(&dag).unwrap();
-    let memo = try_populate_pass1(dag, &inventory).unwrap();
+    let memo = try_add_pass1_candidates(dag, &inventory, &[], &Default::default()).unwrap();
     asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups(memo)
         .unwrap()
 }
@@ -90,7 +90,7 @@ fn key_only_having_of_one_reader_stays_above_a_shared_aggregate() {
     let grouped = format!("SELECT city, {P90} FROM metrics WHERE {WINDOW} GROUP BY city");
     let memo = expanded(&[format!("{grouped} HAVING city = 'NYC'"), grouped]);
     let shared = memo
-        .groups
+        .groups()
         .iter()
         .position(|group| {
             group.reached_by_queries() == [0, 1]
@@ -100,7 +100,7 @@ fn key_only_having_of_one_reader_stays_above_a_shared_aggregate() {
                 )
         })
         .expect("the aggregate both queries read");
-    let expected = cities(&memo.groups[shared].logical_node);
+    let expected = cities(&memo.groups()[shared].logical_node);
     assert_eq!(expected.len(), 3, "{expected:?}");
     assert!(memo.candidates[shared].len() > 1, "summary candidates too");
     for candidate in &memo.candidates[shared] {