diff --git a/crates/integration-tests/src/lib.rs b/crates/integration-tests/src/lib.rs index 390675e97..0e913db53 100644 --- a/crates/integration-tests/src/lib.rs +++ b/crates/integration-tests/src/lib.rs @@ -539,28 +539,19 @@ pub mod planning { 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). + // Each query's own summary candidate, with the candidates on a + // path of groups from it down to the state (Pass 1 and the + // sharing rules add the consumers between them as groups). let reading: Option> = (0..queries) .map(|q| { - let reader = own_choices(memo, q, |c| { - c.input_groups.iter().any(|child| child.0 as usize == state) + own_builds[q].iter().find_map(|&(group, index)| { + let mut path = path_to_state(memo, group, index, state)?; + path.push((group, index)); + Some(path) }) - .into_iter() - .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::>>() - .map(|pairs| pairs.into_iter().flatten().collect()); + .map(|paths| paths.into_iter().flatten().collect()); if let Some(choice) = reading { plans.push(( format!("shared{state}"), @@ -604,6 +595,45 @@ pub mod planning { (roots, separate, cheapest(&plans[separate_plans..])) } + /// The candidates, one per group, on a path from candidate `index` of + /// `group` down to group `state` (excluded, and the start excluded): + /// each reads the next one's group. `None` when no path reaches `state`. + fn path_to_state( + memo: &asap_types::stages::Memo, + group: usize, + index: usize, + state: usize, + ) -> Option> { + fn below( + memo: &asap_types::stages::Memo, + group: usize, + state: usize, + seen: &mut std::collections::HashSet, + ) -> Option> { + if group == state { + return Some(Vec::new()); + } + if !seen.insert(group) { + return None; + } + memo.candidates[group] + .iter() + .enumerate() + .find_map(|(index, candidate)| { + candidate.input_groups.iter().find_map(|input| { + let mut path = below(memo, input.0 as usize, state, seen)?; + path.push((group, index)); + Some(path) + }) + }) + } + let mut seen = std::collections::HashSet::from([group]); + memo.candidates[group][index] + .input_groups + .iter() + .find_map(|input| below(memo, input.0 as usize, state, &mut seen)) + } + /// The summary builds and rollups (regrouping merges) `roots` read, each /// once. pub fn builds_and_rollups(roots: &[Rc]) -> (usize, usize) { diff --git a/crates/logical-asap-optimization/src/added_groups.rs b/crates/logical-asap-optimization/src/added_groups.rs index 103d3cd19..3e9570219 100644 --- a/crates/logical-asap-optimization/src/added_groups.rs +++ b/crates/logical-asap-optimization/src/added_groups.rs @@ -10,17 +10,17 @@ 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 { +pub(crate) struct GroupsReadThroughAddedGroups { /// The added groups passed through, in the order first reached. - pub added_groups: Vec, + pub(crate) added_groups: Vec, /// The Stage 0 groups reached, in the order first reached. - pub stage0_groups: Vec, + pub(crate) 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( +pub(crate) fn groups_read_through_added_groups( memo: &Memo, inputs: &[MemoGroupId], ) -> GroupsReadThroughAddedGroups { 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 9f3f333af..9ec92f647 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs @@ -1,138 +1,19 @@ -//! Shared boundary handling for producer-level reuse rules. Rule outputs retain -//! canonical input identities after guarantee recomposition rebuilds their DAGs. +//! What the producer-level reuse rules share: which candidates they read +//! ([`eligible`]), and how they add what they derive, as one operator over +//! groups ([`add_candidate_over_groups`], [`find_or_add_producer`]). use crate::find_or_add_group::find_or_add_group; use crate::operator_over_groups::{operator_over_groups, without_exact_guarantees_below}; -use crate::pass2::memo::state_groups::{boundary_groups, Pass2RuleError}; +use crate::pass2::memo::state_groups::Pass2RuleError; use asap_types::ir::OperatorNode; +use asap_types::stages::MemoGroupId; use asap_types::stages::{CandidateForMemoGroup, Memo, OriginOfCandidateForMemoGroup}; -use asap_types::stages::{MemoGroup, MemoGroupId}; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::rc::Rc; pub(in crate::pass2::memo) fn eligible(candidate: &CandidateForMemoGroup) -> bool { candidate.accuracy_budget_parts > 0 } -pub(in crate::pass2::memo) fn restore_inputs( - node: &Rc, - memo: &Memo, - inputs: &[MemoGroupId], -) -> Rc { - // These are the same declared logical input groups, rebound after a rule - // cloned the DAG. Value guarantees are recomposed from selected producers - // during reconstruction; state inputs additionally require the same error - // 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)) - .collect(); - fn visit( - node: &Rc, - canonical: &[Rc], - done: &mut HashMap<*const OperatorNode, Rc>, - ) -> Rc { - if let Some(found) = done.get(&Rc::as_ptr(node)) { - return Rc::clone(found); - } - let operator = node - .operator - .map_children(|child| visit(child, canonical, done)); - let rebuilt = Rc::new(OperatorNode { - operator, - ..(**node).clone() - }); - let result = canonical - .iter() - .find(|known| { - same_node( - known, - &rebuilt, - known.result_kind == asap_types::ir::OperatorResultKind::State, - ) - }) - .cloned() - .unwrap_or(rebuilt); - done.insert(Rc::as_ptr(node), Rc::clone(&result)); - result - } - visit(node, &canonical, &mut HashMap::new()) -} - -pub(in crate::pass2::memo) fn same_node( - left: &OperatorNode, - right: &OperatorNode, - guarantees: bool, -) -> bool { - left.operator.map_children(|_| ()) == right.operator.map_children(|_| ()) - && left.schema == right.schema - && left.result_kind == right.result_kind - && left.timing == right.timing - && left.materialization == right.materialization - && left.children().len() == right.children().len() - && left - .children() - .iter() - .zip(right.children()) - .all(|(a, b)| Rc::ptr_eq(a, b)) - && serde_json::to_string(&( - left.operator.map_children(|_| ()), - &left.schema, - left.result_kind, - left.timing, - left.materialization, - )) - .ok() - == serde_json::to_string(&( - right.operator.map_children(|_| ()), - &right.schema, - right.result_kind, - right.timing, - right.materialization, - )) - .ok() - && (!guarantees || left.guarantee == right.guarantee) -} - -pub(in crate::pass2::memo) fn append_producer( - memo: &mut Memo, - node: Rc, - queries: &[usize], - origin: OriginOfCandidateForMemoGroup, - accuracy_budget_parts: usize, -) -> Result { - // 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)| { - !root_groups.contains(&MemoGroupId(*index as u32)) - && memo.candidates[*index].len() == 1 - && memo.candidates[*index][0].origin == origin - && memo.candidates[*index][0].accuracy_budget_parts == accuracy_budget_parts - && same_node(&group.logical_node, &node, true) - }) { - return Ok(MemoGroupId(index as u32)); - } - // 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 children = boundary_groups(memo, &node, |_| true); - 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, - }], - )) -} - /// Whether `actual` has `expected`'s output: the same result kind, fields, /// time index and unique keys, and a closed label set where `expected`'s is /// closed (a projection may enumerate every output of an open schema). @@ -219,31 +100,6 @@ pub(in crate::pass2::memo) fn find_or_add_producer( find_or_add_group(memo, &node, queries, origin, 1) } -pub(in crate::pass2::memo) fn append_reader( - memo: &mut Memo, - owner: usize, - original: &CandidateForMemoGroup, - node: Rc, - origin: OriginOfCandidateForMemoGroup, -) -> Result { - check_output_contract(original, &node, &origin)?; - let children = boundary_groups(memo, &node, |group| group.0 as usize != owner); - if memo.candidates[owner].iter().any(|known| { - known.input_groups == children - && known.accuracy_budget_parts == original.accuracy_budget_parts - && *known.root_node == *node - && serde_json::to_string(&known.root_node).ok() == serde_json::to_string(&node).ok() - }) { - return Ok(false); - } - let mut candidate = original.clone(); - candidate.root_node = node; - candidate.input_groups = children; - candidate.origin = origin; - memo.candidates[owner].push(candidate); - Ok(true) -} - pub(in crate::pass2::memo) fn replace( node: &Rc, replacements: &HashMap<*const OperatorNode, Rc>, @@ -278,16 +134,3 @@ pub(in crate::pass2::memo) fn replace( } visit(node, replacements, &mut HashMap::new()) } - -pub(in crate::pass2::memo) fn combined_inputs( - left: &CandidateForMemoGroup, - right: &CandidateForMemoGroup, -) -> Vec { - let mut seen = HashSet::new(); - left.input_groups - .iter() - .chain(&right.input_groups) - .copied() - .filter(|id| seen.insert(*id)) - .collect() -} 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 9cd78145b..fb6fd0183 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/rollup.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/rollup.rs @@ -1,14 +1,12 @@ //! Producer-level exact and sketch rollups, using the existing capability rules. -use super::reuse::{ - append_producer, append_reader, combined_inputs, eligible, replace, restore_inputs, -}; +use super::reuse::{add_candidate_over_groups, eligible, find_or_add_producer, replace}; use crate::pass2::asap_aware_cse::{rollup_candidate_for, rollup_pass2_candidates}; 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, }; -use asap_types::stages::{Memo, OriginOfCandidateForMemoGroup}; +use asap_types::stages::{Memo, MemoGroupId, OriginOfCandidateForMemoGroup}; use std::collections::HashMap; use std::rc::Rc; @@ -77,27 +75,24 @@ pub(crate) fn add_rollup_candidates(memo: &mut Memo) -> Result::try_from(rewritten) else { unreachable!("two operator roots") }; - let inputs = combined_inputs(coarse, fine); - let coarse_node = restore_inputs(coarse_node, memo, &inputs); - let fine_node = restore_inputs(fine_node, memo, &inputs); if is_sketch || state_rollup { // Concrete resized states become ordinary state groups after all // producer-level configurations have been expanded. - added += usize::from(append_reader( + added += usize::from(add_candidate_over_groups( memo, - *coarse_group, + MemoGroupId(*coarse_group as u32), coarse, coarse_node, OriginOfCandidateForMemoGroup::Rollup, )?); - added += usize::from(append_reader( + added += usize::from(add_candidate_over_groups( memo, - *fine_group, + MemoGroupId(*fine_group as u32), fine, fine_node, OriginOfCandidateForMemoGroup::Rollup, @@ -143,9 +138,9 @@ pub(crate) fn add_rollup_candidates(memo: &mut Memo) -> Result Result Result Result::try_from(nodes) else { continue; }; - if !OperatorNode::reachable(coarse_node).iter().any(|node| { + if !OperatorNode::reachable(&coarse_node).iter().any(|node| { matches!( node.asap(), Some(ASAPOp::SummaryMerge { @@ -290,19 +284,16 @@ pub(crate) fn add_window_rollup_candidates(memo: &mut Memo) -> Result bool { + matches!( + node.asap(), + Some(ASAPOp::SummaryMerge { + group_by: Some(_), + .. + }) + ) + } + // Sketch rollup exposes the resized finer state through an ordinary group // dependency, while each output retains its own grouping and quantile. #[test] @@ -1045,23 +1056,17 @@ mod tests { ("quantile by (job) (0.9, m)", 0.02), ]); let expanded = expand(original).unwrap(); - assert!(expanded - .memo - .candidates - .iter() - .flatten() - .any( - |candidate| candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild - && OperatorNode::reachable(&candidate.root_node) - .iter() - .any(|node| matches!( - node.asap(), - Some(ASAPOp::SummaryMerge { - group_by: Some(_), - .. - }) - )) - )); + // The regrouping merge is a group of the rollup, over the finer build's group. + assert!(expanded.memo.candidates.iter().flatten().any(|candidate| { + candidate.origin == OriginOfCandidateForMemoGroup::Rollup + && regroups(&candidate.root_node) + && candidate.input_groups.iter().any(|group| { + matches!( + expanded.memo.groups()[group.0 as usize].logical_node.asap(), + Some(ASAPOp::SummaryAgg { .. }) + ) + }) + })); assert!(expanded.memo.topological_order().is_some()); } @@ -1208,30 +1213,23 @@ mod tests { &[10_000, 10_000], ); let expanded = expand(memo).unwrap(); - assert!(expanded - .memo - .candidates - .iter() - .flatten() - .any( - |candidate| candidate.origin == OriginOfCandidateForMemoGroup::SharedBuild - && OperatorNode::reachable(&candidate.root_node) - .iter() - .any(|node| matches!( + // The regrouping merge is a group of the window rollup, reading the + // finer query's chain group (through the merge of its newest windows). + assert!(expanded.memo.candidates.iter().flatten().any(|candidate| { + candidate.origin == OriginOfCandidateForMemoGroup::Rollup + && regroups(&candidate.root_node) + && OperatorNode::reachable(&candidate.root_node) + .iter() + .any(|node| { + matches!( node.asap(), - Some(ASAPOp::SummaryMerge { - group_by: Some(_), + Some(ASAPOp::SummaryTumblingWindowChain { + tumbling_window_count: 120, .. }) - )) - && candidate.input_groups.iter().any(|group| matches!( - expanded.memo.groups()[group.0 as usize].logical_node.asap(), - Some(ASAPOp::SummaryTumblingWindowChain { - tumbling_window_count: 120, - .. - }) - )) - )); + ) && expanded.memo.group_whose_node_is(node).is_some() + }) + })); } // Private identity projections retain their full typed contract, but no @@ -1365,7 +1363,7 @@ mod tests { FieldDataType, }; 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 (owner, base, build) = memo.candidates.iter().enumerate().find_map(|(owner, candidates)| candidates.iter().find_map(|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(); @@ -1419,9 +1417,7 @@ mod tests { let QueryRoot::Operator(root) = &roots[0] else { unreachable!() }; - let root = - crate::pass2::memo::rules::reuse::restore_inputs(root, &memo, &base.input_groups); - without_private_guarantees(&root, &boundaries) + without_private_guarantees(root, &boundaries) }; let weak = rewritten(0.001); let strict = rewritten(0.0001); @@ -1429,17 +1425,21 @@ mod tests { *weak, *strict, "ordinary equality omits the probability metadata" ); - base.root_node = weak; - memo.candidates[owner] = vec![base.clone()]; - assert!(crate::pass2::memo::rules::reuse::append_reader( - &mut memo, - owner, - &base, - strict, - OriginOfCandidateForMemoGroup::Rollup - ) - .unwrap()); - assert_eq!(memo.candidates[owner].len(), 2); + let before = memo.candidates[owner].len(); + let mut add = |node: &Rc| { + crate::pass2::memo::rules::reuse::add_candidate_over_groups( + &mut memo, + MemoGroupId(owner as u32), + &base, + Rc::clone(node), + OriginOfCandidateForMemoGroup::Rollup, + ) + .unwrap() + }; + assert!(add(&weak)); + assert!(add(&strict)); + assert!(!add(&weak), "an equal candidate is not added again"); + assert_eq!(memo.candidates[owner].len(), before + 2); } // Equal private producers can be shared within one composite result; they diff --git a/crates/logical-asap-optimization/src/test_support.rs b/crates/logical-asap-optimization/src/test_support.rs index 78abe347c..92d2f8db8 100644 --- a/crates/logical-asap-optimization/src/test_support.rs +++ b/crates/logical-asap-optimization/src/test_support.rs @@ -7,7 +7,6 @@ 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::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,