diff --git a/crates/devtools/src/bin/asap_planner.rs b/crates/devtools/src/bin/asap_planner.rs index 4e231d027..1a2e04f98 100644 --- a/crates/devtools/src/bin/asap_planner.rs +++ b/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 { let ids = |groups: &[MemoGroupId]| groups.iter().map(|g| g.0).collect::>(); - 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, ) }) diff --git a/crates/devtools/src/memo_pricing.rs b/crates/devtools/src/memo_pricing.rs index 128cc4911..f02126d20 100644 --- a/crates/devtools/src/memo_pricing.rs +++ b/crates/devtools/src/memo_pricing.rs @@ -367,7 +367,7 @@ fn group_of( node: &Rc, query: usize, ) -> Result { - 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)) diff --git a/crates/devtools/src/planning_stages.rs b/crates/devtools/src/planning_stages.rs index b40f36dea..152ffc183 100644 --- a/crates/devtools/src/planning_stages.rs +++ b/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 { diff --git a/crates/devtools/src/planning_stages_text.rs b/crates/devtools/src/planning_stages_text.rs index f68332018..50f6e005a 100644 --- a/crates/devtools/src/planning_stages_text.rs +++ b/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> = memo - .groups + .groups() .iter() .enumerate() .map(|(id, group)| { @@ -525,9 +525,12 @@ fn root_label(node: &Rc) -> String { fn chain_label(memo: &Memo, root: &Rc, 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}"); diff --git a/crates/devtools/tests/asap_planner.rs b/crates/devtools/tests/asap_planner.rs index 5f655d3b2..e1f765334 100644 --- a/crates/devtools/tests/asap_planner.rs +++ b/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 diff --git a/crates/devtools/tests/memo_pricing.rs b/crates/devtools/tests/memo_pricing.rs index 1bdc9e850..7a5fb0e0f 100644 --- a/crates/devtools/tests/memo_pricing.rs +++ b/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(), } } diff --git a/crates/integration-tests/src/lib.rs b/crates/integration-tests/src/lib.rs index 49534547c..390675e97 100644 --- a/crates/integration-tests/src/lib.rs +++ b/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> = (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::>>() + .map(|pairs| pairs.into_iter().flatten().collect()); if let Some(choice) = reading { plans.push(( format!("shared{state}"), diff --git a/crates/integration-tests/tests/finer_aggregate_reuse.rs b/crates/integration-tests/tests/finer_aggregate_reuse.rs index ed8a1b06e..854b85a4d 100644 --- a/crates/integration-tests/tests/finer_aggregate_reuse.rs +++ b/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, ); diff --git a/crates/integration-tests/tests/having_shared_aggregate.rs b/crates/integration-tests/tests/having_shared_aggregate.rs index 554e63d38..864a85a8b 100644 --- a/crates/integration-tests/tests/having_shared_aggregate.rs +++ b/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] { diff --git a/crates/integration-tests/tests/memo_enumeration.rs b/crates/integration-tests/tests/memo_enumeration.rs index 6566b08e1..626c070e0 100644 --- a/crates/integration-tests/tests/memo_enumeration.rs +++ b/crates/integration-tests/tests/memo_enumeration.rs @@ -320,16 +320,22 @@ fn ratio_of_quantiles_production_memo() { // model that prunes none and accepts no output, so its summaries stay // and its Pass 2 memo is recorded. let enforced = production_run(&ratio_of_quantiles()); - assert_eq!(candidate_count(pass1_memo(&enforced)), 7); + // 7 groups with one candidate each; the `avg`'s exact accumulator reads + // 6 more groups Pass 1 adds (its two states, their finalizes, a `Project` + // and the division), one candidate each. + assert_eq!(candidate_count(pass1_memo(&enforced)), 13); let models = PlanningModels::builtin().with_accuracy(&AcceptsNoOutput); let run = production_run_with(&ratio_of_quantiles(), models); - // Stage 0 (Pass 1's groups): series identity makes the two + // Stage 0 (Pass 1's first 7 groups): series identity makes the two // `quantile 0.5` one group. // Pass 1: the two quantile groups have 5 candidates each, 3 for the // whole budget (exact, DDSketch, KLL sized for the whole result) and 2 - // for half of it; the other 5 groups one each. - assert_eq!(pass1_memo(&run).groups.len(), 7); - assert_eq!(candidate_count(pass1_memo(&run)), 15); + // for half of it; the other 5 groups one each. Each summary candidate + // reads its estimate's group, over its build's group: the 4 builds are + // shared by the two quantiles, the 8 estimates are not; with the `avg`'s + // 6, Pass 1 adds 18 groups, one candidate each. + assert_eq!(pass1_memo(&run).groups().len(), 25); + assert_eq!(candidate_count(pass1_memo(&run)), 33); // 5 × 5 choices, 3 × 3 of them for the whole budget. The inventory has // 3 × 3 assignments (one target per quantile group and measure, so the // two `quantile 0.5` share one, 3 choices each): the 2 × 2 that take no @@ -343,8 +349,8 @@ fn ratio_of_quantiles_production_memo() { ); // The last step, over the complete tumbling-window-width domain. let last = pass2_memo(&run); - assert_eq!(last.groups.len(), 255); - assert_eq!(candidate_count(last), 1243); + assert_eq!(last.groups().len(), 265); + assert_eq!(candidate_count(last), 1237); let enumeration = enumerate(last, 2_000); assert!(!enumeration.complete); assert_eq!(enumeration.all_chosen_candidates.len(), 2_000); @@ -355,47 +361,46 @@ fn ratio_of_quantiles_production_memo() { .unwrap_or_else(|e| panic!("G{group} candidate {index}: {e}")); } } - // The state-configuration rule rewrites a candidate's KLL to the size - // the readers of a state share; the rewrite is a copy of the candidate - // (`derive` in `pass2/memo/rules/shared_configuration_readers.rs`) with - // the `accuracy_budget_parts` its new size was made for. In each quantile - // group the `accuracy_budget_parts` 1 KLL `k` 567 (sized for the whole - // result) gets copies at `k` 567 (parts 1) and 1184 (parts 2, the size of - // the half-budget KLL); the `accuracy_budget_parts` 2 KLL `k` 1184 keeps - // 2. The `k` 1184 copy composes the same DAG as the original - // `accuracy_budget_parts` 2 candidate: one sketch, sized for half the - // budget. + // The state-configuration rule rewrites a reader of a KLL build group to + // read the build at the size the readers of a state share; the rewrite + // is a copy of the reader (`derive` in + // `pass2/memo/rules/shared_configuration_readers.rs`) with the + // `accuracy_budget_parts` its new size was made for. Each quantile + // group's `accuracy_budget_parts` 1 KLL candidate (`k` 567, sized for + // the whole result) reads an estimate group, which gets a copy reading + // the `k` 1184 build (parts 2, the size of the half-budget KLL): the + // build group the half-budget candidate's estimate reads, so one sketch, + // sized for half the budget. let pass1 = pass1_memo(&run); - let quantile_groups: Vec<_> = (0..pass1.groups.len()) + let quantile_groups: Vec<_> = (0..pass1.groups().len()) .filter(|&group| pass1.candidates[group].len() == 5) .collect(); assert_eq!(quantile_groups.len(), 2); + let sizes = |candidates: &[CandidateForMemoGroup]| { + candidates + .iter() + .filter(|c| c.origin == OriginOfCandidateForMemoGroup::Pass1) + .filter_map(|c| kll(c).first().map(|&(k, _)| (k, c.accuracy_budget_parts))) + .collect::>() + }; for group in quantile_groups { - let sizes = |candidates: &[CandidateForMemoGroup]| { - candidates + assert_eq!(sizes(&pass1.candidates[group]), [(567, 1), (1184, 2)]); + let estimate_read_by = |k: u32| { + let candidate = pass1.candidates[group] .iter() - .filter(|c| c.origin == OriginOfCandidateForMemoGroup::Pass1) - .filter_map(|c| kll(c).first().map(|&(k, _)| (k, c.accuracy_budget_parts))) - .collect::>() + .find(|c| kll(c).first().is_some_and(|&(size, _)| size == k)) + .unwrap(); + let [estimate] = candidate.input_groups[..] else { + panic!("one estimate group") + }; + estimate.0 as usize }; - assert_eq!(sizes(&pass1.candidates[group]), [(567, 1), (1184, 2)]); - let candidates = &last.candidates[group]; + let (whole, half) = (estimate_read_by(567), estimate_read_by(1184)); + assert_eq!(sizes(&last.candidates[whole]), [(567, 1), (1184, 2)]); assert_eq!( - sizes(candidates), - [(567, 1), (1184, 2), (567, 1), (1184, 2)] + last.candidates[whole][1].input_groups, + last.candidates[half][0].input_groups ); - let at_k: Vec = candidates - .iter() - .enumerate() - .filter(|(_, c)| { - c.origin == OriginOfCandidateForMemoGroup::Pass1 - && kll(c).first().is_some_and(|&(k, _)| k == 1184) - }) - .map(|(index, _)| index) - .collect(); - assert_eq!(at_k.len(), 2); - let alone = |index: usize| composed(last, &[(MemoGroupId(group as u32), index)].into()); - assert_eq!(alone(at_k[0]), alone(at_k[1])); } } @@ -455,7 +460,7 @@ fn choices_whose_candidates_read_each_other_are_left_out() { build_logical_asap_dags_from_chosen_candidates(&memo, choice).unwrap(); } - let (group, index, child) = (0..memo.groups.len()) + let (group, index, child) = (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] @@ -466,7 +471,7 @@ fn choices_whose_candidates_read_each_other_are_left_out() { .expect("a candidate reading another group"); let mut cyclic = memo.clone(); let mut back = cyclic.candidates[child.0 as usize][0].clone(); - back.root_node = Rc::clone(&cyclic.groups[group.0 as usize].logical_node); + back.root_node = Rc::clone(&cyclic.groups()[group.0 as usize].logical_node); back.input_groups = vec![group]; cyclic.candidates[child.0 as usize].push(back); let back_index = cyclic.candidates[child.0 as usize].len() - 1; diff --git a/crates/integration-tests/tests/planner_layering_example2.rs b/crates/integration-tests/tests/planner_layering_example2.rs index 9c09da9fd..0b31978b7 100644 --- a/crates/integration-tests/tests/planner_layering_example2.rs +++ b/crates/integration-tests/tests/planner_layering_example2.rs @@ -347,20 +347,28 @@ fn memo_computations(workload: &PlanningWorkload, roots: &[QueryRoot]) -> Vec> = group + // Each query's consumer of the state, and its own UnivMon candidate + // above, which reads the consumer's group (Pass 1 adds the + // estimate beneath the query's aggregate). + let reading: Option> = group .reached_by_queries() .iter() .map(|&q| { - own(q, &|c| { + let reader = own(q, &|c| { c.input_groups.iter().any(|child| child.0 as usize == state) - }) + })?; + let top = own(q, &|c| { + c.origin == OriginOfCandidateForMemoGroup::Pass1 + && OperatorNode::reachable(&c.root_node).iter().any(is_univmon) + })?; + Some([top, reader]) }) .collect(); - choices.extend(reading); + choices.extend(reading.map(|pairs| pairs.into_iter().flatten().collect())); } choices .iter() diff --git a/crates/integration-tests/tests/reuse_rules.rs b/crates/integration-tests/tests/reuse_rules.rs index 2c8d5a4ad..21c9a212e 100644 --- a/crates/integration-tests/tests/reuse_rules.rs +++ b/crates/integration-tests/tests/reuse_rules.rs @@ -172,7 +172,7 @@ fn candidates_of<'a>( .enumerate() .filter(move |(group, _)| { let is_root = root_groups.contains(&MemoGroupId(*group as u32)); - let group = &memo.groups[*group]; + let group = &memo.groups()[*group]; is_root || matches!( group.logical_node.non_asap(), @@ -262,7 +262,7 @@ fn own_candidates<'a>( origins: &'a [OriginOfCandidateForMemoGroup], ) -> impl Iterator)> + 'a { candidates_of(memo, origins) - .filter(|(group, node)| !Rc::ptr_eq(node, &memo.groups[*group].logical_node)) + .filter(|(group, node)| !Rc::ptr_eq(node, &memo.groups()[*group].logical_node)) } /// `sum by (job) (sum_over_time(m[5m]))`: the nested sums and their merge @@ -388,7 +388,7 @@ fn every_candidate_answers_alike(query: &str) { for (group, node) in own_candidates(&memo, &[OriginOfCandidateForMemoGroup::Pass1]) { assert_eq!( execute_implemented(node, &entry), - execute_implemented(&memo.groups[group].logical_node, &entry), + execute_implemented(&memo.groups()[group].logical_node, &entry), "{query}: a candidate of group {group} differs" ); checked += 1; @@ -426,7 +426,7 @@ fn reuse_answers_alike_in( OriginOfCandidateForMemoGroup::TopKReuse, ], ) { - if memo.groups[group] + if memo.groups()[group] .reached_by_queries() .iter() .any(|&q| q >= first_approximate_query) @@ -435,7 +435,7 @@ fn reuse_answers_alike_in( } assert_eq!( execute_implemented(node, &entry), - execute_implemented(&memo.groups[group].logical_node, &entry), + execute_implemented(&memo.groups()[group].logical_node, &entry), "{queries:?}: a reuse candidate of group {group} differs" ); checked += 1; @@ -520,7 +520,7 @@ fn a_rollup_over_a_shared_summary_candidate_answers_as_without_it() { ]; let memo = memo(&queries, repeating()); let shared = memo - .groups + .groups() .iter() .find(|group| { group.reached_by_queries().len() > 1 @@ -536,7 +536,7 @@ fn a_rollup_over_a_shared_summary_candidate_answers_as_without_it() { memo.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() @@ -555,7 +555,7 @@ fn a_rollup_over_a_shared_summary_candidate_answers_as_without_it() { let without = compose_memo_plan(&memo, queries.len(), &sharing); let (group, index) = choices(&memo, &[OriginOfCandidateForMemoGroup::Rollup]) .into_iter() - .find(|(group, _)| memo.groups[*group].reached_by_queries() == [0]) + .find(|(group, _)| memo.groups()[*group].reached_by_queries() == [0]) .expect("a rollup of the coarse sum"); sharing.insert(group, index); let with = compose_memo_plan(&memo, queries.len(), &sharing); diff --git a/crates/integration-tests/tests/sharing_issue_regressions.rs b/crates/integration-tests/tests/sharing_issue_regressions.rs index 04555299e..7775c0921 100644 --- a/crates/integration-tests/tests/sharing_issue_regressions.rs +++ b/crates/integration-tests/tests/sharing_issue_regressions.rs @@ -196,7 +196,7 @@ fn sql_count_topk_has_a_fused_raw_row_heap() { #[test] fn exact_rollup_catalog_crosses_identical_per_series_average() { 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::stages::OriginOfCandidateForMemoGroup; let roots: Vec<_> = Planner::new(FuzzSettings::default()) @@ -209,7 +209,7 @@ fn exact_rollup_catalog_crosses_identical_per_series_average() { &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(); let expanded = asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups(memo) .unwrap(); @@ -230,7 +230,7 @@ use crate::physical_common; fn filtered_group_preserves_sql_empty_and_null_semantics() { use asap_executor::values::Value; 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::{operator::AggIntent, NonASAPOp}; use asap_types::stages::OriginOfCandidateForMemoGroup; @@ -244,7 +244,7 @@ fn filtered_group_preserves_sql_empty_and_null_semantics() { &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(); let expanded = asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups(memo) .unwrap(); @@ -369,7 +369,7 @@ fn filtered_group_preserves_sql_empty_and_null_semantics() { ], ] { let before = physical_common::execute_raw_rows_at( - &expanded.groups[owner].logical_node, + &expanded.groups()[owner].logical_node, rows.clone(), 1000, ); @@ -384,7 +384,7 @@ fn filtered_group_preserves_sql_empty_and_null_semantics() { fn sql_count_rollup_executes_with_integer_results() { use asap_executor::values::Value; 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; use asap_types::stages::OriginOfCandidateForMemoGroup; @@ -407,7 +407,7 @@ fn sql_count_rollup_executes_with_integer_results() { &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(); let expanded = asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups( memo, @@ -449,7 +449,7 @@ fn sql_count_rollup_executes_with_integer_results() { ], ]; let before = - physical_common::execute_raw_rows(&expanded.groups[owner].logical_node, rows.clone()); + physical_common::execute_raw_rows(&expanded.groups()[owner].logical_node, rows.clone()); let after = physical_common::execute_raw_rows(&candidate.root_node, rows); let sorted = |rows: Vec>| { let mut rows: Vec<_> = rows.iter().map(|row| format!("{row:?}")).collect(); @@ -461,7 +461,7 @@ fn sql_count_rollup_executes_with_integer_results() { .all(|row| matches!(row.last(), Some(Value::Int64(_))))); assert_eq!(sorted(before), sorted(after)); let before_empty = - physical_common::execute_raw_rows(&expanded.groups[owner].logical_node, vec![]); + physical_common::execute_raw_rows(&expanded.groups()[owner].logical_node, vec![]); let after_empty = physical_common::execute_raw_rows(&candidate.root_node, vec![]); assert_eq!(sorted(before_empty), sorted(after_empty)); } @@ -510,7 +510,7 @@ fn sql_sum_shares_nested_windows_and_slower_readers() { #[test] fn sketch_rollup_crosses_inner_range_function() { 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::stages::OriginOfCandidateForMemoGroup; let mut w = workload(923); @@ -529,13 +529,18 @@ fn sketch_rollup_crosses_inner_range_function() { &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(); let expanded = asap_logical_asap_optimization::test_support::sharing_rules_and_shared_state_groups(memo) .unwrap(); - let owner = expanded.query_roots[1].as_operator().unwrap().0 as usize; - assert!(expanded.candidates[owner] + // The rule rewrites the reader of the coarse state, a group Pass 1 + // added beneath the query's root. + assert!(expanded + .groups() .iter() + .zip(&expanded.candidates) + .filter(|(group, _)| group.reached_by_queries() == [1]) + .flat_map(|(_, candidates)| candidates) .any(|c| c.origin == OriginOfCandidateForMemoGroup::Rollup)); } diff --git a/crates/logical-asap-optimization/src/added_groups.rs b/crates/logical-asap-optimization/src/added_groups.rs new file mode 100644 index 000000000..103d3cd19 --- /dev/null +++ b/crates/logical-asap-optimization/src/added_groups.rs @@ -0,0 +1,47 @@ +//! Interim Stage 1 helpers that tell Stage 0's groups (the groups Stage 0 +//! produced from the queries' DAG, [`Memo::is_stage0_group`]) from the +//! groups Pass 1 and Pass 2 add beneath them (#1203 2a). They go once no +//! rule needs to tell them apart. + +use std::collections::HashSet; + +use asap_types::stages::{Memo, MemoGroupId}; + +/// The groups some inputs read through the groups added beneath Stage 0's +/// groups ([`groups_read_through_added_groups`]). +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct GroupsReadThroughAddedGroups { + /// The added groups passed through, in the order first reached. + pub added_groups: Vec, + /// The Stage 0 groups reached, in the order first reached. + pub stage0_groups: Vec, +} + +/// The groups `inputs` read, walking through every added group (each read +/// through all its candidates' input groups) and stopping at Stage 0's +/// groups. +pub fn groups_read_through_added_groups( + memo: &Memo, + inputs: &[MemoGroupId], +) -> GroupsReadThroughAddedGroups { + let mut read = GroupsReadThroughAddedGroups::default(); + let mut pending: Vec = inputs.iter().rev().copied().collect(); + let mut seen = HashSet::new(); + while let Some(group) = pending.pop() { + if !seen.insert(group) { + continue; + } + if memo.is_stage0_group(group) { + read.stage0_groups.push(group); + continue; + } + read.added_groups.push(group); + let below: Vec = memo.candidates[group.0 as usize] + .iter() + .flat_map(|candidate| &candidate.input_groups) + .copied() + .collect(); + pending.extend(below.into_iter().rev()); + } + read +} diff --git a/crates/logical-asap-optimization/src/find_or_add_group.rs b/crates/logical-asap-optimization/src/find_or_add_group.rs new file mode 100644 index 000000000..a259c8e30 --- /dev/null +++ b/crates/logical-asap-optimization/src/find_or_add_group.rs @@ -0,0 +1,110 @@ +//! Finding or adding the memo group of a node (Stage 1 only): the group +//! whose node equals it by Stage 0's equality +//! ([`Memo::group_with_equal_node`]), else a new group over its children's +//! groups, each found or added the same way. + +use std::collections::HashMap; +use std::rc::Rc; + +use asap_types::ir::OperatorNode; +use asap_types::stages::{ + CandidateForMemoGroup, Memo, MemoGroup, MemoGroupId, OriginOfCandidateForMemoGroup, +}; + +/// Returns the existing group whose node equals `node` +/// ([`Memo::group_with_equal_node`], Stage 0's equality), after adding +/// `queries` to it and to every group below it ([`Memo::add_queries`]). +/// Otherwise adds a group for `node` over its children's groups +/// ([`add_group`]) and returns it. It never copies a group. +pub(crate) fn find_or_add_group( + memo: &mut Memo, + node: &Rc, + queries: &[usize], + origin: &OriginOfCandidateForMemoGroup, + accuracy_budget_parts: usize, +) -> MemoGroupId { + match memo.group_with_equal_node(node) { + Some(group) => { + memo.add_queries(group, queries); + group + } + None => add_group(memo, node, queries, origin, accuracy_budget_parts), + } +} + +/// Adds a group for `node`, which no group's node equals, over its +/// children's groups ([`over_child_groups`]). The group's logical node is +/// `node` over its children's groups' nodes, `queries` reach it, and its one +/// candidate is that node over those groups, made by `origin` and sized for +/// `accuracy_budget_parts`. +pub(crate) fn add_group( + memo: &mut Memo, + node: &Rc, + queries: &[usize], + origin: &OriginOfCandidateForMemoGroup, + accuracy_budget_parts: usize, +) -> MemoGroupId { + let (logical_node, child_groups) = + over_child_groups(memo, node, queries, origin, accuracy_budget_parts); + let mut queries = queries.to_vec(); + queries.sort_unstable(); + queries.dedup(); + let candidate = CandidateForMemoGroup { + fused_groups: Vec::new(), + root_node: Rc::clone(&logical_node), + input_groups: child_groups.clone(), + origin: origin.clone(), + accuracy_budget_parts, + }; + memo.push_group( + MemoGroup { + logical_node, + child_groups, + parents: Vec::new(), + queries, + }, + vec![candidate], + ) +} + +/// `node`'s operator over the groups of its children, each found or added +/// with every node below it ([`find_or_add_group`], same arguments), and +/// those groups by child position: group `i` is child `i`'s, so an operator +/// reading one group twice (`x / x`) lists it twice. `node` itself when +/// every child is its group's node. +pub(crate) fn over_child_groups( + memo: &mut Memo, + node: &Rc, + queries: &[usize], + origin: &OriginOfCandidateForMemoGroup, + accuracy_budget_parts: usize, +) -> (Rc, Vec) { + let mut group_of_child: HashMap<*const OperatorNode, MemoGroupId> = HashMap::new(); + let mut child_groups = Vec::new(); + for child in node.children() { + let group = match group_of_child.get(&Rc::as_ptr(child)) { + Some(&group) => group, + None => { + let group = find_or_add_group(memo, child, queries, origin, accuracy_budget_parts); + group_of_child.insert(Rc::as_ptr(child), group); + group + } + }; + child_groups.push(group); + } + let group_node = |child: &Rc| { + Rc::clone(&memo.groups()[group_of_child[&Rc::as_ptr(child)].0 as usize].logical_node) + }; + if node + .children() + .into_iter() + .all(|child| Rc::ptr_eq(child, &group_node(child))) + { + return (Rc::clone(node), child_groups); + } + let rebuilt = Rc::new(OperatorNode { + operator: node.operator.map_children(group_node), + ..(**node).clone() + }); + (rebuilt, child_groups) +} diff --git a/crates/logical-asap-optimization/src/lib.rs b/crates/logical-asap-optimization/src/lib.rs index 3864772dc..2a6c14468 100644 --- a/crates/logical-asap-optimization/src/lib.rs +++ b/crates/logical-asap-optimization/src/lib.rs @@ -60,8 +60,11 @@ //! readers; plan selection (Stage 3) checks each selected candidate's //! guarantee against its query's `AccuracyTarget` and never resizes. +mod added_groups; /// The deployment's capabilities: supported summaries and what it retains. mod deployment_capabilities; +mod find_or_add_group; +mod own_candidate_sizing; mod pass1; mod pass2; /// Stage 0: the memo groups of the workload's query roots. @@ -88,7 +91,7 @@ pub use pass1::{ PerComputationLogicalMemoCandidatesError, }; // Pass 2: its entry and its error. -pub use pass1::memo_population::CandidatePopulationError; +pub use pass1::add_pass1_candidates::AddPass1CandidatesError; pub use pass2::logical_memo_candidates::{ generate_logical_asap_memo_candidates, pass2_after_local_window_composition, pass2_local_window_composition, diff --git a/crates/logical-asap-optimization/src/own_candidate_sizing.rs b/crates/logical-asap-optimization/src/own_candidate_sizing.rs new file mode 100644 index 000000000..72ef96e4c --- /dev/null +++ b/crates/logical-asap-optimization/src/own_candidate_sizing.rs @@ -0,0 +1,49 @@ +//! The accuracy budget parts of a group's own candidate, which Pass 1 and +//! Pass 2's sharing rules raise when a consumer sized for more parts finds the +//! group ([`crate::find_or_add_group::find_or_add_group`]); the memo itself +//! knows nothing of accuracy. + +use std::rc::Rc; + +use asap_types::ir::OperatorNode; +use asap_types::stages::{Memo, MemoGroupId, OriginOfCandidateForMemoGroup}; + +use crate::find_or_add_group::add_group; + +/// Sizes `group`'s own candidate (its logical node over its child groups, +/// the one `find_or_add_group` adds) for at least `accuracy_budget_parts`: +/// a consumer sized for that many parts reads the group, and the node is the +/// same, so it meets both sizings. +pub(crate) fn size_own_candidate_for( + memo: &mut Memo, + group: MemoGroupId, + accuracy_budget_parts: usize, +) { + let own = Rc::clone(&memo.groups()[group.0 as usize].logical_node); + for candidate in &mut memo.candidates[group.0 as usize] { + if Rc::ptr_eq(&candidate.root_node, &own) { + candidate.accuracy_budget_parts = + candidate.accuracy_budget_parts.max(accuracy_budget_parts); + } + } +} + +/// [`crate::find_or_add_group::find_or_add_group`], with the group found +/// sized for at least `accuracy_budget_parts` ([`size_own_candidate_for`]): +/// a consumer sized for more parts may find a group added for fewer. +pub(crate) fn find_or_add_sized_group( + memo: &mut Memo, + node: &Rc, + queries: &[usize], + origin: &OriginOfCandidateForMemoGroup, + accuracy_budget_parts: usize, +) -> MemoGroupId { + match memo.group_with_equal_node(node) { + Some(group) => { + memo.add_queries(group, queries); + size_own_candidate_for(memo, group, accuracy_budget_parts); + group + } + None => add_group(memo, node, queries, origin, accuracy_budget_parts), + } +} diff --git a/crates/logical-asap-optimization/src/pass1.rs b/crates/logical-asap-optimization/src/pass1.rs index 2f115053b..ced5d60f2 100644 --- a/crates/logical-asap-optimization/src/pass1.rs +++ b/crates/logical-asap-optimization/src/pass1.rs @@ -4,11 +4,11 @@ //! ([`PerComputationLogicalMemoCandidates`](asap_types::stages::PerComputationLogicalMemoCandidates)) //! Pass 2 extends. +/// Pass 1's candidates on Stage 0's memo groups (`try_add_pass1_candidates`). +pub(crate) mod add_pass1_candidates; /// Stage 1, Pass 1's entry and output: the per-computation candidates on /// the memo groups ([`PerComputationLogicalMemoCandidates`](asap_types::stages::PerComputationLogicalMemoCandidates)). pub(crate) mod logical_memo_candidates; -/// Pass 1's candidates on Stage 0's memo groups (`try_populate_pass1`). -pub(crate) mod memo_population; /// Stage 1, Pass 1's computation targets, their candidates and their /// composition. No execution timing is assigned: that is a Stage 2 /// materialization decision. Its target list diff --git a/crates/logical-asap-optimization/src/pass1/memo_population.rs b/crates/logical-asap-optimization/src/pass1/add_pass1_candidates.rs similarity index 70% rename from crates/logical-asap-optimization/src/pass1/memo_population.rs rename to crates/logical-asap-optimization/src/pass1/add_pass1_candidates.rs index b3b411f1a..362ace932 100644 --- a/crates/logical-asap-optimization/src/pass1/memo_population.rs +++ b/crates/logical-asap-optimization/src/pass1/add_pass1_candidates.rs @@ -12,12 +12,15 @@ //! A Pass 1 target's aggregate group gets one candidate per combination of //! its measures' Pass 1 candidates: the aggregate alone composed by //! [`super::per_computation_candidates::compose_logical_candidate`] with that choice, every target beneath it -//! exact. Its local nodes are the composed nodes above the first node -//! equal to one of the group's descendant groups, which are its children. -//! The local nodes read each child group through its node (the same -//! `Rc`), and the candidate outputs the group's schema, so a parent reads -//! it in the group's place (a query's top-k root, which no parent reads, -//! may output the ranked value beside the key). +//! exact, its summaries sized for its queries' failure probability. The +//! candidate is the composed root's operator alone: every node beneath it +//! is found or added as a group ([`find_or_add_group`](crate::find_or_add_group::find_or_add_group), by Stage 0's +//! node equality), so `Project ← SummaryEstimate ← SummaryAgg` is three +//! groups, and equal builds of different queries are one group. +//! The operator reads each child group through its node (the same `Rc`), +//! and the candidate outputs the group's schema, so a parent reads it in +//! the group's place (a query's top-k root, which no parent reads, may +//! output the ranked value beside the key). //! //! A fused computation is one more candidate of the outer group, whatever //! else reads the groups it absorbs: those keep their own candidates for @@ -41,41 +44,63 @@ use std::rc::Rc; use super::per_computation_candidates::{ compose_logical_candidate_with_accuracy_budget_parts, PerComputationCandidates, }; +use crate::added_groups::groups_read_through_added_groups; +use crate::find_or_add_group::over_child_groups; +use crate::own_candidate_sizing::size_own_candidate_for; +use asap_accuracy::sizing::{ + per_input_failure_probability, with_summaries_sized_for_failure_probability, +}; +use asap_types::ir::properties::AccuracyTarget; use asap_types::ir::schema::FieldDataType; -use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot}; +use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot, SchemaDerivationError}; +use asap_types::workload::DataWorkload; use asap_types::stages::{ CandidateForMemoGroup, DagOfLogicalMemoGroups, Memo, MemoGroup, MemoGroupId, OriginOfCandidateForMemoGroup, }; -/// A Pass 1 inventory does not belong to the initial logical memo. -#[derive(Debug, thiserror::Error, PartialEq, Eq)] -pub enum CandidatePopulationError { +/// Why adding Pass 1's candidates to the memo failed. +#[derive(Debug, thiserror::Error)] +pub enum AddPass1CandidatesError { + /// The inventory walked other roots than the memo's: bad input (an + /// inventory of another workload). #[error("Pass 1 root dependencies differ from the initial logical memo")] RootDependenciesChanged, + /// A target names no measure of an initial aggregate group: bad input. #[error("Pass 1 target {target} is not a valid measure of an initial aggregate group")] InvalidTarget { target: usize }, + /// A fused computation absorbs a target that does not exist: bad input. #[error("Pass 1 target {target} references missing absorbed target {absorbed}")] InvalidAbsorption { target: usize, absorbed: usize }, + /// A candidate's summaries, resized for its readers' failure + /// probability, did not rebuild: an internal bug. + #[error(transparent)] + Sizing(#[from] SchemaDerivationError), } -/// Populate implementation candidates without changing initial group identities. -/// The inventory must have walked `initial`'s groups -/// (`PerComputationCandidates::walked`); its targets name them by id. -pub fn try_populate_pass1( +/// Adds Pass 1's candidates to the memo of `initial`'s groups, without +/// changing the initial groups' identities. The inventory must have walked +/// `initial`'s groups (`PerComputationCandidates::walked`); its targets name +/// them by id. Each composed candidate's summaries are sized for its +/// queries' failure probability ([`size_for_failure_probability`]): query +/// `q`'s target is `accuracy_targets[q]`, over `data`; a query without a +/// target gives no failure probability. +pub fn try_add_pass1_candidates( initial: DagOfLogicalMemoGroups, per_computation: &PerComputationCandidates, -) -> Result { + accuracy_targets: &[AccuracyTarget], + data: &DataWorkload, +) -> Result { if !per_computation.walked(&initial) { - return Err(CandidatePopulationError::RootDependenciesChanged); + return Err(AddPass1CandidatesError::RootDependenciesChanged); } for (index, target) in per_computation.targets.iter().enumerate() { if !matches!( per_computation.aggregate(target).non_asap(), Some(asap_types::ir::NonASAPOp::Aggregate { measures, .. }) if target.measure_index < measures.len() ) { - return Err(CandidatePopulationError::InvalidTarget { target: index }); + return Err(AddPass1CandidatesError::InvalidTarget { target: index }); } if let Some(fused) = &target.fused_computation { if let Some(&absorbed) = fused @@ -83,35 +108,27 @@ pub fn try_populate_pass1( .iter() .find(|&&absorbed| absorbed >= per_computation.targets.len()) { - return Err(CandidatePopulationError::InvalidAbsorption { + return Err(AddPass1CandidatesError::InvalidAbsorption { target: index, absorbed, }); } } } - let candidates = pass1_candidates(per_computation, &initial.groups); - Ok(Memo { - groups: initial.groups, - query_roots: initial.query_roots, - candidates, - }) -} - -/// `g` and every group below it. -fn reach(groups: &[MemoGroup], g: MemoGroupId, reached: &mut HashSet) { - if reached.insert(g) { - for &child in &groups[g.0 as usize].child_groups { - reach(groups, child, reached); - } - } + let mut memo = Memo::new(initial); + add_candidates_to_initial_groups(per_computation, accuracy_targets, data, &mut memo)?; + Ok(memo) } -/// Every group's Pass 1 candidates, by group index. -fn pass1_candidates( +/// Gives every group of the initial memo its Pass 1 candidates, finding or +/// adding the groups their operators read. +fn add_candidates_to_initial_groups( per_computation: &PerComputationCandidates, - groups: &[MemoGroup], -) -> Vec> { + accuracy_targets: &[AccuracyTarget], + data: &DataWorkload, + memo: &mut Memo, +) -> Result<(), SchemaDerivationError> { + let initial_groups = memo.groups().len(); // Each target aggregate group's targets, one per measure: Pass 1 walks // the groups, so an aggregate several queries read has one target per // measure. @@ -123,14 +140,15 @@ fn pass1_candidates( .push(t); } // The groups' own nodes, which the composed candidates read. - let logical_nodes: Rc<[Rc]> = groups + let logical_nodes: Rc<[Rc]> = memo + .groups() .iter() .map(|group| Rc::clone(&group.logical_node)) .collect(); - let own = |g: usize| CandidateForMemoGroup { + let own = |group: &MemoGroup| CandidateForMemoGroup { fused_groups: Vec::new(), - root_node: Rc::clone(&groups[g].logical_node), - input_groups: groups[g].child_groups.clone(), + root_node: Rc::clone(&group.logical_node), + input_groups: group.child_groups.clone(), origin: OriginOfCandidateForMemoGroup::Pass1, accuracy_budget_parts: 1, }; @@ -150,60 +168,62 @@ fn pass1_candidates( let mut most_approximate_inputs = vec![0usize; per_computation.roots.len()]; for (g, targets) in &targets_of { let summaries = targets.iter().filter(|&&t| offers_a_summary(t)).count(); - for &q in &groups[g.0 as usize].queries { + for &q in &memo.groups()[g.0 as usize].queries { most_approximate_inputs[q] += summaries; } } - (0..groups.len()) - .map(|g| match targets_of.get(&MemoGroupId(g as u32)) { - None => vec![own(g)], - Some(targets) => target_candidates( - per_computation, - groups, - &logical_nodes, - MemoGroupId(g as u32), - targets, - groups[g] + let mut structural_ids = NodeToStructuralIdCache::default(); + for g in 0..initial_groups { + let candidates = match targets_of.get(&MemoGroupId(g as u32)) { + None => None, + Some(targets) => { + let max_accuracy_budget_parts = memo.groups()[g] .queries .iter() .map(|&q| most_approximate_inputs[q]) .max() .unwrap_or(1) - .max(1), - ) - .unwrap_or_else(|| vec![own(g)]), - }) - .collect() + .max(1); + target_candidates( + per_computation, + accuracy_targets, + data, + memo, + &mut structural_ids, + &logical_nodes, + MemoGroupId(g as u32), + targets, + max_accuracy_budget_parts, + )? + } + }; + memo.candidates[g] = candidates.unwrap_or_else(|| vec![own(&memo.groups()[g])]); + } + memo.rebuild_parents(); + Ok(()) } /// The candidates of target aggregate group `g`, whose measures' targets /// are `targets`: one per combination of their Pass 1 candidates (the last /// target varying fastest), each the aggregate composed alone with that -/// choice. `None` when none composes. +/// choice, its summaries sized for its queries' failure probability (query +/// `q`'s target `accuracy_targets[q]`, over `data`); the groups beneath it +/// are found or added beneath Stage 0's groups. `None` when none composes. +#[allow(clippy::too_many_arguments)] fn target_candidates( per_computation: &PerComputationCandidates, - groups: &[MemoGroup], + accuracy_targets: &[AccuracyTarget], + data: &DataWorkload, + memo: &mut Memo, + structural_ids: &mut NodeToStructuralIdCache, logical_nodes: &Rc<[Rc]>, g: MemoGroupId, targets: &[usize], max_accuracy_budget_parts: usize, -) -> Option> { - let aggregate = Rc::clone(&groups[g.0 as usize].logical_node); - // The groups a candidate's private nodes may bottom out at, by key. - let mut below = HashSet::new(); - for &child in &groups[g.0 as usize].child_groups { - reach(groups, child, &mut below); - } - let mut structural_ids = NodeToStructuralIdCache::default(); - let mut descendant: HashMap = HashMap::new(); - let mut ordered: Vec = below.into_iter().collect(); - ordered.sort_unstable(); - for d in ordered { - descendant - .entry(structural_ids.structural_id(&groups[d.0 as usize].logical_node)) - .or_insert(d); - } - let own_key = structural_ids.structural_id(&groups[g.0 as usize].logical_node); +) -> Result>, SchemaDerivationError> { + let aggregate = Rc::clone(&memo.groups()[g.0 as usize].logical_node); + let queries = memo.groups()[g.0 as usize].queries.clone(); + let own_key = structural_ids.structural_id(&aggregate); let mut candidates = Vec::new(); // Each candidate's identity in the memo: its node's structural id, its // input groups, the groups it fuses and its accuracy budget parts. A candidate @@ -242,26 +262,38 @@ fn target_candidates( } let (node, children) = if structural_ids.structural_id(&root) == own_key { ( - Rc::clone(&groups[g.0 as usize].logical_node), - groups[g.0 as usize].child_groups.clone(), + Rc::clone(&aggregate), + memo.groups()[g.0 as usize].child_groups.clone(), ) } else { - let mut children = Vec::new(); - let node = over_groups( + let root = size_for_failure_probability( &root, - groups, - &mut structural_ids, - &descendant, - &mut children, - &mut HashMap::new(), + &queries, + accuracy_budget_parts, + &per_computation.roots, + accuracy_targets, + data, + )?; + let root = without_exact_guarantees_below(&root); + let (node, children) = over_child_groups( + memo, + &root, + &queries, + &OriginOfCandidateForMemoGroup::Pass1, + accuracy_budget_parts, ); + // A consumer sized for more parts may find a group added for + // fewer. + for group in groups_read_through_added_groups(memo, &children).added_groups { + size_own_candidate_for(memo, group, accuracy_budget_parts); + } // A parent reads the candidate in the group's place, so it // outputs the group's schema when it has the same columns, // with the group's metadata (composition may derive a closed // label set where the frontend's node leaves it open). A // query's top-k root may output other columns (the ranked // value beside the key): it keeps its own. - let group_node = &groups[g.0 as usize].logical_node; + let group_node = &aggregate; let same_columns = node.schema.fields == group_node.schema.fields && node.schema.time_index == group_node.schema.time_index; let node = if !same_columns || node.schema == group_node.schema { @@ -304,7 +336,53 @@ fn target_candidates( break; } } - (!candidates.is_empty()).then_some(candidates) + Ok((!candidates.is_empty()).then_some(candidates)) +} + +/// `root`, a candidate of a group `queries` reach, with each approximate +/// input sized for `accuracy_budget_parts` parts of its queries' budget +/// (query `q`'s target `accuracy_targets[q]`, over `data`): the +/// smallest [`per_input_failure_probability`] `δ_q / (n · V_q)` among +/// `queries`, `V_q` the values `q`'s target covers (`G` of its root, +/// `roots[q]`, for the whole result) +/// ([`with_summaries_sized_for_failure_probability`]). These are the +/// candidates' final sizes: Pass 2 inherits them or takes the strictest +/// shared consumer's +/// ([`size_for_all_readers`](asap_accuracy::estimators::size_for_all_readers)), +/// and Stage 3 only checks the composed accuracy. `root` is unchanged when +/// no query gives a failure probability (exact, per-value with no `δ`, or +/// a whole-result target over an unbound `G`). +fn size_for_failure_probability( + root: &Rc, + queries: &[usize], + accuracy_budget_parts: usize, + roots: &[QueryRoot], + accuracy_targets: &[AccuracyTarget], + data: &DataWorkload, +) -> Result, SchemaDerivationError> { + let Some(failure_probability) = queries + .iter() + .filter_map(|&query| { + per_input_failure_probability( + accuracy_targets.get(query)?, + accuracy_budget_parts, + roots.get(query)?, + data, + ) + }) + .reduce(f64::min) + else { + return Ok(Rc::clone(root)); + }; + let sized = with_summaries_sized_for_failure_probability( + &[QueryRoot::Operator(Rc::clone(root))], + failure_probability, + data, + )?; + let [QueryRoot::Operator(node)] = &sized[..] else { + unreachable!("one operator root in, one out") + }; + Ok(Rc::clone(node)) } /// `aggregate` composed alone with `choice[t]` for each of its measures' @@ -366,55 +444,59 @@ fn compose_alone( ) } -/// `node` with each first node on a path down from its children that -/// equals a group in `descendant` replaced by that group's node (the same -/// `Rc`), and those groups pushed to `children` in the order met: the -/// candidate's private nodes over its child groups. -fn over_groups( - node: &Rc, - groups: &[MemoGroup], - structural_ids: &mut NodeToStructuralIdCache, - descendant: &HashMap, - children: &mut Vec, - rebuilt: &mut HashMap<*const OperatorNode, Rc>, -) -> Rc { - let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); - for child in node.children() { - let key = Rc::as_ptr(child); - if replaced.contains_key(&key) { - continue; +/// `root` with the exact guarantee removed from every node below it. Stage +/// 0's nodes carry no guarantee, and composition gives each exact node it +/// rebuilds the exact one ("exact over exact inputs"); that difference alone +/// must not split a composed node from the equal Stage 0 group +/// ([`find_or_add_group`](crate::find_or_add_group::find_or_add_group) compares guarantees). An approximate +/// guarantee is kept, so builds that differ only in their failure +/// probability or guarantee stay different groups. `root`'s own guarantee +/// is the candidate's and is kept. +fn without_exact_guarantees_below(root: &Rc) -> Rc { + fn without_exact( + node: &Rc, + done: &mut HashMap<*const OperatorNode, Rc>, + ) -> Rc { + if let Some(found) = done.get(&Rc::as_ptr(node)) { + return Rc::clone(found); } - let new = match descendant.get(&structural_ids.structural_id(child)) { - Some(&group) => { - if !children.contains(&group) { - children.push(group); - } - Rc::clone(&groups[group.0 as usize].logical_node) - } - None => match rebuilt.get(&key) { - Some(done) => Rc::clone(done), - None => { - let done = - over_groups(child, groups, structural_ids, descendant, children, rebuilt); - rebuilt.insert(key, Rc::clone(&done)); - done - } - }, + let operator = node + .operator + .map_children(|child| without_exact(child, done)); + let exact = node.guarantee.as_ref().is_some_and(|g| g.is_exact()); + let changed = exact + || node + .children() + .iter() + .zip(operator.children()) + .any(|(a, b)| !Rc::ptr_eq(a, b)); + let rebuilt = if changed { + Rc::new(OperatorNode { + operator, + guarantee: if exact { None } else { node.guarantee.clone() }, + ..(**node).clone() + }) + } else { + Rc::clone(node) }; - replaced.insert(key, new); + done.insert(Rc::as_ptr(node), Rc::clone(&rebuilt)); + rebuilt } - if node + let mut done = HashMap::new(); + let operator = root + .operator + .map_children(|child| without_exact(child, &mut done)); + if root .children() - .into_iter() - .all(|child| Rc::ptr_eq(child, &replaced[&Rc::as_ptr(child)])) + .iter() + .zip(operator.children()) + .all(|(a, b)| Rc::ptr_eq(a, b)) { - return Rc::clone(node); + return Rc::clone(root); } Rc::new(OperatorNode { - operator: node - .operator - .map_children(|child| Rc::clone(&replaced[&Rc::as_ptr(child)])), - ..(**node).clone() + operator, + ..(**root).clone() }) } @@ -520,7 +602,7 @@ mod tests { initial: DagOfLogicalMemoGroups, inventory: &PerComputationCandidates, ) -> Memo { - try_populate_pass1(initial, inventory).unwrap() + try_add_pass1_candidates(initial, inventory, &[], &Default::default()).unwrap() } fn memo_of(inventory: &PerComputationCandidates) -> Memo { @@ -553,8 +635,8 @@ mod tests { generate_per_computation_candidates(&generate_dag_of_logical_memo_groups(&other_roots)) .unwrap(); assert!(matches!( - try_populate_pass1(initial, &inventory), - Err(CandidatePopulationError::RootDependenciesChanged) + try_add_pass1_candidates(initial, &inventory, &[], &Default::default()), + Err(AddPass1CandidatesError::RootDependenciesChanged) )); } @@ -572,8 +654,8 @@ mod tests { let mut invalid_measure = inventory.clone(); invalid_measure.targets[0].measure_index = usize::MAX; assert!(matches!( - try_populate_pass1(initial.clone(), &invalid_measure), - Err(CandidatePopulationError::InvalidTarget { target: 0 }) + try_add_pass1_candidates(initial.clone(), &invalid_measure, &[], &Default::default()), + Err(AddPass1CandidatesError::InvalidTarget { target: 0 }) )); let mut invalid_absorption = inventory; let missing = invalid_absorption.targets.len(); @@ -582,8 +664,10 @@ mod tests { absorbed_targets: vec![missing], ways_to_compute: Vec::new(), }); - assert!(matches!(try_populate_pass1(initial, &invalid_absorption), - Err(CandidatePopulationError::InvalidAbsorption { target: 0, absorbed }) if absorbed == missing)); + assert!( + matches!(try_add_pass1_candidates(initial, &invalid_absorption, &[], &Default::default()), + Err(AddPass1CandidatesError::InvalidAbsorption { target: 0, absorbed }) if absorbed == missing) + ); } /// Capability pruning must change implementations, not logical identities. @@ -604,11 +688,16 @@ mod tests { let restricted_inventory = inventory.without_unsupported_summaries(Some(&[])); let restricted = populate_pass1(initial, &restricted_inventory); for memo in [&expanded, &restricted] { - assert_eq!(memo.groups.len(), original_groups.len()); - for (before, after) in original_groups.iter().zip(&memo.groups) { + // Pass 1 appends the groups it adds; the initial groups keep + // their ids and nodes, and gain only those groups as parents. + assert!(memo.groups().len() >= original_groups.len()); + for (before, after) in original_groups.iter().zip(memo.groups()) { assert!(Rc::ptr_eq(&before.logical_node, &after.logical_node)); assert_eq!(before.child_groups, after.child_groups); - assert_eq!(before.parents, after.parents); + assert_eq!(before.parents, after.parents[..before.parents.len()]); + assert!(after.parents[before.parents.len()..] + .iter() + .all(|parent| parent.0 as usize >= original_groups.len())); assert_eq!(before.queries, after.queries); } assert_eq!(memo.query_roots, original_roots); @@ -640,12 +729,18 @@ mod tests { /// walking down the groups' first children. fn group_of_kind(memo: &Memo, q: usize, kind: &str) -> usize { let mut g = memo.query_roots[q].as_operator().unwrap().0 as usize; - while memo.groups[g].logical_node.operator.kind_name() != kind { - g = memo.groups[g].child_groups[0].0 as usize; + while memo.groups()[g].logical_node.operator.kind_name() != kind { + g = memo.groups()[g].child_groups[0].0 as usize; } g } + /// The Stage 0 groups `candidate` reads, directly or through the + /// groups Pass 1 added beneath them, in the order first reached. + fn stage0_inputs(memo: &Memo, candidate: &CandidateForMemoGroup) -> Vec { + groups_read_through_added_groups(memo, &candidate.input_groups).stage0_groups + } + /// The summary algorithms a candidate's private nodes build. fn algorithms(candidate: &CandidateForMemoGroup) -> Vec { OperatorNode::reachable(&candidate.root_node) @@ -657,6 +752,102 @@ mod tests { .collect() } + /// Builds of equal `k` but a different failure probability, or a + /// different guarantee, are different groups; an equal build is the same + /// group. Removing exact guarantees before the lookup keeps a build's + /// approximate one. + #[test] + fn equal_k_builds_with_another_failure_probability_or_guarantee_are_other_groups() { + use asap_types::ir::schema::{SketchAlgorithm, SketchKind, SketchParams}; + use asap_types::ir::Operator; + let mut memo = memo( + &["quantile_over_time(0.9, m[5m])"], + AccuracyTarget::max_error(0.01).failure_probability(0.01), + ); + let build = (0..memo.groups().len()) + .find(|&g| algorithms_of(&memo.groups()[g].logical_node).contains("Kll")) + .map(|g| MemoGroupId(g as u32)) + .expect("a KLL build group"); + let node = Rc::clone(&memo.groups()[build.0 as usize].logical_node); + // The estimate over the build carries the approximate guarantee. + let estimate = (0..memo.groups().len()) + .find(|&g| { + memo.groups()[g].child_groups == [build] + && matches!( + memo.groups()[g].logical_node.asap(), + Some(ASAPOp::SummaryEstimate { .. }) + ) + }) + .map(|g| MemoGroupId(g as u32)) + .expect("an estimate over the build"); + let estimate_node = Rc::clone(&memo.groups()[estimate.0 as usize].logical_node); + let guarantee = estimate_node + .guarantee + .clone() + .expect("the estimate's guarantee"); + assert!(!guarantee.is_exact()); + let groups_before = memo.groups().len(); + let mut add = |node: OperatorNode| { + crate::find_or_add_group::find_or_add_group( + &mut memo, + &Rc::new(node), + &[0], + &OriginOfCandidateForMemoGroup::Pass1, + 1, + ) + }; + assert_eq!(add((*node).clone()), build); + let mut other_failure_probability = (*node).clone(); + let Operator::ASAP(ASAPOp::SummaryAgg { + family: FieldDataType::Sketch(kind, _), + .. + }) = &mut other_failure_probability.operator + else { + unreachable!("a sketch build") + }; + let SketchParams::Kll { k, .. } = kind.params().clone() else { + unreachable!("a KLL build") + }; + *kind = SketchKind::new( + SketchAlgorithm::Kll, + SketchParams::Kll { + k, + failure_probability: Some(0.123), + }, + ); + let added = add(other_failure_probability); + assert!(added != build && added.0 as usize >= groups_before); + assert_eq!(add((*estimate_node).clone()), estimate); + let other_guarantee = OperatorNode { + guarantee: None, + ..(*estimate_node).clone() + }; + let added = add(other_guarantee); + assert!(added != estimate && added.0 as usize >= groups_before); + // A consumer of the estimate keeps the estimate's approximate guarantee. + let reader = OperatorNode::new_shared(asap_types::ir::Operator::NonASAP( + asap_types::ir::NonASAPOp::Project { + cols: vec![], + qualifier: None, + child: Rc::clone(&estimate_node), + }, + )) + .unwrap(); + let stripped = without_exact_guarantees_below(&reader); + assert_eq!(stripped.children()[0].guarantee, Some(guarantee)); + } + + /// The summary algorithms `node` and its descendants build. + fn algorithms_of(node: &Rc) -> String { + OperatorNode::reachable(node) + .iter() + .filter_map(|node| match node.asap() { + Some(ASAPOp::SummaryAgg { family, .. }) => Some(format!("{family:?}")), + _ => None, + }) + .collect() + } + /// The per-job rate sum both queries compute is one group, reached by /// both, the root of the first and read by the second's arithmetic. #[test] @@ -670,7 +861,7 @@ mod tests { ); let sum = group_of_kind(&memo, 0, "Aggregate"); assert_eq!(group_of_kind(&memo, 1, "Aggregate"), sum); - let group = &memo.groups[sum]; + let group = &memo.groups()[sum]; assert_eq!(group.queries, vec![0, 1]); assert_eq!( asap_types::stages::queries_rooted_at(&memo.query_roots, MemoGroupId(sum as u32)), @@ -679,11 +870,11 @@ mod tests { let product = memo.query_roots[1].as_operator().unwrap().0 as usize; assert_eq!(group.parents, vec![MemoGroupId(product as u32)]); assert_eq!( - memo.groups[product].child_groups, + memo.groups()[product].child_groups, vec![MemoGroupId(sum as u32)] ); // Children come before parents. - for (g, group) in memo.groups.iter().enumerate() { + for (g, group) in memo.groups().iter().enumerate() { assert!(group.child_groups.iter().all(|c| (c.0 as usize) < g)); } } @@ -698,21 +889,43 @@ mod tests { AccuracyTarget::max_error(0.01), ); let quantile = group_of_kind(&memo, 0, "Aggregate"); - let input = memo.groups[quantile].child_groups.clone(); + let input = memo.groups()[quantile].child_groups.clone(); let candidates = &memo.candidates[quantile]; assert_eq!(candidates.len(), 3); assert!(Rc::ptr_eq( &candidates[0].root_node, - &memo.groups[quantile].logical_node + &memo.groups()[quantile].logical_node )); assert!(algorithms(&candidates[0]).is_empty()); assert!(algorithms(&candidates[1])[0].contains("Kll")); assert!(algorithms(&candidates[2])[0].contains("DDSketch")); + assert_eq!(candidates[0].input_groups, input); for candidate in candidates { assert!(candidate.fused_groups.is_empty()); - assert_eq!(candidate.input_groups, input); + assert_eq!(stage0_inputs(&memo, candidate), input); assert_eq!(candidate.origin, OriginOfCandidateForMemoGroup::Pass1); } + // Each sketch candidate is one operator over groups: `Project` over + // the estimate's group, over the build's group, over the input. + for candidate in &candidates[1..] { + let [estimate] = candidate.input_groups[..] else { + panic!("one input group"); + }; + let estimate = &memo.groups()[estimate.0 as usize]; + assert!(matches!( + estimate.logical_node.asap(), + Some(ASAPOp::SummaryEstimate { .. }) + )); + let [build] = estimate.child_groups[..] else { + panic!("one input group"); + }; + let build = &memo.groups()[build.0 as usize]; + assert!(matches!( + build.logical_node.asap(), + Some(ASAPOp::SummaryAgg { .. }) + )); + assert_eq!(build.child_groups, input); + } } /// The whole-expression top-k heap over the per-job sums implements both @@ -725,8 +938,8 @@ mod tests { AccuracyTarget::max_error(0.01).failure_probability(0.001), ); let topk = group_of_kind(&memo, 0, "Aggregate"); - let sums = memo.groups[topk].child_groups[0]; - let sums_input = memo.groups[sums.0 as usize].child_groups.clone(); + let sums = memo.groups()[topk].child_groups[0]; + let sums_input = memo.groups()[sums.0 as usize].child_groups.clone(); let candidates = &memo.candidates[topk]; assert_eq!(candidates[0].input_groups, vec![sums]); let fused: Vec<_> = candidates @@ -736,7 +949,7 @@ mod tests { assert!(!fused.is_empty()); for candidate in fused { assert_eq!(candidate.fused_groups, vec![sums]); - assert_eq!(candidate.input_groups, sums_input); + assert_eq!(stage0_inputs(&memo, candidate), sums_input); assert!(algorithms(candidate)[0].contains("Heap")); } } @@ -766,12 +979,12 @@ mod tests { .map(Rc::as_ptr) .collect(); for child in &candidate.input_groups { - let node = &memo.groups[child.0 as usize].logical_node; + let node = &memo.groups()[child.0 as usize].logical_node; assert!(reached.contains(&Rc::as_ptr(node)), "{queries:?} g{g}"); } assert_eq!( candidate.root_node.schema, - memo.groups[g].logical_node.schema + memo.groups()[g].logical_node.schema ); } } @@ -845,8 +1058,8 @@ mod tests { let target = AccuracyTarget::max_error(0.01); let memo = memo(&queries, target.clone()); let ratio = memo.query_roots[0].as_operator().unwrap().0 as usize; - let operand = memo.groups[ratio].child_groups[0].0 as usize; - for operand in &memo.groups[ratio].child_groups { + let operand = memo.groups()[ratio].child_groups[0].0 as usize; + for operand in &memo.groups()[ratio].child_groups { let shares: Vec = sketch_shares(&memo, operand.0 as usize) .into_iter() .map(|(_, share)| share) @@ -872,7 +1085,7 @@ mod tests { let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); let memo = memo(&queries, target.clone()); let difference = memo.query_roots[0].as_operator().unwrap().0 as usize; - let operand = memo.groups[difference].child_groups[0].0 as usize; + let operand = memo.groups()[difference].child_groups[0].0 as usize; let half = sketch_shares(&memo, operand) .into_iter() .find(|(family, share)| family.contains("Cms") && *share == 2) @@ -889,16 +1102,21 @@ mod tests { &["quantile_over_time(0.9, m[5m])"], AccuracyTarget::max_error(0.01), ); - assert_eq!( - memo.topological_order().unwrap(), - (0..memo.groups.len() as u32) - .map(MemoGroupId) - .collect::>() - ); + // Pass 1's added groups come after the group whose candidate + // reads them, so the order is not the index order. + let order = memo.topological_order().unwrap(); + let position = |g: &MemoGroupId| order.iter().position(|x| x == g).unwrap(); + for (g, group) in memo.groups().iter().enumerate() { + for read in group + .child_groups + .iter() + .chain(memo.candidates[g].iter().flat_map(|c| &c.input_groups)) + { + assert!(position(read) < position(&MemoGroupId(g as u32))); + } + } let quantile = group_of_kind(&memo, 0, "Aggregate"); - let appended = MemoGroupId(memo.groups.len() as u32); - memo.groups.push(memo.groups[1].clone()); - memo.candidates.push(memo.candidates[1].clone()); + let appended = memo.push_group(memo.groups()[1].clone(), memo.candidates[1].clone()); let mut reader = memo.candidates[quantile][0].clone(); reader.input_groups = vec![appended]; memo.candidates[quantile].push(reader); @@ -1065,7 +1283,7 @@ mod tests { let memo = self.memo; if let Some(child) = children .iter() - .find(|c| Rc::ptr_eq(&memo.groups[c.0 as usize].logical_node, node)) + .find(|c| Rc::ptr_eq(&memo.groups()[c.0 as usize].logical_node, node)) { return self.group(child.0 as usize, composed); } @@ -1104,7 +1322,7 @@ mod tests { let memo = memo_of(&inventory); // Each target's own and fused candidates, absorbed or not. let group_of = |t: usize| { - memo.groups + memo.groups() .iter() .position(|group| { *group.logical_node == **inventory.aggregate(&inventory.targets[t]) @@ -1136,7 +1354,7 @@ mod tests { missing.push(format!("{gauges:?} seed {seed} target {t} candidates")); } } - for (g, group) in memo.groups.iter().enumerate() { + for (g, group) in memo.groups().iter().enumerate() { for candidate in &memo.candidates[g] { if !group.parents.is_empty() && candidate.root_node.schema != group.logical_node.schema @@ -1223,9 +1441,15 @@ mod tests { &generate_per_computation_candidates(&generate_dag_of_logical_memo_groups(&roots)) .unwrap(), ); - assert_eq!(memo.groups.len(), 3); - assert_eq!(memo.groups[0].logical_node.operator.kind_name(), "Scan"); - assert_eq!(memo.groups[0].queries, vec![0, 1]); - assert_eq!(memo.groups[0].parents, vec![MemoGroupId(1), MemoGroupId(2)]); + // Three initial groups; Pass 1's added groups come after them. + assert_eq!(memo.groups()[0].logical_node.operator.kind_name(), "Scan"); + assert_eq!(memo.groups()[1].child_groups, vec![MemoGroupId(0)]); + assert_eq!(memo.groups()[2].child_groups, vec![MemoGroupId(0)]); + assert_eq!(memo.groups()[0].queries, vec![0, 1]); + assert_eq!( + memo.groups()[0].parents[..2], + [MemoGroupId(1), MemoGroupId(2)] + ); + assert!(memo.groups()[0].parents[2..].iter().all(|p| p.0 >= 3)); } } diff --git a/crates/logical-asap-optimization/src/pass1/logical_memo_candidates.rs b/crates/logical-asap-optimization/src/pass1/logical_memo_candidates.rs index 1d20e5e79..a9a00cd0f 100644 --- a/crates/logical-asap-optimization/src/pass1/logical_memo_candidates.rs +++ b/crates/logical-asap-optimization/src/pass1/logical_memo_candidates.rs @@ -2,25 +2,18 @@ //! memo groups of the workload's hash-consed DAG //! ([`PerComputationLogicalMemoCandidates`]), which Pass 2 extends. -use std::rc::Rc; - use asap_types::ir::properties::AccuracyTarget; -use asap_types::ir::{QueryRoot, SchemaDerivationError}; use asap_types::workload::DataWorkload; -use asap_accuracy::sizing::{ - per_input_failure_probability, with_summaries_sized_for_failure_probability, -}; use asap_accuracy::{AccuracyEvidenceProvider, AccuracyModel}; -use super::memo_population::try_populate_pass1; -use super::memo_population::CandidatePopulationError; +use super::add_pass1_candidates::{try_add_pass1_candidates, AddPass1CandidatesError}; use crate::deployment_capabilities::SupportedSummary; use crate::pass1::per_computation_candidates::{ generate_per_computation_candidates, PerComputationCandidateError, }; use asap_types::stages::{ - DagOfLogicalMemoGroups, Memo, PerComputationLogicalMemoCandidates, PrunedCandidate, + DagOfLogicalMemoGroups, PerComputationLogicalMemoCandidates, PrunedCandidate, }; /// Why [`generate_per_computation_logical_memo_candidates`] failed. @@ -30,13 +23,11 @@ pub enum PerComputationLogicalMemoCandidatesError { /// annotations, or a measure's accuracy target is invalid. #[error(transparent)] PerComputation(#[from] PerComputationCandidateError), - /// The targets do not belong to the DAG of logical memo groups. - #[error(transparent)] - Population(#[from] CandidatePopulationError), - /// A candidate's summaries sized for its readers' failure probability - /// did not rebuild. + /// Adding Pass 1's candidates to the memo failed + /// ([`AddPass1CandidatesError`]): the targets do not belong to the DAG of + /// logical memo groups, or a candidate's summaries did not size. #[error(transparent)] - Sizing(#[from] SchemaDerivationError), + AddPass1Candidates(#[from] AddPass1CandidatesError), } /// Stage 1, Pass 1: the per-computation candidates on the memo groups of @@ -66,8 +57,12 @@ pub fn generate_per_computation_logical_memo_candidates( per_computation = per_computation.without_summaries_no_query_can_use(accuracy_targets, data, accuracy); } - let mut memo = try_populate_pass1(dag_of_logical_memo_groups, &per_computation)?; - size_for_failure_probability(&mut memo, &per_computation.roots, accuracy_targets, data)?; + let memo = try_add_pass1_candidates( + dag_of_logical_memo_groups, + &per_computation, + accuracy_targets, + data, + )?; let pruned_candidates = per_computation .pruned_ways_to_compute .into_iter() @@ -83,55 +78,6 @@ pub fn generate_per_computation_logical_memo_candidates( }) } -/// Sizes every candidate of `memo` for its readers' failure probabilities, -/// in place: each approximate input of a candidate with `n` -/// accuracy budget parts (`accuracy_budget_parts`) gets the smallest -/// [`per_input_failure_probability`] `δ_q / (n · V_q)` among the queries `q` -/// reaching its group, `V_q` the values `q`'s target covers (`G` of its -/// root, `roots[q]`, for the whole result) -/// ([`with_summaries_sized_for_failure_probability`]). These are the candidates' final -/// sizes: Pass 2 inherits them or takes the strictest shared reader's -/// ([`size_for_all_readers`](asap_accuracy::estimators::size_for_all_readers)), -/// and Stage 3 only checks the composed accuracy. A candidate no reader -/// gives a failure probability (exact, per-value with no `δ`, or a -/// whole-result target over an unbound `G`) keeps its Pass 1 size. -fn size_for_failure_probability( - memo: &mut Memo, - roots: &[QueryRoot], - accuracy_targets: &[AccuracyTarget], - data: &DataWorkload, -) -> Result<(), PerComputationLogicalMemoCandidatesError> { - for (group, candidates) in memo.groups.iter().zip(&mut memo.candidates) { - for candidate in candidates.iter_mut() { - let Some(failure_probability) = group - .reached_by_queries() - .iter() - .filter_map(|&query| { - per_input_failure_probability( - accuracy_targets.get(query)?, - candidate.accuracy_budget_parts, - roots.get(query)?, - data, - ) - }) - .reduce(f64::min) - else { - continue; - }; - let sized = with_summaries_sized_for_failure_probability( - &[QueryRoot::Operator(Rc::clone(&candidate.root_node))], - failure_probability, - data, - )?; - let [QueryRoot::Operator(node)] = &sized[..] else { - unreachable!("one operator root in, one out") - }; - candidate.root_node = Rc::clone(node); - } - } - Ok(()) -} - #[cfg(test)] mod tests { use super::*; @@ -142,8 +88,10 @@ mod tests { use asap_accuracy::{IgnoreAccuracyGuarantees, NoAccuracyEvidence}; use asap_types::ir::properties::{AccuracyMetric, AccuracyScope, AccuracyTarget}; use asap_types::ir::schema::{FieldDataType, SketchAlgorithm, SketchParams}; + use asap_types::ir::QueryRoot; use asap_types::ir::{ASAPOp, OperatorNode}; use asap_types::stages::CandidateForMemoGroup; + use asap_types::stages::Memo; use asap_types::workload::DurationMs; /// Rank error 0.01 at δ = 0.01 over `scope`. @@ -298,7 +246,7 @@ mod tests { let found = kll_failure_probabilities( shared_summary_state_groups(&memo) .iter() - .filter(|group| memo.groups[group.0 as usize].reached_by_queries() == [0, 1]) + .filter(|group| memo.groups()[group.0 as usize].reached_by_queries() == [0, 1]) .flat_map(|group| &memo.candidates[group.0 as usize]), ); assert!(!found.is_empty()); diff --git a/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs b/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs index ea2b891ca..307ae1d89 100644 --- a/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs @@ -223,7 +223,7 @@ impl PerComputationCandidates { /// and each query's root over the same groups. Its targets' group ids /// then name `dag`'s groups. pub(crate) fn walked(&self, dag: &DagOfLogicalMemoGroups) -> bool { - use super::memo_population::NodeToStructuralIdCache; + use super::add_pass1_candidates::NodeToStructuralIdCache; let mut structural_ids = NodeToStructuralIdCache::default(); let group_of: HashMap<*const OperatorNode, MemoGroupId> = self .logical_nodes diff --git a/crates/logical-asap-optimization/src/pass2/logical_memo_candidates.rs b/crates/logical-asap-optimization/src/pass2/logical_memo_candidates.rs index 678ba5c0d..471d886e4 100644 --- a/crates/logical-asap-optimization/src/pass2/logical_memo_candidates.rs +++ b/crates/logical-asap-optimization/src/pass2/logical_memo_candidates.rs @@ -7,7 +7,7 @@ use crate::pass2::memo::rules::{ summary_results, topk, }; use crate::pass2::memo::state_groups::{ - add_shared_summary_state_groups, check_unexpanded_memo, Pass2RuleError, + add_shared_summary_state_groups, check_group_references_valid_and_acyclic, Pass2RuleError, }; use crate::pass2::memo::window_candidates::add_local_window_composition; use asap_types::stages::Memo; @@ -50,7 +50,7 @@ pub fn pass2_after_local_window_composition( /// Pass 2's rules 2–9, in order, on a memo after local window composition. pub(crate) fn apply_sharing_rules(memo: &mut Memo) -> Result<(), Pass2RuleError> { - check_unexpanded_memo(memo)?; + check_group_references_valid_and_acyclic(memo)?; // 2. Scan widening. scan_labels::add_scan_label_candidates(memo)?; // 3. Summary inputs. diff --git a/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs b/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs index d849a5e75..cb2e56eac 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs @@ -160,7 +160,7 @@ pub fn build_logical_asap_dags_from_chosen_candidates( .collect(); let mut replacements = HashMap::new(); for &child in &candidate.input_groups { - let old = &memo.groups[child.0 as usize].logical_node; + let old = &memo.groups()[child.0 as usize].logical_node; if !own_nodes.contains(&Rc::as_ptr(old)) { return Err( ErrorBuildingLogicalASAPDagsFromChosenCandidates::ChildNotInCandidate { diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/finer_aggregate.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/finer_aggregate.rs index 78aee9654..64da2407d 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/finer_aggregate.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/finer_aggregate.rs @@ -72,7 +72,7 @@ pub(crate) fn add_finer_aggregate_candidates(memo: &mut Memo) -> Result> = HashMap::new(); - for (group, entry) in memo.groups.iter().enumerate() { + for (group, entry) in memo.groups().iter().enumerate() { if let Some(aggregate) = Aggregate::of(&entry.logical_node) { buckets .entry(format!("{:?}", aggregate.measures)) @@ -96,8 +96,8 @@ pub(crate) fn add_finer_aggregate_candidates(memo: &mut Memo) -> Result Result { - let query_node = Rc::clone(&memo.groups[owner].logical_node); - let finer_node = Rc::clone(&memo.groups[finer].logical_node); + let query_node = Rc::clone(&memo.groups()[owner].logical_node); + let finer_node = Rc::clone(&memo.groups()[finer].logical_node); let (Some(query), Some(fine)) = (Aggregate::of(&query_node), Aggregate::of(&finer_node)) else { return Ok(0); }; diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs index 5eac3feb1..15f8c6c32 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs @@ -22,7 +22,7 @@ pub(in crate::pass2::memo) fn restore_inputs( // contract because the intermediate state's guarantee is part of its API. let canonical: Vec<_> = inputs .iter() - .flat_map(|id| OperatorNode::reachable(&memo.groups[id.0 as usize].logical_node)) + .flat_map(|id| OperatorNode::reachable(&memo.groups()[id.0 as usize].logical_node)) .collect(); fn visit( node: &Rc, @@ -101,7 +101,7 @@ pub(in crate::pass2::memo) fn append_producer( // Reuse only dedicated producer groups. A logical result group can contain // other implementations and cannot enforce this producer's exact contract. let root_groups = asap_types::stages::root_groups(&memo.query_roots); - if let Some((index, _)) = memo.groups.iter().enumerate().find(|(index, group)| { + if let Some((index, _)) = memo.groups().iter().enumerate().find(|(index, group)| { !root_groups.contains(&MemoGroupId(*index as u32)) && memo.candidates[*index].len() == 1 && memo.candidates[*index][0].origin == origin @@ -110,30 +110,31 @@ pub(in crate::pass2::memo) fn append_producer( }) { return Ok(MemoGroupId(index as u32)); } - let id = - MemoGroupId(u32::try_from(memo.groups.len()).map_err(|_| Pass2RuleError::TooManyGroups)?); + // The new group's id must fit a `MemoGroupId`. + u32::try_from(memo.groups().len()).map_err(|_| Pass2RuleError::TooManyGroups)?; let node = Rc::new((*node).clone()); let boundaries = memo - .groups + .groups() .iter() .enumerate() .map(|(id, group)| (Rc::as_ptr(&group.logical_node), MemoGroupId(id as u32))) .collect(); let children = boundary_groups(&node, &boundaries); - memo.groups.push(MemoGroup { - logical_node: Rc::clone(&node), - child_groups: children.clone(), - parents: Vec::new(), - queries: queries.to_vec(), - }); - memo.candidates.push(vec![CandidateForMemoGroup { - fused_groups: Vec::new(), - root_node: node, - input_groups: children, - origin, - accuracy_budget_parts, - }]); - Ok(id) + Ok(memo.push_group( + MemoGroup { + logical_node: Rc::clone(&node), + child_groups: children.clone(), + parents: Vec::new(), + queries: queries.to_vec(), + }, + vec![CandidateForMemoGroup { + fused_groups: Vec::new(), + root_node: node, + input_groups: children, + origin, + accuracy_budget_parts, + }], + )) } pub(in crate::pass2::memo) fn append_reader( @@ -155,7 +156,7 @@ pub(in crate::pass2::memo) fn append_reader( ))); } let boundaries = memo - .groups + .groups() .iter() .enumerate() .filter(|(id, _)| *id != owner) diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/rollup.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/rollup.rs index 1f684e04f..9cd78145b 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/rollup.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/rollup.rs @@ -3,7 +3,7 @@ use super::reuse::{ append_producer, append_reader, combined_inputs, eligible, replace, restore_inputs, }; use crate::pass2::asap_aware_cse::{rollup_candidate_for, rollup_pass2_candidates}; -use crate::pass2::memo::state_groups::{simple_build, Pass2RuleError}; +use crate::pass2::memo::state_groups::{build_of, BuildPlace, Pass2RuleError}; use asap_types::ir::schema::FieldDataType; use asap_types::ir::{ ASAPOp, NonASAPOp, Operator, OperatorNode, ProjectItem, QueryRoot, ScalarExpr, @@ -23,7 +23,10 @@ pub(crate) fn add_rollup_candidates(memo: &mut Memo) -> Result Result Result Result Result Result Result Result { let mut owners = Vec::new(); let mut roots = Vec::new(); + // Pass 1's intermediate groups (a summary build, its readout) are + // widened inside their consumers' candidates, down to Stage 0's groups. for (group, candidates) in memo.candidates.iter().enumerate() { + if !memo.is_stage0_group(MemoGroupId(group as u32)) { + continue; + } for (index, candidate) in candidates.iter().enumerate() { // Source-group candidates expose scan reuse to every downstream // operator, including exact/non-summary computations. Their public @@ -26,7 +34,10 @@ pub(crate) fn add_scan_label_candidates(memo: &mut Memo) -> Result Result Result Result = rewritten .children() .iter() .map(|child| { - MemoGroupId( - memo.groups - .iter() - .position(|group| Rc::ptr_eq(child, &group.logical_node)) - .expect("source-chain children were interned") - as u32, - ) + memo.group_whose_node_is(child) + .expect("source-chain children were interned") }) .filter(|input| seen.insert(*input)) .collect(); - memo.groups.push(MemoGroup { - logical_node: Rc::clone(&rewritten), - child_groups: inputs.clone(), - parents: Vec::new(), - queries: queries.to_vec(), - }); - memo.candidates.push(vec![CandidateForMemoGroup { - fused_groups: Vec::new(), - root_node: Rc::clone(&rewritten), - input_groups: inputs, - origin: OriginOfCandidateForMemoGroup::WidenedScans, - accuracy_budget_parts: 1, - }]); + memo.push_group( + MemoGroup { + logical_node: Rc::clone(&rewritten), + child_groups: inputs.clone(), + parents: Vec::new(), + queries: queries.to_vec(), + }, + vec![CandidateForMemoGroup { + fused_groups: Vec::new(), + root_node: Rc::clone(&rewritten), + input_groups: inputs, + origin: OriginOfCandidateForMemoGroup::WidenedScans, + accuracy_budget_parts: 1, + }], + ); rewritten } } else { diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/shared_configuration_readers.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/shared_configuration_readers.rs index 5d45eac0d..7746fdd32 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/shared_configuration_readers.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/shared_configuration_readers.rs @@ -1,6 +1,6 @@ //! Expand compatible readers at finite shared sizes and tumbling-window horizons before //! interning concrete states. Local candidates remain available. -use crate::stage0::cse::same_value; +use asap_types::stages::hash_consing::same_value; use std::collections::{HashMap, HashSet}; use std::rc::Rc; @@ -11,15 +11,16 @@ use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot}; use super::shared_producer::{ add_shared_producer_candidates, spans_groups, Reader, SharedProducer, }; +use crate::own_candidate_sizing::find_or_add_sized_group; use crate::pass2::asap_aware_cse::{with_resized_build, with_tumbling_window_chain}; use crate::pass2::memo::shared_configurations::{ shared_size, state_size_configurations, tumbling_window_configurations, TumblingWindowConfiguration, TumblingWindowRequirement, }; use crate::pass2::memo::state_groups::{ - boundary_ptrs, same_range_input, simple_build, Pass2RuleError, + boundary_ptrs, build_of, same_range_input, BuildPlace, Pass2RuleError, }; -use asap_types::stages::Memo; +use asap_types::stages::{Memo, OriginOfCandidateForMemoGroup}; /// `resized` (a rewrite of `original` with the same shape) with each node /// that is a child group's node in `original` put back, and the nodes above @@ -65,6 +66,8 @@ fn with_child_nodes( struct StateNeed { candidate: usize, build: Rc, + /// Where the candidate reads its build. + place: BuildPlace, merge: Option>, family: FieldDataType, /// The reader candidate's `accuracy_budget_parts`: the shares its build @@ -129,7 +132,9 @@ impl SharedProducer for SummaryStates { let mut sets: Vec>> = Vec::new(); for (group, candidates) in memo.candidates.iter().enumerate() { for (candidate, entry) in candidates.iter().enumerate() { - let Some((build, _)) = simple_build(memo, entry) else { + let Some((build, place @ (BuildPlace::Private | BuildPlace::InputGroup(_)))) = + build_of(memo, entry) + else { continue; }; // Configuration expansion precedes Stage 2. The rewrite helpers @@ -169,6 +174,7 @@ impl SharedProducer for SummaryStates { need: StateNeed { candidate, build, + place, merge, family, accuracy_budget_parts: entry.accuracy_budget_parts, @@ -296,13 +302,80 @@ impl SharedProducer for SummaryStates { with_resized_build(&roots, &need.build, family.clone()) } }; - let boundaries = boundary_ptrs(memo, &original); + let mut boundaries = boundary_ptrs(memo, &original); + // A build read as an input group is the node resized: it is not put + // back, and the resized build reads the build group's own inputs. + if let BuildPlace::InputGroup(build_group) = need.place { + boundaries.remove(&Rc::as_ptr(&need.build)); + boundaries.extend( + memo.groups()[build_group.0 as usize] + .child_groups + .iter() + .map(|child| Rc::as_ptr(&memo.groups()[child.0 as usize].logical_node)), + ); + } let roots = result.map_err(|error| Pass2RuleError::StateConfiguration(error.to_string()))?; let [QueryRoot::Operator(node)] = roots.as_slice() else { unreachable!("operator rewrite") }; let node = with_child_nodes(&original.root_node, node, &boundaries, &mut HashMap::new()); + // The rewritten candidate's sized parts: the covering's, at least the + // reader's own. + let accuracy_budget_parts = match covering { + StateCovering::TumblingWindows { + accuracy_budget_parts, + .. + } + | StateCovering::Size { + accuracy_budget_parts, + .. + } => (*accuracy_budget_parts).max(original.accuracy_budget_parts), + }; + // A build read as an input group stays one: the resized build's group + // is found or added, and the consumer reads it. + let (node, input_groups) = match need.place { + BuildPlace::Private | BuildPlace::BeneathGroups => { + (node, original.input_groups.clone()) + } + BuildPlace::InputGroup(build_group) => { + let resized_build = original + .root_node + .children() + .into_iter() + .zip(node.children()) + .find(|(before, _)| Rc::ptr_eq(before, &need.build)) + .map(|(_, after)| Rc::clone(after)) + .expect("the consumer reads its build group directly"); + let queries = memo.groups()[reader.group].queries.clone(); + // Rule 9 has no origin of its own; its added build groups + // are shared builds. + let resized_group = find_or_add_sized_group( + memo, + &resized_build, + &queries, + &OriginOfCandidateForMemoGroup::SharedBuild, + accuracy_budget_parts, + ); + let group_node = Rc::clone(&memo.groups()[resized_group.0 as usize].logical_node); + let node = super::reuse::replace( + &node, + &HashMap::from([(Rc::as_ptr(&resized_build), group_node)]), + ); + let mut input_groups = Vec::new(); + for &input in &original.input_groups { + let input = if input == build_group { + resized_group + } else { + input + }; + if !input_groups.contains(&input) { + input_groups.push(input); + } + } + (node, input_groups) + } + }; // Resizing affects internal state types, never public output fields. if original.root_node.schema.fields != node.schema.fields || original.root_node.schema.time_index != node.schema.time_index @@ -323,18 +396,8 @@ impl SharedProducer for SummaryStates { previous.push(Rc::clone(&node)); let mut candidate = original; candidate.root_node = node; - // The rewritten build is sized for the covering's shares, at least - // the reader's own. - candidate.accuracy_budget_parts = match covering { - StateCovering::TumblingWindows { - accuracy_budget_parts, - .. - } - | StateCovering::Size { - accuracy_budget_parts, - .. - } => (*accuracy_budget_parts).max(candidate.accuracy_budget_parts), - }; + candidate.input_groups = input_groups; + candidate.accuracy_budget_parts = accuracy_budget_parts; memo.candidates[reader.group].push(candidate); Ok(1) } @@ -498,7 +561,7 @@ fn compatible(left: &StateNeed, right: &StateNeed) -> bool { #[cfg(test)] mod tests { use super::*; - use crate::pass1::memo_population::try_populate_pass1; + use crate::pass1::add_pass1_candidates::try_add_pass1_candidates; use crate::pass1::per_computation_candidates::generate_per_computation_candidates; use crate::stage0::generate_dag_of_logical_memo_groups; use crate::test_support::{entry_with, lower_promql}; @@ -523,7 +586,8 @@ mod tests { let inventory = generate_per_computation_candidates(&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(); let entries = vec![ entry_with( target, @@ -606,7 +670,8 @@ mod tests { let inventory = generate_per_computation_candidates(&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(); // Pass 1's parts per KLL size. let mut parts_of_size: HashMap = HashMap::new(); for candidate in memo.candidates.iter().flatten() { diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/summary_inputs.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/summary_inputs.rs index e2bb0e84f..ba1267965 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/summary_inputs.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/summary_inputs.rs @@ -35,7 +35,7 @@ pub(crate) fn factor_summary_inputs(memo: &mut Memo) -> Result<(), Pass2RuleErro } } } - let queries = memo.groups[owner].queries.clone(); + let queries = memo.groups()[owner].queries.clone(); let mut replacements = HashMap::new(); for input in inputs { canonical_preprocessing(memo, &input, &boundaries, &queries, &mut replacements)?; @@ -62,7 +62,7 @@ fn canonical_preprocessing( // Logical TimeRange groups can remain distinct when their inputs // differ structurally. The existing state-sharing rule nevertheless // permits identical range descriptors over the same canonical source. - if let Some(group) = memo.groups.iter().find(|group| { + if let Some(group) = memo.groups().iter().find(|group| { crate::pass2::memo::state_groups::same_range_input(&group.logical_node, node) }) { replacements.insert(Rc::as_ptr(node), Rc::clone(&group.logical_node)); @@ -94,7 +94,7 @@ fn canonical_preprocessing( )?; replacements.insert( Rc::as_ptr(node), - Rc::clone(&memo.groups[group.0 as usize].logical_node), + Rc::clone(&memo.groups()[group.0 as usize].logical_node), ); Ok(true) } diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/summary_results.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/summary_results.rs index 49d84f991..42327db2d 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/summary_results.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/summary_results.rs @@ -52,7 +52,7 @@ pub(crate) fn factor_summary_results(memo: &mut Memo) -> Result Result Result>>>, Pass2RuleError> { let mut owners = HashMap::new(); let mut roots = Vec::new(); - for (group, entry) in memo.groups.iter().enumerate() { + for (group, entry) in memo.groups().iter().enumerate() { let [input] = entry.logical_node.children()[..] else { continue; }; @@ -101,7 +101,7 @@ impl SharedProducer for TopKRankings { ) -> Result { // Interning the ranking is idempotent, so every reader of the set // reads the producer the first one interned. - let Some(input_group) = memo.groups[covering.reader_groups[0]] + let Some(input_group) = memo.groups()[covering.reader_groups[0]] .child_groups .first() .copied() @@ -112,11 +112,11 @@ impl SharedProducer for TopKRankings { else { return Ok(0); }; - let input = Rc::clone(&memo.groups[exact_input.0 as usize].logical_node); + let input = Rc::clone(&memo.groups()[exact_input.0 as usize].logical_node); let queries: Vec<_> = covering .reader_groups .iter() - .flat_map(|&group| memo.groups[group].queries.iter().copied()) + .flat_map(|&group| memo.groups()[group].queries.iter().copied()) .collect::>() .into_iter() .collect(); @@ -130,7 +130,7 @@ impl SharedProducer for TopKRankings { OriginOfCandidateForMemoGroup::TopKReuse, 1, )?; - let sort = Rc::clone(&memo.groups[sort_group.0 as usize].logical_node); + let sort = Rc::clone(&memo.groups()[sort_group.0 as usize].logical_node); let ranking = covering .ranking .with_new_children(|_| Rc::clone(&sort)) @@ -142,9 +142,9 @@ impl SharedProducer for TopKRankings { OriginOfCandidateForMemoGroup::TopKReuse, 1, )?; - let ranking = Rc::clone(&memo.groups[ranking_group.0 as usize].logical_node); + let ranking = Rc::clone(&memo.groups()[ranking_group.0 as usize].logical_node); let owner = reader.group; - let logical = Rc::clone(&memo.groups[owner].logical_node); + let logical = Rc::clone(&memo.groups()[owner].logical_node); let Some(NonASAPOp::Aggregate { reduction: Reduction::Reduce(keys), measures, @@ -221,8 +221,8 @@ fn exact_input_group( break; }; replacements.insert( - Rc::as_ptr(&memo.groups[child.0 as usize].logical_node), - Rc::clone(&memo.groups[exact_child.0 as usize].logical_node), + Rc::as_ptr(&memo.groups()[child.0 as usize].logical_node), + Rc::clone(&memo.groups()[exact_child.0 as usize].logical_node), ); children.push(exact_child); } @@ -250,7 +250,7 @@ fn exact_input_group( let node = Rc::new((*candidates[0].root_node).clone().with_guarantee(Some( AccuracyGuarantee::exact("only exact implementations and exact inputs"), ))); - let queries = memo.groups[group.0 as usize].queries.clone(); + let queries = memo.groups()[group.0 as usize].queries.clone(); let id = append_producer( memo, node, diff --git a/crates/logical-asap-optimization/src/pass2/memo/shared_configurations.rs b/crates/logical-asap-optimization/src/pass2/memo/shared_configurations.rs index 9eb61a2db..5944b5306 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/shared_configurations.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/shared_configurations.rs @@ -1,10 +1,10 @@ //! Finite state configurations shared by memo rules and candidate selection. -use crate::stage0::cse::same_value; use asap_accuracy::estimators::size_for_all_readers; use asap_types::ir::schema::{ sketch::{SketchKind, SketchParams}, FieldDataType, }; +use asap_types::stages::hash_consing::same_value; use asap_types::workload::greatest_common_divisor; use std::collections::BTreeSet; diff --git a/crates/logical-asap-optimization/src/pass2/memo/state_groups.rs b/crates/logical-asap-optimization/src/pass2/memo/state_groups.rs index e0c5bf94c..44e238b19 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/state_groups.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/state_groups.rs @@ -10,8 +10,9 @@ use std::rc::Rc; use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode}; -use asap_types::stages::{CandidateForMemoGroup, Memo, OriginOfCandidateForMemoGroup}; -use asap_types::stages::{MemoGroup, MemoGroupId}; +use crate::own_candidate_sizing::find_or_add_sized_group; +use asap_types::stages::hash_consing::HashConsingTable; +use asap_types::stages::{CandidateForMemoGroup, Memo, MemoGroupId, OriginOfCandidateForMemoGroup}; /// Why Pass 2's rules failed on a memo: its precondition check, a rule /// (2–9), or the shared summary-state groups (10). @@ -25,8 +26,6 @@ pub enum Pass2RuleError { TooManyGroups, #[error("memo dependencies contain a cycle")] CyclicDependencies, - #[error("explicit summary-state groups have already been introduced")] - AlreadyExpanded, #[error("scan-label candidate expansion: {0}")] ScanLabelWidening(String), #[error("state configuration expansion: {0}")] @@ -35,73 +34,49 @@ pub enum Pass2RuleError { ReuseRule(String), } +/// A private summary build and its occurrences: equal builds (Stage 0's +/// equality) are one. struct BuildReaders { build: Rc, - children: Vec, + /// Each occurrence: the reader's group, its candidate's index there, and + /// the build node in that candidate. readers: Vec<(usize, usize, Rc)>, } /// Rejects a memo the Pass 2 rules cannot extend: mismatched or -/// out-of-range group references, shared summary-state groups already -/// added, or cyclic dependencies. Pass 2 checks this before its rules. -pub(crate) fn check_unexpanded_memo(memo: &Memo) -> Result<(), Pass2RuleError> { +/// out-of-range group references, or cyclic dependencies. Pass 2 checks +/// this before its rules. +pub(crate) fn check_group_references_valid_and_acyclic(memo: &Memo) -> Result<(), Pass2RuleError> { validate_references(memo)?; - if memo.groups.iter().any(|group| { - matches!( - group.logical_node.asap(), - Some(ASAPOp::SummaryAgg { .. } | ASAPOp::SummaryTumblingWindowChain { .. }) - ) - }) { - return Err(Pass2RuleError::AlreadyExpanded); - } if memo.topological_order().is_none() { return Err(Pass2RuleError::CyclicDependencies); } Ok(()) } -/// Pass 2's last step, after its rules: give each summary build that two or -/// more candidates make identically its own group, and add to each reader -/// group a candidate reading that group. Every existing candidate stays -/// (unsupported candidate shapes unchanged). Then rebuilds every group's -/// parents and the queries it may serve. On error, discard `memo`: it may -/// hold partial work. +/// Pass 2's last step, after its rules: give each private summary build that +/// two or more candidates make equally (Stage 0's equality) its own group, +/// and add to each reader group a candidate reading that group; a build +/// equal to an existing group's node is read from that group instead, its +/// own candidate sized for the most budget parts of its readers. Every +/// existing candidate stays (unsupported candidate shapes unchanged). Then +/// rebuilds every group's parents and the queries it may serve. On error, +/// discard `memo`: it may hold partial work. pub(crate) fn add_shared_summary_state_groups(memo: &mut Memo) -> Result<(), Pass2RuleError> { - add_shared_summary_state_groups_with_index(memo, true) -} - -/// `indexed` looks builds up by fingerprint; otherwise by a linear scan -/// (the reference the index is tested against). -fn add_shared_summary_state_groups_with_index( - memo: &mut Memo, - indexed: bool, -) -> Result<(), Pass2RuleError> { let mut builds: Vec = Vec::new(); - let mut build_index: HashMap, Vec> = HashMap::new(); + let mut equal_builds: HashConsingTable = HashConsingTable::default(); for (group, candidates) in memo.candidates.iter().enumerate() { for (index, candidate) in candidates.iter().enumerate() { - for (build, children) in direct_input_builds(memo, candidate) { - let fingerprint = build_fingerprint(&build); - let bucket = build_index.entry(fingerprint).or_default(); - let matching = if indexed { - bucket - .iter() - .copied() - .find(|&index| same_concrete_build(&builds[index].build, &build)) - } else { - builds - .iter() - .position(|entry| same_concrete_build(&entry.build, &build)) - }; - if let Some(build_index) = matching { - builds[build_index].readers.push((group, index, build)); - } else { - bucket.push(builds.len()); - builds.push(BuildReaders { - build: Rc::clone(&build), - children, - readers: vec![(group, index, build)], - }); + for (build, _) in direct_input_builds(memo, candidate) { + match equal_builds.find(&build) { + Some(entry) => builds[entry].readers.push((group, index, build)), + None => { + equal_builds.insert(Rc::clone(&build), builds.len()); + builds.push(BuildReaders { + build: Rc::clone(&build), + readers: vec![(group, index, build)], + }); + } } } } @@ -110,27 +85,6 @@ fn add_shared_summary_state_groups_with_index( let mut replacements: BTreeMap<(usize, usize), HashMap<*const OperatorNode, Rc>> = BTreeMap::new(); for build in builds { - // Repeated private builds within one result are reusable too. The - // reader list records concrete build occurrences, not just owners. - if build.readers.len() < 2 { - continue; - } - let state_group = MemoGroupId( - u32::try_from(memo.groups.len()).map_err(|_| Pass2RuleError::TooManyGroups)?, - ); - let mut queries: Vec<_> = build - .readers - .iter() - .flat_map(|(group, _, _)| memo.groups[*group].queries.iter().copied()) - .collect(); - queries.sort_unstable(); - queries.dedup(); - memo.groups.push(MemoGroup { - logical_node: Rc::clone(&build.build), - child_groups: build.children.clone(), - parents: Vec::new(), - queries, - }); // The build is its readers' own build, so it is sized for their // shares of the budget (equal builds have equal sizes). let accuracy_budget_parts = build @@ -139,20 +93,37 @@ fn add_shared_summary_state_groups_with_index( .map(|(group, index, _)| memo.candidates[*group][*index].accuracy_budget_parts) .max() .unwrap_or(1); - memo.candidates.push(vec![CandidateForMemoGroup { - fused_groups: Vec::new(), - root_node: Rc::clone(&build.build), - input_groups: build.children, - origin: OriginOfCandidateForMemoGroup::SharedBuild, + // A build no group's node equals is given a group when two or more + // candidates make it. Repeated private builds within one result are + // reusable too: the reader list records concrete build occurrences, + // not just owners. + if build.readers.len() < 2 && memo.group_with_equal_node(&build.build).is_none() { + continue; + } + let mut queries: Vec<_> = build + .readers + .iter() + .flat_map(|(group, _, _)| memo.groups()[*group].queries.iter().copied()) + .collect(); + queries.sort_unstable(); + queries.dedup(); + let state_group = find_or_add_sized_group( + memo, + &build.build, + &queries, + &OriginOfCandidateForMemoGroup::SharedBuild, accuracy_budget_parts, - }]); + ); + let node = Rc::clone(&memo.groups()[state_group.0 as usize].logical_node); for (group, index, own_build) in build.readers { replacements .entry((group, index)) .or_default() - .insert(Rc::as_ptr(&own_build), Rc::clone(&build.build)); + .insert(Rc::as_ptr(&own_build), Rc::clone(&node)); + } + if !state_groups.contains(&state_group) { + state_groups.push(state_group); } - state_groups.push(state_group); } for ((group, index), replacements) in replacements { // Factor every available producer boundary together. A producer used by @@ -166,7 +137,7 @@ fn add_shared_summary_state_groups_with_index( .input_groups .iter() .chain(&state_groups) - .map(|id| (Rc::as_ptr(&memo.groups[id.0 as usize].logical_node), *id)) + .map(|id| (Rc::as_ptr(&memo.groups()[id.0 as usize].logical_node), *id)) .collect(); reader.root_node = super::rules::reuse::replace(&reader.root_node, &replacements); reader.input_groups = boundary_groups(&reader.root_node, &boundaries); @@ -174,70 +145,35 @@ fn add_shared_summary_state_groups_with_index( memo.candidates[group].push(reader); } // Parents describe every possible group/candidate dependency, not a selected plan. - for group in &mut memo.groups { - group.parents.clear(); - } - for group in 0..memo.groups.len() { - let dependencies: HashSet<_> = memo.groups[group] - .child_groups - .iter() - .chain( - memo.candidates[group] - .iter() - .flat_map(|candidate| &candidate.input_groups), - ) - .copied() - .collect(); - for child in dependencies { - memo.groups[child.0 as usize] - .parents - .push(MemoGroupId(group as u32)); - } - } - if memo.topological_order().is_none() { - return Err(Pass2RuleError::CyclicDependencies); - } + memo.rebuild_parents(); // The chosen representative input may now serve additional queries through - // state reuse; propagate potential reach along the expanded dependency DAG. + // state reuse; propagate potential reach along the expanded dependency + // DAG, readers first. for group in memo .topological_order() .ok_or(Pass2RuleError::CyclicDependencies)? .into_iter() .rev() { - let queries = memo.groups[group.0 as usize].queries.clone(); - let children: HashSet<_> = memo.groups[group.0 as usize] - .child_groups - .iter() - .chain( - memo.candidates[group.0 as usize] - .iter() - .flat_map(|candidate| &candidate.input_groups), - ) - .copied() - .collect(); - for child in children { - memo.groups[child.0 as usize].queries.extend(&queries); - memo.groups[child.0 as usize].queries.sort_unstable(); - memo.groups[child.0 as usize].queries.dedup(); - } + let queries = memo.groups()[group.0 as usize].queries.clone(); + memo.add_queries(group, &queries); } Ok(()) } fn validate_references(memo: &Memo) -> Result<(), Pass2RuleError> { - if memo.groups.len() != memo.candidates.len() { + if memo.groups().len() != memo.candidates.len() { return Err(Pass2RuleError::MismatchedGroups); } - if memo.groups.len() > u32::MAX as usize { + if memo.groups().len() > u32::MAX as usize { return Err(Pass2RuleError::TooManyGroups); } - for group in &memo.groups { + for group in memo.groups() { if group .child_groups .iter() .chain(&group.parents) - .any(|id| id.0 as usize >= memo.groups.len()) + .any(|id| id.0 as usize >= memo.groups().len()) { return Err(Pass2RuleError::InvalidGroupReference); } @@ -247,7 +183,7 @@ fn validate_references(memo: &Memo) -> Result<(), Pass2RuleError> { .input_groups .iter() .chain(&candidate.fused_groups) - .any(|id| id.0 as usize >= memo.groups.len()) + .any(|id| id.0 as usize >= memo.groups().len()) { return Err(Pass2RuleError::InvalidGroupReference); } @@ -255,27 +191,114 @@ fn validate_references(memo: &Memo) -> Result<(), Pass2RuleError> { Ok(()) } -pub(super) fn simple_build( +/// Where a candidate reads its one summary build ([`build_of`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum BuildPlace { + /// One of the candidate's private nodes. + Private, + /// An input group of the candidate (a build Pass 1 added), which its + /// root reads directly. + InputGroup(MemoGroupId), + /// Beneath the groups the candidate reads, down to Stage 0's groups. + BeneathGroups, +} + +/// The one summary build `candidate` reads, and where, the first found of: +/// - its one private build ([`direct_input_builds`]), when the sharing steps +/// read its builds by origin; +/// - with no private build, its one input group whose node is a +/// `SummaryAgg`, which its root reads directly; +/// - the one build beneath its root down to Stage 0's groups, through the +/// groups added beneath them ([`build_down_to_stage0_groups`]). +/// +/// `None` when there is none, or more than one where it looks. +pub(super) fn build_of( + memo: &Memo, + candidate: &CandidateForMemoGroup, +) -> Option<(Rc, BuildPlace)> { + if direct_input_builds_eligible(candidate) { + let mut builds = direct_input_builds(memo, candidate); + match builds.len() { + 1 => return Some((builds.remove(0).0, BuildPlace::Private)), + 0 => { + if let Some(found) = input_group_build(memo, candidate) { + return Some(found); + } + } + _ => {} + } + } + let (build, _) = build_down_to_stage0_groups(memo, candidate)?; + Some((build, BuildPlace::BeneathGroups)) +} + +/// The only summary build beneath `candidate`'s root down to Stage 0's +/// groups, through its private nodes and the groups added beneath them +/// ([`builds_beneath`]), with the Stage 0 groups it reads; `None` unless +/// there is exactly one and it reads only Stage 0 groups. Unlike +/// [`build_of`]'s first look, a build over another private node counts: +/// scan widening rewrites the whole candidate, so it needs the only build. +pub(super) fn build_down_to_stage0_groups( memo: &Memo, candidate: &CandidateForMemoGroup, ) -> Option<(Rc, Vec)> { - let mut builds = direct_input_builds(memo, candidate); - (builds.len() == 1).then(|| builds.remove(0)) + let [(build, children)]: [_; 1] = builds_beneath(&candidate.root_node, &|node| { + memo.group_whose_node_is(node) + .filter(|&group| memo.is_stage0_group(group)) + }) + .try_into() + .ok()?; + Some((build, children?)) } -pub(super) fn direct_input_builds( +/// `candidate`'s one input group whose node is a `SummaryAgg`, which its +/// root reads directly, with that node. +fn input_group_build( memo: &Memo, candidate: &CandidateForMemoGroup, -) -> Vec<(Rc, Vec)> { - if !matches!( +) -> Option<(Rc, BuildPlace)> { + let mut builds = candidate.input_groups.iter().filter(|group| { + matches!( + memo.groups()[group.0 as usize].logical_node.asap(), + Some(ASAPOp::SummaryAgg { .. }) + ) + }); + let group = *builds.next()?; + let build = &memo.groups()[group.0 as usize].logical_node; + if builds.next().is_some() + || !candidate + .root_node + .children() + .iter() + .any(|child| Rc::ptr_eq(child, build)) + { + return None; + } + Some((Rc::clone(build), BuildPlace::InputGroup(group))) +} + +/// Whether the sharing steps read `candidate`'s builds: one made by Pass 1 +/// or a rule that composes its own build, sized for some budget parts. +fn direct_input_builds_eligible(candidate: &CandidateForMemoGroup) -> bool { + matches!( candidate.origin, OriginOfCandidateForMemoGroup::Pass1 | OriginOfCandidateForMemoGroup::TumblingWindows | OriginOfCandidateForMemoGroup::WidenedScans | OriginOfCandidateForMemoGroup::Rollup | OriginOfCandidateForMemoGroup::SummaryResult - ) || candidate.accuracy_budget_parts == 0 - { + ) && candidate.accuracy_budget_parts != 0 +} + +/// The summary builds `candidate` makes beneath its root, each with the +/// input groups it reads ([`builds_beneath`] down to its input groups): a +/// build reading another of its private nodes is left out. Empty unless +/// the candidate is one the sharing steps read by origin. +pub(super) fn direct_input_builds( + memo: &Memo, + candidate: &CandidateForMemoGroup, +) -> Vec<(Rc, Vec)> { + if !direct_input_builds_eligible(candidate) { return Vec::new(); } let boundaries: HashMap<_, _> = candidate @@ -283,33 +306,50 @@ pub(super) fn direct_input_builds( .iter() .map(|&group| { ( - Rc::as_ptr(&memo.groups[group.0 as usize].logical_node), + Rc::as_ptr(&memo.groups()[group.0 as usize].logical_node), group, ) }) .collect(); - let mut stack = vec![Rc::clone(&candidate.root_node)]; + builds_beneath(&candidate.root_node, &|node| { + boundaries.get(&Rc::as_ptr(node)).copied() + }) + .into_iter() + .filter_map(|(build, children)| Some((build, children?))) + .collect() +} + +/// The summary builds (`SummaryAgg`, `SummaryTumblingWindowChain`) beneath +/// `root`, down to the nodes `boundary` gives a group (excluded), each +/// once, with the boundary groups it reads (repeats next to each other +/// removed), or `None` when it reads a node that is not a boundary. +pub(super) fn builds_beneath( + root: &Rc, + boundary: &dyn Fn(&Rc) -> Option, +) -> Vec<(Rc, Option>)> { + let mut stack: Vec> = root.children().into_iter().cloned().collect(); let mut seen = HashSet::new(); let mut builds = Vec::new(); while let Some(node) = stack.pop() { - if boundaries.contains_key(&Rc::as_ptr(&node)) || !seen.insert(Rc::as_ptr(&node)) { + if boundary(&node).is_some() || !seen.insert(Rc::as_ptr(&node)) { continue; } - if !Rc::ptr_eq(&node, &candidate.root_node) - && matches!( - node.asap(), - Some(ASAPOp::SummaryAgg { .. } | ASAPOp::SummaryTumblingWindowChain { .. }) - ) - { + if matches!( + node.asap(), + Some(ASAPOp::SummaryAgg { .. } | ASAPOp::SummaryTumblingWindowChain { .. }) + ) { let children: Option> = node .children() .iter() - .map(|child| boundaries.get(&Rc::as_ptr(child)).copied()) + .map(|child| boundary(child)) .collect(); - if let Some(mut children) = children { - children.dedup(); - builds.push((Rc::clone(&node), children)); - } + builds.push(( + Rc::clone(&node), + children.map(|mut children| { + children.dedup(); + children + }), + )); } stack.extend(node.children().into_iter().cloned()); } @@ -346,7 +386,7 @@ pub(super) fn boundary_ptrs( candidate .input_groups .iter() - .map(|group| Rc::as_ptr(&memo.groups[group.0 as usize].logical_node)) + .map(|group| Rc::as_ptr(&memo.groups()[group.0 as usize].logical_node)) .collect() } @@ -369,61 +409,6 @@ pub(super) fn nodes_within( nodes } -// Hashing narrows comparisons only; full strict equality below remains the -// authority, including probability metadata and floating-point corner cases. -#[derive(Hash, PartialEq, Eq)] -struct BuildFingerprint { - local: String, - inputs: Vec<(usize, Option)>, -} - -fn local_fingerprint(node: &OperatorNode) -> Option { - serde_json::to_string(&( - node.operator.map_children(|_| ()), - &node.schema, - node.result_kind, - &node.guarantee, - node.timing, - node.materialization, - )) - .ok() -} - -fn build_fingerprint(node: &OperatorNode) -> Option { - let inputs = node - .children() - .into_iter() - .map(|input| { - match input.non_asap() { - Some(NonASAPOp::TimeRange { child, .. }) => { - // Exactly the existing same_range_input contract: identical - // local range semantics over the same canonical source. - Some((Rc::as_ptr(child) as usize, Some(local_fingerprint(input)?))) - } - _ => Some((Rc::as_ptr(input) as usize, None)), - } - }) - .collect::>>()?; - Some(BuildFingerprint { - local: local_fingerprint(node)?, - inputs, - }) -} - -// Immediate inputs are canonical memo nodes. Compare local fields without -// recursively walking their potentially shared DAGs; serialization additionally -// distinguishes signed-zero parameters, matching conservative CSE semantics. -fn same_concrete_build(left: &OperatorNode, right: &OperatorNode) -> bool { - let left_inputs = left.children(); - let right_inputs = right.children(); - left_inputs.len() == right_inputs.len() - && left_inputs - .iter() - .zip(right_inputs) - .all(|(left, right)| Rc::ptr_eq(left, right) || same_range_input(left, right)) - && same_local_fields(left, right) -} - // TimeRange lacks a unique row key and therefore is not generally memo-interned. // Equal ranges over the same canonical input nevertheless feed identical data // into a summary build, as the existing ASAP-aware summary reuse rule allows. @@ -471,19 +456,18 @@ fn replace_node( #[cfg(test)] mod tests { use super::*; - use crate::pass1::memo_population::try_populate_pass1; + use crate::pass1::add_pass1_candidates::try_add_pass1_candidates; use crate::pass1::per_computation_candidates::generate_per_computation_candidates; use crate::stage0::generate_dag_of_logical_memo_groups; use asap_types::ir::properties::AccuracyTarget; use asap_types::ir::QueryRoot; - /// An expanded memo with what its shared summary-state step added. + /// An expanded memo with its shared summary-state groups. struct Expanded { memo: Memo, - /// The groups the shared summary-state step added. + /// The shared summary-state groups + /// ([`crate::test_support::shared_summary_state_groups`]). state_groups: Vec, - /// The candidates it added to reader groups. - added_reader_candidates: usize, } /// Pass 2 after local window composition, over the complete tumbling-window-width @@ -494,71 +478,12 @@ mod tests { } fn expanded(memo: Memo) -> Expanded { - let state_groups = crate::test_support::shared_summary_state_groups(&memo); - // Before this step no candidate is a shared build; after it, each - // state group has one and each added reader candidate is one. - let shared_builds = memo - .candidates - .iter() - .flatten() - .filter(|candidate| candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild) - .count(); Expanded { - added_reader_candidates: shared_builds - state_groups.len(), - state_groups, + state_groups: crate::test_support::shared_summary_state_groups(&memo), memo, } } - // Hash indexing is only an acceleration: independently run the original - // linear state lookup and compare every generated candidate and state. - #[test] - fn indexed_state_interning_preserves_linear_catalog() { - let memo = with_local_tumbling_windows( - workload(&[ - ("quantile_over_time(0.95, latency[5m])", 0.01), - ("quantile_over_time(0.99, latency[15m])", 0.02), - ("quantile_over_time(0.9, latency{job=\"other\"}[5m])", 0.01), - ]), - &[6, 6, 6], - ); - let mut memo = memo; - crate::pass2::logical_memo_candidates::apply_sharing_rules(&mut memo).unwrap(); - let mut linear = memo.clone(); - let mut indexed = memo; - add_shared_summary_state_groups_with_index(&mut indexed, true).unwrap(); - add_shared_summary_state_groups_with_index(&mut linear, false).unwrap(); - let (indexed, linear) = (expanded(indexed), expanded(linear)); - assert_eq!(indexed.state_groups, linear.state_groups); - assert_eq!(indexed.memo.groups.len(), linear.memo.groups.len()); - for (left, right) in indexed.memo.candidates.iter().zip(&linear.memo.candidates) { - let signatures = |candidates: &[CandidateForMemoGroup]| { - candidates - .iter() - .map(|candidate| { - ( - candidate.input_groups.clone(), - candidate.accuracy_budget_parts, - serde_json::to_string( - &asap_types::ir::flat::flatten(&[QueryRoot::Operator(Rc::clone( - &candidate.root_node, - ))]) - .0, - ) - .unwrap(), - ) - }) - .collect::>() - }; - assert_eq!(signatures(left), signatures(right)); - } - for (left, right) in indexed.memo.groups.iter().zip(&linear.memo.groups) { - assert_eq!(left.child_groups, right.child_groups); - assert_eq!(left.parents, right.parents); - assert_eq!(left.queries, right.queries); - } - } - fn workload(queries: &[(&str, f64)]) -> Memo { workload_with_identity(queries, false) } @@ -584,29 +509,29 @@ mod tests { let inventory = generate_per_computation_candidates(&generate_dag_of_logical_memo_groups(&roots)) .unwrap(); - try_populate_pass1(initial, &inventory).unwrap() + try_add_pass1_candidates(initial, &inventory, &[], &Default::default()).unwrap() } - // Equal-sized quantiles share one canonical state while local choices remain. + // Equal-sized quantiles share one canonical state (Pass 1 adds it as a group) + // while local choices remain. #[test] fn quantile_readers_share_concrete_state_and_retain_local_candidates() { let original = workload(&[ ("quantile_over_time(0.95, latency[5m])", 0.01), ("quantile_over_time(0.99, latency[5m])", 0.01), ]); - let original_count = original.groups.len(); + let original_count = original.groups().len(); let expanded = expand(original.clone()).unwrap(); assert!(!expanded.state_groups.is_empty()); - assert!(expanded.added_reader_candidates >= 2); let order = expanded.memo.topological_order().unwrap(); for state in &expanded.state_groups { - let group = &expanded.memo.groups[state.0 as usize]; + let group = &expanded.memo.groups()[state.0 as usize]; assert_eq!(group.queries, vec![0, 1]); for input in &group.child_groups { - assert_eq!(expanded.memo.groups[input.0 as usize].queries, vec![0, 1]); + assert_eq!(expanded.memo.groups()[input.0 as usize].queries, vec![0, 1]); assert!(group.logical_node.children().iter().any(|node| Rc::ptr_eq( node, - &expanded.memo.groups[input.0 as usize].logical_node + &expanded.memo.groups()[input.0 as usize].logical_node ))); } assert_eq!(group.parents.len(), 2); @@ -635,10 +560,6 @@ mod tests { assert_eq!(before.input_groups, after.input_groups); } } - assert!(matches!( - expand(expanded.memo), - Err(Pass2RuleError::AlreadyExpanded) - )); } // A stricter concrete configuration can serve both readers; local sizes remain. @@ -660,7 +581,7 @@ mod tests { assert!(expanded .state_groups .iter() - .any(|id| expanded.memo.groups[id.0 as usize].queries == vec![0, 1])); + .any(|id| expanded.memo.groups()[id.0 as usize].queries == vec![0, 1])); } // Resizing never makes different source predicates compatible. @@ -670,10 +591,10 @@ mod tests { ("quantile_over_time(0.95, latency{job=\"a\"}[5m])", 0.01), ("quantile_over_time(0.99, latency{job=\"b\"}[5m])", 0.02), ]); - let count = original.groups.len(); + let count = original.groups().len(); let expanded = expand(original).unwrap(); assert!(expanded.state_groups.is_empty()); - assert_eq!(expanded.memo.groups.len(), count); + assert_eq!(expanded.memo.groups().len(), count); } fn with_local_tumbling_windows(mut memo: Memo, grids: &[u64]) -> Memo { @@ -682,13 +603,17 @@ mod tests { }; let boundaries: HashMap<_, _> = memo - .groups + .groups() .iter() .enumerate() .map(|(id, group)| (Rc::as_ptr(&group.logical_node), MemoGroupId(id as u32))) .collect(); - for group in 0..memo.groups.len() { - let grid = grids[memo.groups[group].queries[0]]; + // As local window composition does: Stage 0's groups only. + for group in 0..memo.groups().len() { + if !memo.is_stage0_group(MemoGroupId(group as u32)) { + continue; + } + let grid = grids[memo.groups()[group].queries[0]]; for candidate in memo.candidates[group].clone() { let roots = vec![QueryRoot::Operator(Rc::clone(&candidate.root_node))]; let tumbling_windows = tumbling_window_candidates(&roots, &[Some(grid)]); @@ -737,7 +662,7 @@ mod tests { ]), &[60_000, 60_000], ); - assert!(!original.groups.iter().any(|group| matches!( + assert!(!original.groups().iter().any(|group| matches!( group.logical_node.non_asap(), Some(NonASAPOp::Scan { .. }) ) && group @@ -767,7 +692,7 @@ mod tests { .state_groups .iter() .find(|id| { - let group = &expanded.memo.groups[id.0 as usize]; + let group = &expanded.memo.groups()[id.0 as usize]; matches!( group.logical_node.asap(), Some(ASAPOp::SummaryTumblingWindowChain { .. }) @@ -782,7 +707,7 @@ mod tests { }) }) .expect("one state over union labels"); - let state = &expanded.memo.groups[shared.0 as usize]; + let state = &expanded.memo.groups()[shared.0 as usize]; assert!(state.parents.len() >= 2); for parent in &state.parents { let reader = expanded.memo.candidates[parent.0 as usize] @@ -791,7 +716,7 @@ mod tests { .unwrap(); assert_eq!( reader.root_node.schema.fields, - expanded.memo.groups[parent.0 as usize] + expanded.memo.groups()[parent.0 as usize] .logical_node .schema .fields @@ -811,7 +736,7 @@ mod tests { assert!(state .child_groups .iter() - .all(|child| expanded.memo.groups[child.0 as usize].queries == vec![0, 1])); + .all(|child| expanded.memo.groups()[child.0 as usize].queries == vec![0, 1])); } // Widening shifts identity's position but must retain its full typed field @@ -835,7 +760,7 @@ mod tests { let state = expanded .state_groups .iter() - .map(|id| &expanded.memo.groups[id.0 as usize]) + .map(|id| &expanded.memo.groups()[id.0 as usize]) .find(|group| { matches!( group.logical_node.asap(), @@ -870,7 +795,7 @@ mod tests { candidate.input_groups.contains(&MemoGroupId( expanded .memo - .groups + .groups() .iter() .position(|group| Rc::ptr_eq(&group.logical_node, &state.logical_node)) .unwrap() as u32, @@ -879,7 +804,9 @@ mod tests { .unwrap(); assert_eq!( reader.root_node.schema, - expanded.memo.groups[parent.0 as usize].logical_node.schema + expanded.memo.groups()[parent.0 as usize] + .logical_node + .schema ); } } @@ -926,7 +853,7 @@ mod tests { ), "{queries:?}: the narrower scan is widened" ); - for group in 0..memo.groups.len() { + for group in 0..memo.groups().len() { assert!( !groups_read_by(&memo, group).contains(&group), "{queries:?}: group {group} reads itself: {:?}", @@ -937,7 +864,7 @@ mod tests { ); } let expanded = expand(workload_with_identity(&queries, true)).unwrap(); - for group in 0..expanded.memo.groups.len() { + for group in 0..expanded.memo.groups().len() { assert!(!groups_read_by(&expanded.memo, group).contains(&group)); } } @@ -947,13 +874,13 @@ mod tests { #[test] fn a_group_reading_itself_is_a_cyclic_memo() { let mut memo = workload(&[("sum by (job) (m)", 0.01)]); - assert!(check_unexpanded_memo(&memo).is_ok()); + assert!(check_group_references_valid_and_acyclic(&memo).is_ok()); let mut reader = memo.candidates[0][0].clone(); reader.input_groups = vec![MemoGroupId(0)]; memo.candidates[0].push(reader); assert_eq!(memo.topological_order(), None); assert!(matches!( - check_unexpanded_memo(&memo), + check_group_references_valid_and_acyclic(&memo), Err(Pass2RuleError::CyclicDependencies) )); } @@ -978,7 +905,7 @@ mod tests { assert!(expanded .state_groups .iter() - .all(|id| expanded.memo.groups[id.0 as usize].queries.len() == 1)); + .all(|id| expanded.memo.groups()[id.0 as usize].queries.len() == 1)); } // Real tumbling window rewrites expose one identical chain to distinct quantile readers. @@ -1004,14 +931,14 @@ mod tests { .iter() .filter(|id| { matches!( - expanded.memo.groups[id.0 as usize].logical_node.asap(), + expanded.memo.groups()[id.0 as usize].logical_node.asap(), Some(ASAPOp::SummaryTumblingWindowChain { .. }) ) }) .collect(); assert!(!chains.is_empty()); for id in chains { - let group = &expanded.memo.groups[id.0 as usize]; + let group = &expanded.memo.groups()[id.0 as usize]; assert_eq!(group.queries, vec![0, 1]); assert_eq!(group.parents.len(), 2); assert!(group.child_groups.iter().all(|input| group @@ -1020,7 +947,7 @@ mod tests { .iter() .any(|node| Rc::ptr_eq( node, - &expanded.memo.groups[input.0 as usize].logical_node + &expanded.memo.groups()[input.0 as usize].logical_node )))); } for (group, candidates) in original.candidates.iter().enumerate() { @@ -1030,10 +957,6 @@ mod tests { .any(|candidate| Rc::ptr_eq(&original.root_node, &candidate.root_node))); } } - assert!(matches!( - expand(expanded.memo), - Err(Pass2RuleError::AlreadyExpanded) - )); } // A finer grid and longer horizon serves a shorter reader through newest_tumbling_windows. @@ -1052,7 +975,7 @@ mod tests { .iter() .find(|id| { matches!( - expanded.memo.groups[id.0 as usize].logical_node.asap(), + expanded.memo.groups()[id.0 as usize].logical_node.asap(), Some(ASAPOp::SummaryTumblingWindowChain { tumbling_window_width_ms: 30_000, tumbling_window_count: 30, @@ -1061,7 +984,7 @@ mod tests { ) }) .expect("shared fine grid covering fifteen minutes"); - assert_eq!(expanded.memo.groups[state.0 as usize].queries, vec![0, 1]); + assert_eq!(expanded.memo.groups()[state.0 as usize].queries, vec![0, 1]); assert!(expanded.memo.candidates.iter().flatten().any(|reader| { reader.input_groups.contains(state) && OperatorNode::reachable(&reader.root_node) @@ -1153,7 +1076,7 @@ mod tests { ("topk by (job) (5, sum_over_time(m[1m]))", 0.01), ]); let expanded = expand(original).unwrap(); - assert!(expanded.memo.groups.iter().any(|group| matches!( + assert!(expanded.memo.groups().iter().any(|group| matches!( group.logical_node.non_asap(), Some(NonASAPOp::Limit { n: Some(5), .. }) ) && group.parents.len() >= 2)); @@ -1177,7 +1100,7 @@ mod tests { let expanded = expand(original.clone()).unwrap(); let union = expanded .memo - .groups + .groups() .iter() .enumerate() .find(|(_, group)| { @@ -1194,7 +1117,7 @@ mod tests { .map(|(id, _)| MemoGroupId(id as u32)) .expect("union scan"); let readers = original - .groups + .groups() .iter() .enumerate() .filter(|(_, group)| { @@ -1210,61 +1133,32 @@ mod tests { assert!(readers >= 2); } - // A composite average can share only its sum state with another query while - // retaining the count state locally; reader budget allocation is preserved. + // A composite average shares only its sum state with another query and + // keeps its count state: Pass 1 adds each state as its own group. #[test] - fn multiple_private_states_can_share_only_the_common_producer() { - let mut original = workload(&[("avg(m)", 0.01), ("sum(m)", 0.01)]); - let average = original.query_roots[0].as_operator().unwrap().0 as usize; - assert!(original.candidates[average] - .iter() - .any(|candidate| direct_input_builds(&original, candidate).len() >= 2)); - for candidates in &mut original.candidates { - for candidate in candidates { - candidate.accuracy_budget_parts = 2; - } - } - let expanded = expand(original.clone()).unwrap(); - assert!(!expanded.state_groups.is_empty()); - let reader = expanded.memo.candidates[average] - .iter() - .find(|candidate| { - candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild - && candidate - .input_groups - .iter() - .any(|group| expanded.state_groups.contains(group)) - }) - .expect("average has shared state input"); - assert_eq!(reader.accuracy_budget_parts, 2); - // The shared build is the readers' own, sized for two shares too. - for state in &expanded.state_groups { - assert_eq!( - expanded.memo.candidates[state.0 as usize][0].accuracy_budget_parts, - 2 - ); - } - assert!( - reader - .input_groups + fn multiple_states_share_only_the_common_state() { + let original = workload(&[("avg(m)", 0.01), ("sum(m)", 0.01)]); + let expanded = expand(original).unwrap(); + let states_of = |queries: &[usize]| { + expanded + .memo + .groups() .iter() - .any(|id| expanded.state_groups.contains(id) - && expanded.memo.groups[id.0 as usize].queries == vec![0, 1]), + .filter(|group| { + matches!(group.logical_node.asap(), Some(ASAPOp::SummaryAgg { .. })) + && group.queries == queries + }) + .count() + }; + assert!( + states_of(&[0, 1]) >= 1, "sum state is reusable by both queries" ); assert!( - reader - .input_groups - .iter() - .any(|id| expanded.state_groups.contains(id) - && expanded.memo.groups[id.0 as usize].queries == vec![0]), + states_of(&[0]) >= 1, "count state belongs only to the average query" ); - assert!(original.candidates[average] - .iter() - .all(|before| expanded.memo.candidates[average] - .iter() - .any(|after| Rc::ptr_eq(&before.root_node, &after.root_node)))); + assert!(expanded.memo.topological_order().is_some()); } // A global SQL-window quantile can consume newest tumbling windows of a longer, @@ -1312,7 +1206,7 @@ mod tests { generate_per_computation_candidates(&generate_dag_of_logical_memo_groups(&roots)) .unwrap(); let memo = with_local_tumbling_windows( - try_populate_pass1(initial, &generated).unwrap(), + try_add_pass1_candidates(initial, &generated, &[], &Default::default()).unwrap(), &[10_000, 10_000], ); let expanded = expand(memo).unwrap(); @@ -1333,7 +1227,7 @@ mod tests { }) )) && candidate.input_groups.iter().any(|group| matches!( - expanded.memo.groups[group.0 as usize].logical_node.asap(), + expanded.memo.groups()[group.0 as usize].logical_node.asap(), Some(ASAPOp::SummaryTumblingWindowChain { tumbling_window_count: 120, .. @@ -1359,10 +1253,15 @@ mod tests { fn check_private_inputs(filter: bool) { use asap_types::ir::scalar::ScalarValue; use asap_types::ir::{Operator, Predicate, ProjectItem, ScalarExpr}; - let mut memo = workload(&[ - ("quantile_over_time(0.95, latency[5m])", 0.01), - ("quantile_over_time(0.99, latency[5m])", 0.01), - ]); + // Pass 1 adds its builds as groups; the local tumbling-window + // candidates keep private builds. + let mut memo = with_local_tumbling_windows( + workload(&[ + ("quantile_over_time(0.95, latency[5m])", 0.01), + ("quantile_over_time(0.99, latency[5m])", 0.01), + ]), + &[60_000, 60_000], + ); let mut injected = 0; for owner in 0..memo.candidates.len() { let originals = memo.candidates[owner].clone(); @@ -1410,7 +1309,7 @@ mod tests { assert!(injected >= 2); let expanded = expand(memo).unwrap(); assert!(expanded.state_groups.iter().any(|state| { - let group = &expanded.memo.groups[state.0 as usize]; + let group = &expanded.memo.groups()[state.0 as usize]; group.queries == vec![0, 1] && group.child_groups.iter().any(|input| { expanded.memo.candidates[input.0 as usize] @@ -1439,13 +1338,11 @@ mod tests { .collect(); assert!(!before.is_empty(), "real generation must offer UnivMon"); crate::pass2::memo::rules::summary_inputs::factor_summary_inputs(&mut memo).unwrap(); + // Pass 1 added the UnivMon build as its own group. assert!(memo - .candidates + .groups() .iter() - .flatten() - .any(|candidate| direct_input_builds(&memo, candidate) - .iter() - .any(|(node, _)| is_univmon(node)))); + .any(|group| is_univmon(&group.logical_node))); for original in before { assert!( memo.candidates @@ -1468,7 +1365,7 @@ mod tests { }; let mut memo = workload(&[("quantile_over_time(0.95, latency[5m])", 0.01)]); let (owner, mut base, build) = memo.candidates.iter().enumerate().find_map(|(owner, candidates)| candidates.iter().find_map(|candidate| { - let (build, _) = simple_build(&memo, candidate)?; + let (build, BuildPlace::InputGroup(_)) = build_of(&memo, candidate)? else { return None }; matches!(build.asap(), Some(ASAPOp::SummaryAgg {family: FieldDataType::Sketch(kind, _), ..}) if kind.algorithm() == &SketchAlgorithm::Kll).then(|| (owner, candidate.clone(), build)) })).unwrap(); let Some(ASAPOp::SummaryAgg { @@ -1484,7 +1381,7 @@ mod tests { let boundaries: HashSet<_> = base .input_groups .iter() - .map(|id| Rc::as_ptr(&memo.groups[id.0 as usize].logical_node)) + .map(|id| Rc::as_ptr(&memo.groups()[id.0 as usize].logical_node)) .collect(); fn without_private_guarantees( node: &Rc, @@ -1548,7 +1445,11 @@ mod tests { // do not need to originate from different memo owner groups. #[test] fn repeated_states_within_one_owner_share_a_producer() { - let mut memo = workload(&[("quantile_over_time(0.95, latency[5m])", 0.01)]); + // A local tumbling-window candidate keeps its build private. + let mut memo = with_local_tumbling_windows( + workload(&[("quantile_over_time(0.95, latency[5m])", 0.01)]), + &[60_000], + ); let (owner, mut candidate) = memo .candidates .iter() @@ -1556,14 +1457,16 @@ mod tests { .find_map(|(group, candidates)| { candidates .iter() - .find(|candidate| simple_build(&memo, candidate).is_some()) + .find(|candidate| { + matches!(build_of(&memo, candidate), Some((_, BuildPlace::Private))) + }) .map(|candidate| (group, candidate.clone())) }) .unwrap(); let boundaries: HashSet<_> = candidate .input_groups .iter() - .map(|id| Rc::as_ptr(&memo.groups[id.0 as usize].logical_node)) + .map(|id| Rc::as_ptr(&memo.groups()[id.0 as usize].logical_node)) .collect(); fn duplicate( node: &Rc, @@ -1586,27 +1489,71 @@ mod tests { })) .unwrap(); candidate.root_node = Rc::clone(&combined); - memo.groups[owner].logical_node = combined; + memo.groups_mut()[owner].logical_node = combined; memo.candidates[owner] = vec![candidate]; let expanded = expand(memo).unwrap(); - assert_eq!(expanded.state_groups.len(), 1); - let state = expanded.state_groups[0]; - assert!(expanded.memo.groups[state.0 as usize] - .parents - .contains(&MemoGroupId(owner as u32))); - assert!(expanded.memo.candidates[owner] - .iter() - .any( - |candidate| candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild - && candidate.input_groups.contains(&state) - )); + assert!(expanded.state_groups.iter().any(|state| { + expanded.memo.groups()[state.0 as usize] + .parents + .contains(&MemoGroupId(owner as u32)) + && expanded.memo.candidates[owner].iter().any(|candidate| { + candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild + && candidate.input_groups.contains(state) + }) + })); + } + + /// A private build equal to a build group is read from that group, and + /// the group's own candidate is sized for the most budget parts of its + /// readers. + #[test] + fn a_private_build_equal_to_a_group_raises_its_budget_parts() { + let mut memo = workload(&[("quantile_over_time(0.95, latency[5m])", 0.01)]); + let (build_group, estimate_group) = (0..memo.groups().len()) + .find_map(|group| { + let [own] = memo.candidates[group].as_slice() else { + return None; + }; + let (_, BuildPlace::InputGroup(build)) = build_of(&memo, own)? else { + return None; + }; + matches!( + memo.groups()[build.0 as usize].logical_node.asap(), + Some(ASAPOp::SummaryAgg { .. }) + ) + .then_some((build, group)) + }) + .expect("an estimate over a build group"); + assert_eq!( + memo.candidates[build_group.0 as usize][0].accuracy_budget_parts, + 1 + ); + // The estimate over a private copy of the build, sized for 2 parts. + let build_node = Rc::clone(&memo.groups()[build_group.0 as usize].logical_node); + let private_build = Rc::new((*build_node).clone()); + let mut reader = memo.candidates[estimate_group][0].clone(); + reader.root_node = replace_node(&reader.root_node, &build_node, &private_build); + reader.input_groups = memo.groups()[build_group.0 as usize].child_groups.clone(); + reader.accuracy_budget_parts = 2; + memo.candidates[estimate_group].push(reader); + add_shared_summary_state_groups(&mut memo).unwrap(); + assert_eq!( + memo.candidates[build_group.0 as usize][0].accuracy_budget_parts, + 2 + ); + assert!(memo.candidates[estimate_group].iter().any(|candidate| { + candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild + && candidate.input_groups == [build_group] + })); } // Public expansion rejects malformed IDs before indexing dependency arrays. #[test] fn malformed_references_fail_before_dependency_traversal() { let mut memo = workload(&[("quantile_over_time(0.95, latency[5m])", 0.01)]); - memo.groups[0].child_groups.push(MemoGroupId(u32::MAX)); + memo.groups_mut()[0] + .child_groups + .push(MemoGroupId(u32::MAX)); assert!(matches!( expand(memo), Err(Pass2RuleError::InvalidGroupReference) diff --git a/crates/logical-asap-optimization/src/pass2/memo/window_candidates.rs b/crates/logical-asap-optimization/src/pass2/memo/window_candidates.rs index 468de54d5..4c3629735 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/window_candidates.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/window_candidates.rs @@ -8,8 +8,8 @@ use crate::pass2::asap_aware_cse::{ }; use asap_types::ir::flat::flatten; use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::stages::MemoGroupId as GroupId; use asap_types::stages::{CandidateForMemoGroup, Memo, OriginOfCandidateForMemoGroup}; -use asap_types::stages::{MemoGroup, MemoGroupId as GroupId}; use asap_types::workload::QueryWorkloadEntry; #[cfg(test)] use std::collections::BTreeSet; @@ -55,9 +55,13 @@ fn local_tumbling_window_candidates( memo: &Memo, entries: &[QueryWorkloadEntry], ) -> Vec<(usize, CandidateForMemoGroup)> { - let group_of = group_of_node(&memo.groups); let mut own: Vec<(usize, CandidateForMemoGroup)> = Vec::new(); - for (g, group) in memo.groups.iter().enumerate() { + for (g, group) in memo.groups().iter().enumerate() { + // Pass 1's intermediate groups (a summary build, its readout) are + // rewritten inside their consumers' candidates instead. + if !memo.is_stage0_group(GroupId(g as u32)) { + continue; + } let grids = tumbling_window_width_domain( group .queries @@ -79,7 +83,7 @@ fn local_tumbling_window_candidates( matches!( node.non_asap(), Some(asap_types::ir::NonASAPOp::Scan { .. }) - ) && group_of.contains_key(&Rc::as_ptr(node)) + ) && memo.group_whose_node_is(node).is_some() }) .collect::>(); let roots = [QueryRoot::Operator(Rc::clone(&candidate.root_node))]; @@ -122,7 +126,7 @@ fn local_tumbling_window_candidates( own.push(( g, rewritten_candidate( - &group_of, + memo, GroupId(g as u32), candidate, node, @@ -171,7 +175,10 @@ fn restore_scan_boundaries( // compatible subset: common widths belong to each reader's divisor domain. // Nongridded potential consumers do not veto a repeated reader's tumbling windows. #[cfg(test)] -fn compatible_reader_grids(group: &MemoGroup, entries: &[QueryWorkloadEntry]) -> Vec { +fn compatible_reader_grids( + group: &asap_types::stages::MemoGroup, + entries: &[QueryWorkloadEntry], +) -> Vec { super::shared_configurations::tumbling_window_width_domain( group .queries @@ -180,17 +187,9 @@ fn compatible_reader_grids(group: &MemoGroup, entries: &[QueryWorkloadEntry]) -> ) } -fn group_of_node(groups: &[MemoGroup]) -> HashMap<*const OperatorNode, usize> { - groups - .iter() - .enumerate() - .map(|(g, group)| (Rc::as_ptr(&group.logical_node), g)) - .collect() -} - /// `from`, a candidate of group `owner`, with root `node` and `origin`. fn rewritten_candidate( - group_of: &HashMap<*const OperatorNode, usize>, + memo: &Memo, owner: GroupId, from: &CandidateForMemoGroup, node: Rc, @@ -199,7 +198,7 @@ fn rewritten_candidate( let implements: Vec = std::iter::once(owner) .chain(from.fused_groups.iter().copied()) .collect(); - let reached = |node: &Rc| child_groups(group_of, node, &implements); + let reached = |node: &Rc| child_groups(memo, node, &implements); let children = if reached(&from.root_node) == from.input_groups { reached(&node) } else { @@ -214,14 +213,10 @@ fn rewritten_candidate( } } -fn child_groups( - group_of: &HashMap<*const OperatorNode, usize>, - node: &Rc, - implements: &[GroupId], -) -> Vec { +fn child_groups(memo: &Memo, node: &Rc, implements: &[GroupId]) -> Vec { fn walk( node: &Rc, - group_of: &HashMap<*const OperatorNode, usize>, + memo: &Memo, implements: &[GroupId], seen: &mut HashSet<*const OperatorNode>, out: &mut Vec, @@ -230,18 +225,18 @@ fn child_groups( if !seen.insert(Rc::as_ptr(child)) { continue; } - match group_of.get(&Rc::as_ptr(child)) { - Some(&g) if !implements.contains(&GroupId(g as u32)) => { - if !out.contains(&GroupId(g as u32)) { - out.push(GroupId(g as u32)); + match memo.group_whose_node_is(child) { + Some(g) if !implements.contains(&g) => { + if !out.contains(&g) { + out.push(g); } } - _ => walk(child, group_of, implements, seen, out), + _ => walk(child, memo, implements, seen, out), } } } let mut out = Vec::new(); - walk(node, group_of, implements, &mut HashSet::new(), &mut out); + walk(node, memo, implements, &mut HashSet::new(), &mut out); out } @@ -252,7 +247,7 @@ fn candidate_signature(candidate: &CandidateForMemoGroup) -> Result; - -/// The operator with every child (operator inputs and the operator nodes -/// referenced from its scalar expressions alike) replaced by `()`. What -/// remains is the node's own data: variant tag, scalar expressions, -/// parameters. -fn own_fields(node: &OperatorNode) -> Operator<()> { - node.operator.map_children(|_| ()) -} - -/// Coarse structural hash used only to bucket [`InternTable::intern`]'s -/// candidate search — never the sharing decision ([`same_node`] is). -/// -/// `OperatorNode` carries `f64`s (`ScalarValue::Float64`, quantile targets, -/// `AccuracyGuarantee` bounds, …), so it cannot derive `std::hash::Hash`. The -/// hash is SipHash over two parts: -/// -/// 1. the canonical JSON of [`own_fields`] plus `result_kind`, `schema`, -/// `guarantee`, `timing` and `materialization` — every field `PartialEq` compares except the -/// children. A scalar expression is serialized as data with each operator -/// node it reads replaced by `()`, so a reference to an -/// interned sub-DAG contributes nothing of its own here; -/// 2. for every child in [`OperatorNode::children`] order (operator inputs, -/// then scalar-referenced nodes), the child's own `structural_hash`, -/// memoized in `cache` by `Rc` pointer identity. -/// -/// Part 2 is what makes equal sub-DAGs hash equal whether they are reached -/// through an operator input or through a `scalar(v)`, and what keeps the -/// pass linear: a node is generally a DAG, and re-serializing a shared -/// descendant once per parent would cost `O(sub-DAG)` per node instead of -/// `O(1)` beyond the children's already-known hashes. A non-finite `f64` -/// serializes as `null`, merely widening one (still equality-checked) bucket. -pub(crate) fn structural_hash(node: &OperatorNode, cache: &mut NodeToStructuralHashCache) -> u64 { - fn child_hash(child: &Rc, cache: &mut NodeToStructuralHashCache) -> u64 { - let ptr = Rc::as_ptr(child); - if let Some(&h) = cache.get(&ptr) { - return h; - } - let h = structural_hash(child, cache); - cache.insert(ptr, h); - h - } - - let mut hasher = DefaultHasher::new(); - let own = ( - own_fields(node), - node.result_kind, - &node.schema, - &node.guarantee, - node.timing, - node.materialization, - ); - serde_json::to_string(&own) - .unwrap_or_default() - .hash(&mut hasher); - for child in node.children() { - child_hash(child, cache).hash(&mut hasher); - } - hasher.finish() -} - -/// Numeric `PartialEq` alone conflates signed zeros. The serialized check is -/// additional evidence, never a replacement for typed equality (JSON maps -/// non-finite floats to `null`). Used for the guarantee, whose bounds are -/// floats a shared node must preserve bit-for-bit, and for summary families, -/// whose `SketchParams` equality omits the recorded error and probability. -pub(crate) fn same_value(left: &T, right: &T) -> bool { - left == right - && match (serde_json::to_string(left), serde_json::to_string(right)) { - (Ok(left), Ok(right)) => left == right, - _ => false, - } -} - -/// Cache of child-pair comparisons already decided by [`same_node`], keyed by -/// pointer pair. Only interned (table-owned, hence alive) nodes are keys. -type NodePairToIsSameNodeCache = HashMap<(*const OperatorNode, *const OperatorNode), bool>; - -/// The sharing decision: typed equality of two nodes. -/// -/// `OperatorNode`'s derived `PartialEq` would recurse into children by value -/// even when both sides hold the same `Rc` (`OperatorNode` is not `Eq`, so -/// `Rc` gets no pointer shortcut), expanding a shared diamond once per path. -/// Children are therefore compared by pointer first; only when the pointers -/// differ (an equal child that was not legal to share) are the values -/// compared, memoized per pair so a diamond is still walked once. -fn same_node( - left: &OperatorNode, - right: &OperatorNode, - cache: &mut NodePairToIsSameNodeCache, -) -> bool { - let (lc, rc) = (left.children(), right.children()); - if lc.len() != rc.len() { - return false; - } - let children_equal = lc.iter().zip(&rc).all(|(a, b)| { - if Rc::ptr_eq(a, b) { - return true; - } - let key = (Rc::as_ptr(a), Rc::as_ptr(b)); - if let Some(&eq) = cache.get(&key) { - return eq; - } - let eq = same_node(a, b, cache); - cache.insert(key, eq); - eq - }); - children_equal - && left.result_kind == right.result_kind - && left.schema == right.schema - && left.timing == right.timing - && left.materialization == right.materialization - && same_value(&left.guarantee, &right.guarantee) - && same_value(&own_fields(left), &own_fields(right)) -} - -/// Bottom-up hash-consing table: structurally-equal, sharing-legal nodes -/// collapse onto one `Rc`. -/// -/// `buckets` is keyed by [`structural_hash`] — a coarse candidate filter -/// only. Every entry within one bucket is a full node kept around for the -/// [`same_node`] comparison that actually decides a match; a hash collision -/// between structurally different nodes just means a harmless linear scan of -/// a few extra candidates. -struct InternTable { - buckets: HashMap>>, - /// Persisted for the table's whole lifetime so hashing is `O(1)` per node - /// beyond its children; every cached node is owned by `buckets`. - structural_hash_cache: NodeToStructuralHashCache, - same_node_cache: NodePairToIsSameNodeCache, -} - -impl InternTable { - fn new() -> Self { - Self { - buckets: HashMap::new(), - structural_hash_cache: HashMap::new(), - same_node_cache: HashMap::new(), - } - } - - /// Intern one node whose children are already interned: look it up by - /// [`structural_hash`], confirm with [`same_node`], and — only when - /// sharing is legal (module doc, "Legality") — return the existing `Rc` - /// instead of allocating a new one. - fn intern(&mut self, node: OperatorNode) -> Rc { - let hash = structural_hash(&node, &mut self.structural_hash_cache); - // A node that is not legal to share is never *returned* as a match - // for something else; it still occupies a fresh slot in the bucket - // (harmless: later scans require legality of the new node too). - let reusable = node.is_asap() - || node - .non_asap() - .is_some_and(NonASAPOp::is_shareable_when_equal); - let bucket = self.buckets.entry(hash).or_default(); - if reusable { - if let Some(existing) = bucket - .iter() - .find(|candidate| same_node(candidate, &node, &mut self.same_node_cache)) - { - return Rc::clone(existing); - } - } - let rc = Rc::new(node); - bucket.push(Rc::clone(&rc)); - rc - } -} +use asap_types::ir::operator::node::OperatorNode; +use asap_types::stages::hash_consing::HashConsingTable; /// Input pointer → (input `Rc`, interned result). The input `Rc` is retained /// so its address cannot be freed and reused by a fresh allocation while the @@ -247,7 +71,7 @@ type NodeToInternedNodeCache = HashMap<*const OperatorNode, (Rc, R /// every child is replaced by an equal node, so each derived property stays /// valid, and the result is `PartialEq`-equal to the input. fn intern_bottom_up( - table: &mut InternTable, + table: &mut HashConsingTable>, visited: &mut NodeToInternedNodeCache, node: &Rc, ) -> Rc { @@ -261,7 +85,14 @@ fn intern_bottom_up( .with_guarantee(node.guarantee.clone()) .with_timing(node.timing) .with_materialization(node.materialization); - let interned = table.intern(rebuilt); + let rebuilt = Rc::new(rebuilt); + let interned = match table.find(&rebuilt) { + Some(existing) => existing, + None => { + table.insert(Rc::clone(&rebuilt), Rc::clone(&rebuilt)); + rebuilt + } + }; visited.insert(Rc::as_ptr(node), (Rc::clone(node), Rc::clone(&interned))); interned } @@ -276,7 +107,7 @@ fn intern_bottom_up( pub fn share_common_sub_dags( roots: Vec<(Id, Rc)>, ) -> Vec<(Id, Rc)> { - let mut table = InternTable::new(); + let mut table = HashConsingTable::default(); let mut visited = NodeToInternedNodeCache::new(); roots .into_iter() @@ -289,6 +120,7 @@ mod tests { use super::*; use asap_types::ir::operator::agg_intent::AggIntent; use asap_types::ir::operator::asap::ASAPOp; + use asap_types::ir::operator::node::Operator; use asap_types::ir::operator::operator_properties::{ BinaryOpKind, GroupKeys, Reduction, Source, }; @@ -577,32 +409,6 @@ mod tests { ); } - #[test] - fn structural_hash_sees_through_a_scalar_reference() { - // Two equal bridges must hash equal whether or not their referenced - // vector is the same Rc — the reference contributes the vector's - // memoized hash, not its identity. - let a = scalar_of_vector(); - let b = scalar_of_vector(); - let mut cache = HashMap::new(); - assert_eq!( - structural_hash(&a, &mut cache), - structural_hash(&b, &mut cache) - ); - assert_eq!( - cache.len(), - 4, - "aggregate + scan cached once per root: {cache:?}" - ); - let other = node(NonASAPOp::PromqlVectorFromScalar( - ScalarExpr::PromqlScalarFromVector(quantile_agg(vec![1], Some(2), 0.5)), - )); - assert_ne!( - structural_hash(&a, &mut cache), - structural_hash(&other, &mut cache) - ); - } - // ── ASAP nodes ────────────────────────────────────────────────────── fn summary_agg(alpha: f64, guarantee: Option) -> Rc { @@ -670,48 +476,6 @@ mod tests { assert!(Rc::ptr_eq(&producer(&p95), &producer(&p99))); } - // ── structural_hash (DAG-aware memoization) ───────────────────────── - - #[test] - fn structural_hash_is_stable_across_cache_states() { - let agg = quantile_agg(vec![1], Some(2), 0.5); - let mut cold = HashMap::new(); - let mut warm = HashMap::new(); - structural_hash(&scan(), &mut warm); - assert_eq!( - structural_hash(&agg, &mut cold), - structural_hash(&agg, &mut warm), - "hash must be independent of unrelated cache state" - ); - } - - #[test] - fn structural_hash_of_an_internally_shared_dag_matches_the_unshared_equivalent() { - let agg = quantile_agg(vec![1], Some(2), 0.5); - let shared_root = compare(Rc::clone(&agg), Rc::clone(&agg)); - let unshared_root = compare( - quantile_agg(vec![1], Some(2), 0.5), - quantile_agg(vec![1], Some(2), 0.5), - ); - assert_eq!( - structural_hash(&shared_root, &mut HashMap::new()), - structural_hash(&unshared_root, &mut HashMap::new()), - ); - } - - #[test] - fn structural_hash_memoizes_a_shared_descendant_exactly_once() { - let agg = quantile_agg(vec![1], Some(2), 0.5); - let root = compare(Rc::clone(&agg), Rc::clone(&agg)); - let mut cache = HashMap::new(); - structural_hash(&root, &mut cache); - assert_eq!( - cache.len(), - 2, - "one entry per unique node in the shared branch (Aggregate + Scan): {cache:?}" - ); - } - /// `DISTINCT` (every column) returns the same rows each time; a `Dedup` /// on some columns keeps an unspecified row per key. #[test] diff --git a/crates/logical-asap-optimization/src/test_support.rs b/crates/logical-asap-optimization/src/test_support.rs index 988cb4210..78abe347c 100644 --- a/crates/logical-asap-optimization/src/test_support.rs +++ b/crates/logical-asap-optimization/src/test_support.rs @@ -7,7 +7,8 @@ use std::rc::Rc; // Stage 1 internals other crates' tests call: Pass 1's working form and its // population step, single-candidate composition, Stage 0 hash-consing, and // the Pass 2 rule entry points those tests build memos with. -pub use crate::pass1::memo_population::try_populate_pass1; +pub use crate::added_groups::groups_read_through_added_groups; +pub use crate::pass1::add_pass1_candidates::try_add_pass1_candidates; pub use crate::pass1::per_computation_candidates::{ compose_logical_candidate, generate_per_computation_candidates, PerComputationCandidates, }; @@ -39,24 +40,41 @@ pub fn sharing_rules_and_shared_state_groups( Ok(memo) } -/// The groups `add_shared_summary_state_groups` -/// added to an expanded memo, in order: each has one candidate, the -/// shared build (origin `SharedBuild`); no other group's candidates are -/// only that. +/// The shared summary-state groups of an expanded memo, in order: each +/// group `add_shared_summary_state_groups` added (one candidate, the shared +/// build, origin `SharedBuild`), and each summary build Pass 1 added +/// (logical node `SummaryAgg` or `SummaryTumblingWindowChain`) whose result +/// candidates of two or more groups read. pub fn shared_summary_state_groups( memo: &asap_types::stages::Memo, ) -> Vec { - memo.candidates - .iter() - .enumerate() - .filter(|(_, candidates)| { - matches!( - candidates.as_slice(), - [only] if only.origin - == asap_types::stages::OriginOfCandidateForMemoGroup::SharedBuild - ) - }) - .map(|(group, _)| asap_types::stages::MemoGroupId(group as u32)) + use asap_types::ir::ASAPOp; + use asap_types::stages::{MemoGroupId, OriginOfCandidateForMemoGroup}; + let added_by_the_shared_step = |group: usize| { + matches!( + memo.candidates[group].as_slice(), + [only] if only.origin == OriginOfCandidateForMemoGroup::SharedBuild + ) + }; + let added_and_shared = |group: usize| { + let id = MemoGroupId(group as u32); + matches!( + memo.groups()[group].logical_node.asap(), + Some(ASAPOp::SummaryAgg { .. } | ASAPOp::SummaryTumblingWindowChain { .. }) + ) && memo + .candidates + .iter() + .filter(|candidates| { + candidates + .iter() + .any(|candidate| candidate.input_groups.contains(&id)) + }) + .count() + >= 2 + }; + (0..memo.groups().len()) + .filter(|&group| added_by_the_shared_step(group) || added_and_shared(group)) + .map(|group| MemoGroupId(group as u32)) .collect() } diff --git a/crates/physical-asap-optimization/src/catalog/inventory.rs b/crates/physical-asap-optimization/src/catalog/inventory.rs index cc8daf559..78f536498 100644 --- a/crates/physical-asap-optimization/src/catalog/inventory.rs +++ b/crates/physical-asap-optimization/src/catalog/inventory.rs @@ -105,12 +105,12 @@ pub fn build_with_compatible_producers_cache( compatible_producers_cache: &mut MemoGroupOutputToCompatibleProducersCache, ) -> Result { let invalid = |message: &str| CandidateInventoryError::InvalidMemo(message.into()); - if memo.groups.len() != memo.candidates.len() { + if memo.groups().len() != memo.candidates.len() { return Err(invalid("group and candidate vector lengths differ")); } if output_groups .iter() - .any(|group| group.0 as usize >= memo.groups.len()) + .any(|group| group.0 as usize >= memo.groups().len()) { return Err(invalid("query output references missing group")); } @@ -118,7 +118,7 @@ pub fn build_with_compatible_producers_cache( operator_nodes_of_candidates_for_memo_groups: Vec::new(), producer_nodes_per_input: Vec::new(), query_output_nodes: Vec::new(), - implementation_roots_of_group: vec![Vec::new(); memo.groups.len()], + implementation_roots_of_group: vec![Vec::new(); memo.groups().len()], }; for (group, candidates) in memo.candidates.iter().enumerate() { for (index, candidate) in candidates.iter().enumerate() { @@ -133,7 +133,7 @@ pub fn build_with_compatible_producers_cache( if implemented.len() != 1 + candidate.fused_groups.len() || implemented .iter() - .any(|id| id.0 as usize >= memo.groups.len()) + .any(|id| id.0 as usize >= memo.groups().len()) { return Err(invalid( "candidate names duplicate or missing absorbed groups", @@ -151,14 +151,14 @@ pub fn build_with_compatible_producers_cache( if candidate .input_groups .iter() - .any(|child| child.0 as usize >= memo.groups.len()) + .any(|child| child.0 as usize >= memo.groups().len()) { return Err(invalid("candidate references missing input group")); } let inputs = candidate .input_groups .iter() - .map(|child| Rc::clone(&memo.groups[child.0 as usize].logical_node)) + .map(|child| Rc::clone(&memo.groups()[child.0 as usize].logical_node)) .collect::>(); if inputs .iter() @@ -195,7 +195,7 @@ pub fn build_with_compatible_producers_cache( for (index, candidate) in candidates.iter().enumerate() { let mut boundaries = HashMap::new(); for &child in &candidate.input_groups { - let pointer = Rc::as_ptr(&memo.groups[child.0 as usize].logical_node); + let pointer = Rc::as_ptr(&memo.groups()[child.0 as usize].logical_node); if boundaries .insert(pointer, child) .is_some_and(|previous| previous != child) @@ -242,7 +242,7 @@ pub fn build_with_compatible_producers_cache( } } for group in output_groups { - let expected = &memo.groups[group.0 as usize].logical_node; + let expected = &memo.groups()[group.0 as usize].logical_node; let mut compatible_candidates = Vec::new(); for &id in &inventory.implementation_roots_of_group[group.0 as usize] { let template = diff --git a/crates/physical-asap-optimization/src/catalog/realize.rs b/crates/physical-asap-optimization/src/catalog/realize.rs index f473709c1..5a9e7426c 100644 --- a/crates/physical-asap-optimization/src/catalog/realize.rs +++ b/crates/physical-asap-optimization/src/catalog/realize.rs @@ -215,7 +215,7 @@ mod tests { use super::*; 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::test_support::{ entry, lower_promql, sharing_rules_and_shared_state_groups, }; @@ -272,7 +272,8 @@ mod tests { roots[0].clone() ])) .unwrap(); - let memo = try_populate_pass1(initial, &per_computation).unwrap(); + let memo = + try_add_pass1_candidates(initial, &per_computation, &[], &Default::default()).unwrap(); let logical = LogicalASAPMemoCandidates { memo: sharing_rules_and_shared_state_groups(memo).unwrap(), }; diff --git a/crates/planner/src/run.rs b/crates/planner/src/run.rs index 5069ca1de..f1ec8b489 100644 --- a/crates/planner/src/run.rs +++ b/crates/planner/src/run.rs @@ -182,7 +182,7 @@ pub fn run_planning( let logical = logical_asap_candidates.insert( logical.map_err(|error| SelectionError::LogicalASAPCandidates(error.to_string()))?, ); - report.memo_groups = logical.memo.groups.len(); + report.memo_groups = logical.memo.groups().len(); report.logical_candidates = logical.memo.candidates.iter().map(Vec::len).sum(); // Stage 2. let started = Instant::now(); diff --git a/crates/types/src/stages/candidates.rs b/crates/types/src/stages/candidates.rs index 8d18a64bd..80e4a6c13 100644 --- a/crates/types/src/stages/candidates.rs +++ b/crates/types/src/stages/candidates.rs @@ -1,22 +1,42 @@ //! Implementation candidates attached to memo groups. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::rc::Rc; -use super::{MemoGroup, MemoGroupId}; +use super::hash_consing::HashConsingTable; +use super::{DagOfLogicalMemoGroups, MemoGroup, MemoGroupId}; use crate::ir::{OperatorNode, QueryRoot}; /// The memo: the memo groups, each query's root over them, and each group's /// candidate implementations. #[derive(Debug, Clone, Default)] pub struct Memo { - pub groups: Vec, + /// The groups, by id ([`Self::groups`]); private so every group added + /// ([`Self::push_group`]) is in [`Self::group_lookup`]. + groups: Vec, /// Each query's root over groups, by index in the workload (the /// initial logical memo's query roots); groups added later are never a /// query's root. pub query_roots: Vec>, /// Each group's implementation candidates, by group index. pub candidates: Vec>, + /// Every group, by its logical node: [`Self::group_whose_node_is`] and + /// [`Self::group_with_equal_node`] look groups up in it. + group_lookup: MemoGroupLookup, + /// The number of groups Stage 0 produced from the queries' DAG + /// ([`Self::new`]): groups `0..stage0_group_count`. Groups added later + /// come after them ([`Self::is_stage0_group`]). + stage0_group_count: usize, +} + +/// Every group of a memo, by the address of its logical node and by node +/// equality ([`HashConsingTable`], Stage 0's equality). +#[derive(Debug, Clone, Default)] +struct MemoGroupLookup { + /// Each group's logical node, by address (the memo keeps it alive). + group_whose_node_is: HashMap<*const OperatorNode, MemoGroupId>, + /// Each group's logical node, by equality. + group_with_equal_node: HashConsingTable, } /// One way to implement a group: its node as Stage 2 implements it, over @@ -47,7 +67,8 @@ pub struct CandidateForMemoGroup { /// Where an implementation candidate came from. #[derive(Debug, Clone, PartialEq, Eq)] pub enum OriginOfCandidateForMemoGroup { - /// Pass 1: exact, exact accumulator, or a summary family and size. + /// Pass 1: exact, exact accumulator, or a summary family and size; or + /// the operator of an intermediate result Pass 1 added as a group. Pass1, /// One build serving several reader groups (summary capability). SharedBuild, @@ -67,6 +88,144 @@ pub enum OriginOfCandidateForMemoGroup { } impl Memo { + /// The memo of Stage 0's groups and query roots, every group without + /// candidates yet. + pub fn new(dag: DagOfLogicalMemoGroups) -> Self { + let mut memo = Self { + groups: Vec::with_capacity(dag.groups.len()), + query_roots: dag.query_roots, + candidates: vec![Vec::new(); dag.groups.len()], + group_lookup: MemoGroupLookup::default(), + stage0_group_count: dag.groups.len(), + }; + for group in dag.groups { + memo.index_and_push(group); + } + memo + } + + /// The groups, by id. + pub fn groups(&self) -> &[MemoGroup] { + &self.groups + } + + /// Whether `group` is one Stage 0 produced from the queries' DAG, not + /// one Pass 1 or Pass 2 added beneath them. + pub fn is_stage0_group(&self, group: MemoGroupId) -> bool { + (group.0 as usize) < self.stage0_group_count + } + + /// The groups, by id, to change in a test (a malformed reference, a + /// different logical node); a changed logical node is not looked up again. + #[cfg(any(test, feature = "test-support"))] + pub fn groups_mut(&mut self) -> &mut [MemoGroup] { + &mut self.groups + } + + /// Adds `group` with its `candidates` and returns its id. + pub fn push_group( + &mut self, + group: MemoGroup, + candidates: Vec, + ) -> MemoGroupId { + let id = self.index_and_push(group); + self.candidates.push(candidates); + id + } + + /// Indexes `group` by its logical node and adds it; its id. + fn index_and_push(&mut self, group: MemoGroup) -> MemoGroupId { + let id = MemoGroupId( + u32::try_from(self.groups.len()).expect("a memo has at most u32::MAX groups"), + ); + let node = Rc::clone(&group.logical_node); + self.group_lookup + .group_whose_node_is + .entry(Rc::as_ptr(&node)) + .or_insert(id); + let equal_nodes = &mut self.group_lookup.group_with_equal_node; + if equal_nodes.find(&node).is_none() { + equal_nodes.insert(node, id); + } + self.groups.push(group); + id + } + + /// The group whose logical node is `node` itself (the same `Rc`). + pub fn group_whose_node_is(&self, node: &Rc) -> Option { + self.group_lookup + .group_whose_node_is + .get(&Rc::as_ptr(node)) + .copied() + } + + /// The group whose logical node is `node`, else the lowest group whose + /// logical node equals `node` by Stage 0's equality when sharing `node` + /// is legal ([`HashConsingTable::find`]). + pub fn group_with_equal_node(&self, node: &Rc) -> Option { + self.group_whose_node_is(node) + .or_else(|| self.group_lookup.group_with_equal_node.find(node)) + } + + /// The groups `group` reads: its logical child groups and its + /// candidates' input groups. + fn groups_read_by(&self, group: MemoGroupId) -> impl Iterator + '_ { + self.groups[group.0 as usize] + .child_groups + .iter() + .chain( + self.candidates + .get(group.0 as usize) + .into_iter() + .flatten() + .flat_map(|candidate| &candidate.input_groups), + ) + .copied() + } + + /// Adds `queries` to `group` and to every group below it, through + /// logical child groups and candidates' input groups. Below `group`, a + /// group that already has every query is not walked further: the groups + /// below it have them already, or get them when this is called for it. + pub fn add_queries(&mut self, group: MemoGroupId, queries: &[usize]) { + let mut pending: Vec = self.groups_read_by(group).collect(); + let own = &mut self.groups[group.0 as usize].queries; + own.extend_from_slice(queries); + own.sort_unstable(); + own.dedup(); + let mut seen = HashSet::from([group]); + while let Some(next) = pending.pop() { + if !seen.insert(next) { + continue; + } + let reached = &mut self.groups[next.0 as usize].queries; + if queries.iter().all(|query| reached.contains(query)) { + continue; + } + reached.extend_from_slice(queries); + reached.sort_unstable(); + reached.dedup(); + pending.extend(self.groups_read_by(next)); + } + } + + /// Sets every group's parents to the groups that read it: through a + /// logical child group or a candidate's input group. + pub fn rebuild_parents(&mut self) { + for group in &mut self.groups { + group.parents.clear(); + } + for group in 0..self.groups.len() { + let consumer = MemoGroupId(group as u32); + let read: HashSet = self.groups_read_by(consumer).collect(); + let mut read: Vec = read.into_iter().collect(); + read.sort_unstable(); + for child in read { + self.groups[child.0 as usize].parents.push(consumer); + } + } + } + /// The groups, each after every group it or one of its candidates /// reads, the lower index first among those ready: a dependency order for /// candidate traversal, whatever order transformations appended diff --git a/crates/types/src/stages/hash_consing.rs b/crates/types/src/stages/hash_consing.rs new file mode 100644 index 000000000..dcf16edfa --- /dev/null +++ b/crates/types/src/stages/hash_consing.rs @@ -0,0 +1,377 @@ +//! The one node equality that makes two nodes one memo group: Stage 0's +//! hash-consing of the frontends' DAG and the memo's group lookup +//! ([`group_with_equal_node`](super::Memo::group_with_equal_node)) both use +//! [`HashConsingTable`]. Two nodes are one when they are equal by +//! [`same_node`] (operator, schema, result kind, guarantee, timing, +//! materialization and children) and sharing them is legal +//! ([`shareable_when_equal`]). + +use std::collections::hash_map::DefaultHasher; +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use crate::ir::operator::node::{Operator, OperatorNode}; +use crate::ir::NonASAPOp; + +/// [`structural_hash`]'s memoization cache: a hashed node's address to the +/// node (kept alive, so its address is not reused) and its hash. +type NodeToStructuralHashCache = HashMap<*const OperatorNode, (Rc, u64)>; + +/// The operator with every child (operator inputs and the operator nodes +/// referenced from its scalar expressions alike) replaced by `()`. What +/// remains is the node's own data: variant tag, scalar expressions, +/// parameters. +fn own_fields(node: &OperatorNode) -> Operator<()> { + node.operator.map_children(|_| ()) +} + +/// Coarse structural hash used only to bucket [`HashConsingTable`]'s +/// candidate search — never the sharing decision ([`same_node`] is). +/// +/// `OperatorNode` carries `f64`s (`ScalarValue::Float64`, quantile targets, +/// `AccuracyGuarantee` bounds, …), so it cannot derive `std::hash::Hash`. The +/// hash is SipHash over two parts: +/// +/// 1. the canonical JSON of [`own_fields`] plus `result_kind`, `schema`, +/// `guarantee`, `timing` and `materialization` — every field `PartialEq` compares except the +/// children. A scalar expression is serialized as data with each operator +/// node it reads replaced by `()`, so a reference to an +/// interned sub-DAG contributes nothing of its own here; +/// 2. for every child in [`OperatorNode::children`] order (operator inputs, +/// then scalar-referenced nodes), the child's own `structural_hash`, +/// memoized in `cache` by `Rc` pointer identity. +/// +/// Part 2 is what makes equal sub-DAGs hash equal whether they are reached +/// through an operator input or through a `scalar(v)`, and what keeps the +/// pass linear: a node is generally a DAG, and re-serializing a shared +/// descendant once per parent would cost `O(sub-DAG)` per node instead of +/// `O(1)` beyond the children's already-known hashes. A non-finite `f64` +/// serializes as `null`, merely widening one (still equality-checked) bucket. +fn structural_hash(node: &OperatorNode, cache: &mut NodeToStructuralHashCache) -> u64 { + structural_hash_with(node, &mut |child| cached_structural_hash(child, cache)) +} + +/// [`structural_hash`] of `node` with each child's hash from `child_hash`. +fn structural_hash_with( + node: &OperatorNode, + child_hash: &mut dyn FnMut(&Rc) -> u64, +) -> u64 { + let mut hasher = DefaultHasher::new(); + let own = ( + own_fields(node), + node.result_kind, + &node.schema, + &node.guarantee, + node.timing, + node.materialization, + ); + serde_json::to_string(&own) + .unwrap_or_default() + .hash(&mut hasher); + for child in node.children() { + child_hash(child).hash(&mut hasher); + } + hasher.finish() +} + +/// [`structural_hash`] of `node`, memoized in `cache`. +fn cached_structural_hash(node: &Rc, cache: &mut NodeToStructuralHashCache) -> u64 { + if let Some((_, hash)) = cache.get(&Rc::as_ptr(node)) { + return *hash; + } + let hash = structural_hash(node, cache); + cache.insert(Rc::as_ptr(node), (Rc::clone(node), hash)); + hash +} + +/// [`structural_hash`] of `node`, reading `cache` but not adding to it: the +/// nodes hashed only for a lookup are memoized in `looked_up`, which lives +/// no longer than the borrowed `node` that keeps them alive. +fn looked_up_structural_hash( + node: &OperatorNode, + cache: &NodeToStructuralHashCache, + looked_up: &mut HashMap<*const OperatorNode, u64>, +) -> u64 { + structural_hash_with(node, &mut |child| { + let pointer = Rc::as_ptr(child); + if let Some((_, hash)) = cache.get(&pointer) { + return *hash; + } + if let Some(hash) = looked_up.get(&pointer) { + return *hash; + } + let hash = looked_up_structural_hash(child, cache, looked_up); + looked_up.insert(pointer, hash); + hash + }) +} + +/// Numeric `PartialEq` alone conflates signed zeros. The serialized check is +/// additional evidence, never a replacement for typed equality (JSON maps +/// non-finite floats to `null`). Used for the guarantee, whose bounds are +/// floats a shared node must preserve bit-for-bit, and for summary families, +/// whose `SketchParams` equality omits the recorded error and probability. +pub fn same_value(left: &T, right: &T) -> bool { + left == right + && match (serde_json::to_string(left), serde_json::to_string(right)) { + (Ok(left), Ok(right)) => left == right, + _ => false, + } +} + +/// Cache of child-pair comparisons already decided by [`same_node`], keyed by +/// pointer pair; both nodes must stay alive while it is used. +type NodePairToIsSameNodeCache = HashMap<(*const OperatorNode, *const OperatorNode), bool>; + +/// The sharing decision: typed equality of two nodes. +/// +/// `OperatorNode`'s derived `PartialEq` would recurse into children by value +/// even when both sides hold the same `Rc` (`OperatorNode` is not `Eq`, so +/// `Rc` gets no pointer shortcut), expanding a shared diamond once per path. +/// Children are therefore compared by pointer first; only when the pointers +/// differ (an equal child that was not legal to share) are the values +/// compared, memoized per pair so a diamond is still walked once. +fn same_node( + left: &OperatorNode, + right: &OperatorNode, + cache: &mut NodePairToIsSameNodeCache, +) -> bool { + let (lc, rc) = (left.children(), right.children()); + if lc.len() != rc.len() { + return false; + } + let children_equal = lc.iter().zip(&rc).all(|(a, b)| { + if Rc::ptr_eq(a, b) { + return true; + } + let key = (Rc::as_ptr(a), Rc::as_ptr(b)); + if let Some(&eq) = cache.get(&key) { + return eq; + } + let eq = same_node(a, b, cache); + cache.insert(key, eq); + eq + }); + children_equal + && left.result_kind == right.result_kind + && left.schema == right.schema + && left.timing == right.timing + && left.materialization == right.materialization + && same_value(&left.guarantee, &right.guarantee) + && same_value(&own_fields(left), &own_fields(right)) +} + +/// Whether two equal copies of `node` return the same rows for the same +/// input rows, so they may be one node: an ASAP node, or a non-ASAP node +/// [`NonASAPOp::is_shareable_when_equal`]. +fn shareable_when_equal(node: &OperatorNode) -> bool { + node.is_asap() + || node + .non_asap() + .is_some_and(NonASAPOp::is_shareable_when_equal) +} + +/// Hash-consing table: each node added with its value (an interned node, a +/// memo group), found again by an equal node legal to share. +/// +/// Buckets are keyed by [`structural_hash`] — a coarse candidate filter +/// only; [`same_node`] decides a match, so a hash collision between +/// different nodes only means a longer scan of one bucket. +#[derive(Clone)] +pub struct HashConsingTable { + buckets: HashMap, V)>>, + /// The added nodes' hashes (and their descendants'), persisted for the + /// table's whole lifetime so hashing is `O(1)` per node beyond its + /// children; a lookup ([`Self::find`]) adds nothing, so it keeps no node + /// alive. + structural_hashes: NodeToStructuralHashCache, +} + +impl Default for HashConsingTable { + fn default() -> Self { + Self { + buckets: HashMap::new(), + structural_hashes: HashMap::new(), + } + } +} + +impl std::fmt::Debug for HashConsingTable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("HashConsingTable") + .field("nodes", &self.buckets.values().map(Vec::len).sum::()) + .finish_non_exhaustive() + } +} + +impl HashConsingTable { + /// The value of a node added before that equals `node` ([`same_node`]), + /// when sharing `node` is legal ([`shareable_when_equal`]); `None` + /// otherwise. + pub fn find(&self, node: &Rc) -> Option { + if !shareable_when_equal(node) { + return None; + } + let hash = match self.structural_hashes.get(&Rc::as_ptr(node)) { + Some((_, hash)) => *hash, + None => looked_up_structural_hash(node, &self.structural_hashes, &mut HashMap::new()), + }; + let mut same_nodes = NodePairToIsSameNodeCache::new(); + self.buckets + .get(&hash)? + .iter() + .find(|(candidate, _)| same_node(candidate, node, &mut same_nodes)) + .map(|(_, value)| value.clone()) + } + + /// Adds `node` with `value`. A node not legal to share is added too; it + /// is never found for another node. + pub fn insert(&mut self, node: Rc, value: V) { + let hash = cached_structural_hash(&node, &mut self.structural_hashes); + self.buckets.entry(hash).or_default().push((node, value)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ir::operator::agg_intent::AggIntent; + use crate::ir::operator::operator_properties::{BinaryOpKind, Reduction, Source}; + use crate::ir::properties::AccuracyTarget; + use crate::ir::scalar::expr_ir::CompareOpKind; + use crate::ir::schema::{DataType, Field, Schema}; + use crate::ir::{BinaryOperator, ScalarExpr}; + + fn node(op: NonASAPOp) -> Rc { + OperatorNode::new_shared(Operator::NonASAP(op)).unwrap() + } + + /// `[ts, service, value, latency]`, no unique key. + fn scan() -> Rc { + node(NonASAPOp::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("service", DataType::Utf8, false), + Field::plain("value", DataType::Float64, false), + Field::plain("latency", DataType::Float64, false), + ], + 0, + vec![], + ), + }) + } + + fn quantile_agg(by: Vec, col: Option, q: f64) -> Rc { + node(NonASAPOp::Aggregate { + reduction: Reduction::by(by), + measures: vec![AggIntent::Quantile { + col, + q, + accuracy: AccuracyTarget::exact(), + }], + output_names: vec![], + filters: vec![], + having: None, + child: scan(), + }) + } + + fn compare(lhs: Rc, rhs: Rc) -> Rc { + node(NonASAPOp::BinaryOp { + operator: BinaryOperator { + checked_relative_division: false, + checked_finite_division: false, + kind: BinaryOpKind::Compare(CompareOpKind::Eq), + vector_match: None, + }, + return_bool: false, + lhs, + rhs, + }) + } + + /// `vector(scalar(sum by (service) (m)))`. + fn scalar_of_vector() -> Rc { + let sum_up = node(NonASAPOp::Aggregate { + reduction: Reduction::by(vec![1]), + measures: vec![AggIntent::Sum { col: Some(2) }], + output_names: vec![], + filters: vec![], + having: None, + child: scan(), + }); + node(NonASAPOp::PromqlVectorFromScalar( + ScalarExpr::PromqlScalarFromVector(sum_up), + )) + } + + #[test] + fn structural_hash_sees_through_a_scalar_reference() { + // Two equal bridges must hash equal whether or not their referenced + // vector is the same Rc: the reference contributes the vector's + // memoized hash, not its identity. + let a = scalar_of_vector(); + let b = scalar_of_vector(); + let mut cache = HashMap::new(); + assert_eq!( + structural_hash(&a, &mut cache), + structural_hash(&b, &mut cache) + ); + assert_eq!( + cache.len(), + 4, + "aggregate + scan cached once per root: {cache:?}" + ); + let other = node(NonASAPOp::PromqlVectorFromScalar( + ScalarExpr::PromqlScalarFromVector(quantile_agg(vec![1], Some(2), 0.5)), + )); + assert_ne!( + structural_hash(&a, &mut cache), + structural_hash(&other, &mut cache) + ); + } + + #[test] + fn structural_hash_is_stable_across_cache_states() { + let agg = quantile_agg(vec![1], Some(2), 0.5); + let mut cold = HashMap::new(); + let mut warm = HashMap::new(); + structural_hash(&scan(), &mut warm); + assert_eq!( + structural_hash(&agg, &mut cold), + structural_hash(&agg, &mut warm), + "hash must be independent of unrelated cache state" + ); + } + + #[test] + fn structural_hash_of_an_internally_shared_dag_matches_the_unshared_equivalent() { + let agg = quantile_agg(vec![1], Some(2), 0.5); + let shared_root = compare(Rc::clone(&agg), Rc::clone(&agg)); + let unshared_root = compare( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(2), 0.5), + ); + assert_eq!( + structural_hash(&shared_root, &mut HashMap::new()), + structural_hash(&unshared_root, &mut HashMap::new()), + ); + } + + #[test] + fn structural_hash_memoizes_a_shared_descendant_exactly_once() { + let agg = quantile_agg(vec![1], Some(2), 0.5); + let root = compare(Rc::clone(&agg), Rc::clone(&agg)); + let mut cache = HashMap::new(); + structural_hash(&root, &mut cache); + assert_eq!( + cache.len(), + 2, + "one entry per unique node in the shared branch (Aggregate + Scan): {cache:?}" + ); + } +} diff --git a/crates/types/src/stages/mod.rs b/crates/types/src/stages/mod.rs index 109ce200e..69e9f1a06 100644 --- a/crates/types/src/stages/mod.rs +++ b/crates/types/src/stages/mod.rs @@ -9,6 +9,7 @@ //! live here. mod candidates; +pub mod hash_consing; mod logical_memo_candidates; mod memo; mod per_computation_logical_memo_candidates;