diff --git a/crates/logical-asap-optimization/src/memo.rs b/crates/logical-asap-optimization/src/memo.rs index 38051a4c9..4803a0bca 100644 --- a/crates/logical-asap-optimization/src/memo.rs +++ b/crates/logical-asap-optimization/src/memo.rs @@ -8,6 +8,8 @@ pub mod candidates; mod groups; mod rules; +#[cfg(test)] +mod shared_producer_snapshot; /// Shared finite state configuration enumeration. pub mod state_configurations; /// Concrete summary-state groups exposed by sharing rules. diff --git a/crates/logical-asap-optimization/src/memo/rules.rs b/crates/logical-asap-optimization/src/memo/rules.rs index 69dd57f49..b639c00a7 100644 --- a/crates/logical-asap-optimization/src/memo/rules.rs +++ b/crates/logical-asap-optimization/src/memo/rules.rs @@ -4,6 +4,7 @@ pub(super) mod filtered_group; pub(super) mod reuse; pub(super) mod rollup; pub(super) mod scan_labels; +pub(super) mod shared_producer; pub(super) mod state_configurations; pub(super) mod summary_inputs; pub(super) mod summary_results; diff --git a/crates/logical-asap-optimization/src/memo/rules/shared_producer.rs b/crates/logical-asap-optimization/src/memo/rules/shared_producer.rs new file mode 100644 index 000000000..ed1e73f17 --- /dev/null +++ b/crates/logical-asap-optimization/src/memo/rules/shared_producer.rs @@ -0,0 +1,77 @@ +//! The shared-producer rule: readers that need one computation at different +//! values (a top-k `k`, a sketch size, a window and step) share one producer +//! built at a need that covers them, and each reader derives its own result +//! from it as one more candidate; its own candidates stay, so plan selection +//! decides. Identical expressions are the trivial case, shared by memo group +//! identity, and add no candidates here. +use crate::memo::candidates::CandidateMemo; +use crate::memo::check_generation_deadline; +use crate::memo::state_groups::StateGroupExpansionError; +use std::time::Instant; + +/// One reader of a kind of shared producer: its memo group and its need. +pub(in crate::memo) struct Reader { + pub group: usize, + pub need: N, +} + +/// A kind of producer several readers can share. +pub(in crate::memo) trait SharedProducer { + /// What one reader requires of the producer. + type Need; + /// A producer's need that serves several readers at once. + type Covering; + /// The memo's readers of this kind, in sets whose members can all read + /// one producer; in memo order. + fn sets( + &mut self, + memo: &CandidateMemo, + deadline: Option, + ) -> Result>>, StateGroupExpansionError>; + /// The coverings to offer the readers of `set`. + fn coverings( + &self, + set: &[Reader], + deadline: Option, + ) -> Result, StateGroupExpansionError>; + /// Adds `reader`'s derivation from the producer at `covering`, if that + /// producer serves it; the number of candidates added. + fn derive( + &mut self, + memo: &mut CandidateMemo, + reader: &Reader, + covering: &Self::Covering, + ) -> Result; +} + +/// Whether `set`'s readers belong to two or more groups; a set inside one +/// group shares nothing. +pub(in crate::memo) fn spans_groups(set: &[Reader]) -> bool { + set.iter().any(|reader| reader.group != set[0].group) +} + +/// Gives every reader of a set spanning two or more groups its derivation +/// from each of the set's coverings; the number of candidates added. +pub(in crate::memo) fn add_shared_producer_candidates( + memo: &mut CandidateMemo, + producer: &mut P, + deadline: Option, +) -> Result { + check_generation_deadline(deadline)?; + let mut added = 0; + for set in producer.sets(memo, deadline)? { + check_generation_deadline(deadline)?; + if !spans_groups(&set) { + continue; + } + let coverings = producer.coverings(&set, deadline)?; + for reader in &set { + for covering in &coverings { + check_generation_deadline(deadline)?; + added += producer.derive(memo, reader, covering)?; + } + } + } + check_generation_deadline(deadline)?; + Ok(added) +} diff --git a/crates/logical-asap-optimization/src/memo/rules/state_configurations.rs b/crates/logical-asap-optimization/src/memo/rules/state_configurations.rs index 225310bee..7cf6bc00b 100644 --- a/crates/logical-asap-optimization/src/memo/rules/state_configurations.rs +++ b/crates/logical-asap-optimization/src/memo/rules/state_configurations.rs @@ -9,11 +9,14 @@ use asap_types::ir::flat::{flatten, FlatDag}; use asap_types::ir::schema::{sketch::SketchKind, FieldDataType}; use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot}; +use super::shared_producer::{ + add_shared_producer_candidates, spans_groups, Reader, SharedProducer, +}; use crate::asap_aware_cse::{with_pane_chain, with_resized_build}; use crate::memo::candidates::CandidateMemo; use crate::memo::state_configurations::{ - pane_configurations_before, shared_size, state_size_configurations, PaneRequirement, - PaneWidthDomain, + pane_configurations_before, shared_size, state_size_configurations, PaneConfiguration, + PaneRequirement, PaneWidthDomain, }; use crate::memo::state_groups::{same_range_input, simple_build, StateGroupExpansionError}; @@ -57,12 +60,39 @@ fn with_child_nodes( out } -struct Reader { - group: usize, +/// A summary build one candidate reads: its own need of the shared producer. +struct StateNeed { candidate: usize, build: Rc, merge: Option>, family: FieldDataType, + /// The reader's normalized pane layout, when its set shares panes: equal + /// layouts of one set have one index, so a configuration rewrites them once. + normalized_layout: Option, +} + +/// A shared build's need: one pane layout (its index among the set's pane +/// configurations, for request deduplication) or one size. +enum StateCovering { + Panes { + configuration_index: usize, + configuration: PaneConfiguration, + }, + Size(FieldDataType), +} + +/// Summary builds of one family over one input at different sizes or pane +/// layouts: each reader is rewritten to a build at each covering size or +/// layout that serves it, and equal concrete builds become one state group +/// later ([`super::super::state_groups`]). +struct SummaryStates { + deduplicate_requests: bool, + domain: PaneWidthDomain, + /// Every candidate's graph signature by group, so no rewrite repeats one. + known: HashMap<(usize, String), Vec>>, + rewrite_count: usize, + /// The (normalized layout, configuration index) rewrites already made. + requests: HashSet<(usize, usize)>, } pub(in crate::memo) fn add_state_configuration_candidates_in_domain( @@ -79,230 +109,292 @@ fn expand_configurations( domain: PaneWidthDomain, deadline: Option, ) -> Result { - check_generation_deadline(deadline)?; - let mut rewrite_count = 0; - let mut buckets: Vec> = Vec::new(); - for (group, candidates) in memo.candidates.iter().enumerate() { - check_generation_deadline(deadline)?; - for (candidate, entry) in candidates.iter().enumerate() { + let mut states = SummaryStates { + deduplicate_requests, + domain, + known: HashMap::new(), + rewrite_count: 0, + requests: HashSet::new(), + }; + add_shared_producer_candidates(memo, &mut states, deadline)?; + Ok(states.rewrite_count) +} + +impl SharedProducer for SummaryStates { + type Need = StateNeed; + type Covering = StateCovering; + + fn sets( + &mut self, + memo: &CandidateMemo, + deadline: Option, + ) -> Result>>, StateGroupExpansionError> { + let mut sets: Vec>> = Vec::new(); + for (group, candidates) in memo.candidates.iter().enumerate() { check_generation_deadline(deadline)?; - let Some((build, _)) = simple_build(memo, entry) else { - continue; - }; - // Configuration expansion precedes Stage 2. The rewrite helpers - // create unassigned state; preserve any already assigned build. - if build.timing.is_some() - || build.materialization != asap_types::physical::Materialization::NotMaterialized - { - continue; - } - let family = match build.asap() { - Some(ASAPOp::SummaryAgg { family, .. } | ASAPOp::PaneChain { family, .. }) => { - family.clone() - } - _ => unreachable!(), - }; - let merge = if matches!(build.asap(), Some(ASAPOp::PaneChain { .. })) { - let merges: Vec<_> = OperatorNode::reachable(&entry.node).into_iter().filter(|node| { - matches!(node.asap(), Some(ASAPOp::SummaryMerge {children, group_by: None, ..}) - if children.len() == 1 && Rc::ptr_eq(&children[0], &build)) - }).collect(); - // The existing rewrite changes one plain merge. Other pane uses - // need their own explicit transformation, not guessed semantics. - if merges.len() != 1 { + for (candidate, entry) in candidates.iter().enumerate() { + check_generation_deadline(deadline)?; + let Some((build, _)) = simple_build(memo, entry) else { + continue; + }; + // Configuration expansion precedes Stage 2. The rewrite helpers + // create unassigned state; preserve any already assigned build. + if build.timing.is_some() + || build.materialization + != asap_types::physical::Materialization::NotMaterialized + { continue; } - Some(Rc::clone(&merges[0])) - } else { - None - }; - let reader = Reader { - group, - candidate, - build, - merge, - family, - }; - if let Some(bucket) = buckets - .iter_mut() - .find(|bucket| compatible(&bucket[0], &reader)) - { - bucket.push(reader); - } else { - buckets.push(vec![reader]); + let family = match build.asap() { + Some(ASAPOp::SummaryAgg { family, .. } | ASAPOp::PaneChain { family, .. }) => { + family.clone() + } + _ => unreachable!(), + }; + let merge = if matches!(build.asap(), Some(ASAPOp::PaneChain { .. })) { + let merges: Vec<_> = OperatorNode::reachable(&entry.node).into_iter().filter(|node| { + matches!(node.asap(), Some(ASAPOp::SummaryMerge {children, group_by: None, ..}) + if children.len() == 1 && Rc::ptr_eq(&children[0], &build)) + }).collect(); + // The existing rewrite changes one plain merge. Other pane uses + // need their own explicit transformation, not guessed semantics. + if merges.len() != 1 { + continue; + } + Some(Rc::clone(&merges[0])) + } else { + None + }; + let reader = Reader { + group, + need: StateNeed { + candidate, + build, + merge, + family, + normalized_layout: None, + }, + }; + match sets + .iter_mut() + .find(|set| compatible(&set[0].need, &reader.need)) + { + Some(set) => set.push(reader), + None => sets.push(vec![reader]), + } } } - } - let mut known: HashMap<(usize, String), Vec>> = HashMap::new(); - for (group, candidates) in memo.candidates.iter().enumerate() { - check_generation_deadline(deadline)?; - for candidate in candidates { + for (group, candidates) in memo.candidates.iter().enumerate() { check_generation_deadline(deadline)?; - known - .entry((group, graph_signature(&candidate.node)?)) - .or_default() - .push(Rc::clone(&candidate.node)); + for candidate in candidates { + check_generation_deadline(deadline)?; + self.known + .entry((group, graph_signature(&candidate.node)?)) + .or_default() + .push(Rc::clone(&candidate.node)); + } + } + let mut layouts: Vec<(usize, usize, FlatDag)> = Vec::new(); + for set in &mut sets { + if spans_groups(set) { + assign_normalized_layouts(memo, set, &mut layouts, deadline)?; + } } + Ok(sets) } - for readers in buckets { - check_generation_deadline(deadline)?; - if !readers + + fn coverings( + &self, + set: &[Reader], + deadline: Option, + ) -> Result, StateGroupExpansionError> { + let pane_requirements: Option> = set .iter() - .any(|reader| reader.group != readers[0].group) - { - continue; - } - let pane_requirements: Option> = readers.iter().map(pane_requirement).collect(); - let pane_configurations = pane_requirements - .as_ref() - .map(|requirements| pane_configurations_before(requirements, domain, deadline)) - .transpose()?; - let families = if pane_configurations.is_none() { - let requirements = readers + .map(|reader| pane_requirement(&reader.need)) + .collect(); + let Some(requirements) = pane_requirements else { + let requirements = set .iter() - .map(|reader| match &reader.family { + .map(|reader| match &reader.need.family { FieldDataType::Sketch(kind, _) => Some(kind.params().clone()), _ => None, }) .collect(); - state_size_configurations(requirements, deadline)? + let first = &set[0].need.family; + return Ok(state_size_configurations(requirements, deadline)? .into_iter() - .map(|configuration| match (&readers[0].family, configuration) { - (FieldDataType::Sketch(kind, grouping), Some(params)) => FieldDataType::Sketch( - SketchKind::new(kind.algorithm().clone(), params), - grouping.clone(), - ), - _ => readers[0].family.clone(), + .map(|configuration| { + StateCovering::Size(match (first, configuration) { + (FieldDataType::Sketch(kind, grouping), Some(params)) => { + FieldDataType::Sketch( + SketchKind::new(kind.algorithm().clone(), params), + grouping.clone(), + ) + } + _ => first.clone(), + }) }) - .collect::>() - } else { - Vec::new() + .collect()); }; - // Different original pane layouts can produce the same requested - // reader/configuration. Normalize once per reader and perform that - // rewrite once, instead of once for every existing pane width. - let mut templates: Vec<(usize, usize, FlatDag)> = Vec::new(); - let mut template_keys: HashMap<(usize, String), Vec> = HashMap::new(); - let mut requests = HashSet::new(); - for reader in &readers { - check_generation_deadline(deadline)?; - let original = memo.candidates[reader.group][reader.candidate].clone(); - let roots = [QueryRoot::Operator(Rc::clone(&original.node))]; - let mut rewritten = Vec::new(); - if let Some(configurations) = &pane_configurations { - let requirement = pane_requirement(reader).expect("pane bucket"); - let template = if let Ok(count) = u32::try_from(requirement.window_ms) { - let normalized = with_pane_chain( - &roots, - reader.merge.as_ref().unwrap(), - reader.family.clone(), - 1, - count, - ) - .map_err(|error| { - StateGroupExpansionError::StateConfiguration(error.to_string()) - })?; - let flat = flatten(&normalized).0; - let signature = serde_json::to_string(&flat).map_err(|error| { - StateGroupExpansionError::StateConfiguration(error.to_string()) - })?; - let indices = template_keys.entry((reader.group, signature)).or_default(); - let existing = indices.iter().copied().find(|&index| { - let (group, candidate, previous) = &templates[index]; - let previous_candidate = &memo.candidates[*group][*candidate]; - previous == &flat - && previous_candidate.origin == original.origin - && previous_candidate.children == original.children - && previous_candidate.implements == original.implements - && previous_candidate.budget_share == original.budget_share - }); - Some(existing.unwrap_or_else(|| { - let index = templates.len(); - templates.push((reader.group, reader.candidate, flat)); - indices.push(index); - index - })) - } else { - // Normalization must not narrow the legal long-window domain. - None - }; - for (configuration_index, configuration) in configurations.iter().enumerate() { - check_generation_deadline(deadline)?; - if !configuration.serves(&requirement) { - continue; - } - let Ok(count) = u32::try_from( - configuration.longest_window_ms / configuration.pane_width_ms, - ) else { - continue; - }; - if deduplicate_requests - && template.is_some_and(|template| { - !requests.insert((template, configuration_index)) - }) - { - continue; - } - rewrite_count += 1; - rewritten.push(with_pane_chain( - &roots, - reader.merge.as_ref().unwrap(), - configuration.family.clone(), - configuration.pane_width_ms, - count, - )); - } - } else { - for family in &families { - check_generation_deadline(deadline)?; - if !shared_size(family, &reader.family) - .is_some_and(|joined| same_family(&joined, family)) - { - continue; - } - rewrite_count += 1; - rewritten.push(with_resized_build(&roots, &reader.build, family.clone())); - } + Ok( + pane_configurations_before(&requirements, self.domain, deadline)? + .into_iter() + .enumerate() + .map( + |(configuration_index, configuration)| StateCovering::Panes { + configuration_index, + configuration, + }, + ) + .collect(), + ) + } + + fn derive( + &mut self, + memo: &mut CandidateMemo, + reader: &Reader, + covering: &StateCovering, + ) -> Result { + let need = &reader.need; + let serves = match covering { + StateCovering::Panes { configuration, .. } => { + configuration.serves(&pane_requirement(need).expect("pane set")) } - let boundaries: HashSet<_> = original - .children - .iter() - .map(|group| Rc::as_ptr(&memo.groups[group.0 as usize].node)) - .collect(); - for result in rewritten { - check_generation_deadline(deadline)?; - let roots = result.map_err(|error| { - StateGroupExpansionError::StateConfiguration(error.to_string()) - })?; - let [QueryRoot::Operator(node)] = roots.as_slice() else { - unreachable!("operator rewrite") + StateCovering::Size(family) => { + shared_size(family, &need.family).is_some_and(|joined| same_family(&joined, family)) + } + }; + if !serves { + return Ok(0); + } + let original = memo.candidates[reader.group][need.candidate].clone(); + let roots = [QueryRoot::Operator(Rc::clone(&original.node))]; + let result = match covering { + StateCovering::Panes { + configuration_index, + configuration, + } => { + let Ok(count) = + u32::try_from(configuration.longest_window_ms / configuration.pane_width_ms) + else { + return Ok(0); }; - let node = with_child_nodes(&original.node, node, &boundaries, &mut HashMap::new()); - // Resizing affects internal state types, never public output fields. - if original.node.schema.fields != node.schema.fields - || original.node.schema.time_index != node.schema.time_index - || original.node.schema.unique_keys != node.schema.unique_keys - || (original.node.schema.closed && !node.schema.closed) + if self.deduplicate_requests + && need + .normalized_layout + .is_some_and(|layout| !self.requests.insert((layout, *configuration_index))) { - return Err(StateGroupExpansionError::StateConfiguration( - "state rewrite changed the reader output contract".into(), - )); - } - let previous = known - .entry((reader.group, graph_signature(&node)?)) - .or_default(); - if previous.iter().any(|known| **known == *node) { - continue; + return Ok(0); } - previous.push(Rc::clone(&node)); - let mut candidate = original.clone(); - candidate.node = node; - memo.candidates[reader.group].push(candidate); + self.rewrite_count += 1; + with_pane_chain( + &roots, + need.merge.as_ref().unwrap(), + configuration.family.clone(), + configuration.pane_width_ms, + count, + ) + } + StateCovering::Size(family) => { + self.rewrite_count += 1; + with_resized_build(&roots, &need.build, family.clone()) } + }; + let boundaries: HashSet<_> = original + .children + .iter() + .map(|group| Rc::as_ptr(&memo.groups[group.0 as usize].node)) + .collect(); + let roots = result + .map_err(|error| StateGroupExpansionError::StateConfiguration(error.to_string()))?; + let [QueryRoot::Operator(node)] = roots.as_slice() else { + unreachable!("operator rewrite") + }; + let node = with_child_nodes(&original.node, node, &boundaries, &mut HashMap::new()); + // Resizing affects internal state types, never public output fields. + if original.node.schema.fields != node.schema.fields + || original.node.schema.time_index != node.schema.time_index + || original.node.schema.unique_keys != node.schema.unique_keys + || (original.node.schema.closed && !node.schema.closed) + { + return Err(StateGroupExpansionError::StateConfiguration( + "state rewrite changed the reader output contract".into(), + )); + } + let previous = self + .known + .entry((reader.group, graph_signature(&node)?)) + .or_default(); + if previous.iter().any(|known| **known == *node) { + return Ok(0); } + previous.push(Rc::clone(&node)); + let mut candidate = original; + candidate.node = node; + memo.candidates[reader.group].push(candidate); + Ok(1) + } +} + +/// Numbers the normalized pane layouts of `set`'s readers, when it shares +/// panes: different original layouts can produce the same requested +/// reader/configuration, so each is normalized once and its rewrite made once, +/// instead of once for every existing pane width. Numbers continue across +/// sets (`layouts`), so a configuration index of one set never matches +/// another's. +fn assign_normalized_layouts( + memo: &CandidateMemo, + set: &mut [Reader], + layouts: &mut Vec<(usize, usize, FlatDag)>, + deadline: Option, +) -> Result<(), StateGroupExpansionError> { + let Some(requirements) = set + .iter() + .map(|reader| pane_requirement(&reader.need)) + .collect::>>() + else { + return Ok(()); + }; + let mut layout_keys: HashMap<(usize, String), Vec> = HashMap::new(); + for (reader, requirement) in set.iter_mut().zip(&requirements) { + check_generation_deadline(deadline)?; + let Ok(count) = u32::try_from(requirement.window_ms) else { + // Normalization must not narrow the legal long-window domain. + continue; + }; + let need = &reader.need; + let original = &memo.candidates[reader.group][need.candidate]; + let roots = [QueryRoot::Operator(Rc::clone(&original.node))]; + let normalized = with_pane_chain( + &roots, + need.merge.as_ref().unwrap(), + need.family.clone(), + 1, + count, + ) + .map_err(|error| StateGroupExpansionError::StateConfiguration(error.to_string()))?; + let flat = flatten(&normalized).0; + let signature = serde_json::to_string(&flat) + .map_err(|error| StateGroupExpansionError::StateConfiguration(error.to_string()))?; + let indices = layout_keys.entry((reader.group, signature)).or_default(); + let existing = indices.iter().copied().find(|&index| { + let (group, candidate, previous) = &layouts[index]; + let previous_candidate = &memo.candidates[*group][*candidate]; + previous == &flat + && previous_candidate.origin == original.origin + && previous_candidate.children == original.children + && previous_candidate.implements == original.implements + && previous_candidate.budget_share == original.budget_share + }); + reader.need.normalized_layout = Some(existing.unwrap_or_else(|| { + let index = layouts.len(); + layouts.push((reader.group, need.candidate, flat)); + indices.push(index); + index + })); } - check_generation_deadline(deadline)?; - Ok(rewrite_count) + Ok(()) } fn graph_signature(node: &Rc) -> Result { @@ -310,7 +402,7 @@ fn graph_signature(node: &Rc) -> Result Option { +fn pane_requirement(reader: &StateNeed) -> Option { let ASAPOp::PaneChain { pane_width_ms, pane_count, @@ -333,7 +425,7 @@ fn same_family(left: &FieldDataType, right: &FieldDataType) -> bool { left == right && serde_json::to_string(left).ok() == serde_json::to_string(right).ok() } -fn compatible(left: &Reader, right: &Reader) -> bool { +fn compatible(left: &StateNeed, right: &StateNeed) -> bool { if shared_size(&left.family, &right.family).is_none() { return false; } @@ -349,7 +441,7 @@ fn compatible(left: &Reader, right: &Reader) -> bool { } // Ignore only the fields this rule rewrites. Build guarantees are derived // from these operators and are recomposed by the existing rewrite helpers. - let signature = |reader: &Reader| { + let signature = |reader: &StateNeed| { let mut operator = reader.build.operator.map_children(|_| ()); match &mut operator { Operator::ASAP(ASAPOp::SummaryAgg { family, .. }) => *family = left.family.clone(), diff --git a/crates/logical-asap-optimization/src/memo/rules/topk.rs b/crates/logical-asap-optimization/src/memo/rules/topk.rs index bbbe88d09..fb3cde63c 100644 --- a/crates/logical-asap-optimization/src/memo/rules/topk.rs +++ b/crates/logical-asap-optimization/src/memo/rules/topk.rs @@ -1,6 +1,7 @@ //! Exact ranking reuse with explicit exact-input groups. No global selection //! flags are required: only exact implementations can supply these inputs. use super::reuse::{append_producer, append_reader, eligible, replace}; +use super::shared_producer::{add_shared_producer_candidates, Reader, SharedProducer}; use crate::asap_aware_cse::topk_reuse_candidates; use crate::memo::candidates::{CandidateMemo, CandidateOrigin}; use crate::memo::check_generation_deadline; @@ -18,60 +19,126 @@ pub(in crate::memo) fn add_topk_candidates( memo: &mut CandidateMemo, deadline: Option, ) -> Result { - check_generation_deadline(deadline)?; - let mut owners = HashMap::new(); - let mut roots = Vec::new(); - for (group, entry) in memo.groups.iter().enumerate() { - check_generation_deadline(deadline)?; - let [input] = entry.node.children()[..] else { - continue; - }; - let exact = Rc::new( - (**input) - .clone() - .with_guarantee(Some(AccuracyGuarantee::exact( - "exact input candidate group", - ))), - ); - let Ok(copy) = entry.node.with_new_children(|_| Rc::clone(&exact)) else { - continue; - }; - let copy = Rc::new(copy); - owners.insert(Rc::as_ptr(©), group); - roots.push(QueryRoot::Operator(copy)); - } - let mut exact_groups = HashMap::new(); - let mut added = 0; - for reuse in topk_reuse_candidates(&roots) { - check_generation_deadline(deadline)?; - let members: Vec<_> = reuse - .members - .iter() - .filter_map(|node| owners.get(&Rc::as_ptr(node)).copied()) - .collect(); - if members.len() < 2 { - continue; + add_shared_producer_candidates( + memo, + &mut TopKRankings { + exact_groups: HashMap::new(), + }, + deadline, + ) +} + +/// Exact top-k readers of one input and partition with different `k`: the +/// covering need is the ranking with the largest `k`, `Limit(k_max)` over +/// `Sort` of an exact input ([`topk_reuse_candidates`]); each reader derives +/// `Limit(k) ▸ Project` from it. +struct TopKRankings { + /// Each input group's exact-only group ([`exact_input_group`]). + exact_groups: HashMap>, +} + +/// One input's top-k readers' shared ranking: `Limit(k_max)` over `Sort`, +/// and the readers' groups, whose queries the producer serves. +struct SharedRanking { + ranking: Rc, + reader_groups: Vec, +} + +impl SharedProducer for TopKRankings { + /// The ranking of the reader's input and partition, at the largest `k` + /// of its set. + type Need = Rc; + type Covering = SharedRanking; + + fn sets( + &mut self, + memo: &CandidateMemo, + deadline: Option, + ) -> Result>>>, StateGroupExpansionError> { + let mut owners = HashMap::new(); + let mut roots = Vec::new(); + for (group, entry) in memo.groups.iter().enumerate() { + check_generation_deadline(deadline)?; + let [input] = entry.node.children()[..] else { + continue; + }; + let exact = Rc::new( + (**input) + .clone() + .with_guarantee(Some(AccuracyGuarantee::exact( + "exact input candidate group", + ))), + ); + let Ok(copy) = entry.node.with_new_children(|_| Rc::clone(&exact)) else { + continue; + }; + let copy = Rc::new(copy); + owners.insert(Rc::as_ptr(©), group); + roots.push(QueryRoot::Operator(copy)); } - let Some(input_group) = memo.groups[members[0]].children.first().copied() else { - continue; + Ok(topk_reuse_candidates(&roots) + .into_iter() + .map(|reuse| { + reuse + .members + .iter() + .filter_map(|node| { + Some(Reader { + group: *owners.get(&Rc::as_ptr(node))?, + need: Rc::clone(&reuse.ranking), + }) + }) + .collect() + }) + .collect()) + } + + fn coverings( + &self, + set: &[Reader>], + _deadline: Option, + ) -> Result, StateGroupExpansionError> { + // The largest `k` holds every smaller top-k, so it serves the set. + Ok(vec![SharedRanking { + ranking: Rc::clone(&set[0].need), + reader_groups: set.iter().map(|reader| reader.group).collect(), + }]) + } + + fn derive( + &mut self, + memo: &mut CandidateMemo, + reader: &Reader>, + covering: &SharedRanking, + ) -> 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]] + .children + .first() + .copied() + else { + return Ok(0); }; - let Some(exact_input) = exact_input_group(memo, input_group, &mut exact_groups)? else { - continue; + let Some(exact_input) = exact_input_group(memo, input_group, &mut self.exact_groups)? + else { + return Ok(0); }; let input = Rc::clone(&memo.groups[exact_input.0 as usize].node); - let queries: Vec<_> = members + 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(); - let sort = reuse.ranking.children()[0] + let sort = covering.ranking.children()[0] .with_new_children(|_| Rc::clone(&input)) .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; let sort_group = append_producer(memo, Rc::new(sort), &queries, CandidateOrigin::TopKReuse, 1)?; let sort = Rc::clone(&memo.groups[sort_group.0 as usize].node); - let ranking = reuse + let ranking = covering .ranking .with_new_children(|_| Rc::clone(&sort)) .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; @@ -83,51 +150,47 @@ pub(in crate::memo) fn add_topk_candidates( 1, )?; let ranking = Rc::clone(&memo.groups[ranking_group.0 as usize].node); - for owner in members { - check_generation_deadline(deadline)?; - let logical = Rc::clone(&memo.groups[owner].node); - let Some(NonASAPOp::Aggregate { - reduction: Reduction::Reduce(keys), - measures, - .. - }) = logical.non_asap() - else { - continue; - }; - let [AggIntent::TopK { k, .. }] = measures.as_slice() else { - continue; - }; - let Some(original) = memo.candidates[owner].iter().find(|candidate| eligible(candidate) && matches!(candidate.node.non_asap(), Some(NonASAPOp::Aggregate {measures, ..}) if matches!(measures.as_slice(), [AggIntent::TopK {..}]))).cloned() else { continue }; - let limited = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Limit { - n: Some(*k), - offset: 0, - partition_by: keys.clone(), - child: Rc::clone(&ranking), - })) - .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; - let columns = asap_types::ir::schema::topk::ranked_row_projection( - &ranking, - &Reduction::Reduce(keys.clone()), - &logical.schema.fields[keys.len()].name, - ) - .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; - let reader = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Project { - cols: columns, - qualifier: None, - child: limited, - })) - .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; - added += usize::from(append_reader( - memo, - owner, - &original, - reader, - CandidateOrigin::TopKReuse, - )?); - } + let owner = reader.group; + let logical = Rc::clone(&memo.groups[owner].node); + let Some(NonASAPOp::Aggregate { + reduction: Reduction::Reduce(keys), + measures, + .. + }) = logical.non_asap() + else { + return Ok(0); + }; + let [AggIntent::TopK { k, .. }] = measures.as_slice() else { + return Ok(0); + }; + let Some(original) = memo.candidates[owner].iter().find(|candidate| eligible(candidate) && matches!(candidate.node.non_asap(), Some(NonASAPOp::Aggregate {measures, ..}) if matches!(measures.as_slice(), [AggIntent::TopK {..}]))).cloned() else { return Ok(0) }; + let limited = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Limit { + n: Some(*k), + offset: 0, + partition_by: keys.clone(), + child: Rc::clone(&ranking), + })) + .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; + let columns = asap_types::ir::schema::topk::ranked_row_projection( + &ranking, + &Reduction::Reduce(keys.clone()), + &logical.schema.fields[keys.len()].name, + ) + .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; + let reader = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Project { + cols: columns, + qualifier: None, + child: limited, + })) + .map_err(|error| StateGroupExpansionError::ReuseRule(error.to_string()))?; + Ok(usize::from(append_reader( + memo, + owner, + &original, + reader, + CandidateOrigin::TopKReuse, + )?)) } - check_generation_deadline(deadline)?; - Ok(added) } fn exact_input_group( diff --git a/crates/logical-asap-optimization/src/memo/shared_producer_snapshot.rs b/crates/logical-asap-optimization/src/memo/shared_producer_snapshot.rs new file mode 100644 index 000000000..1a2755de6 --- /dev/null +++ b/crates/logical-asap-optimization/src/memo/shared_producer_snapshot.rs @@ -0,0 +1,165 @@ +//! Pins the candidates the shared-producer cases generate, recorded before +//! they moved onto one generation path: each workload's expanded memo +//! (groups, candidates, their origins and children) hashes to a fixed value. +use super::candidates::{try_populate_pass1, CandidateMemo, CandidateOrigin}; +use super::state_configurations::PaneWidthDomain; +use super::state_groups::expand_concrete_summary_state_groups_in_domain; +use super::window_candidates::add_local_window_composition; +use super::LogicalMemo; +use crate::per_computation_candidates::generate_per_computation_candidates; +use crate::test_support::{entry_with, lower_promql}; +use asap_types::ir::flat::flatten; +use asap_types::ir::properties::AccuracyTarget; +use asap_types::ir::QueryRoot; +use asap_types::workload::{QueryRecurrence, RepeatedDemand, RepetitionInterval}; +use std::rc::Rc; + +/// The expanded memo of `queries` (each with its error bound), every one +/// repeated every `step_ms` milliseconds, with local pane candidates. +fn expanded(queries: &[(&str, f64)], step_ms: u32) -> CandidateMemo { + let targets: Vec<_> = queries + .iter() + .map(|&(_, error)| { + if error == 0.0 { + AccuracyTarget::exact() + } else { + AccuracyTarget::max_error(error).failure_probability(0.001) + } + }) + .collect(); + let roots: Vec<_> = queries + .iter() + .zip(&targets) + .enumerate() + .map(|(id, ((query, _), target))| { + (id, QueryRoot::Operator(lower_promql(query, target.clone()))) + }) + .collect(); + let initial = LogicalMemo::from_roots(&roots); + let inventory = generate_per_computation_candidates(roots).unwrap(); + let mut memo = try_populate_pass1(initial, &inventory).unwrap(); + let entries: Vec<_> = targets + .into_iter() + .map(|target| { + entry_with( + target, + QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval( + step_ms, + ))), + ) + }) + .collect(); + add_local_window_composition(&mut memo, &entries); + expand_concrete_summary_state_groups_in_domain(memo, PaneWidthDomain::SeedGrids) + .unwrap() + .memo +} + +fn signature(node: &Rc) -> String { + serde_json::to_string(&flatten(&[QueryRoot::Operator(Rc::clone(node))]).0).unwrap() +} + +/// FNV-1a over every group (node, edges, queries, roots) and its candidates +/// (origin, budget share, implemented and child groups, node), in memo +/// order: the search's tie-breaks depend on that order too. +fn fingerprint(memo: &CandidateMemo) -> (usize, u64) { + let mut entries: Vec = Vec::new(); + for (group, candidates) in memo.groups.iter().zip(&memo.candidates) { + let own: Vec = candidates + .iter() + .map(|candidate| { + format!( + "{:?}|{}|{:?}|{:?}|{}", + candidate.origin, + candidate.budget_share, + candidate.implements, + candidate.children, + signature(&candidate.node) + ) + }) + .collect(); + entries.push(format!( + "{}|{:?}|{:?}|{:?}|{:?}=>{:?}", + signature(&group.node), + group.children, + group.parents, + group.queries, + group.roots_of, + own + )); + } + let count = memo.candidates.iter().map(Vec::len).sum(); + let mut hash: u64 = 0xcbf29ce484222325; + for byte in entries.join("\n").bytes() { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x100000001b3); + } + (count, hash) +} + +/// Top-k readers of one ranking with different `k`. +#[test] +fn top_k_reuse_candidates_are_unchanged() { + let memo = expanded( + &[ + ("topk(3, sum_over_time(m_gauge[2m]))", 0.0), + ("topk(10, sum_over_time(m_gauge[2m]))", 0.0), + ], + 60_000, + ); + assert!(memo + .candidates + .iter() + .flatten() + .any(|candidate| candidate.origin == CandidateOrigin::TopKReuse)); + assert_eq!(fingerprint(&memo), (313, 16940643818053118362)); +} + +/// Quantile readers of one input at different sizes. +#[test] +fn shared_size_candidates_are_unchanged() { + let memo = expanded( + &[ + ("quantile_over_time(0.95, latency[5m])", 0.01), + ("quantile_over_time(0.5, latency[5m])", 0.05), + ], + 6_000, + ); + assert_eq!(fingerprint(&memo), (508, 7457272222287953080)); +} + +/// Nested windows of one sum, sharing a pane chain. +#[test] +fn shared_pane_candidates_are_unchanged() { + let memo = expanded( + &[ + ("sum(sum_over_time(m_gauge[2m]))", 0.0), + ("sum(sum_over_time(m_gauge[1m]))", 0.0), + ], + 60_000, + ); + assert!(memo.candidates.iter().flatten().any(|candidate| { + asap_types::ir::OperatorNode::reachable(&candidate.node) + .iter() + .any(|node| matches!(node.asap(), Some(asap_types::ir::ASAPOp::PaneChain { .. }))) + })); + assert_eq!(fingerprint(&memo), (729, 15532247561651645380)); +} + +/// Two top-k sets over different inputs, next to a pane set: pins the order +/// sets are expanded in. +#[test] +fn several_top_k_sets_are_unchanged() { + let memo = expanded( + &[ + ("topk(3, sum_over_time(m_gauge[2m]))", 0.0), + ("topk(10, sum_over_time(m_gauge[2m]))", 0.0), + ("topk(5, sum_over_time(m_gauge[2m]))", 0.0), + ("topk(2, sum_over_time(other[2m]))", 0.0), + ("topk(4, sum_over_time(other[2m]))", 0.0), + ("sum(sum_over_time(other[1m]))", 0.0), + ], + 60_000, + ); + assert_eq!(fingerprint(&memo), (994, 9213568444657927922)); +}