diff --git a/crates/logical-asap-optimization/src/operator_over_groups.rs b/crates/logical-asap-optimization/src/operator_over_groups.rs index dd5ea61a0..f6a6a7a1d 100644 --- a/crates/logical-asap-optimization/src/operator_over_groups.rs +++ b/crates/logical-asap-optimization/src/operator_over_groups.rs @@ -16,7 +16,8 @@ use crate::own_candidate_sizing::size_own_candidate_for; /// `root`, a candidate of a group `queries` reach, as one operator over /// groups: the exact guarantees below it removed -/// ([`without_exact_guarantees_below`]), each child found or added with +/// ([`without_exact_guarantees`]; its own guarantee is the candidate's and +/// is kept), each child found or added with /// every node below it ([`over_child_groups`]: an added group is made /// by `origin`, sized for `accuracy_budget_parts`), and every added group it /// reads through sized for at least `accuracy_budget_parts` (a reader sized @@ -29,7 +30,20 @@ pub(crate) fn operator_over_groups( origin: &OriginOfCandidateForMemoGroup, accuracy_budget_parts: usize, ) -> (Rc, Vec) { - let root = without_exact_guarantees_below(root); + let stripped = without_exact_guarantees(root); + let root = if stripped + .children() + .iter() + .zip(root.children()) + .all(|(a, b)| Rc::ptr_eq(a, b)) + { + Rc::clone(root) + } else { + Rc::new(OperatorNode { + guarantee: root.guarantee.clone(), + ..(*stripped).clone() + }) + }; let (node, children) = over_child_groups(memo, &root, queries, origin, accuracy_budget_parts); for group in groups_read_through_added_groups(memo, &children).added_groups { size_own_candidate_for(memo, group, accuracy_budget_parts); @@ -37,16 +51,15 @@ pub(crate) fn operator_over_groups( (node, children) } -/// `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 +/// `node` with every exact guarantee removed, its own included. 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. -pub(crate) fn without_exact_guarantees_below(root: &Rc) -> Rc { +/// compares guarantees). An approximate guarantee is kept, so builds that +/// differ only in their failure probability or guarantee stay different +/// groups. +pub(crate) fn without_exact_guarantees(node: &Rc) -> Rc { fn without_exact( node: &Rc, done: &mut HashMap<*const OperatorNode, Rc>, @@ -76,20 +89,5 @@ pub(crate) fn without_exact_guarantees_below(root: &Rc) -> Rc MemoGroupId { - let below = without_exact_guarantees_below(node); - let node = if below.guarantee.as_ref().is_some_and(|g| g.is_exact()) { - Rc::new(OperatorNode { - guarantee: None, - ..(*below).clone() - }) - } else { - below - }; - find_or_add_group(memo, &node, queries, origin, 1) + find_or_add_group(memo, &without_exact_guarantees(node), queries, origin, 1) } pub(in crate::pass2::memo) fn replace( diff --git a/crates/logical-asap-optimization/src/pass2/memo/rules/scan_labels.rs b/crates/logical-asap-optimization/src/pass2/memo/rules/scan_labels.rs index 4fdb5319d..04e5f1de5 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/rules/scan_labels.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/rules/scan_labels.rs @@ -1,19 +1,16 @@ //! Scan-label widening for explicit memo state sharing. Reuses the existing //! Pass 2 widening rule; only its source-chain boundaries are interned here. -use std::collections::{HashMap, HashSet}; use std::rc::Rc; use asap_accuracy::propagation::recompose_guarantees; use asap_types::ir::operator::Reduction; use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode, QueryRoot}; -use super::reuse::same_output_contract; +use super::reuse::{add_candidate_over_groups, same_output_contract}; use crate::added_groups::groups_read_through_added_groups; -use crate::pass2::memo::state_groups::{ - boundary_groups, build_down_to_stage0_groups, Pass2RuleError, -}; -use asap_types::stages::{CandidateForMemoGroup, Memo, OriginOfCandidateForMemoGroup}; -use asap_types::stages::{MemoGroup, MemoGroupId}; +use crate::operator_over_groups::without_exact_guarantees; +use crate::pass2::memo::state_groups::{build_down_to_stage0_groups, Pass2RuleError}; +use asap_types::stages::{Memo, MemoGroupId, OriginOfCandidateForMemoGroup}; /// Preserve local implementations and append readers over union-label inputs. /// Parent/query reachability is rebuilt by the enclosing state-group expansion. @@ -83,53 +80,49 @@ pub(crate) fn add_scan_label_candidates(memo: &mut Memo) -> Result bool { _ => false, } } - -fn same_source_node(left: &OperatorNode, right: &OperatorNode) -> bool { - // These operators preserve exact source rows; recomposition's exact - // guarantee annotation does not change their logical/execution contract. - 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() - .iter() - .zip(right.children()) - .all(|(a, b)| Rc::ptr_eq(a, b)) - && left.children().len() == right.children().len() -} - -fn intern_inputs( - node: &Rc, - memo: &mut Memo, - queries: &[usize], - done: &mut HashMap<*const OperatorNode, Rc>, -) -> Result, Pass2RuleError> { - if let Some(known) = done.get(&Rc::as_ptr(node)) { - return Ok(Rc::clone(known)); - } - let mut children = HashMap::new(); - for child in node.children() { - children.insert( - Rc::as_ptr(child), - intern_inputs(child, memo, queries, done)?, - ); - } - let rewritten = Rc::new(OperatorNode { - operator: node - .operator - .map_children(|child| Rc::clone(&children[&Rc::as_ptr(child)])), - ..(**node).clone() - }); - let result = if is_source_chain(&rewritten) { - if let Some(existing) = memo - .groups() - .iter() - .find(|group| same_source_node(&group.logical_node, &rewritten)) - { - Rc::clone(&existing.logical_node) - } else { - // The new group's id must fit a `MemoGroupId`. - u32::try_from(memo.groups().len()).map_err(|_| Pass2RuleError::TooManyGroups)?; - let mut seen = HashSet::new(); - let inputs: Vec<_> = rewritten - .children() - .iter() - .map(|child| { - memo.group_whose_node_is(child) - .expect("source-chain children were interned") - }) - .filter(|input| seen.insert(*input)) - .collect(); - 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 { - rewritten - }; - done.insert(Rc::as_ptr(node), Rc::clone(&result)); - Ok(result) -} 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 abe49ab85..93b0ca498 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 @@ -11,15 +11,14 @@ use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot}; use super::shared_producer::{ add_shared_producer_candidates, spans_groups, Reader, SharedProducer, }; +use crate::operator_over_groups::operator_over_groups; 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, build_of, same_range_input, BuildPlace, Pass2RuleError, -}; +use crate::pass2::memo::state_groups::{boundary_ptrs, build_of, BuildPlace, Pass2RuleError}; use asap_types::stages::{Memo, MemoGroupId, OriginOfCandidateForMemoGroup}; /// `resized` (a rewrite of `original` with the same shape) with each node @@ -102,6 +101,9 @@ struct SummaryStates { rewrite_count: usize, /// The (normalized layout, configuration index) rewrites already made. requests: HashSet<(usize, usize)>, + /// The nodes of Stage 0's groups, which a rewrite of a candidate whose + /// build is beneath groups puts back ([`BuildPlace::BeneathGroups`]). + stage0_group_nodes: HashSet<*const OperatorNode>, } pub(crate) fn add_state_configuration_candidates(memo: &mut Memo) -> Result<(), Pass2RuleError> { @@ -116,6 +118,7 @@ fn expand_configurations( deduplicate_requests, rewrite_count: 0, requests: HashSet::new(), + stage0_group_nodes: HashSet::new(), }; add_shared_producer_candidates(memo, &mut states)?; Ok(states.rewrite_count) @@ -127,13 +130,31 @@ impl SharedProducer for SummaryStates { fn sets(&mut self, memo: &Memo) -> Result>>, Pass2RuleError> { let mut sets: Vec>> = Vec::new(); + self.stage0_group_nodes = memo + .groups() + .iter() + .enumerate() + .filter(|&(id, _)| memo.is_stage0_group(MemoGroupId(id as u32))) + .map(|(_, group)| Rc::as_ptr(&group.logical_node)) + .collect(); for (group, candidates) in memo.candidates.iter().enumerate() { for (candidate, entry) in candidates.iter().enumerate() { - let Some((build, place @ (BuildPlace::Private | BuildPlace::InputGroup(_)))) = - build_of(memo, entry) - else { + let Some((build, place)) = build_of(memo, entry) else { continue; }; + // A build beneath groups (a tumbling-window chain a rule + // added as groups) is rewritten whole, from a Stage 0 + // group's candidate only: per-width merge groups below would + // defeat the per-group layout deduplication. + if place == BuildPlace::BeneathGroups + && !(memo.is_stage0_group(MemoGroupId(group as u32)) + && matches!( + build.asap(), + Some(ASAPOp::SummaryTumblingWindowChain { .. }) + )) + { + continue; + } // Configuration expansion precedes Stage 2. The rewrite helpers // create unassigned state; preserve any already assigned build. if build.timing.is_some() @@ -291,7 +312,12 @@ impl SharedProducer for SummaryStates { with_resized_build(&roots, &need.build, family.clone()) } }; - let mut boundaries = boundary_ptrs(memo, &original); + let mut boundaries = match need.place { + // The rewrite reaches below the input groups: only Stage 0's + // groups' nodes are put back. + BuildPlace::BeneathGroups => self.stage0_group_nodes.clone(), + BuildPlace::Private | BuildPlace::InputGroup(_) => 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 { @@ -324,9 +350,16 @@ impl SharedProducer for SummaryStates { // 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::Private => (node, original.input_groups.clone()), + // Rule 9 has no origin of its own; the groups it adds are shared + // builds. + BuildPlace::BeneathGroups => operator_over_groups( + memo, + &node, + &memo.groups()[reader.group].queries.clone(), + &OriginOfCandidateForMemoGroup::SharedBuild, + accuracy_budget_parts, + ), BuildPlace::InputGroup(build_group) => { let resized_build = original .root_node @@ -337,8 +370,6 @@ impl SharedProducer for SummaryStates { .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, @@ -503,7 +534,7 @@ fn compatible(left: &StateNeed, right: &StateNeed) -> bool { || !left_inputs .iter() .zip(right_inputs) - .all(|(a, b)| Rc::ptr_eq(a, b) || same_range_input(a, b)) + .all(|(a, b)| Rc::ptr_eq(a, b)) { return false; } 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 c438d448c..8fa3e7006 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/state_groups.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/state_groups.rs @@ -8,7 +8,7 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::rc::Rc; -use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode}; +use asap_types::ir::{ASAPOp, OperatorNode}; use crate::own_candidate_sizing::find_or_add_sized_group; use asap_types::stages::hash_consing::HashConsingTable; @@ -413,41 +413,6 @@ pub(super) fn nodes_within( nodes } -// 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. -pub(super) fn same_range_input(left: &OperatorNode, right: &OperatorNode) -> bool { - let ( - Some(NonASAPOp::TimeRange { - child: left_child, .. - }), - Some(NonASAPOp::TimeRange { - child: right_child, .. - }), - ) = (left.non_asap(), right.non_asap()) - else { - return false; - }; - Rc::ptr_eq(left_child, right_child) && same_local_fields(left, right) -} - -fn same_local_fields(left: &OperatorNode, right: &OperatorNode) -> bool { - let signature = |node: &OperatorNode| { - ( - node.operator.map_children(|_| ()), - node.result_kind, - node.schema.clone(), - node.guarantee.clone(), - node.timing, - node.materialization, - ) - }; - let left = signature(left); - let right = signature(right); - left == right - && matches!((serde_json::to_string(&left), serde_json::to_string(&right)), (Ok(left), Ok(right)) if left == right) -} - #[cfg(test)] fn replace_node( node: &Rc, @@ -464,6 +429,7 @@ mod tests { 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::NonASAPOp; use asap_types::ir::QueryRoot; /// An expanded memo with its shared summary-state groups. @@ -648,6 +614,77 @@ mod tests { // Different label projections can consume one wider tumbling-window state through // typed reader projections; no original scan needs to contain the union. + /// Rollup gap check after scan widening became one operator over groups: + /// a widened tumbling-window reader's chain now sits beneath groups + /// (`BuildPlace::BeneathGroups`), which rule 6 does not take. Before, + /// the chain was the reader's private build, which rule 6 took. No + /// rollup is lost: rule 6's rewrite finds none for these readers, with + /// either query as the coarse one. + #[test] + fn rollups_lose_nothing_from_widened_readers_whose_chain_is_beneath_groups() { + use crate::pass2::asap_aware_cse::rollup_candidate_for; + let mut memo = with_local_tumbling_windows( + workload(&[ + ("sum by (job) (sum_over_time(latency[5m]))", 0.01), + ("sum by (job, region) (sum_over_time(latency[5m]))", 0.01), + ]), + &[60_000, 60_000], + ); + super::super::rules::scan_labels::add_scan_label_candidates(&mut memo).unwrap(); + let memo = &memo; + let beneath: Vec<_> = (0..memo.groups().len()) + .flat_map(|group| { + memo.candidates[group] + .iter() + .filter(|candidate| { + candidate.origin == OriginOfCandidateForMemoGroup::WidenedScans + }) + .filter_map(move |candidate| match build_of(memo, candidate)? { + (build, BuildPlace::BeneathGroups) => Some((group, candidate, build)), + _ => None, + }) + }) + .collect(); + assert!( + beneath.iter().any(|(_, _, build)| matches!( + build.asap(), + Some(ASAPOp::SummaryTumblingWindowChain { .. }) + )), + "a widened reader whose chain is beneath groups" + ); + // Rule 6's other inputs: Pass 1 and widened candidates whose build + // it takes. + let inputs: Vec<_> = (0..memo.groups().len()) + .flat_map(|group| { + memo.candidates[group].iter().filter_map(move |candidate| { + match build_of(memo, candidate)? { + (_, BuildPlace::Private | BuildPlace::InputGroup(_)) => { + Some((group, candidate)) + } + _ => None, + } + }) + }) + .chain( + beneath + .iter() + .map(|&(group, candidate, _)| (group, candidate)), + ) + .collect(); + for (coarse_group, coarse, build) in &beneath { + for (fine_group, fine) in &inputs { + if coarse_group == fine_group { + continue; + } + let roots = [ + QueryRoot::Operator(Rc::clone(&coarse.root_node)), + QueryRoot::Operator(Rc::clone(&fine.root_node)), + ]; + assert!(rollup_candidate_for(&roots, Rc::as_ptr(build)).is_none()); + } + } + } + #[test] fn disjoint_scan_labels_share_a_union_tumbling_window_state_with_typed_readers() { let original = with_local_tumbling_windows( @@ -686,9 +723,9 @@ mod tests { }) .collect(); let expanded = expand(original).unwrap(); - let shared = expanded - .state_groups - .iter() + // The union-label chain is a group both queries reach. + let shared = (0..expanded.memo.groups().len()) + .map(|id| MemoGroupId(id as u32)) .find(|id| { let group = &expanded.memo.groups()[id.0 as usize]; matches!( @@ -706,24 +743,33 @@ mod tests { }) .expect("one state over union labels"); 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] - .iter() - .find(|candidate| candidate.input_groups.contains(shared)) - .unwrap(); - assert_eq!( - reader.root_node.schema.fields, - expanded.memo.groups()[parent.0 as usize] - .logical_node - .schema - .fields - ); - assert!(OperatorNode::reachable(&reader.root_node) - .iter() - .any(|node| matches!(node.non_asap(), Some(NonASAPOp::Project { .. })))); - reader.root_node.validate_structure().unwrap(); + // Each query reads it through a widened-scan candidate of a Stage 0 + // group that projects back to the group's own labels. + let mut reading_queries = Vec::new(); + for (group, candidates) in expanded.memo.candidates.iter().enumerate() { + if !expanded.memo.is_stage0_group(MemoGroupId(group as u32)) { + continue; + } + for reader in candidates.iter().filter(|candidate| { + candidate.origin == OriginOfCandidateForMemoGroup::WidenedScans + && OperatorNode::reachable(&candidate.root_node) + .iter() + .any(|node| Rc::ptr_eq(node, &state.logical_node)) + }) { + assert_eq!( + reader.root_node.schema.fields, + expanded.memo.groups()[group].logical_node.schema.fields + ); + assert!(OperatorNode::reachable(&reader.root_node) + .iter() + .any(|node| matches!(node.non_asap(), Some(NonASAPOp::Project { .. })))); + reader.root_node.validate_structure().unwrap(); + reading_queries.extend(expanded.memo.groups()[group].queries.iter().copied()); + } } + reading_queries.sort_unstable(); + reading_queries.dedup(); + assert_eq!(reading_queries, vec![0, 1]); for (group, locals) in local_nodes.into_iter().enumerate() { assert!(locals .iter() @@ -755,10 +801,11 @@ mod tests { &[60_000, 60_000], ); let expanded = expand(memo).unwrap(); + // The union-label chain is a group both queries reach. let state = expanded - .state_groups + .memo + .groups() .iter() - .map(|id| &expanded.memo.groups()[id.0 as usize]) .find(|group| { matches!( group.logical_node.asap(),