From 465bd4eab98b61b4de9d1641beb73c6be2a91912 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:02:18 +0000 Subject: [PATCH 01/25] plan-selection: sub-expression memo: groups and Pass 1 candidates, capability, window/rollup/top-k transformations, per-group pricing (behind no flag yet) --- .../src/accuracy/estimators/mod.rs | 2 +- .../src/asap_aware_cse.rs | 74 ++ crates/plan-selection/src/lib.rs | 50 +- crates/plan-selection/src/pass/mod.rs | 1 + .../src/pass/subexpr/capability.rs | 402 ++++++++ .../plan-selection/src/pass/subexpr/groups.rs | 517 ++++++++++ crates/plan-selection/src/pass/subexpr/mod.rs | 138 +++ .../src/pass/subexpr/pricing.rs | 716 +++++++++++++ .../plan-selection/src/pass/subexpr/search.rs | 1 + .../src/pass/subexpr/windows.rs | 942 ++++++++++++++++++ 10 files changed, 2841 insertions(+), 2 deletions(-) create mode 100644 crates/plan-selection/src/pass/subexpr/capability.rs create mode 100644 crates/plan-selection/src/pass/subexpr/groups.rs create mode 100644 crates/plan-selection/src/pass/subexpr/mod.rs create mode 100644 crates/plan-selection/src/pass/subexpr/pricing.rs create mode 100644 crates/plan-selection/src/pass/subexpr/search.rs create mode 100644 crates/plan-selection/src/pass/subexpr/windows.rs diff --git a/crates/logical-asap-optimization/src/accuracy/estimators/mod.rs b/crates/logical-asap-optimization/src/accuracy/estimators/mod.rs index afc1501c5..21209c851 100644 --- a/crates/logical-asap-optimization/src/accuracy/estimators/mod.rs +++ b/crates/logical-asap-optimization/src/accuracy/estimators/mod.rs @@ -235,7 +235,7 @@ pub(crate) fn binomial_upper_tail(n: u32, p: f64, k: u32) -> f64 { /// shrinks, so each reader's `k` or heap is first resized to the shared /// (smallest) failure probability at no larger error ([`kll::shared_k`], /// [`univmon::shared_heap`]), and the shared build meets every reader. -pub(crate) fn size_for_all_readers(params: &[&SketchParams]) -> Option { +pub fn size_for_all_readers(params: &[&SketchParams]) -> Option { use SketchParams as P; let (first, rest) = params.split_first()?; let shared = rest.iter().try_fold((*first).clone(), |a, b| { diff --git a/crates/logical-asap-optimization/src/asap_aware_cse.rs b/crates/logical-asap-optimization/src/asap_aware_cse.rs index 0d7fc2dac..cf97221b1 100644 --- a/crates/logical-asap-optimization/src/asap_aware_cse.rs +++ b/crates/logical-asap-optimization/src/asap_aware_cse.rs @@ -1704,6 +1704,80 @@ fn in_panes( recompose_guarantees(&out) } +/// `roots` with the `SummaryMerge` `merge` (over one [`ASAPOp::PaneChain`] +/// of `roots`) reading the newest panes of another chain over the same +/// input: `pane_count` panes of `pane_width_ms` of family `family` (a size +/// that serves every reader of the chain). The merge reads the newest +/// panes covering its own chain's window, which `pane_width_ms` divides +/// and the new chain covers; every guarantee is composed again. +pub fn with_pane_chain( + roots: &[QueryRoot], + merge: &Rc, + family: FieldDataType, + pane_width_ms: u64, + pane_count: u32, +) -> Result, PerComputationCandidateError> { + let Some(ASAPOp::SummaryMerge { children, .. }) = merge.asap() else { + unreachable!("a merge over a pane chain"); + }; + let own = &children[0]; + let Some(ASAPOp::PaneChain { + child, + input, + reduction, + grouping, + filter, + pane_width_ms: own_width, + pane_count: own_count, + .. + }) = own.asap() + else { + unreachable!("a merge over a pane chain"); + }; + let window_ms = own_width * u64::from(*own_count); + let newest = + u32::try_from(window_ms / pane_width_ms).expect("a reader's window is at most the chain's"); + let mut chain = OperatorNode::new(Operator::ASAP(ASAPOp::PaneChain { + child: Rc::clone(child), + family, + input: input.clone(), + reduction: reduction.clone(), + grouping: grouping.clone(), + filter: filter.clone(), + pane_width_ms, + pane_count, + }))?; + for (field, kept) in chain.schema.fields.iter_mut().zip(&own.schema.fields) { + field.name = kept.name.clone(); + field.table = kept.table.clone(); + } + let read = OperatorNode::new(Operator::ASAP(ASAPOp::SummaryMerge { + children: vec![Rc::new(chain)], + newest_panes: (newest < pane_count).then_some(newest), + group_by: None, + }))?; + let replaced = HashMap::from([(Rc::as_ptr(merge), Rc::new(read))]); + let mut rebuilt = HashMap::new(); + let mut out = Vec::new(); + for root in roots { + out.push(match root { + QueryRoot::Operator(node) => { + QueryRoot::Operator(rebuild(node, &replaced, &mut rebuilt)?) + } + QueryRoot::Scalar(expr) => { + let mut nodes = HashMap::new(); + for node in expr.operator_refs() { + nodes.insert(Rc::as_ptr(node), rebuild(node, &replaced, &mut rebuilt)?); + } + QueryRoot::Scalar( + expr.map_operator_refs(&mut |node| nodes[&Rc::as_ptr(node)].clone()), + ) + } + }); + } + recompose_guarantees(&out) +} + /// The [`ASAPOp::PaneChain`] of `candidate`'s panes, with the build's /// output names. fn pane_chain(candidate: &PaneCandidate) -> Result, PerComputationCandidateError> { diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 2b591ce70..28146e320 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1185,6 +1185,54 @@ pub struct PlanSelector<'a> { runner_up: Option, } +/// Stage 3's per-node pricing of `candidate`, a part of a workload whose +/// roots are read by the queries of `entries` (one entry per root), sized +/// for their failure probabilities when that resizes, with the formulas' +/// CPU time scaled by `formula_scale`: its cost and calibration fit. +/// `None` when it fails a check no summary size repairs: a per-series +/// aggregate the executor cannot run, or raw data the deployment does not +/// keep. The memo over sub-expressions prices each group's sub-DAG this +/// way; accuracy and latency are checked on the whole plan. +pub(crate) fn price_part( + candidate: &PhysicalCandidate, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, + formula_scale: f64, +) -> Result, SelectionError> { + let targets = accuracy_targets(entries); + let sized = sized_for_failure_probability(candidate, &targets, data) + .unwrap_or_else(|_| candidate.clone()); + if per_entity_aggregate_violation(&sized.dag).is_some() { + return Ok(None); + } + let nodes = node_map(&sized.dag); + let kept_chains = kept_pane_chains(&sized.dag); + if !models.deployment_capabilities.raw_data_retained + && raw_data_violation(&sized.dag, &nodes, &kept_chains).is_some() + { + return Ok(None); + } + let recurrence: Vec = entries.iter().map(|e| e.recurrence.clone()).collect(); + let range_evaluations: Vec<_> = entries.iter().map(|e| e.range_evaluation).collect(); + price_at_scale( + &sized.dag, + &nodes, + &kept_chains, + &recurrence, + &range_evaluations, + data, + models.plan_selection, + formula_scale, + ) + .map(Some) + .map_err(|(node, error)| SelectionError::Cost { + candidate: candidate.id.clone(), + node, + error, + }) +} + /// Prices within this fraction of each other are the same sum of charges /// rounded differently (Stage 3's tie tolerance is 1e-9). const ROUNDING: f64 = 1e-12; @@ -4064,7 +4112,7 @@ impl CalibrationFit { /// Measured ÷ formula time, 1 without measured nodes (or with only /// nodes the formulas price free). - fn factor(self) -> f64 { + pub(crate) fn factor(self) -> f64 { match self.formula_ms_per_second > 0.0 && self.measured_ms_per_second > 0.0 { true => self.measured_ms_per_second / self.formula_ms_per_second, false => 1.0, diff --git a/crates/plan-selection/src/pass/mod.rs b/crates/plan-selection/src/pass/mod.rs index d95cbef92..16053a432 100644 --- a/crates/plan-selection/src/pass/mod.rs +++ b/crates/plan-selection/src/pass/mod.rs @@ -14,6 +14,7 @@ mod components; mod interchangeable; +mod subexpr; mod query_candidates; mod stage1_to_stage3; diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs new file mode 100644 index 000000000..f67dd3511 --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -0,0 +1,402 @@ +//! The identical-expression and summary-capability transformations of +//! Pass 2 over the memo. +//! +//! **Identical expression.** The memo hash-conses the workload's logical +//! DAG, so identical sub-expressions of several queries are already one +//! group, read by each of their parents: there is nothing to add. Whether +//! the shared group is computed once is the group's single choice of +//! implementation and materialization (Stage 2's `NotMaterialized` or +//! `MatAtQuery` for a group with several parents), which the search makes +//! once for all of its parents. Leaving identical sub-expressions apart is +//! never cheaper under the stages' per-node pricing, since every parent +//! reads the same implementation either way. +//! +//! **Summary capability.** Reader groups whose Pass 1 summary candidates +//! build one family (one sketch algorithm and grouping at any size, or one +//! exact accumulator) over one input group, with the same reduction, +//! grouping, item and weight, can read one build. Each such reader +//! candidate gets a [`CandidateOrigin::SharedBuild`] twin carrying a +//! [`CouplingKind::SharedBuildSize`] site. The twin's node keeps the +//! reader's own size: the shared build's size is the site's threshold, +//! chosen by the search among [`size_thresholds`], and a reader may read it +//! iff [`may_read`]. + +use std::collections::hash_map::DefaultHasher; +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use asap_logical_asap_optimization::accuracy::estimators::size_for_all_readers; +use asap_types::ir::operator::Reduction; +use asap_types::ir::schema::sketch::{GroupingStrategy, SketchParams}; +use asap_types::ir::schema::FieldDataType; +use asap_types::ir::{ASAPOp, OperatorNode}; + +use super::{CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo}; + +/// One size a shared build may take: the per-parameter maximum +/// ([`size_for_all_readers`]) of the requirements of the readers it admits. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct SizeThreshold { + /// The build's size parameters; `None` for an exact accumulator, which + /// has no size and admits every reader. + pub params: Option, +} + +/// Add the summary-capability transformation's candidates to `memo`: for +/// every set of two or more reader groups whose Pass 1 summary candidates +/// build one family over one input group, a +/// [`CandidateOrigin::SharedBuild`] copy of each such candidate with that +/// set's [`CouplingKind::SharedBuildSize`] site added. +pub(crate) fn add_shared_builds(memo: &mut Memo) { + // Per site key, the (group, candidate index) of each reader candidate. + let mut readers: HashMap> = HashMap::new(); + let mut order: Vec = Vec::new(); + for (group, candidates) in memo.candidates.iter().enumerate() { + for (index, candidate) in candidates.iter().enumerate() { + if candidate.origin != CandidateOrigin::Pass1 { + continue; + } + let Some(key) = site_key(memo, candidate) else { + continue; + }; + let entry = readers.entry(key).or_default(); + if entry.is_empty() { + order.push(key); + } + entry.push((group, index)); + } + } + for key in order { + let members = &readers[&key]; + let first_group = members[0].0; + if members.iter().all(|&(group, _)| group == first_group) { + continue; + } + let site = CouplingSite { + key, + kind: CouplingKind::SharedBuildSize, + }; + for &(group, index) in members { + let mut shared = memo.candidates[group][index].clone(); + shared.origin = CandidateOrigin::SharedBuild; + shared.coupling_sites.push(site); + memo.candidates[group].push(shared); + } + } +} + +/// The sizes the shared build of `site` may take: the per-parameter +/// maximum of every set of its readers' requirements (each reader's alone +/// among them), each once, without a threshold that admits no more +/// readers than a smaller one. Empty for a site no reader carries. +pub(crate) fn size_thresholds(memo: &Memo, site: &CouplingSite) -> Vec { + let requirements: Vec> = memo + .candidates + .iter() + .flatten() + .filter(|c| c.origin == CandidateOrigin::SharedBuild && c.coupling_sites.contains(site)) + .filter_map(|c| reader_build(c).map(|(_, params)| params)) + .collect(); + if requirements.iter().any(Option::is_none) { + // Exact accumulators: one build, no size. + return vec![SizeThreshold { params: None }]; + } + let requirements: Vec = requirements.into_iter().flatten().collect(); + // Every join of a set of requirements, each once. + let mut joins: Vec = Vec::new(); + for requirement in &requirements { + let mut grown: Vec = joins + .iter() + .filter_map(|join| size_for_all_readers(&[join, requirement])) + .collect(); + grown.push(requirement.clone()); + for join in grown { + if !joins.iter().any(|known| same_size(known, &join)) { + joins.push(join); + } + } + } + // A threshold larger than the join of the readers it admits admits no + // more readers than that join. + joins + .into_iter() + .filter(|threshold| { + let admitted: Vec<&SketchParams> = requirements + .iter() + .filter(|r| admits(threshold, r)) + .collect(); + size_for_all_readers(&admitted).is_some_and(|join| same_size(&join, threshold)) + }) + .map(|params| SizeThreshold { + params: Some(params), + }) + .collect() +} + +/// Whether `reader` (a shared-build reader candidate) may read the shared +/// build at `threshold`: its own requirement is no stricter, so the build +/// sized for both is the threshold itself. +pub(crate) fn may_read(threshold: &SizeThreshold, reader: &ImplementationCandidate) -> bool { + match (&threshold.params, reader_build(reader)) { + (_, None) => false, + (None, Some((_, params))) => params.is_none(), + (Some(threshold), Some((_, Some(params)))) => admits(threshold, ¶ms), + (Some(_), Some((_, None))) => false, + } +} + +/// Whether a build at `threshold` serves a reader that requires `required`. +fn admits(threshold: &SketchParams, required: &SketchParams) -> bool { + size_for_all_readers(&[threshold, required]).is_some_and(|join| same_size(&join, threshold)) +} + +/// Equal sizes, the recorded failure probabilities and errors included +/// (which `SketchParams` equality ignores). +fn same_size(a: &SketchParams, b: &SketchParams) -> bool { + a == b && format!("{a:?}") == format!("{b:?}") +} + +/// The site key of `candidate`'s summary build, when it is a shareable +/// one over a group of `memo`: a hash of the family (algorithm and +/// grouping, or the exact accumulator), the input group, the reduction, +/// the grouping, the filter, and the item and weight. +fn site_key(memo: &Memo, candidate: &ImplementationCandidate) -> Option { + let (build, _) = reader_build(candidate)?; + let Some(ASAPOp::SummaryAgg { + child, + family, + input, + reduction, + grouping, + filter, + }) = build.asap() + else { + return None; + }; + let input_group = candidate.children.iter().find(|&&GroupId(g)| { + let node = &memo.groups[g as usize].node; + Rc::ptr_eq(node, child) || **node == **child + })?; + let family_key = match family { + FieldDataType::Sketch(kind, strategy) => format!("{:?}/{strategy:?}", kind.algorithm()), + exact => format!("{exact:?}"), + }; + let mut hasher = DefaultHasher::new(); + ( + family_key, + input_group, + format!("{reduction:?}/{grouping:?}/{filter:?}/{input:?}"), + ) + .hash(&mut hasher); + Some(hasher.finish()) +} + +/// `candidate`'s summary build, nearest its root, that the summary +/// capability rule can share, with its size (`None` for an exact +/// accumulator): a per-subpopulation sketch read by an estimate, or a +/// per-series exact accumulator read by its finalization. +fn reader_build( + candidate: &ImplementationCandidate, +) -> Option<(Rc, Option)> { + let mut queue = std::collections::VecDeque::from([Rc::clone(&candidate.node)]); + while let Some(node) = queue.pop_front() { + let build = match node.asap() { + Some(ASAPOp::SummaryEstimate { summary_input, .. }) => Some(summary_input), + Some(ASAPOp::FinalizeExactAccumulator { child }) => Some(child), + _ => None, + }; + if let Some(build) = build { + match build.asap() { + Some(ASAPOp::SummaryAgg { + family: FieldDataType::Sketch(kind, GroupingStrategy::PerSubpopulationInstance), + .. + }) => return Some((Rc::clone(build), Some(kind.params().clone()))), + Some(ASAPOp::SummaryAgg { + family: FieldDataType::ExactAggregate(..), + reduction: Reduction::PerEntity, + .. + }) => return Some((Rc::clone(build), None)), + _ => {} + } + } + queue.extend(node.children().into_iter().cloned()); + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_logical_asap_optimization::test_support::{metric_scan, time_range}; + use asap_types::ir::scalar::ColumnRef; + use asap_types::ir::schema::sketch::{SketchAlgorithm, SketchKind, SketchStatistic}; + use asap_types::ir::schema::{Field, Schema, SummaryUpdate}; + use asap_types::ir::Operator; + use std::time::Duration; + + /// A per-series sketch build of `algorithm` with `params` over `child`. + fn build( + algorithm: SketchAlgorithm, + params: SketchParams, + child: &Rc, + ) -> Rc { + let family = FieldDataType::Sketch( + SketchKind::new(algorithm, params), + GroupingStrategy::PerSubpopulationInstance, + ); + let schema = Schema::lifted(vec![Field::new("state", family.clone(), false)], None); + Rc::new(OperatorNode::with_schema( + Operator::ASAP(ASAPOp::SummaryAgg { + child: Rc::clone(child), + family, + input: SummaryUpdate::column(ColumnRef::SampleValue), + reduction: Reduction::PerEntity, + grouping: GroupingStrategy::PerSubpopulationInstance, + filter: None, + }), + schema, + )) + } + + /// The 0.9 quantile of `build`. + fn estimate(build: Rc) -> Rc { + Rc::new(OperatorNode::with_schema( + Operator::ASAP(ASAPOp::SummaryEstimate { + summary_input: build, + query: SketchStatistic::Quantile { q: 0.9 }, + }), + Schema::lifted(Vec::new(), None), + )) + } + + fn kll(k: u32) -> (SketchAlgorithm, SketchParams) { + ( + SketchAlgorithm::Kll, + SketchParams::Kll { + k, + failure_probability: None, + }, + ) + } + + fn ddsketch(alpha: f64) -> (SketchAlgorithm, SketchParams) { + (SketchAlgorithm::DDSketch, SketchParams::DDSketch { alpha }) + } + + /// A memo of a scan (group 0), its 5-minute window (group 1), and one + /// reader group per entry of `readers` over the window, each with one + /// Pass 1 candidate: the 0.9 quantile of that sketch over the window. + fn memo(readers: &[(SketchAlgorithm, SketchParams)]) -> Memo { + let scan = metric_scan(&[]); + let window = time_range(Duration::from_secs(300), Rc::clone(&scan)); + let group = |node: &Rc, children: Vec| super::super::MemoGroup { + node: Rc::clone(node), + children, + parents: Vec::new(), + queries: Vec::new(), + roots_of: Vec::new(), + }; + let mut memo = Memo { + groups: vec![group(&scan, vec![]), group(&window, vec![GroupId(0)])], + candidates: vec![Vec::new(), Vec::new()], + }; + for (q, (algorithm, params)) in readers.iter().enumerate() { + let node = estimate(build(algorithm.clone(), params.clone(), &window)); + memo.groups.push(super::super::MemoGroup { + roots_of: vec![q], + queries: vec![q], + ..group(&node, vec![GroupId(1)]) + }); + memo.candidates.push(vec![ImplementationCandidate { + implements: vec![GroupId(memo.groups.len() as u32 - 1)], + node, + children: vec![GroupId(1)], + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + }]); + } + add_shared_builds(&mut memo); + memo + } + + /// The shared-build candidates of `memo`, and their sites. + fn shared(memo: &Memo) -> (Vec<&ImplementationCandidate>, Vec) { + let readers: Vec<&ImplementationCandidate> = memo + .candidates + .iter() + .flatten() + .filter(|c| c.origin == CandidateOrigin::SharedBuild) + .collect(); + let mut sites: Vec = readers + .iter() + .flat_map(|c| c.coupling_sites.iter().copied()) + .collect(); + sites.sort(); + sites.dedup(); + (readers, sites) + } + + fn k_of(threshold: &SizeThreshold) -> u32 { + match threshold.params { + Some(SketchParams::Kll { k, .. }) => k, + ref other => panic!("not a KLL size: {other:?}"), + } + } + + /// Two KLL readers of different `k` over one window: one site, a + /// shared-build candidate per reader, and one threshold per reader's + /// `k`. + #[test] + fn two_kll_readers_give_one_site_with_two_thresholds() { + let memo = memo(&[kll(200), kll(400)]); + let (readers, sites) = shared(&memo); + assert_eq!(readers.len(), 2); + assert_eq!(sites.len(), 1); + assert_eq!(sites[0].kind, CouplingKind::SharedBuildSize); + let mut ks: Vec = size_thresholds(&memo, &sites[0]).iter().map(k_of).collect(); + ks.sort_unstable(); + assert_eq!(ks, vec![200, 400]); + } + + /// The larger threshold admits both readers; the smaller only the + /// reader that requires it. + #[test] + fn the_larger_threshold_admits_both_readers_the_smaller_one() { + let memo = memo(&[kll(200), kll(400)]); + let (readers, sites) = shared(&memo); + let thresholds = size_thresholds(&memo, &sites[0]); + let admitted = |k: u32| { + let threshold = thresholds.iter().find(|t| k_of(t) == k).unwrap(); + readers.iter().filter(|r| may_read(threshold, r)).count() + }; + assert_eq!(admitted(400), 2); + assert_eq!(admitted(200), 1); + } + + /// A DDSketch and a KLL reader over one window are different families: + /// no site, no shared-build candidate. + #[test] + fn ddsketch_and_kll_readers_do_not_share_a_site() { + let memo = memo(&[ddsketch(0.01), kll(200)]); + let (readers, sites) = shared(&memo); + assert!(readers.is_empty()); + assert!(sites.is_empty()); + } + + /// Three DDSketch readers, two of one α: thresholds are the distinct + /// α's, the strictest admitting all three. + #[test] + fn equal_requirements_give_one_threshold() { + let memo = memo(&[ddsketch(0.02), ddsketch(0.01), ddsketch(0.02)]); + let (readers, sites) = shared(&memo); + assert_eq!(readers.len(), 3); + let thresholds = size_thresholds(&memo, &sites[0]); + assert_eq!(thresholds.len(), 2); + let strictest = thresholds + .iter() + .find(|t| t.params == Some(SketchParams::DDSketch { alpha: 0.01 })) + .unwrap(); + assert!(readers.iter().all(|r| may_read(strictest, r))); + } +} diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs new file mode 100644 index 000000000..4928e8675 --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -0,0 +1,517 @@ +//! The hash-consed group DAG and Pass 1's implementation candidates per +//! group (design steps 1 and 2, Pass 1 part). +//! +//! Groups are the workload's logical DAG interned by +//! [`share_common_sub_dags`], the identical-expression rule's legality +//! included: structurally equal nodes are one group when sharing them is +//! legal (an ASAP node, or one whose schema has a unique key), so the +//! unkeyed scans of two queries stay two groups while a keyed aggregate +//! over them is one. +//! +//! A Pass 1 target's aggregate group gets one candidate per combination of +//! its measures' Pass 1 candidates: the aggregate alone composed by +//! [`compose_logical_candidate`] with that choice, every target beneath it +//! exact. Its private nodes are the composed nodes above the first node +//! equal to one of the group's descendant groups, which are its children. +//! Summaries are sized for their measure's whole budget, as composing the +//! aggregate alone does. + +use std::collections::hash_map::DefaultHasher; +use std::collections::{HashMap, HashSet}; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use asap_logical_asap_optimization::per_computation_candidates::{ + compose_logical_candidate, PerComputationCandidates, +}; +use asap_types::ir::cse::share_common_sub_dags; +use asap_types::ir::{OperatorNode, QueryRoot}; + +use super::{CandidateOrigin, GroupId, ImplementationCandidate, Memo, MemoGroup}; + +/// The memo of `per_computation`'s workload: its roots hash-consed into +/// groups, children before parents, and each group's Pass 1 candidates. +/// `per_computation` is the workload's Pass 1 inventory, pruned for the +/// deployment's capabilities. +pub(crate) fn memo_of(per_computation: &PerComputationCandidates) -> Memo { + let (groups, group_of) = hash_consed(&per_computation.roots); + let candidates = pass1_candidates(per_computation, &groups, &group_of); + Memo { groups, candidates } +} + +/// Each original node (of `roots`), by address, to its group. +type GroupOf = HashMap<*const OperatorNode, GroupId>; + +/// The groups of `roots`, children before parents, and each original node's +/// group. A scalar query's roots are the operator nodes its expression +/// reads. +fn hash_consed(roots: &[(Id, QueryRoot)]) -> (Vec, GroupOf) { + let originals: Vec<(usize, Rc)> = roots + .iter() + .enumerate() + .flat_map(|(query, (_, root))| { + operator_roots(root) + .into_iter() + .map(move |node| (query, node)) + }) + .collect(); + let interned = share_common_sub_dags(originals.clone()); + // Number the interned nodes children first. + let mut groups: Vec = Vec::new(); + let mut id_of: HashMap<*const OperatorNode, GroupId> = HashMap::new(); + for (_, root) in &interned { + number(root, &mut groups, &mut id_of); + } + for (query, root) in &interned { + let g = id_of[&Rc::as_ptr(root)]; + let roots_of = &mut groups[g.0 as usize].roots_of; + if !roots_of.contains(query) { + roots_of.push(*query); + } + let mut reached = HashSet::new(); + reach(&groups, g, &mut reached); + for reached in reached { + let queries = &mut groups[reached.0 as usize].queries; + if !queries.contains(query) { + queries.push(*query); + } + } + } + for group in &mut groups { + group.queries.sort_unstable(); + } + // Interning rebuilds every node over its children in order, so the two + // DAGs are walked side by side. + let mut group_of = GroupOf::new(); + for ((_, original), (_, root)) in originals.iter().zip(&interned) { + pair(original, root, &id_of, &mut group_of); + } + (groups, group_of) +} + +/// The operator nodes a query root reads. +fn operator_roots(root: &QueryRoot) -> Vec> { + match root { + QueryRoot::Operator(node) => vec![Rc::clone(node)], + QueryRoot::Scalar(expr) => expr.operator_refs().into_iter().cloned().collect(), + } +} + +/// Number `node` (interned) and its descendants, children first. +fn number( + node: &Rc, + groups: &mut Vec, + id_of: &mut HashMap<*const OperatorNode, GroupId>, +) -> GroupId { + if let Some(&id) = id_of.get(&Rc::as_ptr(node)) { + return id; + } + let children: Vec = node + .children() + .into_iter() + .map(|child| number(child, groups, id_of)) + .collect(); + let id = GroupId(groups.len() as u32); + for &child in &children { + let parents = &mut groups[child.0 as usize].parents; + if !parents.contains(&id) { + parents.push(id); + } + } + groups.push(MemoGroup { + node: Rc::clone(node), + children, + parents: Vec::new(), + queries: Vec::new(), + roots_of: Vec::new(), + }); + id_of.insert(Rc::as_ptr(node), id); + id +} + +/// `g` and every group below it. +fn reach(groups: &[MemoGroup], g: GroupId, reached: &mut HashSet) { + if reached.insert(g) { + for &child in &groups[g.0 as usize].children { + reach(groups, child, reached); + } + } +} + +/// Map `original` and its descendants to the groups of `interned`, its +/// interned copy. +fn pair( + original: &Rc, + interned: &Rc, + id_of: &HashMap<*const OperatorNode, GroupId>, + group_of: &mut GroupOf, +) { + if group_of + .insert(Rc::as_ptr(original), id_of[&Rc::as_ptr(interned)]) + .is_some() + { + return; + } + for (child, interned_child) in original.children().into_iter().zip(interned.children()) { + pair(child, interned_child, id_of, group_of); + } +} + +/// Every group's Pass 1 candidates, by group index. +fn pass1_candidates( + per_computation: &PerComputationCandidates, + groups: &[MemoGroup], + group_of: &GroupOf, +) -> Vec> { + // Each target aggregate group's targets, one per measure (identical + // aggregates of several queries are one group with one target each). + let mut targets_of: HashMap> = HashMap::new(); + for (t, target) in per_computation.targets.iter().enumerate() { + let g = group_of[&Rc::as_ptr(&target.aggregate)]; + let targets = targets_of.entry(g).or_default(); + if targets + .iter() + .all(|&other| per_computation.targets[other].measure != target.measure) + { + targets.push(t); + } + } + let own = |g: usize| ImplementationCandidate { + implements: vec![GroupId(g as u32)], + node: Rc::clone(&groups[g].node), + children: groups[g].children.clone(), + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + }; + (0..groups.len()) + .map(|g| match targets_of.get(&GroupId(g as u32)) { + None => vec![own(g)], + Some(targets) => target_candidates( + per_computation, + groups, + group_of, + GroupId(g as u32), + targets, + ) + .unwrap_or_else(|| vec![own(g)]), + }) + .collect() +} + +/// 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. +fn target_candidates( + per_computation: &PerComputationCandidates, + groups: &[MemoGroup], + group_of: &GroupOf, + g: GroupId, + targets: &[usize], +) -> Option> { + let aggregate = Rc::clone(&per_computation.targets[targets[0]].aggregate); + // 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].children { + reach(groups, child, &mut below); + } + let mut keys = NodeKeys::default(); + let mut descendant: HashMap = HashMap::new(); + let mut ordered: Vec = below.into_iter().collect(); + ordered.sort_unstable(); + for d in ordered { + descendant + .entry(keys.key(&groups[d.0 as usize].node)) + .or_insert(d); + } + let own_key = keys.key(&groups[g.0 as usize].node); + let mut candidates = Vec::new(); + let mut choice = vec![0; per_computation.targets.len()]; + loop { + if let Ok(composed) = compose_alone(per_computation, &aggregate, targets, &choice) { + if let Some((_, QueryRoot::Operator(root))) = composed.into_iter().next() { + let mut implements = vec![g]; + for &t in targets { + if let Some((_, Some(absorbs))) = + per_computation.targets[t].computation_for(choice[t]) + { + for &a in absorbs { + let absorbed = + group_of[&Rc::as_ptr(&per_computation.targets[a].aggregate)]; + if !implements.contains(&absorbed) { + implements.push(absorbed); + } + } + } + } + let candidate = if keys.key(&root) == own_key { + ImplementationCandidate { + implements, + node: Rc::clone(&groups[g.0 as usize].node), + children: groups[g.0 as usize].children.clone(), + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + } + } else { + let mut children = Vec::new(); + boundary( + &root, + &mut keys, + &descendant, + &mut children, + &mut HashSet::new(), + ); + ImplementationCandidate { + implements, + node: root, + children, + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + } + }; + candidates.push(candidate); + } + } + // Mixed-radix over this group's targets only. + let mut advanced = false; + for &t in targets.iter().rev() { + choice[t] += 1; + if choice[t] < per_computation.targets[t].choice_count() { + advanced = true; + break; + } + choice[t] = 0; + } + if !advanced { + break; + } + } + (!candidates.is_empty()).then_some(candidates) +} + +/// `aggregate` composed alone with `choice[t]` for each of its measures' +/// `targets`, every other target left exact: the inventory holds only +/// `targets`, and the targets a chosen fused computation absorbs (not +/// computed then). +fn compose_alone( + per_computation: &PerComputationCandidates, + aggregate: &Rc, + targets: &[usize], + choice: &[usize], +) -> Result< + Vec<(usize, QueryRoot)>, + asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError, +> { + let mut kept: Vec = targets.to_vec(); + for &t in targets { + if let Some((_, Some(absorbs))) = per_computation.targets[t].computation_for(choice[t]) { + for &a in absorbs { + if !kept.contains(&a) { + kept.push(a); + } + } + } + } + let position = |t: usize| kept.iter().position(|&k| k == t); + let alone = PerComputationCandidates { + roots: vec![(0usize, QueryRoot::Operator(Rc::clone(aggregate)))], + targets: kept + .iter() + .map(|&t| { + let mut target = per_computation.targets[t].clone(); + if let Some(fused) = &mut target.fused { + // Absorbed targets not kept are never absorbed here: the + // fusion is kept only when chosen, with all of them. + match fused.absorbs.iter().map(|&a| position(a)).collect() { + Some(absorbs) => fused.absorbs = absorbs, + None => target.fused = None, + } + } + target + }) + .collect(), + }; + let alone_choice: Vec = kept + .iter() + .map(|&t| if targets.contains(&t) { choice[t] } else { 0 }) + .collect(); + compose_logical_candidate(&alone, &alone_choice) +} + +/// The groups `node`'s private nodes read: walking down from `node`'s +/// children, the first node on each path equal to a group in `descendant`. +fn boundary( + node: &Rc, + keys: &mut NodeKeys, + descendant: &HashMap, + children: &mut Vec, + seen: &mut HashSet<*const OperatorNode>, +) { + for child in node.children() { + if !seen.insert(Rc::as_ptr(child)) { + continue; + } + match descendant.get(&keys.key(child)) { + Some(&group) => { + if !children.contains(&group) { + children.push(group); + } + } + None => boundary(child, keys, descendant, children, seen), + } + } +} + +/// Structural hashes of nodes: the operator over its children's hashes, +/// without the schema, guarantee, timing and materialization composition +/// derives. Memoized per node (the memo keeps each node alive, so its +/// address is not reused). +#[derive(Default)] +struct NodeKeys { + memo: HashMap<*const OperatorNode, (Rc, u64)>, +} + +impl NodeKeys { + fn key(&mut self, node: &Rc) -> u64 { + if let Some((_, key)) = self.memo.get(&Rc::as_ptr(node)) { + return *key; + } + let children: HashMap<*const OperatorNode, u64> = node + .children() + .into_iter() + .map(|child| (Rc::as_ptr(child), self.key(child))) + .collect(); + let operator = node.operator.map_children(|c| children[&Rc::as_ptr(c)]); + let mut hasher = DefaultHasher::new(); + serde_json::to_string(&operator) + .unwrap_or_default() + .hash(&mut hasher); + let key = hasher.finish(); + self.memo.insert(Rc::as_ptr(node), (Rc::clone(node), key)); + key + } +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + + use asap_types::ir::ASAPOp; + + /// The memo of PromQL `queries`, each with accuracy target `target`. + fn memo(queries: &[&str], target: AccuracyTarget) -> Memo { + let roots = queries + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, target.clone()))) + .enumerate() + .collect(); + memo_of(&generate_per_computation_candidates(roots).unwrap()) + } + + /// The group of the first node of `kind` reached from query `q`'s root, + /// walking down the groups' first children. + fn group_of_kind(memo: &Memo, q: usize, kind: &str) -> usize { + let mut g = memo + .groups + .iter() + .position(|group| group.roots_of.contains(&q)) + .unwrap(); + while memo.groups[g].node.operator.kind_name() != kind { + g = memo.groups[g].children[0].0 as usize; + } + g + } + + /// The summary algorithms a candidate's private nodes build. + fn algorithms(candidate: &ImplementationCandidate) -> Vec { + OperatorNode::reachable(&candidate.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] + fn identical_sub_expressions_of_two_queries_are_one_group() { + let memo = memo( + &[ + "sum by (job) (rate(m[5m]))", + "sum by (job) (rate(m[5m])) * 2", + ], + AccuracyTarget::exact(), + ); + let sum = group_of_kind(&memo, 0, "Aggregate"); + assert_eq!(group_of_kind(&memo, 1, "Aggregate"), sum); + let group = &memo.groups[sum]; + assert_eq!(group.queries, vec![0, 1]); + assert_eq!(group.roots_of, vec![0]); + let product = memo + .groups + .iter() + .position(|group| group.roots_of == vec![1]) + .unwrap(); + assert_eq!(group.parents, vec![GroupId(product as u32)]); + assert_eq!(memo.groups[product].children, vec![GroupId(sum as u32)]); + // Children come before parents. + for (g, group) in memo.groups.iter().enumerate() { + assert!(group.children.iter().all(|c| (c.0 as usize) < g)); + } + } + + /// A quantile's aggregate group has an exact candidate (its own node over + /// its input's group) and a KLL and a DDSketch candidate over the same + /// input group. + #[test] + fn a_quantile_target_group_has_exact_kll_and_ddsketch_candidates() { + let memo = memo( + &["quantile_over_time(0.9, m[5m])"], + AccuracyTarget::max_error(0.01), + ); + let quantile = group_of_kind(&memo, 0, "Aggregate"); + let input = memo.groups[quantile].children.clone(); + let candidates = &memo.candidates[quantile]; + assert_eq!(candidates.len(), 3); + assert!(Rc::ptr_eq(&candidates[0].node, &memo.groups[quantile].node)); + assert!(algorithms(&candidates[0]).is_empty()); + assert!(algorithms(&candidates[1])[0].contains("Kll")); + assert!(algorithms(&candidates[2])[0].contains("DDSketch")); + for candidate in candidates { + assert_eq!(candidate.implements, vec![GroupId(quantile as u32)]); + assert_eq!(candidate.children, input); + assert_eq!(candidate.origin, CandidateOrigin::Pass1); + } + } + + /// The whole-expression top-k heap over the per-job sums implements both + /// aggregates' groups, and reads the sums' input group; the exact top-k + /// reads the sums' group. + #[test] + fn a_fused_candidate_lists_both_groups_it_implements() { + let memo = memo( + &["topk by (job) (10, sum_over_time(m[1m]))"], + AccuracyTarget::max_error(0.01).failure_probability(0.001), + ); + let topk = group_of_kind(&memo, 0, "Aggregate"); + let sums = memo.groups[topk].children[0]; + let sums_input = memo.groups[sums.0 as usize].children.clone(); + let candidates = &memo.candidates[topk]; + assert_eq!(candidates[0].children, vec![sums]); + let fused: Vec<_> = candidates + .iter() + .filter(|c| c.implements.len() > 1) + .collect(); + assert!(!fused.is_empty()); + for candidate in fused { + assert_eq!(candidate.implements, vec![GroupId(topk as u32), sums]); + assert_eq!(candidate.children, sums_input); + assert!(algorithms(candidate)[0].contains("Heap")); + } + } +} diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs new file mode 100644 index 000000000..8eb747abb --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -0,0 +1,138 @@ +//! Memo over logical sub-expressions (joint-planning-search design, +//! "Memo-group search"): the workload's logical DAG hash-consed into +//! groups, each group's implementation candidates generated by Pass 1 and +//! the Pass 2 transformations, priced per (group, candidate, context) by +//! Stage 3, and searched by enumerating coupling thresholds with bounds and +//! a dynamic program over the rest. +//! +//! Behind [`crate::pass::Stage1ToStage3Search::memo_groups`] until it +//! selects the exhaustive search's price on every equivalence test. +//! +//! - [`groups`] (W1): the hash-consed group DAG and Pass 1's +//! implementation candidates per group; +//! - [`capability`] (W3): identical-expression and summary-capability +//! transformations; +//! - [`windows`] (W4): window-composition, rollup and top-k reuse +//! transformations; +//! - [`pricing`] (W2): local prices per (group, candidate, context), +//! cached, and frontiers; +//! - [`search`] (integrator): clusters, coupling thresholds with branch and +//! bound, the dynamic program, build-late. +#![allow(dead_code)] + +pub(crate) mod capability; +pub(crate) mod groups; +pub(crate) mod pricing; +pub(crate) mod search; +pub(crate) mod windows; + +use std::rc::Rc; + +use asap_types::ir::OperatorNode; +use asap_types::physical::Materialization; + +/// A group, by index in [`Memo::groups`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub(crate) struct GroupId(pub u32); + +/// One distinct logical sub-expression of the workload. +#[derive(Debug, Clone)] +pub(crate) struct MemoGroup { + /// Its logical node as the frontends (and Pass 1 targets) see it; its + /// children are other groups' nodes. + pub node: Rc, + /// Its child groups, in the node's child order. + pub children: Vec, + /// The groups that read it. + pub parents: Vec, + /// The queries whose roots reach it, by index in the workload. + pub queries: Vec, + /// The queries whose root it is. + pub roots_of: Vec, +} + +/// The hash-consed workload: its groups, children before parents. +#[derive(Debug, Clone, Default)] +pub(crate) struct Memo { + pub groups: Vec, + /// Each group's implementation candidates, by group index. + pub candidates: Vec>, +} + +/// One way to implement a group: its node as Stage 2 implements it, over +/// child groups. +#[derive(Debug, Clone)] +pub(crate) struct ImplementationCandidate { + /// The groups it implements: one, or several for a fused computation + /// (the topmost first). + pub implements: Vec, + /// The implementation's own node(s): its root node, whose descendants + /// down to `children` are its private nodes. + pub node: Rc, + /// The groups it reads, whose implementations are chosen separately. + pub children: Vec, + /// How it was made: by Pass 1, or by which Pass 2 transformation. + pub origin: CandidateOrigin, + /// The coupling sites its cost depends on. + pub coupling_sites: Vec, +} + +/// Where an implementation candidate came from. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CandidateOrigin { + /// Pass 1: exact, exact accumulator, or a summary family and size. + Pass1, + /// One build serving several reader groups (summary capability). + SharedBuild, + /// Tumbling panes, or the newest panes of another group's chain. + Panes, + /// A coarse group regrouping a finer group. + Rollup, + /// A top-k read from another group's ranking. + TopKReuse, +} + +/// A choice several groups' costs depend on; equal keys are one site. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub(crate) struct CouplingSite { + pub key: u64, + pub kind: CouplingKind, +} + +/// What a coupling site's threshold is. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub(crate) enum CouplingKind { + /// A shared summary build's size, set by its strictest reader. + SharedBuildSize, + /// A shared pane chain: grid width, longest window, size. + PaneChain, + /// A metric's scan label set. + ScanLabels, + /// A source's raw-data horizon, the longest raw scan. + RawHorizon, +} + +/// The thresholds chosen for a cluster's coupling sites, by site key: the +/// context a candidate is priced in. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub(crate) struct Context { + /// Per site key, the index of the chosen threshold among the site's + /// enumerated thresholds. + pub thresholds: Vec<(u64, u32)>, + /// The materialization chosen for each shared group in the context. + pub materialization: Vec<(GroupId, Materialization)>, +} + +/// One candidate's usage priced in a context (Stage 3's per-node pricing). +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct LocalUsage { + pub cost_per_second: Option, + pub cpu_ops_per_second: f64, + pub scanned_bytes_per_second: f64, + pub average_memory_bytes: f64, + pub peak_memory_bytes: u64, + pub peak_cpu_cores: u32, + pub disk_bytes: u64, + /// Each query's latency the candidate adds per evaluation, in ms. + pub latency_ms: Vec<(usize, f64)>, +} diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs new file mode 100644 index 000000000..6c0e3f63c --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -0,0 +1,716 @@ +//! Local prices per (group, candidate, context), cached, the frozen +//! calibration factor, and frontiers (design "Memo-group search": Local +//! price, Frontier, Coupled charges). +//! +//! A candidate's local price is the Stage 3 usage of its private nodes: +//! from its node down to, but excluding, its child groups' nodes. It is +//! priced as the difference of two Stage 3 prices of whole sub-DAGs, each +//! built by Stage 2 and priced by [`crate::price_part`] with the readers' +//! entries: +//! +//! - the candidate's node with everything below it, its children being the +//! child groups' nodes, so its private nodes see their children's output +//! statistics and the readers' rates; +//! - minus each child group's node with everything below it, priced alike. +//! +//! Additive usage (price under cost model 1, CPU, scans, average memory, +//! peak memory, disk) is the difference; peak cores are the candidate's +//! whole sub-DAG's (frontiers combine them by maximum); each reader's +//! latency is what the candidate adds to its longest child's. Summed over a +//! plan's groups, the differences telescope to Stage 3's price of the whole +//! plan, coupled charges included (memory held for a query's whole latency, +//! a source's longest raw scan), which the root's sub-DAG carries. +//! +//! Materialization: the context's materializations, in order, are the +//! options of the sub-DAG's builds in Stage 2's discovery order; builds +//! past the list are `NotMaterialized`. A child's builds are the trailing +//! ones of its parent's sub-DAG. + +use std::collections::HashMap; + +use asap_physical_asap_optimization::{optimize_physical_asap_assigned, PhysicalCandidate}; +use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::physical::Materialization; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; +use std::rc::Rc; + +use super::{Context, GroupId, ImplementationCandidate, LocalUsage, Memo}; +use crate::{ + price_part, CalibrationFit, CandidateCost, ClusterLimits, CostModel, PlanningModels, + SelectionError, +}; + +/// The cached prices of one search: valid for one workload, data workload +/// and set of planning models, priced with one frozen calibration factor. +#[derive(Debug)] +pub(crate) struct PriceCache { + /// The frozen calibration factor ([`frozen_calibration`]). + formula_scale: f64, + /// Local prices by (group, candidate's node, context, readers). + local: HashMap>, + /// Whole sub-DAG prices by (node, readers, materializations). + trees: HashMap>, + /// Local prices found in the cache. + pub hits: u64, + /// Local prices computed. + pub misses: u64, +} + +impl PriceCache { + /// An empty cache pricing with the calibration factor `formula_scale`. + pub(crate) fn new(formula_scale: f64) -> Self { + Self { + formula_scale, + local: HashMap::new(), + trees: HashMap::new(), + hits: 0, + misses: 0, + } + } +} + +/// A node by its structure: its serialized tree (a hash alone does not +/// prove two nodes equal). +type NodeKey = String; +type LocalKey = (GroupId, NodeKey, Context, Vec); +type TreeKey = (NodeKey, Vec, Vec); + +/// One whole sub-DAG's Stage 3 price. +#[derive(Debug, Clone, PartialEq)] +struct TreePrice { + usage: LocalUsage, + /// The builds Stage 2 found in it, in discovery order. + builds: usize, +} + +fn node_key(node: &OperatorNode) -> NodeKey { + serde_json::to_string(node).expect("an operator node serializes") +} + +/// `candidate` of `group` priced in `context` for the queries `readers` (by +/// index in the workload, whose `entries` they are): the Stage 3 usage of +/// its private nodes (module doc). `None` when it fails a check sharing +/// cannot repair, or Stage 2 does not build it with the context's +/// materializations. +#[allow(clippy::too_many_arguments)] +pub(crate) fn local_price( + memo: &Memo, + group: GroupId, + candidate: &ImplementationCandidate, + context: &Context, + readers: &[usize], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + cache: &mut PriceCache, +) -> Result, SelectionError> { + let key = ( + group, + node_key(&candidate.node), + context.clone(), + readers.to_vec(), + ); + if let Some(price) = cache.local.get(&key) { + cache.hits += 1; + return Ok(price.clone()); + } + cache.misses += 1; + let reader_entries: Vec = + readers.iter().map(|&q| entries[q].clone()).collect(); + let materializations: Vec = + context.materialization.iter().map(|(_, m)| *m).collect(); + let pricing = TreePricing { + readers, + entries: &reader_entries, + data, + models: &models, + }; + let Some(whole) = pricing.price(&candidate.node, &materializations, cache)? else { + cache.local.insert(key, None); + return Ok(None); + }; + let mut children = Vec::with_capacity(candidate.children.len()); + // Each child's builds trail its parent's, in child order. + let mut end = whole.builds; + let child_builds: Vec = candidate + .children + .iter() + .map(|c| pricing.builds(&memo.groups[c.0 as usize].node)) + .collect(); + let mut offsets = vec![0; child_builds.len()]; + for (i, builds) in child_builds.iter().enumerate().rev() { + end = end.saturating_sub(*builds); + offsets[i] = end; + } + for (child, (&offset, &builds)) in candidate + .children + .iter() + .zip(offsets.iter().zip(&child_builds)) + { + let node = &memo.groups[child.0 as usize].node; + let own = materializations + .get(offset..(offset + builds).min(materializations.len())) + .unwrap_or(&[]); + // A child that does not price alone (e.g. a raw scan only its + // parent's ingestion build makes valid) is charged inside its + // parent: nothing is subtracted for it. + if let Some(price) = pricing.price(node, own, cache)? { + children.push(price.usage); + } + } + let local = difference(&whole.usage, &children, &models.plan_selection.cost_model); + cache.local.insert(key, Some(local.clone())); + Ok(Some(local)) +} + +/// Pricing whole sub-DAGs for one set of readers. +struct TreePricing<'a> { + readers: &'a [usize], + entries: &'a [QueryWorkloadEntry], + data: &'a DataWorkload, + models: &'a PlanningModels<'a>, +} + +impl TreePricing<'_> { + /// Stage 2's sub-DAG of `node`, one root per reader, its builds given + /// `materializations` in order (`NotMaterialized` past them), and the + /// number of builds; `None` when Stage 2 does not build it so. + fn built( + &self, + node: &Rc, + materializations: &[Materialization], + ) -> Option<(PhysicalCandidate, usize)> { + let roots = vec![QueryRoot::Operator(Rc::clone(node)); self.readers.len()]; + let mut next = materializations.iter().copied(); + let mut builds = 0; + let mut choose = |_: &[usize]| { + builds += 1; + Some(next.next().unwrap_or(Materialization::NotMaterialized)) + }; + let built = optimize_physical_asap_assigned( + "S", + &roots, + self.entries, + self.data, + self.models.deployment_capabilities, + &mut choose, + ); + let physical = built.ok().flatten()?; + Some((physical, builds)) + } + + /// The builds Stage 2 finds in `node`'s sub-DAG. + fn builds(&self, node: &Rc) -> usize { + self.built(node, &[]).map_or(0, |(_, builds)| builds) + } + + /// `node`'s whole sub-DAG priced by Stage 3, cached. + fn price( + &self, + node: &Rc, + materializations: &[Materialization], + cache: &mut PriceCache, + ) -> Result, SelectionError> { + let key = ( + node_key(node), + self.readers.to_vec(), + materializations.to_vec(), + ); + if let Some(price) = cache.trees.get(&key) { + return Ok(price.clone()); + } + let price = match self.built(node, materializations) { + None => None, + Some((physical, builds)) => price_part( + &physical, + self.entries, + self.data, + self.models, + cache.formula_scale, + )? + .map(|(cost, _)| TreePrice { + usage: usage_of(&cost, self.readers), + builds, + }), + }; + cache.trees.insert(key, price.clone()); + Ok(price) + } +} + +/// `cost` as a usage, each root's latency that of the reader at its index. +fn usage_of(cost: &CandidateCost, readers: &[usize]) -> LocalUsage { + let usage = &cost.usage; + LocalUsage { + cost_per_second: cost.cost_per_second, + cpu_ops_per_second: usage.cpu_ops_per_second, + scanned_bytes_per_second: usage.scanned_bytes_per_second, + average_memory_bytes: usage.average_memory_bytes, + peak_memory_bytes: usage.peak_memory_bytes, + peak_cpu_cores: usage.peak_cpu_cores, + disk_bytes: usage.disk_bytes, + latency_ms: readers + .iter() + .zip(&usage.latency_per_query) + .map(|(&q, latency)| (q, latency.total_ms())) + .collect(), + } +} + +/// `whole` minus `children`: additive components subtracted, peak cores +/// `whole`'s, each reader's latency over its longest child's. Under cost +/// model 2 the price is recomputed from the difference's peaks. +fn difference(whole: &LocalUsage, children: &[LocalUsage], cost_model: &CostModel) -> LocalUsage { + let mut local = whole.clone(); + for child in children { + local.cost_per_second = match (local.cost_per_second, child.cost_per_second) { + (Some(a), Some(b)) => Some(a - b), + (a, _) => a, + }; + local.cpu_ops_per_second -= child.cpu_ops_per_second; + local.scanned_bytes_per_second -= child.scanned_bytes_per_second; + local.average_memory_bytes -= child.average_memory_bytes; + local.peak_memory_bytes = local + .peak_memory_bytes + .saturating_sub(child.peak_memory_bytes); + local.disk_bytes = local.disk_bytes.saturating_sub(child.disk_bytes); + } + for (query, latency) in &mut local.latency_ms { + let longest_child = children + .iter() + .flat_map(|c| &c.latency_ms) + .filter(|(q, _)| q == query) + .map(|(_, l)| *l) + .fold(0.0, f64::max); + *latency -= longest_child; + } + if let CostModel::ProvisionedCapacity(_) = cost_model { + local.cost_per_second = provisioned_cost(&local, cost_model); + } + local +} + +/// The calibration factor of the formulas' CPU time to a measured cost +/// table, fitted once over every candidate of every group, each with all +/// its group's queries as readers, before the search, and frozen (design +/// "Coupled charges"); 1 without a measured table. +pub(crate) fn frozen_calibration( + memo: &Memo, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result { + if models + .plan_selection + .measured_operation_costs + .operators + .is_empty() + { + return Ok(1.0); + } + let mut fit = CalibrationFit::default(); + for (group, candidates) in memo.groups.iter().zip(&memo.candidates) { + let reader_entries: Vec = + group.queries.iter().map(|&q| entries[q].clone()).collect(); + let pricing = TreePricing { + readers: &group.queries, + entries: &reader_entries, + data, + models: &models, + }; + for candidate in candidates { + if let Some((physical, _)) = pricing.built(&candidate.node, &[]) { + if let Some((_, candidate_fit)) = + price_part(&physical, &reader_entries, data, &models, 1.0)? + { + fit.add(candidate_fit); + } + } + } + } + Ok(fit.factor()) +} + +/// Under cost model 2, `usage`'s provisioned price; otherwise its own. +fn provisioned_cost(usage: &LocalUsage, cost_model: &CostModel) -> Option { + match cost_model { + CostModel::ProvisionedCapacity(prices) => Some( + prices.cost_per_provisioned_memory_byte_second * usage.peak_memory_bytes as f64 + + prices.cost_per_provisioned_cpu_core_second * f64::from(usage.peak_cpu_cores) + + prices.cost_per_disk_byte_second * usage.disk_bytes as f64, + ), + _ => usage.cost_per_second, + } +} + +/// The candidates of a subproblem no other dominates in usage: under cost +/// model 1 the cheapest (Stage 3's tie tolerance; ties keep the earlier), +/// under cost models 2 and 3 the Pareto front over (price, peak memory, +/// peak cores, disk, worst latency), exact ties kept. Under cost model 3, +/// usage over the cluster's limits is never kept. +#[derive(Debug, Clone)] +pub(crate) struct Frontier { + cost_model: CostModel, + /// The kept usages with what each stands for, in insertion order. + pub entries: Vec<(LocalUsage, T)>, +} + +impl Frontier { + pub(crate) fn new(cost_model: &CostModel) -> Self { + Self { + cost_model: cost_model.clone(), + entries: Vec::new(), + } + } + + /// Add `usage` standing for `item` unless a kept entry dominates it, + /// dropping the kept entries it dominates; whether it was kept. + pub(crate) fn insert(&mut self, usage: LocalUsage, item: T) -> bool { + if let CostModel::ClusterCapacity(limits) = &self.cost_model { + if over_limits(&usage, limits) { + return false; + } + } + if let CostModel::UsageOverTime(_) = self.cost_model { + let price = usage.cost_per_second.unwrap_or(f64::INFINITY); + match self.entries.first() { + Some((best, _)) + if !cheaper(price, best.cost_per_second.unwrap_or(f64::INFINITY)) => + { + return false + } + _ => { + self.entries.clear(); + self.entries.push((usage, item)); + return true; + } + } + } + let point = front_point(&usage); + if self + .entries + .iter() + .any(|(kept, _)| dominates(&front_point(kept), &point)) + { + return false; + } + self.entries + .retain(|(kept, _)| !dominates(&point, &front_point(kept))); + self.entries.push((usage, item)); + true + } + + /// Every pair of `self`'s and `other`'s entries combined, two disjoint + /// parts of a plan: additive usage summed, peak memory summed (Stage 3 + /// adds every node's peak), peak cores by maximum, each query's latency + /// by maximum, the price summed (recomputed from the peaks under cost + /// model 2); `join` combines the items. Kept as [`Self::insert`] keeps. + pub(crate) fn combine( + &self, + other: &Frontier, + join: impl Fn(&T, &U) -> V, + ) -> Frontier { + let mut combined = Frontier::new(&self.cost_model); + for (a, x) in &self.entries { + for (b, y) in &other.entries { + combined.insert(sum(a, b, &self.cost_model), join(x, y)); + } + } + combined + } +} + +/// `a` and `b` of two disjoint parts of a plan, combined ([`Frontier::combine`]). +fn sum(a: &LocalUsage, b: &LocalUsage, cost_model: &CostModel) -> LocalUsage { + let mut latency_ms = a.latency_ms.clone(); + for &(query, latency) in &b.latency_ms { + match latency_ms.iter_mut().find(|(q, _)| *q == query) { + Some((_, l)) => *l = l.max(latency), + None => latency_ms.push((query, latency)), + } + } + latency_ms.sort_by_key(|(q, _)| *q); + let mut usage = LocalUsage { + cost_per_second: match (a.cost_per_second, b.cost_per_second) { + (Some(x), Some(y)) => Some(x + y), + _ => None, + }, + cpu_ops_per_second: a.cpu_ops_per_second + b.cpu_ops_per_second, + scanned_bytes_per_second: a.scanned_bytes_per_second + b.scanned_bytes_per_second, + average_memory_bytes: a.average_memory_bytes + b.average_memory_bytes, + peak_memory_bytes: a.peak_memory_bytes.saturating_add(b.peak_memory_bytes), + peak_cpu_cores: a.peak_cpu_cores.max(b.peak_cpu_cores), + disk_bytes: a.disk_bytes.saturating_add(b.disk_bytes), + latency_ms, + }; + usage.cost_per_second = provisioned_cost(&usage, cost_model); + usage +} + +fn over_limits(usage: &LocalUsage, limits: &ClusterLimits) -> bool { + usage.peak_memory_bytes > limits.memory_bytes + || usage.disk_bytes > limits.disk_bytes + || limits + .cpu_cores + .is_some_and(|cores| usage.peak_cpu_cores > cores) +} + +/// Stage 3's price comparison: `cost` is cheaper than `best` beyond its +/// tie tolerance (1e-9). +fn cheaper(cost: f64, best: f64) -> bool { + cost < best - best.abs() * 1e-9 +} + +/// The Pareto front's coordinates: price (when there is one), peak memory, +/// peak cores, disk, worst latency. +fn front_point(usage: &LocalUsage) -> Vec { + let mut point: Vec = usage.cost_per_second.into_iter().collect(); + point.extend([ + usage.peak_memory_bytes as f64, + f64::from(usage.peak_cpu_cores), + usage.disk_bytes as f64, + usage.latency_ms.iter().map(|(_, l)| *l).fold(0.0, f64::max), + ]); + point +} + +/// Whether `a` is no greater than `b` everywhere and smaller somewhere +/// (points of different shapes are incomparable). +fn dominates(a: &[f64], b: &[f64]) -> bool { + a.len() == b.len() + && a.iter().zip(b).all(|(x, y)| x <= y) + && a.iter().zip(b).any(|(x, y)| x < y) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pass::subexpr::{CandidateOrigin, MemoGroup}; + use crate::{ClusterObjective, ResourceUnitCost}; + use asap_logical_asap_optimization::test_support::{entry_with, lower_promql}; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::workload::{ + DataArrival, QueryRecurrence, RepeatedDemand, RepetitionInterval, SourceFacts, + }; + + fn data() -> DataWorkload { + DataWorkload { + arrival: DataArrival::Ingesting { + rows_per_second: 10_000.0, + }, + sources: [( + "m".to_string(), + SourceFacts { + series_count: Some(10_000), + ..SourceFacts::default() + }, + )] + .into_iter() + .collect(), + } + } + + fn entries() -> Vec { + let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); + vec![entry_with( + target, + QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval(60_000))), + )] + } + + /// A one-query workload's selected plan split into two groups, its root + /// node over the rest, and the context of the selected + /// materializations. + fn split_plan(models: PlanningModels<'_>) -> (Memo, Context, Option) { + let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); + let root = lower_promql("quantile_over_time(0.99, m[5m])", target); + let run = crate::pass::run_stage1_to_stage3( + vec![(0, QueryRoot::Operator(root))], + &entries(), + &data(), + models, + ) + .unwrap(); + let selection = run.stage3_plan_selection.as_ref().unwrap(); + let price = selection.costs[&selection.selected].cost_per_second; + let assignment = run.stage3_selected.as_ref().unwrap().assignment.clone(); + let roots = run + .stage1_selected + .as_ref() + .unwrap() + .roots + .as_ref() + .unwrap(); + let QueryRoot::Operator(top) = &roots[0] else { + panic!("an operator root") + }; + let below = Rc::clone(top.children()[0]); + let group = |node: &Rc, children: Vec| MemoGroup { + node: Rc::clone(node), + children, + parents: Vec::new(), + queries: vec![0], + roots_of: Vec::new(), + }; + let candidate = + |id: u32, node: &Rc, children: Vec| ImplementationCandidate { + implements: vec![GroupId(id)], + node: Rc::clone(node), + children, + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + }; + let memo = Memo { + groups: vec![group(&below, vec![]), group(top, vec![GroupId(0)])], + candidates: vec![ + vec![candidate(0, &below, vec![])], + vec![candidate(1, top, vec![GroupId(0)])], + ], + }; + let context = Context { + thresholds: Vec::new(), + materialization: assignment.into_iter().map(|m| (GroupId(0), m)).collect(), + }; + (memo, context, price) + } + + #[test] + fn local_prices_add_up_to_the_whole_plans_price() { + let models = PlanningModels::builtin(); + let (memo, context, expected) = split_plan(models); + let (entries, data) = (entries(), data()); + let mut cache = + PriceCache::new(frozen_calibration(&memo, &entries, &data, models).unwrap()); + let mut total = 0.0; + for (g, candidates) in memo.candidates.iter().enumerate() { + let local = local_price( + &memo, + GroupId(g as u32), + &candidates[0], + &context, + &[0], + &entries, + &data, + models, + &mut cache, + ) + .unwrap() + .unwrap(); + // Each group carries a part of the price, not the whole. + let price = local.cost_per_second.unwrap(); + assert!( + price > 0.0 && price < expected.unwrap(), + "group {g}: {price}" + ); + total += price; + } + let expected = expected.unwrap(); + assert!( + (total - expected).abs() <= 1e-9 * expected.abs(), + "{total} vs {expected}" + ); + } + + #[test] + fn the_cache_hits_on_a_repeated_context() { + let models = PlanningModels::builtin(); + let (memo, context, _) = split_plan(models); + let (entries, data) = (entries(), data()); + let mut cache = PriceCache::new(1.0); + let mut price = || { + local_price( + &memo, + GroupId(1), + &memo.candidates[1][0], + &context, + &[0], + &entries, + &data, + models, + &mut cache, + ) + .unwrap() + }; + let first = price(); + let second = price(); + assert!(first.is_some()); + assert_eq!(first, second); + assert_eq!((cache.misses, cache.hits), (1, 1)); + } + + fn usage(cost: f64, peak_memory_bytes: u64, peak_cpu_cores: u32, latency: f64) -> LocalUsage { + LocalUsage { + cost_per_second: Some(cost), + cpu_ops_per_second: 1.0, + scanned_bytes_per_second: 1.0, + average_memory_bytes: 1.0, + peak_memory_bytes, + peak_cpu_cores, + disk_bytes: 1, + latency_ms: vec![(0, latency)], + } + } + + fn items(frontier: &Frontier) -> Vec { + frontier.entries.iter().map(|(_, t)| t.clone()).collect() + } + + #[test] + fn cost_model_1_keeps_the_cheapest_and_the_earlier_tie() { + let mut frontier = + Frontier::new(&CostModel::UsageOverTime(ResourceUnitCost::illustrative())); + assert!(frontier.insert(usage(2.0, 1, 1, 1.0), "b")); + assert!(frontier.insert(usage(1.0, 9, 9, 9.0), "a")); + assert!(!frontier.insert(usage(1.0, 1, 1, 1.0), "tie")); + assert!(!frontier.insert(usage(3.0, 0, 0, 0.0), "c")); + assert_eq!(items(&frontier), vec!["a"]); + } + + #[test] + fn cost_model_2_keeps_the_pareto_front() { + let mut frontier = Frontier::new(&CostModel::ProvisionedCapacity( + ResourceUnitCost::illustrative(), + )); + assert!(frontier.insert(usage(2.0, 10, 1, 5.0), "cheap memory")); + assert!(frontier.insert(usage(1.0, 20, 1, 5.0), "cheap price")); + // Dominated by "cheap price". + assert!(!frontier.insert(usage(1.0, 20, 1, 6.0), "slower")); + // An exact tie is kept. + assert!(frontier.insert(usage(1.0, 20, 1, 5.0), "tie")); + // Dominates "cheap memory". + assert!(frontier.insert(usage(2.0, 10, 1, 4.0), "faster")); + assert_eq!(items(&frontier), vec!["cheap price", "tie", "faster"]); + } + + #[test] + fn combining_sums_usage_and_takes_the_largest_peak_cores_and_latency() { + let cluster = CostModel::ClusterCapacity(ClusterLimits { + memory_bytes: 25, + disk_bytes: u64::MAX, + cpu_cores: Some(4), + objective: ClusterObjective::ParetoFront, + }); + let mut a = Frontier::new(&cluster); + let none = |mut u: LocalUsage| { + u.cost_per_second = None; + u + }; + a.insert(none(usage(0.0, 10, 2, 5.0)), 1); + a.insert(none(usage(0.0, 5, 3, 7.0)), 2); + let mut b = Frontier::new(&cluster); + b.insert(none(usage(0.0, 12, 4, 6.0)), 10); + let combined = a.combine(&b, |x, y| x + y); + // 10 + 12 and 5 + 12 bytes both fit under 25. + assert_eq!(items(&combined), vec![11, 12]); + let (first, _) = &combined.entries[0]; + assert_eq!(first.peak_memory_bytes, 22); + assert_eq!(first.peak_cpu_cores, 4); + assert_eq!(first.latency_ms, vec![(0, 6.0)]); + assert_eq!(first.cpu_ops_per_second, 2.0); + assert_eq!(first.disk_bytes, 2); + // Over the memory limit: dropped. + let mut c = Frontier::new(&cluster); + c.insert(none(usage(0.0, 21, 1, 1.0)), 100); + assert!(a.combine(&c, |x, y| x + y).entries.is_empty()); + } +} diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs new file mode 100644 index 000000000..e6f41fc1e --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -0,0 +1 @@ +//! See the module doc of [`super`]. diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs new file mode 100644 index 000000000..194071e4c --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -0,0 +1,942 @@ +//! The window-composition, rollup and exact top-k reuse transformations +//! (joint-planning-search design, "Memo-group search", stage interfaces): +//! Pass 2 rules that add implementation candidates to the memo's groups. +//! +//! - **Window composition** ([`add_window_composition`]): a group whose +//! candidate builds a summary over a range `TimeRange(W)`, read only by +//! queries that repeat on a grid, gets the candidate with that build in +//! tumbling panes (the width rules of `pane_candidates`). When another +//! group builds the same family over the same input, it also gets the +//! candidate that reads the newest panes of a shared chain, with a +//! [`CouplingKind::PaneChain`] site: the chain's grid, longest window and +//! size are the site's threshold ([`chain_thresholds`], [`may_read`], +//! [`read_at_threshold`]). +//! - **Rollup** ([`add_rollup`]): a coarse exact accumulator answered by +//! regrouping a finer group's values (`rollup_candidates`' conditions). +//! - **Exact top-k reuse** ([`add_topk_reuse`]): a top-k answered from a +//! ranking group, `Limit(k_max)` over a descending `Sort` of the input, +//! added to the memo when missing. +//! +//! The joint search never reuses a top-k ranking in a plan with panes +//! ([`combinable`]). + +use std::collections::hash_map::DefaultHasher; +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use asap_logical_asap_optimization::accuracy::estimators::size_for_all_readers; +use asap_logical_asap_optimization::asap_aware_cse::{ + pane_candidates, rollup_candidates, topk_reuse_candidates, window_composition_pass2_candidates, + with_pane_chain, +}; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; +use asap_types::ir::flat::flatten; +use asap_types::ir::schema::sketch::SketchKind; +use asap_types::ir::schema::FieldDataType; +use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot}; +use asap_types::workload::{greatest_common_divisor, QueryWorkloadEntry}; + +use super::{ + CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo, MemoGroup, +}; + +/// One threshold of a shared pane chain: its grid, the longest window it +/// covers and the size it is built at. A reader no stricter than it may +/// read it ([`may_read`]). +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ChainThreshold { + /// The width of one pane. + pub pane_width_ms: u64, + /// The window the chain covers: the longest window among its readers. + pub longest_window_ms: u64, + /// The family and size the chain is built at: the size for its + /// strictest reader (`size_for_all_readers`). + pub family: FieldDataType, +} + +/// A candidate's read of a shared pane chain: its own panes, window and +/// size, and the `SummaryMerge` that reads them. +#[derive(Debug, Clone)] +pub(crate) struct ChainReader { + /// The chain's [`CouplingSite::key`]. + pub key: u64, + /// The width of the reader's own panes: the gcd of its window and its + /// readers' grids. + pub pane_width_ms: u64, + /// The reader's window. + pub window_ms: u64, + /// The family and size the reader needs. + pub family: FieldDataType, + /// The merge over the reader's own chain, which a threshold replaces. + pub merge: Rc, +} + +/// Every transformation of this module, in order: window composition, +/// rollup, exact top-k reuse. `entries[q]` is workload query `q`'s entry. +pub(crate) fn add_pass2_candidates(memo: &mut Memo, entries: &[QueryWorkloadEntry]) { + add_window_composition(memo, entries); + add_rollup(memo); + add_topk_reuse(memo); +} + +/// Whether the joint search can choose `chosen` together: never an exact +/// top-k reuse with panes (Pass 2 never applies window composition after +/// top-k reuse, nor top-k reuse after it). +pub(crate) fn combinable<'a>( + chosen: impl IntoIterator, +) -> bool { + let (mut panes, mut topk) = (false, false); + for candidate in chosen { + panes |= candidate.origin == CandidateOrigin::Panes; + topk |= candidate.origin == CandidateOrigin::TopKReuse; + } + !(panes && topk) +} + +/// Window composition: for each group whose every query repeats on a grid, +/// each Pass 1 (or shared-build) candidate with a summary build over a +/// range window gets a [`CandidateOrigin::Panes`] candidate per non-empty +/// set of its builds in panes (`window_composition_pass2_candidates` on +/// the candidate alone). Each build another group also builds (same +/// family at any size, same input below the range, grouping, item and +/// weight, filter and names) is also offered reading a shared chain: the +/// same node with a [`CouplingKind::PaneChain`] site per such build, whose +/// chain the site's threshold sets ([`read_at_threshold`]). +pub(crate) fn add_window_composition(memo: &mut Memo, entries: &[QueryWorkloadEntry]) { + let group_of = group_of_node(&memo.groups); + let mut own: Vec<(usize, ImplementationCandidate)> = Vec::new(); + for (g, group) in memo.groups.iter().enumerate() { + let Some(grid) = group_grid_ms(group, entries) else { + continue; + }; + for candidate in &memo.candidates[g] { + if !matches!( + candidate.origin, + CandidateOrigin::Pass1 | CandidateOrigin::SharedBuild + ) { + continue; + } + let roots = [QueryRoot::Operator(Rc::clone(&candidate.node))]; + let panes = pane_candidates(&roots, &[Some(grid)]); + if panes.is_empty() { + continue; + } + let Ok(pass2_candidates) = window_composition_pass2_candidates(&roots, &panes) else { + continue; + }; + for pass2_candidate in pass2_candidates { + let Ok(rewritten) = pass2_candidate.roots else { + continue; + }; + let [QueryRoot::Operator(node)] = &rewritten[..] else { + continue; + }; + own.push(( + g, + rewritten_candidate( + &group_of, + candidate, + Rc::clone(node), + CandidateOrigin::Panes, + ), + )); + } + } + } + // The chain keys each group builds. + let mut builders: HashMap> = HashMap::new(); + for (g, candidate) in &own { + for reader in chain_readers(&candidate.node) { + builders.entry(reader.key).or_default().insert(*g); + } + } + for (g, candidate) in own { + let mut keys: Vec = chain_readers(&candidate.node) + .into_iter() + .map(|reader| reader.key) + .filter(|key| builders[key].len() >= 2) + .collect(); + keys.sort_unstable(); + keys.dedup(); + if !keys.is_empty() { + let mut shared = candidate.clone(); + shared + .coupling_sites + .extend(keys.into_iter().map(|key| CouplingSite { + key, + kind: CouplingKind::PaneChain, + })); + push_new(memo, g, shared); + } + push_new(memo, g, candidate); + } +} + +/// The reads of a pane chain in `node`: each `SummaryMerge` over an +/// [`ASAPOp::PaneChain`], with the chain's key. +pub(crate) fn chain_readers(node: &Rc) -> Vec { + let mut seen = HashSet::new(); + let mut readers = Vec::new(); + for merge in OperatorNode::reachable(node) { + if !seen.insert(Rc::as_ptr(&merge)) { + continue; + } + let Some(ASAPOp::SummaryMerge { children, .. }) = merge.asap() else { + continue; + }; + let [chain] = &children[..] else { + continue; + }; + let Some(ASAPOp::PaneChain { + child, + family, + input, + reduction, + grouping, + filter, + pane_width_ms, + pane_count, + }) = chain.asap() + else { + continue; + }; + let below = serde_json::to_string(&flatten(&[QueryRoot::Operator(Rc::clone(child))]).0) + .unwrap_or_default(); + let names: Vec<(&String, &Option)> = chain + .schema + .fields + .iter() + .map(|f| (&f.name, &f.table)) + .collect(); + let key = hash_of( + &serde_json::json!({ + "family": family_without_size(family), + "below": below, + "input": input, + "reduction": reduction, + "grouping": grouping, + "filter": filter, + "names": names, + }) + .to_string(), + ); + readers.push(ChainReader { + key, + pane_width_ms: *pane_width_ms, + window_ms: pane_width_ms * u64::from(*pane_count), + family: family.clone(), + merge: Rc::clone(&merge), + }); + } + readers +} + +/// The thresholds of the shared pane chain `site`: its grids are the +/// closure of its readers' pane widths under gcd; for each grid, every +/// distinct (longest window, strictest size) pair that some set of the +/// readers whose widths the grid divides has, as the readers no stricter +/// than it ([`may_read`]). Ordered by grid, window, then size. +pub(crate) fn chain_thresholds(memo: &Memo, site: &CouplingSite) -> Vec { + let mut readers: Vec = Vec::new(); + for candidate in memo.candidates.iter().flatten() { + if !candidate.coupling_sites.contains(site) { + continue; + } + for reader in chain_readers(&candidate.node) { + let known = readers.iter().any(|r| { + r.pane_width_ms == reader.pane_width_ms + && r.window_ms == reader.window_ms + && r.family == reader.family + }); + if reader.key == site.key && !known { + readers.push(reader); + } + } + } + let mut grids: BTreeSet = readers.iter().map(|r| r.pane_width_ms).collect(); + loop { + let more: Vec = grids + .iter() + .flat_map(|&a| grids.iter().map(move |&b| greatest_common_divisor(a, b))) + .filter(|w| !grids.contains(w)) + .collect(); + if more.is_empty() { + break; + } + grids.extend(more); + } + let mut thresholds = Vec::new(); + for &pane_width_ms in &grids { + let eligible: Vec<&ChainReader> = readers + .iter() + .filter(|r| r.pane_width_ms.is_multiple_of(pane_width_ms)) + .collect(); + let windows: BTreeSet = eligible.iter().map(|r| r.window_ms).collect(); + // The sizes: the readers' own, closed under the shared size. + let mut sizes: Vec = Vec::new(); + for reader in &eligible { + if !sizes.contains(&reader.family) { + sizes.push(reader.family.clone()); + } + } + let mut i = 0; + while i < sizes.len() { + for j in 0..i { + if let Some(joined) = shared_size(&sizes[i], &sizes[j]) { + if !sizes.contains(&joined) { + sizes.push(joined); + } + } + } + i += 1; + } + for &longest_window_ms in &windows { + for family in &sizes { + let threshold = ChainThreshold { + pane_width_ms, + longest_window_ms, + family: family.clone(), + }; + let admitted: Vec<&&ChainReader> = eligible + .iter() + .filter(|r| may_read(&threshold, r)) + .collect(); + let window = admitted.iter().map(|r| r.window_ms).max(); + let size = admitted.iter().try_fold(None::, |acc, r| { + Some(Some(match acc { + None => r.family.clone(), + Some(acc) => shared_size(&acc, &r.family)?, + })) + }); + if window == Some(longest_window_ms) && size == Some(Some(family.clone())) { + thresholds.push(threshold); + } + } + } + } + thresholds +} + +/// Whether `reader` may read a chain at `threshold`: its panes are +/// multiples of the chain's, its window is no longer than the chain's, and +/// its size is no stricter (the chain's size serves it). +pub(crate) fn may_read(threshold: &ChainThreshold, reader: &ChainReader) -> bool { + reader.pane_width_ms.is_multiple_of(threshold.pane_width_ms) + && reader.window_ms <= threshold.longest_window_ms + && shared_size(&threshold.family, &reader.family).as_ref() == Some(&threshold.family) +} + +/// `candidate`'s node with each of its chain reads of `site`'s key reading +/// the newest panes of the chain at `threshold`, which each may read +/// ([`may_read`]); every guarantee composed again. +pub(crate) fn read_at_threshold( + candidate: &ImplementationCandidate, + site: &CouplingSite, + threshold: &ChainThreshold, +) -> Result, PerComputationCandidateError> { + let pane_count = u32::try_from(threshold.longest_window_ms / threshold.pane_width_ms) + .expect("a chain's pane count fits"); + let mut roots = vec![QueryRoot::Operator(Rc::clone(&candidate.node))]; + loop { + // Each rewrite rebuilds the node: find the next unread merge again. + let QueryRoot::Operator(node) = &roots[0] else { + unreachable!("an operator root"); + }; + let next = chain_readers(node).into_iter().find(|reader| { + reader.key == site.key + && !(reader.pane_width_ms == threshold.pane_width_ms + && matches!(reader.merge.asap(), Some(ASAPOp::SummaryMerge { children, .. }) + if matches!(children[0].asap(), Some(ASAPOp::PaneChain { pane_count: n, family, .. }) + if *n == pane_count && *family == threshold.family))) + }); + let Some(reader) = next else { + break; + }; + roots = with_pane_chain( + &roots, + &reader.merge, + threshold.family.clone(), + threshold.pane_width_ms, + pane_count, + )?; + } + let [QueryRoot::Operator(node)] = &roots[..] else { + unreachable!("an operator root"); + }; + Ok(Rc::clone(node)) +} + +/// Rollup: each coarse exact accumulator of a group's Pass 1 (or +/// shared-build) candidate that another group's candidate can answer by +/// regrouping its finalized values (`rollup_candidates`), as a +/// [`CandidateOrigin::Rollup`] candidate that reads that group: the finer +/// value is the finer group's output, so the finer candidate must be that +/// group's whole implementation. Sketch rollups, which merge a finer +/// build's states rather than a group's output, are not offered. +pub(crate) fn add_rollup(memo: &mut Memo) { + let group_of = group_of_node(&memo.groups); + let mut added: Vec<(usize, ImplementationCandidate)> = Vec::new(); + let base = |c: &ImplementationCandidate| { + matches!( + c.origin, + CandidateOrigin::Pass1 | CandidateOrigin::SharedBuild + ) + }; + for (coarse_group, coarse_candidates) in memo.candidates.iter().enumerate() { + for coarse in coarse_candidates.iter().filter(|c| base(c)) { + let reached: HashSet<*const OperatorNode> = OperatorNode::reachable(&coarse.node) + .iter() + .map(Rc::as_ptr) + .collect(); + for (finer_group, finer_candidates) in memo.candidates.iter().enumerate() { + if finer_group == coarse_group { + continue; + } + for finer in finer_candidates.iter().filter(|c| base(c)) { + let roots = [ + QueryRoot::Operator(Rc::clone(&coarse.node)), + QueryRoot::Operator(Rc::clone(&finer.node)), + ]; + for rollup in rollup_candidates(&roots) { + let exact = matches!( + rollup.finer.asap(), + Some(ASAPOp::FinalizeExactAccumulator { .. }) + ); + if !exact + || !Rc::ptr_eq(&rollup.finer, &finer.node) + || !reached.contains(&Rc::as_ptr(&rollup.coarse)) + { + continue; + } + // As `rollup_pass2_candidates` rewrites it, with the finer + // value read from the finer group: every finalized value + // of the coarse build replaced by the rollup. + let finer_output = &memo.groups[finer_group].node; + let to_group = HashMap::from([(Rc::as_ptr(&finer.node), finer_output)]); + let Ok(rolled) = replace(&rollup.rollup, &to_group, &mut HashMap::new()) + else { + continue; + }; + let finalized: Vec> = OperatorNode::reachable( + &coarse.node, + ) + .into_iter() + .filter(|n| { + matches!(n.asap(), Some(ASAPOp::FinalizeExactAccumulator { child }) + if Rc::ptr_eq(child, &rollup.coarse)) + }) + .collect(); + let by_rollup: HashMap<*const OperatorNode, &Rc> = + finalized.iter().map(|n| (Rc::as_ptr(n), &rolled)).collect(); + let Ok(node) = replace(&coarse.node, &by_rollup, &mut HashMap::new()) + else { + continue; + }; + added.push(( + coarse_group, + rewritten_candidate(&group_of, coarse, node, CandidateOrigin::Rollup), + )); + } + } + } + } + } + for (g, candidate) in added { + push_new(memo, g, candidate); + } +} + +/// Exact top-k reuse: the groups whose Pass 1 candidates are exact top-k +/// computations of one input and partition with at least two different +/// `k` (`topk_reuse_candidates` over the candidates' nodes, which carry +/// their guarantees) each get a [`CandidateOrigin::TopKReuse`] candidate +/// reading the ranking group, `Limit(k_max)` over a descending `Sort` of +/// the input group (the top-k group's child). The ranking's two groups +/// are added when the memo has no equal ones, each with its one candidate +/// (also of origin `TopKReuse`). The rule reads an exact input: the input +/// group's chosen implementation must be exact. +pub(crate) fn add_topk_reuse(memo: &mut Memo) { + let mut group_of_candidate: HashMap<*const OperatorNode, usize> = HashMap::new(); + let mut roots: Vec = Vec::new(); + for (g, candidates) in memo.candidates.iter().enumerate() { + for candidate in candidates { + if candidate.origin == CandidateOrigin::Pass1 { + group_of_candidate.insert(Rc::as_ptr(&candidate.node), g); + roots.push(QueryRoot::Operator(Rc::clone(&candidate.node))); + } + } + } + for candidate in topk_reuse_candidates(&roots) { + let mut members: Vec = candidate + .members + .iter() + .filter_map(|m| group_of_candidate.get(&Rc::as_ptr(m)).copied()) + .collect(); + members.sort_unstable(); + members.dedup(); + if members.len() < 2 { + continue; + } + let queries: Vec = members + .iter() + .flat_map(|&g| memo.groups[g].queries.iter().copied()) + .collect::>() + .into_iter() + .collect(); + // The ranking over the input group's node. + let Some(&input) = memo.groups[members[0]].node.children().first() else { + continue; + }; + let input = Rc::clone(input); + let sort = candidate.ranking.children()[0]; + let Ok(sort) = sort.with_new_children(|_| Rc::clone(&input)) else { + continue; + }; + let sort = Rc::new(sort); + let input_group = find_group(&memo.groups, &input); + let sort_group = group_for(memo, &sort, input_group.into_iter().collect(), &queries); + let sort = Rc::clone(&memo.groups[sort_group.0 as usize].node); + let ranking = match candidate.ranking.with_new_children(|_| Rc::clone(&sort)) { + Ok(ranking) => Rc::new(ranking), + Err(_) => continue, + }; + let ranking_group = group_for(memo, &ranking, vec![sort_group], &queries); + let ranking = Rc::clone(&memo.groups[ranking_group.0 as usize].node); + for &g in &members { + let Ok(node) = memo.groups[g] + .node + .with_new_children(|_| Rc::clone(&ranking)) + else { + continue; + }; + let reuse = ImplementationCandidate { + implements: vec![GroupId(g as u32)], + node: Rc::new(node), + children: vec![ranking_group], + origin: CandidateOrigin::TopKReuse, + coupling_sites: Vec::new(), + }; + push_new(memo, g, reuse); + } + } +} + +/// The group whose node equals `node`, or a new one over `children`, read +/// by `queries`, with `node` itself as its one candidate. +fn group_for( + memo: &mut Memo, + node: &Rc, + children: Vec, + queries: &[usize], +) -> GroupId { + if let Some(found) = find_group(&memo.groups, node) { + return found; + } + let id = GroupId(memo.groups.len() as u32); + for child in &children { + memo.groups[child.0 as usize].parents.push(id); + } + memo.groups.push(MemoGroup { + node: Rc::clone(node), + children: children.clone(), + parents: Vec::new(), + queries: queries.to_vec(), + roots_of: Vec::new(), + }); + memo.candidates.push(vec![ImplementationCandidate { + implements: vec![id], + node: Rc::clone(node), + children, + origin: CandidateOrigin::TopKReuse, + coupling_sites: Vec::new(), + }]); + id +} + +/// The group whose node is `node` (the same node, or an equal one). +fn find_group(groups: &[MemoGroup], node: &Rc) -> Option { + groups + .iter() + .position(|g| Rc::ptr_eq(&g.node, node)) + .or_else(|| groups.iter().position(|g| *g.node == **node)) + .map(|g| GroupId(g as u32)) +} + +/// Each group's node, by address. +fn group_of_node(groups: &[MemoGroup]) -> HashMap<*const OperatorNode, usize> { + groups + .iter() + .enumerate() + .map(|(g, group)| (Rc::as_ptr(&group.node), g)) + .collect() +} + +/// `from` rewritten to `node`: the groups it implements and its sites +/// kept, its child groups those `node` reaches (the groups' nodes below +/// its root), or `from`'s when `from`'s node does not reach its children +/// by their nodes. +fn rewritten_candidate( + group_of: &HashMap<*const OperatorNode, usize>, + from: &ImplementationCandidate, + node: Rc, + origin: CandidateOrigin, +) -> ImplementationCandidate { + let reached = |node: &Rc| child_groups(group_of, node, &from.implements); + let children = if reached(&from.node) == from.children { + reached(&node) + } else { + from.children.clone() + }; + ImplementationCandidate { + implements: from.implements.clone(), + node, + children, + origin, + coupling_sites: from.coupling_sites.clone(), + } +} + +/// The groups other than `implements` whose nodes `node` reaches below its +/// root without passing through another group's node, in discovery order. +fn child_groups( + group_of: &HashMap<*const OperatorNode, usize>, + node: &Rc, + implements: &[GroupId], +) -> Vec { + fn walk( + node: &Rc, + group_of: &HashMap<*const OperatorNode, usize>, + implements: &[GroupId], + seen: &mut HashSet<*const OperatorNode>, + out: &mut Vec, + ) { + for child in node.children() { + 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)); + } + } + _ => walk(child, group_of, implements, seen, out), + } + } + } + let mut out = Vec::new(); + walk(node, group_of, implements, &mut HashSet::new(), &mut out); + out +} + +/// Add `candidate` to group `g` unless an equal one (same node, sites and +/// origin) is there. +fn push_new(memo: &mut Memo, g: usize, candidate: ImplementationCandidate) { + let flat = |node: &Rc| flatten(&[QueryRoot::Operator(Rc::clone(node))]).0; + let new = flat(&candidate.node); + let known = memo.candidates[g].iter().any(|c| { + c.origin == candidate.origin + && c.coupling_sites == candidate.coupling_sites + && flat(&c.node) == new + }); + if !known { + memo.candidates[g].push(candidate); + } +} + +/// `node` with each node of `replaced` (by address) swapped. +fn replace( + node: &Rc, + replaced: &HashMap<*const OperatorNode, &Rc>, + done: &mut HashMap<*const OperatorNode, Rc>, +) -> Result, asap_types::ir::SchemaDerivationError> { + if let Some(found) = replaced.get(&Rc::as_ptr(node)) { + return Ok(Rc::clone(found)); + } + if let Some(found) = done.get(&Rc::as_ptr(node)) { + return Ok(Rc::clone(found)); + } + let mut children = HashMap::new(); + let mut changed = false; + for child in node.children() { + let new = replace(child, replaced, done)?; + changed |= !Rc::ptr_eq(&new, child); + children.insert(Rc::as_ptr(child), new); + } + let out = if changed { + Rc::new(node.with_new_children(|c| Rc::clone(&children[&Rc::as_ptr(c)]))?) + } else { + Rc::clone(node) + }; + done.insert(Rc::as_ptr(node), Rc::clone(&out)); + Ok(out) +} + +/// The grid every query of `group` is evaluated on: the gcd of their +/// grids; `None` when one does not repeat at a fixed interval. +fn group_grid_ms(group: &MemoGroup, entries: &[QueryWorkloadEntry]) -> Option { + group + .queries + .iter() + .map(|&q| entries[q].evaluation_grid_ms()) + .try_fold(0, |grid, query| Some(greatest_common_divisor(grid, query?))) + .filter(|&grid| grid > 0) +} + +/// The size that serves readers of families `a` and `b`: either, for one +/// exact accumulator; a sketch of their algorithm sized for both; `None` +/// when they cannot be one build. +fn shared_size(a: &FieldDataType, b: &FieldDataType) -> Option { + match (a, b) { + (FieldDataType::Sketch(a, grouping), FieldDataType::Sketch(b, grouping_b)) + if a.algorithm() == b.algorithm() && grouping == grouping_b => + { + let params = size_for_all_readers(&[a.params(), b.params()])?; + Some(FieldDataType::Sketch( + SketchKind::new(a.algorithm().clone(), params), + grouping.clone(), + )) + } + (FieldDataType::Sketch(..), _) | (_, FieldDataType::Sketch(..)) => None, + _ => (a == b).then(|| a.clone()), + } +} + +/// `family` without its size: a sketch's algorithm and grouping, or the +/// family itself. +fn family_without_size(family: &FieldDataType) -> String { + match family { + FieldDataType::Sketch(kind, grouping) => format!("{:?}/{grouping:?}", kind.algorithm()), + other => serde_json::to_string(other).unwrap_or_default(), + } +} + +fn hash_of(value: &T) -> u64 { + let mut hasher = DefaultHasher::new(); + value.hash(&mut hasher); + hasher.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::test_support::{entry_with, lower_promql}; + use asap_logical_asap_optimization::workload_candidates::stage1_logical_asap_optimization; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::ir::schema::schema_support::with_promql_series_identity; + use asap_types::ir::NonASAPOp; + use asap_types::workload::{QueryRecurrence, RepeatedDemand, RepetitionInterval}; + + fn repeating(interval_ms: u32) -> QueryWorkloadEntry { + entry_with( + AccuracyTarget::exact(), + QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval( + interval_ms, + ))), + ) + } + + /// A memo by hand: one group per query (`queries[q]` repeating every + /// `intervals_ms[q]`), its node the lowered query, its candidates every + /// Pass 1 candidate of the query alone, each wholly private. + fn memo(queries: &[&str], intervals_ms: &[u32]) -> (Memo, Vec) { + let entries: Vec = + intervals_ms.iter().map(|&ms| repeating(ms)).collect(); + let mut memo = Memo::default(); + for (q, query) in queries.iter().enumerate() { + let node = lower_promql(query, AccuracyTarget::exact()); + let node = with_promql_series_identity(&node).unwrap(); + let inventory = generate_per_computation_candidates(vec![( + q, + QueryRoot::Operator(Rc::clone(&node)), + )]) + .unwrap(); + let id = GroupId(q as u32); + let candidates = stage1_logical_asap_optimization(&inventory, &entries[q..=q]) + .filter(|c| c.pass2.is_none()) + .filter_map(|c| match c.roots.ok()?.pop()? { + QueryRoot::Operator(node) => Some(node), + QueryRoot::Scalar(_) => None, + }) + .map(|node| ImplementationCandidate { + implements: vec![id], + node, + children: Vec::new(), + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + }) + .collect(); + memo.groups.push(MemoGroup { + node, + children: Vec::new(), + parents: Vec::new(), + queries: vec![q], + roots_of: vec![q], + }); + memo.candidates.push(candidates); + } + (memo, entries) + } + + /// The pane-chain sites of group `g`'s candidates. + fn chain_sites(memo: &Memo, g: usize) -> BTreeSet { + memo.candidates[g] + .iter() + .flat_map(|c| c.coupling_sites.iter().copied()) + .filter(|s| s.kind == CouplingKind::PaneChain) + .collect() + } + + /// Sums over 5 m and 1 h, both every 60 s: each group gets its own pane + /// candidate and one reading a shared chain, the two at one site, + /// whose only grid is 60 s; the 1-hour threshold serves both, and a + /// chain read at it merges the newest 5 of 60 panes. + #[test] + fn nested_windows_share_one_chain_site_on_one_grid() { + let (mut memo, entries) = memo( + &["sum_over_time(m[5m])", "sum_over_time(m[1h])"], + &[60_000, 60_000], + ); + add_window_composition(&mut memo, &entries); + for g in 0..2 { + assert!(memo.candidates[g] + .iter() + .any(|c| c.origin == CandidateOrigin::Panes && c.coupling_sites.is_empty())); + } + let site = chain_sites(&memo, 0); + assert_eq!(site.len(), 1); + assert_eq!(site, chain_sites(&memo, 1)); + let site = *site.first().unwrap(); + let thresholds = chain_thresholds(&memo, &site); + let grids: BTreeSet = thresholds.iter().map(|t| t.pane_width_ms).collect(); + assert_eq!(grids, BTreeSet::from([60_000])); + let hour = thresholds + .iter() + .find(|t| t.longest_window_ms == 3_600_000) + .unwrap(); + let five_minutes = memo.candidates[0] + .iter() + .find(|c| c.coupling_sites.contains(&site)) + .unwrap(); + let reader = &chain_readers(&five_minutes.node)[0]; + assert!(may_read(hour, reader)); + let read = read_at_threshold(five_minutes, &site, hour).unwrap(); + let newest: Vec<(Option, u32)> = OperatorNode::reachable(&read) + .iter() + .filter_map(|n| match n.asap() { + Some(ASAPOp::SummaryMerge { + children, + newest_panes, + .. + }) => match children[0].asap() { + Some(ASAPOp::PaneChain { pane_count, .. }) => { + Some((*newest_panes, *pane_count)) + } + _ => None, + }, + _ => None, + }) + .collect(); + assert_eq!(newest, vec![(Some(5), 60)]); + } + + /// A 5-minute sum every 60 s (panes of 60 s) and a 1-hour sum every + /// 300 s (panes of 300 s): grids 60 s and 300 s, and only at 60 s does + /// a threshold serve both. + #[test] + fn readers_at_60_and_300_seconds_meet_on_the_60_second_grid() { + let (mut memo, entries) = memo( + &["sum_over_time(m[5m])", "sum_over_time(m[1h])"], + &[60_000, 300_000], + ); + add_window_composition(&mut memo, &entries); + let site = *chain_sites(&memo, 0).first().unwrap(); + let thresholds = chain_thresholds(&memo, &site); + let grids: BTreeSet = thresholds.iter().map(|t| t.pane_width_ms).collect(); + assert_eq!(grids, BTreeSet::from([60_000, 300_000])); + let readers: Vec = (0..2) + .flat_map(|g| memo.candidates[g].iter()) + .filter(|c| c.coupling_sites.contains(&site)) + .flat_map(|c| chain_readers(&c.node)) + .collect(); + let both = |t: &ChainThreshold| readers.iter().all(|r| may_read(t, r)); + assert!(thresholds + .iter() + .any(|t| t.pane_width_ms == 60_000 && both(t))); + assert!(!thresholds + .iter() + .any(|t| t.pane_width_ms == 300_000 && both(t))); + } + + /// Sums over 5 m by `(job, instance)` and by `job`: the per-job group + /// gets a rollup candidate that reads the finer group's output. + #[test] + fn a_rollup_candidate_reads_the_finer_group() { + let (mut memo, entries) = memo( + &[ + "sum by (job, instance) (sum_over_time(m[5m]))", + "sum by (job) (sum_over_time(m[5m]))", + ], + &[60_000, 60_000], + ); + add_pass2_candidates(&mut memo, &entries); + let rollups: Vec<&ImplementationCandidate> = memo.candidates[1] + .iter() + .filter(|c| c.origin == CandidateOrigin::Rollup) + .collect(); + assert!(!rollups.is_empty()); + for rollup in rollups { + assert_eq!(rollup.children, vec![GroupId(0)]); + assert!(OperatorNode::reachable(&rollup.node) + .iter() + .any(|n| Rc::ptr_eq(n, &memo.groups[0].node))); + } + assert!(memo.candidates[0] + .iter() + .all(|c| c.origin != CandidateOrigin::Rollup)); + } + + /// Exact top-5 and top-10 of one per-job sum: the ranking (`Limit(10)` + /// over a `Sort`) is added as a group, and each top-k gets a candidate + /// reading it; it cannot be chosen with a pane candidate. + #[test] + fn a_topk_reuse_candidate_reads_the_ranking() { + let (mut memo, entries) = memo( + &[ + "topk(5, sum by (job) (sum_over_time(m[5m])))", + "topk(10, sum by (job) (sum_over_time(m[5m])))", + ], + &[60_000, 60_000], + ); + add_pass2_candidates(&mut memo, &entries); + let ranking = memo + .groups + .iter() + .position(|g| { + matches!( + g.node.non_asap(), + Some(NonASAPOp::Limit { n: Some(10), .. }) + ) + }) + .expect("a ranking group"); + for g in 0..2 { + let reuse = memo.candidates[g] + .iter() + .find(|c| c.origin == CandidateOrigin::TopKReuse) + .expect("a top-k reuse candidate"); + assert_eq!(reuse.children, vec![GroupId(ranking as u32)]); + assert!(Rc::ptr_eq( + reuse.node.children()[0], + &memo.groups[ranking].node + )); + let panes = ImplementationCandidate { + origin: CandidateOrigin::Panes, + ..reuse.clone() + }; + assert!(!combinable([reuse, &panes])); + assert!(combinable([reuse])); + } + } +} From 2606697a804161229349ebf853b6b86210380eae Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:08:54 +0000 Subject: [PATCH 02/25] plan-selection: sub-expression memo candidates read child group nodes, budget-share variants, topological order Each Pass 1 candidate's private nodes read their child groups through the groups' own nodes and output the group's schema. A sketch candidate is offered once per share of its measure's budget a reader may give it (budget_share), composed by the new compose_logical_candidate_with_budget_share. Memo::topological_order orders groups after every group they or their candidates read. --- .../src/per_computation_candidates.rs | 35 +- crates/plan-selection/src/pass/mod.rs | 2 +- .../src/pass/subexpr/capability.rs | 1 + .../plan-selection/src/pass/subexpr/groups.rs | 444 +++++++++++++++--- crates/plan-selection/src/pass/subexpr/mod.rs | 6 + .../src/pass/subexpr/pricing.rs | 1 + .../src/pass/subexpr/windows.rs | 4 + 7 files changed, 437 insertions(+), 56 deletions(-) diff --git a/crates/logical-asap-optimization/src/per_computation_candidates.rs b/crates/logical-asap-optimization/src/per_computation_candidates.rs index ad63a20ca..e9c2f4a1e 100644 --- a/crates/logical-asap-optimization/src/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/per_computation_candidates.rs @@ -1157,6 +1157,30 @@ fn with_raw_data_samples_requirement( pub fn compose_logical_candidate( inventory: &PerComputationCandidates, choice: &[usize], +) -> Result, PerComputationCandidateError> { + compose_with_budget_share(inventory, choice, None) +} + +/// [`compose_logical_candidate`] with every sketch summary sized for one of +/// `share` equal shares of its measure's budget, whatever the roots +/// compose: the size a summary gets when it is one of `share` approximate +/// inputs of a query ([`budget_shares`]). `share` 1 sizes for the whole +/// budget. The sub-expression memo composes one aggregate alone this way, +/// once per share its readers may give it. +pub fn compose_logical_candidate_with_budget_share( + inventory: &PerComputationCandidates, + choice: &[usize], + share: usize, +) -> Result, PerComputationCandidateError> { + compose_with_budget_share(inventory, choice, Some(share)) +} + +/// [`compose_logical_candidate`], each summary sized for `budget_share` +/// shares of its budget, or as [`budget_shares`] splits it when `None`. +fn compose_with_budget_share( + inventory: &PerComputationCandidates, + choice: &[usize], + budget_share: Option, ) -> Result, PerComputationCandidateError> { if choice.len() != inventory.targets.len() { return Err(PerComputationCandidateError::InvalidChoice); @@ -1189,7 +1213,16 @@ pub fn compose_logical_candidate( // again with each sized for an equal share ([`budget_shares`]). let whole = Sizing::default(); let roots = compose_roots(inventory, &chosen, &pushed_below, &whole)?; - let shares = budget_shares(&roots, &whole.built.into_inner()); + let shares = match budget_share { + None => budget_shares(&roots, &whole.built.into_inner()), + Some(share) if share > 1 => whole + .built + .into_inner() + .into_iter() + .map(|(_, key)| (key, share)) + .collect(), + Some(_) => HashMap::new(), + }; if shares.is_empty() { return Ok(roots); } diff --git a/crates/plan-selection/src/pass/mod.rs b/crates/plan-selection/src/pass/mod.rs index 16053a432..060b4bfde 100644 --- a/crates/plan-selection/src/pass/mod.rs +++ b/crates/plan-selection/src/pass/mod.rs @@ -14,9 +14,9 @@ mod components; mod interchangeable; -mod subexpr; mod query_candidates; mod stage1_to_stage3; +mod subexpr; use std::collections::BTreeMap; use std::rc::Rc; diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index f67dd3511..e8481b4da 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -314,6 +314,7 @@ mod tests { children: vec![GroupId(1)], origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + budget_share: 1, }]); } add_shared_builds(&mut memo); diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index 4928e8675..c0c1e0c64 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -13,8 +13,19 @@ //! [`compose_logical_candidate`] with that choice, every target beneath it //! exact. Its private nodes are the composed nodes above the first node //! equal to one of the group's descendant groups, which are its children. -//! Summaries are sized for their measure's whole budget, as composing the -//! aggregate alone does. +//! The private 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 whose value combines `n` approximate inputs splits its budget +//! into `n` equal shares, so a summary's size depends on its readers: each +//! sketch candidate is offered once per share `1..=n`, `n` the most +//! approximate inputs a query reaching the group can have +//! ([`ImplementationCandidate::budget_share`]). A plan is valid with the +//! share composition gives: the most approximate inputs among the queries +//! that read the summary, 1 when such a query's value carries no +//! guarantee. A smaller share fails the reader's accuracy check, a larger +//! one costs more. use std::collections::hash_map::DefaultHasher; use std::collections::{HashMap, HashSet}; @@ -22,10 +33,11 @@ use std::hash::{Hash, Hasher}; use std::rc::Rc; use asap_logical_asap_optimization::per_computation_candidates::{ - compose_logical_candidate, PerComputationCandidates, + compose_logical_candidate_with_budget_share, ComputationCandidate, PerComputationCandidates, }; use asap_types::ir::cse::share_common_sub_dags; -use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::ir::schema::FieldDataType; +use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot}; use super::{CandidateOrigin, GroupId, ImplementationCandidate, Memo, MemoGroup}; @@ -182,7 +194,26 @@ fn pass1_candidates( children: groups[g].children.clone(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + budget_share: 1, + }; + // Each query's most approximate inputs: its targets that offer a + // summary, each one sketch state at most. + let offers_a_summary = |t: usize| { + let target = &per_computation.targets[t]; + let summary = |c: &ComputationCandidate| matches!(c, ComputationCandidate::Summary(_)); + target.choices.candidates.iter().any(summary) + || target + .fused + .as_ref() + .is_some_and(|fused| fused.candidates.iter().any(summary)) }; + 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 { + most_approximate_inputs[q] += summaries; + } + } (0..groups.len()) .map(|g| match targets_of.get(&GroupId(g as u32)) { None => vec![own(g)], @@ -192,6 +223,13 @@ fn pass1_candidates( group_of, GroupId(g as u32), targets, + groups[g] + .queries + .iter() + .map(|&q| most_approximate_inputs[q]) + .max() + .unwrap_or(1) + .max(1), ) .unwrap_or_else(|| vec![own(g)]), }) @@ -208,6 +246,7 @@ fn target_candidates( group_of: &GroupOf, g: GroupId, targets: &[usize], + max_budget_share: usize, ) -> Option> { let aggregate = Rc::clone(&per_computation.targets[targets[0]].aggregate); // The groups a candidate's private nodes may bottom out at, by key. @@ -228,49 +267,73 @@ fn target_candidates( let mut candidates = Vec::new(); let mut choice = vec![0; per_computation.targets.len()]; loop { - if let Ok(composed) = compose_alone(per_computation, &aggregate, targets, &choice) { - if let Some((_, QueryRoot::Operator(root))) = composed.into_iter().next() { - let mut implements = vec![g]; - for &t in targets { - if let Some((_, Some(absorbs))) = - per_computation.targets[t].computation_for(choice[t]) - { - for &a in absorbs { - let absorbed = - group_of[&Rc::as_ptr(&per_computation.targets[a].aggregate)]; - if !implements.contains(&absorbed) { - implements.push(absorbed); - } - } + let mut implements = vec![g]; + for &t in targets { + if let Some((_, Some(absorbs))) = per_computation.targets[t].computation_for(choice[t]) + { + for &a in absorbs { + let absorbed = group_of[&Rc::as_ptr(&per_computation.targets[a].aggregate)]; + if !implements.contains(&absorbed) { + implements.push(absorbed); } } - let candidate = if keys.key(&root) == own_key { - ImplementationCandidate { - implements, - node: Rc::clone(&groups[g.0 as usize].node), - children: groups[g.0 as usize].children.clone(), - origin: CandidateOrigin::Pass1, - coupling_sites: Vec::new(), - } + } + } + for budget_share in 1..=max_budget_share { + let Ok(composed) = + compose_alone(per_computation, &aggregate, targets, &choice, budget_share) + else { + continue; + }; + let Some((_, QueryRoot::Operator(root))) = composed.into_iter().next() else { + continue; + }; + if budget_share > 1 && !builds_a_sketch(&root) { + break; + } + let (node, children) = if keys.key(&root) == own_key { + ( + Rc::clone(&groups[g.0 as usize].node), + groups[g.0 as usize].children.clone(), + ) + } else { + let mut children = Vec::new(); + let node = over_groups( + &root, + groups, + &mut keys, + &descendant, + &mut children, + &mut HashMap::new(), + ); + // A parent reads the candidate in the group's place, so it + // outputs the group's schema: the same columns, with the + // group's metadata (composition may derive a closed label + // set where the frontend's node leaves it open). + let group_node = &groups[g.0 as usize].node; + if node.schema.fields != group_node.schema.fields + || node.schema.time_index != group_node.schema.time_index + { + continue; + } + let node = if node.schema == group_node.schema { + node } else { - let mut children = Vec::new(); - boundary( - &root, - &mut keys, - &descendant, - &mut children, - &mut HashSet::new(), - ); - ImplementationCandidate { - implements, - node: root, - children, - origin: CandidateOrigin::Pass1, - coupling_sites: Vec::new(), - } + Rc::new(OperatorNode { + schema: group_node.schema.clone(), + ..(*node).clone() + }) }; - candidates.push(candidate); - } + (node, children) + }; + candidates.push(ImplementationCandidate { + implements: implements.clone(), + node, + children, + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + budget_share, + }); } // Mixed-radix over this group's targets only. let mut advanced = false; @@ -298,6 +361,7 @@ fn compose_alone( aggregate: &Rc, targets: &[usize], choice: &[usize], + budget_share: usize, ) -> Result< Vec<(usize, QueryRoot)>, asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError, @@ -335,31 +399,76 @@ fn compose_alone( .iter() .map(|&t| if targets.contains(&t) { choice[t] } else { 0 }) .collect(); - compose_logical_candidate(&alone, &alone_choice) + compose_logical_candidate_with_budget_share(&alone, &alone_choice, budget_share) } -/// The groups `node`'s private nodes read: walking down from `node`'s -/// children, the first node on each path equal to a group in `descendant`. -fn boundary( +/// `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], keys: &mut NodeKeys, descendant: &HashMap, children: &mut Vec, - seen: &mut HashSet<*const OperatorNode>, -) { + rebuilt: &mut HashMap<*const OperatorNode, Rc>, +) -> Rc { + let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); for child in node.children() { - if !seen.insert(Rc::as_ptr(child)) { + let key = Rc::as_ptr(child); + if replaced.contains_key(&key) { continue; } - match descendant.get(&keys.key(child)) { + let new = match descendant.get(&keys.key(child)) { Some(&group) => { if !children.contains(&group) { children.push(group); } + Rc::clone(&groups[group.0 as usize].node) } - None => boundary(child, keys, descendant, children, seen), - } + None => match rebuilt.get(&key) { + Some(done) => Rc::clone(done), + None => { + let done = over_groups(child, groups, keys, descendant, children, rebuilt); + rebuilt.insert(key, Rc::clone(&done)); + done + } + }, + }; + replaced.insert(key, new); } + if node + .children() + .into_iter() + .all(|child| Rc::ptr_eq(child, &replaced[&Rc::as_ptr(child)])) + { + return Rc::clone(node); + } + Rc::new(OperatorNode { + operator: node + .operator + .map_children(|child| Rc::clone(&replaced[&Rc::as_ptr(child)])), + ..(**node).clone() + }) +} + +/// Whether `node`'s private nodes build a sketch, which a budget share sizes. +fn builds_a_sketch(node: &Rc) -> bool { + OperatorNode::reachable(node).iter().any(|n| { + matches!( + n.asap(), + Some( + ASAPOp::SummaryAgg { + family: FieldDataType::Sketch(..), + .. + } | ASAPOp::PaneChain { + family: FieldDataType::Sketch(..), + .. + } + ) + ) + }) } /// Structural hashes of nodes: the operator over its children's hashes, @@ -392,10 +501,57 @@ impl NodeKeys { } } +impl Memo { + /// The groups, each after every group it or one of its candidates + /// reads, the lower index first among those ready: the order the + /// dynamic program visits them, whatever order transformations appended + /// groups in. `None` when those reads form a cycle. + pub(crate) fn topological_order(&self) -> Option> { + let count = self.groups.len(); + let mut reads: Vec> = vec![HashSet::new(); count]; + for (g, group) in self.groups.iter().enumerate() { + let candidates = self.candidates.get(g).map_or(&[][..], Vec::as_slice); + for child in group + .children + .iter() + .chain(candidates.iter().flat_map(|c| &c.children)) + { + if child.0 as usize != g { + reads[g].insert(child.0 as usize); + } + } + } + let mut readers: Vec> = vec![Vec::new(); count]; + let mut waiting: Vec = reads.iter().map(HashSet::len).collect(); + for (g, read) in reads.iter().enumerate() { + for &child in read { + readers[child].push(g); + } + } + let mut ready: std::collections::BinaryHeap> = (0..count) + .filter(|&g| waiting[g] == 0) + .map(std::cmp::Reverse) + .collect(); + let mut order = Vec::with_capacity(count); + while let Some(std::cmp::Reverse(g)) = ready.pop() { + order.push(GroupId(g as u32)); + for &reader in &readers[g] { + waiting[reader] -= 1; + if waiting[reader] == 0 { + ready.push(std::cmp::Reverse(reader)); + } + } + } + (order.len() == count).then_some(order) + } +} + #[cfg(test)] mod tests { use super::*; - use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::per_computation_candidates::{ + compose_logical_candidate, generate_per_computation_candidates, + }; use asap_logical_asap_optimization::test_support::lower_promql; use asap_types::ir::properties::AccuracyTarget; @@ -514,4 +670,184 @@ mod tests { assert!(algorithms(candidate)[0].contains("Heap")); } } + + /// Every candidate reads its child groups through their nodes (the same + /// `Rc`), and outputs its group's schema, so a parent can read it in + /// the group's place. + #[test] + fn candidates_read_child_group_nodes_and_output_the_group_schema() { + for (queries, target) in [ + ( + vec!["quantile_over_time(0.9, m[5m]) / quantile_over_time(0.5, m[5m])"], + AccuracyTarget::max_error(0.01), + ), + ( + vec!["topk by (job) (10, sum_over_time(m[1m]))"], + AccuracyTarget::max_error(0.01).failure_probability(0.001), + ), + (vec!["sum by (job) (rate(m[5m]))"], AccuracyTarget::exact()), + ] { + let memo = memo(&queries, target); + for (g, candidates) in memo.candidates.iter().enumerate() { + for candidate in candidates { + let reached: HashSet<*const OperatorNode> = + OperatorNode::reachable(&candidate.node) + .iter() + .map(Rc::as_ptr) + .collect(); + for child in &candidate.children { + let node = &memo.groups[child.0 as usize].node; + assert!(reached.contains(&Rc::as_ptr(node)), "{queries:?} g{g}"); + } + assert_eq!(candidate.node.schema, memo.groups[g].node.schema); + } + } + } + } + + /// The sketch summaries of the first query of `queries` composed with + /// every target at its `algorithm` summary, by [`compose_logical_candidate`]: + /// as the exhaustive search sizes them. + fn composed_sketches( + queries: &[&str], + target: &AccuracyTarget, + algorithm: &str, + ) -> Vec { + let roots = queries + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, target.clone()))) + .enumerate() + .collect(); + let inventory = generate_per_computation_candidates(roots).unwrap(); + let choice: Vec = inventory + .targets + .iter() + .map(|t| { + t.choices + .candidates + .iter() + .position(|c| { + matches!(c, ComputationCandidate::Summary(a) if format!("{a:?}") == algorithm) + }) + .unwrap() + }) + .collect(); + let composed = compose_logical_candidate(&inventory, &choice).unwrap(); + let QueryRoot::Operator(root) = &composed[0].1 else { + unreachable!("an operator root") + }; + OperatorNode::reachable(root) + .iter() + .filter_map(|n| match n.asap() { + Some(ASAPOp::SummaryAgg { family, .. }) => Some(format!("{family:?}")), + _ => None, + }) + .collect() + } + + /// The sketch candidates of `memo`'s group `g`: each one's summary and + /// budget share. + fn sketch_shares(memo: &Memo, g: usize) -> Vec<(String, usize)> { + memo.candidates[g] + .iter() + .filter(|c| !algorithms(c).is_empty()) + .map(|c| (algorithms(c)[0].clone(), c.budget_share)) + .collect() + } + + /// `q0.9 / q0.5` beside a plain `q0.9`: each of the ratio's quantile + /// groups offers every sketch for the whole budget and for half of it, + /// since its query has two inputs that may be approximate; the plain + /// query's quantile group (its own: its unkeyed window is not shared) + /// offers only the whole budget. A ratio of quantiles carries no + /// guarantee, so composition does not split its budget: the + /// whole-budget candidate is the exhaustive search's sizing, and the + /// half-budget one, larger, is never the cheaper valid one. + #[test] + fn a_ratio_operand_has_a_candidate_per_budget_share() { + let queries = [ + "quantile_over_time(0.9, m[5m]) / quantile_over_time(0.5, m[5m])", + "quantile_over_time(0.9, m[5m])", + ]; + let target = AccuracyTarget::max_error(0.01); + let memo = memo(&queries, target.clone()); + let ratio = memo + .groups + .iter() + .position(|group| group.roots_of == vec![0]) + .unwrap(); + let operand = memo.groups[ratio].children[0].0 as usize; + for operand in &memo.groups[ratio].children { + let shares: Vec = sketch_shares(&memo, operand.0 as usize) + .into_iter() + .map(|(_, share)| share) + .collect(); + assert_eq!(shares, vec![1, 2, 1, 2]); + } + let plain = group_of_kind(&memo, 1, "Aggregate"); + let plain_shares: Vec = sketch_shares(&memo, plain) + .into_iter() + .map(|(_, share)| share) + .collect(); + assert_eq!(plain_shares, vec![1, 1]); + let whole_kll = &sketch_shares(&memo, operand)[0]; + assert_eq!(whole_kll.1, 1); + assert_eq!( + composed_sketches(&queries[..1], &target, "Kll"), + vec![whole_kll.0.clone(); 2] + ); + } + + /// A difference of two Count-Min counts splits its budget: each + /// operand's half-budget candidate is the Count-Min composition sizes + /// for the difference with both counts sketched. + #[test] + fn a_split_budget_matches_its_half_budget_candidate() { + let queries = ["count by (job) (a) - count by (job) (b)"]; + let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); + let memo = memo(&queries, target.clone()); + let difference = memo + .groups + .iter() + .position(|group| group.roots_of == vec![0]) + .unwrap(); + let operand = memo.groups[difference].children[0].0 as usize; + let half = sketch_shares(&memo, operand) + .into_iter() + .find(|(family, share)| family.contains("Cms") && *share == 2) + .unwrap(); + assert_eq!(composed_sketches(&queries, &target, "Cms"), vec![half.0; 2]); + } + + /// A group appended after its readers (as a transformation appends a + /// ranking group) is ordered before every group whose candidate reads + /// it. + #[test] + fn topological_order_puts_appended_groups_before_their_readers() { + let mut memo = memo( + &["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(GroupId) + .collect::>() + ); + let quantile = group_of_kind(&memo, 0, "Aggregate"); + let appended = GroupId(memo.groups.len() as u32); + memo.groups.push(memo.groups[1].clone()); + memo.candidates.push(memo.candidates[1].clone()); + let mut reader = memo.candidates[quantile][0].clone(); + reader.children = vec![appended]; + memo.candidates[quantile].push(reader); + let order = memo.topological_order().unwrap(); + let position = |g: GroupId| order.iter().position(|&x| x == g).unwrap(); + assert!(position(appended) < position(GroupId(quantile as u32))); + // A cycle has no order. + memo.candidates[1][0] + .children + .push(GroupId(quantile as u32)); + assert_eq!(memo.topological_order(), None); + } } diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 8eb747abb..cedff6346 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -75,6 +75,12 @@ pub(crate) struct ImplementationCandidate { pub origin: CandidateOrigin, /// The coupling sites its cost depends on. pub coupling_sites: Vec, + /// The number of equal shares of its measures' accuracy budget its + /// summaries are sized for: 1 for the whole budget, `n` when it is one + /// of `n` approximate inputs of a query (a ratio's operand). Valid in a + /// plan only when `n` is the largest number of approximate inputs among + /// the queries that read it. + pub budget_share: usize, } /// Where an implementation candidate came from. diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index 6c0e3f63c..bc2743f0f 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -559,6 +559,7 @@ mod tests { children, origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + budget_share: 1, }; let memo = Memo { groups: vec![group(&below, vec![]), group(top, vec![GroupId(0)])], diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index 194071e4c..bab8362fc 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -516,6 +516,7 @@ pub(crate) fn add_topk_reuse(memo: &mut Memo) { children: vec![ranking_group], origin: CandidateOrigin::TopKReuse, coupling_sites: Vec::new(), + budget_share: 1, }; push_new(memo, g, reuse); } @@ -550,6 +551,7 @@ fn group_for( children, origin: CandidateOrigin::TopKReuse, coupling_sites: Vec::new(), + budget_share: 1, }]); id } @@ -594,6 +596,7 @@ fn rewritten_candidate( children, origin, coupling_sites: from.coupling_sites.clone(), + budget_share: from.budget_share, } } @@ -766,6 +769,7 @@ mod tests { children: Vec::new(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + budget_share: 1, }) .collect(); memo.groups.push(MemoGroup { From ac5270d0e5f63714998a7800ecca1e47042142cd Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:10:07 +0000 Subject: [PATCH 03/25] plan-selection: sub-expression memo scan widening and raw-horizon sites ScanLabels sites per metric whose scans project different labels, with WidenedScans twins, label thresholds and widening at a threshold; and RawHorizon sites per source whose raw windows two or more groups read. Identical expression tested as group identity. --- .../src/pass/subexpr/capability.rs | 40 ++ .../plan-selection/src/pass/subexpr/groups.rs | 4 +- crates/plan-selection/src/pass/subexpr/mod.rs | 6 + .../plan-selection/src/pass/subexpr/scans.rs | 596 ++++++++++++++++++ 4 files changed, 644 insertions(+), 2 deletions(-) create mode 100644 crates/plan-selection/src/pass/subexpr/scans.rs diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index f67dd3511..3162b0da9 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -228,6 +228,7 @@ fn reader_build( #[cfg(test)] mod tests { use super::*; + use crate::pass::subexpr::MemoGroup; use asap_logical_asap_optimization::test_support::{metric_scan, time_range}; use asap_types::ir::scalar::ColumnRef; use asap_types::ir::schema::sketch::{SketchAlgorithm, SketchKind, SketchStatistic}; @@ -399,4 +400,43 @@ mod tests { .unwrap(); assert!(readers.iter().all(|r| may_read(strictest, r))); } + + /// Identical expression is group identity: the per-job rate sum of two + /// queries is one group, reached by both and read by each query's + /// arithmetic, and the summary-capability transformation adds nothing + /// for it. + #[test] + fn identical_sub_expressions_are_one_group_with_two_parents() { + use crate::pass::subexpr::groups::memo_of; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::ir::QueryRoot; + let roots = [ + "sum by (job) (rate(m[5m])) * 2", + "sum by (job) (rate(m[5m])) / 2", + ] + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, AccuracyTarget::exact()))) + .enumerate() + .collect(); + let mut memo = memo_of(&generate_per_computation_candidates(roots).unwrap()); + let before: Vec = memo.candidates.iter().map(Vec::len).collect(); + add_shared_builds(&mut memo); + let after: Vec = memo.candidates.iter().map(Vec::len).collect(); + assert_eq!(before, after); + let sums: Vec<&MemoGroup> = memo + .groups + .iter() + .filter(|g| g.node.operator.kind_name() == "Aggregate" && g.queries == vec![0, 1]) + .collect(); + let sum = sums.last().expect("the shared sum"); + assert_eq!(sum.parents.len(), 2); + let parents_roots: Vec> = sum + .parents + .iter() + .map(|p| memo.groups[p.0 as usize].roots_of.clone()) + .collect(); + assert_eq!(parents_roots, vec![vec![0], vec![1]]); + } } diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index 4928e8675..35ddc582d 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -367,12 +367,12 @@ fn boundary( /// derives. Memoized per node (the memo keeps each node alive, so its /// address is not reused). #[derive(Default)] -struct NodeKeys { +pub(super) struct NodeKeys { memo: HashMap<*const OperatorNode, (Rc, u64)>, } impl NodeKeys { - fn key(&mut self, node: &Rc) -> u64 { + pub(super) fn key(&mut self, node: &Rc) -> u64 { if let Some((_, key)) = self.memo.get(&Rc::as_ptr(node)) { return *key; } diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 8eb747abb..8bfaf2492 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -12,6 +12,8 @@ //! implementation candidates per group; //! - [`capability`] (W3): identical-expression and summary-capability //! transformations; +//! - [`scans`] (W3): scan widening to a metric's shared label set, and +//! each source's raw-data horizon; //! - [`windows`] (W4): window-composition, rollup and top-k reuse //! transformations; //! - [`pricing`] (W2): local prices per (group, candidate, context), @@ -23,6 +25,7 @@ pub(crate) mod capability; pub(crate) mod groups; pub(crate) mod pricing; +pub(crate) mod scans; pub(crate) mod search; pub(crate) mod windows; @@ -90,6 +93,9 @@ pub(crate) enum CandidateOrigin { Rollup, /// A top-k read from another group's ranking. TopKReuse, + /// Its metric's scans widened to a shared label set, read through a + /// `Project` (scan widening). + WidenedScans, } /// A choice several groups' costs depend on; equal keys are one site. diff --git a/crates/plan-selection/src/pass/subexpr/scans.rs b/crates/plan-selection/src/pass/subexpr/scans.rs new file mode 100644 index 000000000..3684e950d --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/scans.rs @@ -0,0 +1,596 @@ +//! Couplings through a source's rows: scan widening and the raw-data +//! horizon (joint-planning-search design, "Memo-group search", step 4). +//! +//! - **Scan widening** ([`add_scan_widening`]): the time-series scans of a +//! metric that project different labels are separate groups (an unkeyed +//! scan is never one group with another). The scans of the candidates +//! that read them can be widened to one label set, as +//! `with_union_scan_labels` does for a whole candidate, so that they are +//! one scan. Each reader candidate gets a [`CandidateOrigin::WidenedScans`] +//! twin with the metric's [`CouplingKind::ScanLabels`] site; the label +//! set is the site's threshold ([`label_thresholds`]), a reader may read +//! it iff its own labels are a subset ([`may_read_labels`]), and its node +//! at a threshold is [`widened_at_threshold`]: the labels passed on by the +//! nodes that pass them (a range or shift, a per-series build and its +//! reads) and narrowed back by a `Project` above them. +//! - **Raw horizon** ([`add_raw_horizons`]): the raw data a source keeps on +//! disk is the longest window a query-time scan reads from it. Every +//! candidate that reads a raw window of a source read by two or more +//! groups gets the source's [`CouplingKind::RawHorizon`] site; the +//! horizon is the site's threshold ([`horizon_thresholds`]), and a reader +//! may read at most it ([`may_read_horizon`]). Whether a read is at query +//! time is the reader's materialization, which the search knows: a read +//! built at ingestion keeps no raw data. + +use std::collections::hash_map::DefaultHasher; +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use asap_logical_asap_optimization::asap_aware_cse::with_union_scan_labels; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; +use asap_types::ir::operator::Source; +use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; + +use super::groups::NodeKeys; +use super::{CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo}; + +/// One label set a metric's scans may be widened to. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) struct LabelThreshold { + /// The column names the widened scan projects. + pub labels: BTreeSet, +} + +/// A candidate's read of a metric's scan: the columns its scan projects. +#[derive(Debug, Clone)] +pub(crate) struct ScanReader { + /// The metric's [`CouplingSite::key`]. + pub key: u64, + /// The column names the reader's scan projects. + pub labels: BTreeSet, + /// The scan. + pub scan: Rc, +} + +/// One horizon a source's raw data may be kept for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) struct HorizonThreshold { + /// The longest window read from the raw data, in ms. + pub horizon_ms: u64, +} + +/// A candidate's read of a raw window of a source. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct RawReader { + /// The source's [`CouplingSite::key`]. + pub key: u64, + /// The window it reads, in ms. + pub horizon_ms: u64, +} + +/// Scan widening: for each metric whose scans read by candidates of two or +/// more groups project different labels, a +/// [`CandidateOrigin::WidenedScans`] twin of each reader candidate with +/// the metric's [`CouplingKind::ScanLabels`] site. The twin keeps the +/// reader's node and children; [`widened_at_threshold`] gives both at a +/// threshold. +pub(crate) fn add_scan_widening(memo: &mut Memo) { + // Per metric key, the readers: (group, candidate index, labels). + let mut readers: HashMap)>> = HashMap::new(); + let mut order = Vec::new(); + for (g, candidates) in memo.candidates.iter().enumerate() { + if passes_rows_on(&memo.groups[g].node) { + // Its readers widen it. + continue; + } + for (index, candidate) in candidates.iter().enumerate() { + if candidate.origin == CandidateOrigin::WidenedScans { + continue; + } + for reader in reads(memo, candidate, scan_reader) { + let entry = readers.entry(reader.key).or_default(); + if entry.is_empty() { + order.push(reader.key); + } + if !entry.iter().any(|(eg, ei, _)| (*eg, *ei) == (g, index)) { + entry.push((g, index, reader.labels)); + } + } + } + } + for key in order { + let members = &readers[&key]; + let groups: HashSet = members.iter().map(|(g, _, _)| *g).collect(); + let label_sets: HashSet<&BTreeSet> = members.iter().map(|(_, _, l)| l).collect(); + if groups.len() < 2 || label_sets.len() < 2 { + continue; + } + let site = CouplingSite { + key, + kind: CouplingKind::ScanLabels, + }; + for (g, index, _) in members.clone() { + let mut twin = memo.candidates[g][index].clone(); + twin.origin = CandidateOrigin::WidenedScans; + twin.coupling_sites.push(site); + memo.candidates[g].push(twin); + } + } +} + +/// The reads of a metric's scans by `candidate`'s private nodes (directly, +/// or through ranges and shifts), each once. +pub(crate) fn scan_readers(memo: &Memo, candidate: &ImplementationCandidate) -> Vec { + reads(memo, candidate, scan_reader) +} + +/// The label sets the scans of `site` may be widened to: each reader's +/// own, and the union of every reader's, each once, smallest first. +pub(crate) fn label_thresholds(memo: &Memo, site: &CouplingSite) -> Vec { + let mut thresholds: BTreeSet> = BTreeSet::new(); + let mut union = BTreeSet::new(); + for reader in site_scan_readers(memo, site) { + union.extend(reader.labels.iter().cloned()); + thresholds.insert(reader.labels.into_iter().collect()); + } + thresholds.insert(union.into_iter().collect()); + let mut thresholds: Vec = thresholds + .into_iter() + .map(|labels| LabelThreshold { + labels: labels.into_iter().collect(), + }) + .collect(); + thresholds.sort_by_key(|t| t.labels.len()); + thresholds +} + +/// Whether `reader` may read a scan widened to `threshold`: its labels are +/// a subset. +pub(crate) fn may_read_labels(threshold: &LabelThreshold, reader: &ScanReader) -> bool { + reader.labels.is_subset(&threshold.labels) +} + +/// `candidate` (a twin with `site`) with its scans of the site's metric +/// widened to `threshold`, which each of its reads may read +/// ([`may_read_labels`]): its node rewritten as `with_union_scan_labels` +/// rewrites a candidate, its children the groups the new node reaches. +/// The candidate itself when its scans already project `threshold`; +/// `None` when the widening is not possible (a node that would change +/// arity, a summary state to narrow, a grouped or filtered pane chain over +/// the scan, or another metric's scans that would change too). +pub(crate) fn widened_at_threshold( + memo: &Memo, + candidate: &ImplementationCandidate, + site: &CouplingSite, + threshold: &LabelThreshold, +) -> Result, PerComputationCandidateError> { + let own: Vec = scan_readers(memo, candidate) + .into_iter() + .filter(|r| r.key == site.key) + .collect(); + if own.iter().all(|r| r.labels == threshold.labels) { + return Ok(Some(candidate.clone())); + } + if !own.iter().all(|r| may_read_labels(threshold, r)) { + return Ok(None); + } + // The site's scans whose labels the threshold covers: their union is + // the threshold when it is a union of readers' label sets. + let mut roots = vec![QueryRoot::Operator(Rc::clone(&candidate.node))]; + let mut covered = BTreeSet::new(); + for reader in site_scan_readers(memo, site) { + if may_read_labels(threshold, &reader) { + covered.extend(reader.labels.iter().cloned()); + roots.push(QueryRoot::Operator(Rc::clone(&reader.scan))); + } + } + if covered != threshold.labels { + return Ok(None); + } + let Some(widened) = with_union_scan_labels(&roots)? else { + return Ok(None); + }; + let QueryRoot::Operator(node) = &widened[0] else { + unreachable!("an operator root"); + }; + // Only the site's metric changes. + let before: HashSet<*const OperatorNode> = OperatorNode::reachable(&candidate.node) + .iter() + .map(Rc::as_ptr) + .collect(); + let metric = metric_of_site(memo, site); + let other_changed = OperatorNode::reachable(node).iter().any(|n| { + matches!(n.non_asap(), Some(NonASAPOp::Scan { source, .. }) + if Some(source) != metric.as_ref()) + && !before.contains(&Rc::as_ptr(n)) + }); + if other_changed { + return Ok(None); + } + Ok(Some(ImplementationCandidate { + implements: candidate.implements.clone(), + node: Rc::clone(node), + children: child_groups(&mut Boundary::of(memo, candidate), node), + origin: CandidateOrigin::WidenedScans, + coupling_sites: candidate.coupling_sites.clone(), + })) +} + +/// Raw horizon: for each source whose raw windows candidates of two or +/// more groups read, the source's [`CouplingKind::RawHorizon`] site added +/// to each of those candidates. +pub(crate) fn add_raw_horizons(memo: &mut Memo) { + let mut readers: HashMap> = HashMap::new(); + let mut order = Vec::new(); + for (g, candidates) in memo.candidates.iter().enumerate() { + for (index, candidate) in candidates.iter().enumerate() { + for reader in reads(memo, candidate, raw_reader) { + let entry = readers.entry(reader.key).or_default(); + if entry.is_empty() { + order.push(reader.key); + } + if !entry.contains(&(g, index)) { + entry.push((g, index)); + } + } + } + } + for key in order { + let members = &readers[&key]; + let groups: HashSet = members.iter().map(|(g, _)| *g).collect(); + if groups.len() < 2 { + continue; + } + let site = CouplingSite { + key, + kind: CouplingKind::RawHorizon, + }; + for &(g, index) in members { + memo.candidates[g][index].coupling_sites.push(site); + } + } +} + +/// The raw windows `candidate`'s private nodes read: each range window +/// directly over a scan, read by one of them. +pub(crate) fn raw_readers(memo: &Memo, candidate: &ImplementationCandidate) -> Vec { + reads(memo, candidate, raw_reader) +} + +/// The horizons the raw data of `site` may be kept for: its readers' +/// windows, each once, shortest first. +pub(crate) fn horizon_thresholds(memo: &Memo, site: &CouplingSite) -> Vec { + let horizons: BTreeSet = memo + .candidates + .iter() + .flatten() + .filter(|c| c.coupling_sites.contains(site)) + .flat_map(|c| raw_readers(memo, c)) + .filter(|r| r.key == site.key) + .map(|r| r.horizon_ms) + .collect(); + horizons + .into_iter() + .map(|horizon_ms| HorizonThreshold { horizon_ms }) + .collect() +} + +/// Whether `reader` may read raw data kept for `threshold`: its window is +/// no longer. +pub(crate) fn may_read_horizon(threshold: &HorizonThreshold, reader: &RawReader) -> bool { + reader.horizon_ms <= threshold.horizon_ms +} + +/// The scan readers of the candidates that carry `site`, one per distinct +/// label set. +fn site_scan_readers(memo: &Memo, site: &CouplingSite) -> Vec { + let mut out: Vec = Vec::new(); + for candidate in memo.candidates.iter().flatten() { + if !candidate.coupling_sites.contains(site) { + continue; + } + for reader in reads(memo, candidate, scan_reader) { + if reader.key == site.key && !out.iter().any(|r| r.labels == reader.labels) { + out.push(reader); + } + } + } + out +} + +/// The metric of a scan-labels site. +fn metric_of_site(memo: &Memo, site: &CouplingSite) -> Option { + site_scan_readers(memo, site) + .first() + .and_then(|r| match r.scan.non_asap() { + Some(NonASAPOp::Scan { source, .. }) => Some(source.clone()), + _ => None, + }) +} + +/// What `read` finds in each node a private node of `candidate` reads +/// (its children, private or a child group's node), each node once. +fn reads( + memo: &Memo, + candidate: &ImplementationCandidate, + read: impl Fn(&Rc) -> Option, +) -> Vec { + let mut boundary = Boundary::of(memo, candidate); + let mut out = Vec::new(); + let mut seen = HashSet::new(); + let mut stack = vec![Rc::clone(&candidate.node)]; + while let Some(node) = stack.pop() { + if !seen.insert(Rc::as_ptr(&node)) { + continue; + } + for child in node.children() { + if let Some(found) = read(child) { + out.push(found); + } + if boundary.group(child).is_none() { + stack.push(Rc::clone(child)); + } + } + } + out +} + +/// The scan of a metric `node` reads, directly or through ranges and +/// shifts, without predicates. +fn scan_reader(node: &Rc) -> Option { + match node.non_asap() { + Some(NonASAPOp::Scan { + source: source @ Source::TimeSeries { .. }, + predicates, + schema, + }) if predicates.is_empty() => Some(ScanReader { + key: hash_of(&("ScanLabels", format!("{source:?}"))), + labels: schema.fields.iter().map(|f| f.name.clone()).collect(), + scan: Rc::clone(node), + }), + Some(NonASAPOp::TimeRange { child, .. } | NonASAPOp::TimeShift { child, .. }) => { + scan_reader(child) + } + _ => None, + } +} + +/// The raw window `node` is, when it is a range window directly over a +/// scan. +fn raw_reader(node: &Rc) -> Option { + let Some(NonASAPOp::TimeRange { range, child, .. }) = node.non_asap() else { + return None; + }; + let Some(NonASAPOp::Scan { source, .. }) = child.non_asap() else { + return None; + }; + Some(RawReader { + key: hash_of(&("RawHorizon", format!("{source:?}"))), + horizon_ms: u64::try_from(range.as_millis()).unwrap_or(u64::MAX), + }) +} + +/// Whether `node` passes its input's rows on as a scan's reader sees them: +/// a scan, range or shift, whose readers widen it. +fn passes_rows_on(node: &OperatorNode) -> bool { + matches!( + node.non_asap(), + Some(NonASAPOp::Scan { .. } | NonASAPOp::TimeRange { .. } | NonASAPOp::TimeShift { .. }) + ) +} + +/// A candidate's child groups by node key ([`NodeKeys`]): its private +/// nodes end at the first node equal to one of them up to what composition +/// derives (schema, guarantee, timing, materialization). +struct Boundary { + keys: NodeKeys, + children: HashMap, +} + +impl Boundary { + fn of(memo: &Memo, candidate: &ImplementationCandidate) -> Self { + let mut keys = NodeKeys::default(); + let children = candidate + .children + .iter() + .map(|&child| (keys.key(&memo.groups[child.0 as usize].node), child)) + .collect(); + Self { keys, children } + } + + /// The child group `node` is, if any. + fn group(&mut self, node: &Rc) -> Option { + let key = self.keys.key(node); + self.children.get(&key).copied() + } +} + +/// The child groups `node` reaches below its root without passing through +/// another child group's node, in discovery order. +fn child_groups(boundary: &mut Boundary, node: &Rc) -> Vec { + let mut out = Vec::new(); + let mut seen = HashSet::new(); + let mut stack = vec![Rc::clone(node)]; + while let Some(node) = stack.pop() { + for child in node.children() { + if !seen.insert(Rc::as_ptr(child)) { + continue; + } + match boundary.group(child) { + Some(g) => { + if !out.contains(&g) { + out.push(g); + } + } + None => stack.push(Rc::clone(child)), + } + } + } + out +} + +fn hash_of(value: &T) -> u64 { + let mut hasher = DefaultHasher::new(); + value.hash(&mut hasher); + hasher.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pass::subexpr::groups::memo_of; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + + /// The memo of exact PromQL `queries`. + fn memo(queries: &[&str]) -> Memo { + let roots = queries + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, AccuracyTarget::exact()))) + .enumerate() + .collect(); + memo_of(&generate_per_computation_candidates(roots).unwrap()) + } + + /// The sites of `kind` the candidates of `memo` carry. + fn sites(memo: &Memo, kind: CouplingKind) -> Vec { + let sites: BTreeSet = memo + .candidates + .iter() + .flatten() + .flat_map(|c| c.coupling_sites.iter().copied()) + .filter(|s| s.kind == kind) + .collect(); + sites.into_iter().collect() + } + + /// The candidates of `memo` that carry `site`. + fn carrying<'a>(memo: &'a Memo, site: &CouplingSite) -> Vec<&'a ImplementationCandidate> { + memo.candidates + .iter() + .flatten() + .filter(|c| c.coupling_sites.contains(site)) + .collect() + } + + fn labels(names: &[&str]) -> BTreeSet { + names.iter().map(|n| n.to_string()).collect() + } + + /// The column names of the metric scans `node` reaches. + fn scan_labels(node: &Rc) -> Vec> { + OperatorNode::reachable(node) + .iter() + .filter_map(|n| match n.non_asap() { + Some(NonASAPOp::Scan { schema, .. }) => { + Some(schema.fields.iter().map(|f| f.name.clone()).collect()) + } + _ => None, + }) + .collect() + } + + /// Per-job and per-instance rates of one metric: its two scans project + /// different labels. One scan-labels site, read by each rate's twin; + /// its thresholds are each reader's labels and their union, which + /// admits both. Widened to the union, each rate reads a scan of every + /// label, through a `Project` back to its own, and the scan and window + /// groups it read become its own nodes. + #[test] + fn scans_of_one_metric_with_different_labels_share_a_widened_scan() { + let mut memo = memo(&[ + "sum by (job) (rate(m[5m]))", + "sum by (instance) (rate(m[5m]))", + ]); + add_scan_widening(&mut memo); + let sites = sites(&memo, CouplingKind::ScanLabels); + assert_eq!(sites.len(), 1); + let twins = carrying(&memo, &sites[0]); + assert_eq!(twins.len(), 2); + assert!(twins + .iter() + .all(|c| c.origin == CandidateOrigin::WidenedScans)); + let thresholds = label_thresholds(&memo, &sites[0]); + let union = labels(&["ts", "value", "job", "instance"]); + assert_eq!( + thresholds + .iter() + .map(|t| t.labels.clone()) + .collect::>(), + vec![ + labels(&["ts", "value", "instance"]), + labels(&["ts", "value", "job"]), + union.clone(), + ] + ); + let widest = &thresholds[2]; + for twin in &twins { + let readers = scan_readers(&memo, twin); + assert!(readers.iter().all(|r| may_read_labels(widest, r))); + // Its own labels admit only itself. + let own = LabelThreshold { + labels: readers[0].labels.clone(), + }; + let other = twins + .iter() + .find(|c| !Rc::ptr_eq(&c.node, &twin.node)) + .unwrap(); + assert!(!scan_readers(&memo, other) + .iter() + .all(|r| may_read_labels(&own, r))); + let widened = widened_at_threshold(&memo, twin, &sites[0], widest) + .unwrap() + .expect("a rate over a range widens"); + assert_eq!(scan_labels(&widened.node), vec![union.clone()]); + assert_eq!(widened.node.schema.fields, twin.node.schema.fields); + assert!(widened.children.is_empty()); + } + } + + /// Scans of two metrics with different labels: no site. + #[test] + fn scans_of_different_metrics_are_not_widened() { + let mut memo = memo(&[ + "sum by (job) (rate(m[5m]))", + "sum by (instance) (rate(x[5m]))", + ]); + add_scan_widening(&mut memo); + assert!(sites(&memo, CouplingKind::ScanLabels).is_empty()); + } + + /// One-minute and five-minute windows of one metric read raw data: one + /// raw-horizon site with both windows as thresholds; the one-minute + /// reader may read either, the five-minute one only the longer. + #[test] + fn raw_windows_of_one_source_give_one_horizon_site() { + let mut memo = memo(&["sum_over_time(m[1m])", "sum_over_time(m[5m])"]); + add_raw_horizons(&mut memo); + let sites = sites(&memo, CouplingKind::RawHorizon); + assert_eq!(sites.len(), 1); + let thresholds = horizon_thresholds(&memo, &sites[0]); + assert_eq!( + thresholds.iter().map(|t| t.horizon_ms).collect::>(), + vec![60_000, 300_000] + ); + let readers: Vec = carrying(&memo, &sites[0]) + .into_iter() + .flat_map(|c| raw_readers(&memo, c)) + .collect(); + let short = readers.iter().find(|r| r.horizon_ms == 60_000).unwrap(); + let long = readers.iter().find(|r| r.horizon_ms == 300_000).unwrap(); + assert!(may_read_horizon(&thresholds[0], short)); + assert!(may_read_horizon(&thresholds[1], short)); + assert!(!may_read_horizon(&thresholds[0], long)); + assert!(may_read_horizon(&thresholds[1], long)); + } + + /// A source read by one group only couples nothing. + #[test] + fn a_source_read_by_one_group_has_no_horizon_site() { + let mut memo = memo(&["sum_over_time(m[1m])", "sum_over_time(x[5m])"]); + add_raw_horizons(&mut memo); + assert!(sites(&memo, CouplingKind::RawHorizon).is_empty()); + } +} From cfdb07f1fa5b05ac43eb5a77160cf7b3beaf2cdb Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:10:26 +0000 Subject: [PATCH 04/25] plan-selection: sub-expression memo: exact-input and rollup requirements, sketch rollup over a finer group's state, groups kept in topological order --- crates/plan-selection/src/pass/mod.rs | 2 +- .../src/pass/subexpr/capability.rs | 2 + .../plan-selection/src/pass/subexpr/groups.rs | 6 + crates/plan-selection/src/pass/subexpr/mod.rs | 25 + .../src/pass/subexpr/pricing.rs | 2 + .../src/pass/subexpr/windows.rs | 809 ++++++++++++++++-- 6 files changed, 762 insertions(+), 84 deletions(-) diff --git a/crates/plan-selection/src/pass/mod.rs b/crates/plan-selection/src/pass/mod.rs index 16053a432..060b4bfde 100644 --- a/crates/plan-selection/src/pass/mod.rs +++ b/crates/plan-selection/src/pass/mod.rs @@ -14,9 +14,9 @@ mod components; mod interchangeable; -mod subexpr; mod query_candidates; mod stage1_to_stage3; +mod subexpr; use std::collections::BTreeMap; use std::rc::Rc; diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index f67dd3511..9418c9b13 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -314,6 +314,8 @@ mod tests { children: vec![GroupId(1)], origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, }]); } add_shared_builds(&mut memo); diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index 4928e8675..bad4c8009 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -182,6 +182,8 @@ fn pass1_candidates( children: groups[g].children.clone(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, }; (0..groups.len()) .map(|g| match targets_of.get(&GroupId(g as u32)) { @@ -251,6 +253,8 @@ fn target_candidates( children: groups[g.0 as usize].children.clone(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, } } else { let mut children = Vec::new(); @@ -267,6 +271,8 @@ fn target_candidates( children, origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, } }; candidates.push(candidate); diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 8eb747abb..b997f699a 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -75,6 +75,31 @@ pub(crate) struct ImplementationCandidate { pub origin: CandidateOrigin, /// The coupling sites its cost depends on. pub coupling_sites: Vec, + /// Groups whose chosen implementation must be exact for this one to + /// be valid (a rollup's finer group, a top-k reuse's input group); a + /// child or a descendant. + pub requires_exact_children: Vec, + /// For a rollup, what it rolls up and from where; `None` otherwise. + pub rollup_of: Option, +} + +/// What a [`CandidateOrigin::Rollup`] candidate reads, and what it +/// rewrites: the search checks that the finer group's chosen +/// implementation is the one the joint rule would roll up from. +#[derive(Debug, Clone)] +pub(crate) struct RollupRead { + /// The finer group it reads. + pub finer_group: GroupId, + /// For a sketch rollup, the finer build whose state it merges, as its + /// node reads it: a boundary node, priced with `finer_group`, which + /// must choose an implementation building that state and exposing it + /// at the same [`CouplingKind::SharedBuildSize`] site. `None` for an + /// exact rollup, which reads `finer_group`'s output. + pub finer_state: Option>, + /// The coarse group's candidate it rewrites. + pub coarse_candidate: Rc, + /// That candidate's coarse build, which it no longer builds. + pub coarse_build: Rc, } /// Where an implementation candidate came from. diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index 6c0e3f63c..1d29ddd07 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -559,6 +559,8 @@ mod tests { children, origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, }; let memo = Memo { groups: vec![group(&below, vec![]), group(top, vec![GroupId(0)])], diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index 194071e4c..edc4e75c3 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -12,13 +12,16 @@ //! size are the site's threshold ([`chain_thresholds`], [`may_read`], //! [`read_at_threshold`]). //! - **Rollup** ([`add_rollup`]): a coarse exact accumulator answered by -//! regrouping a finer group's values (`rollup_candidates`' conditions). +//! regrouping a finer group's values, or a coarse sketch by merging a +//! finer group's build states (`rollup_candidates`' conditions). //! - **Exact top-k reuse** ([`add_topk_reuse`]): a top-k answered from a //! ranking group, `Limit(k_max)` over a descending `Sort` of the input, //! added to the memo when missing. //! -//! The joint search never reuses a top-k ranking in a plan with panes -//! ([`combinable`]). +//! What the joint rules require across groups, the search checks on a +//! choice with [`requirements_met`]: no top-k reuse with panes, exact +//! inputs where a rule reads exact values, and each rollup reading the +//! build the joint rule would choose. use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeSet, HashMap, HashSet}; @@ -27,11 +30,12 @@ use std::rc::Rc; use asap_logical_asap_optimization::accuracy::estimators::size_for_all_readers; use asap_logical_asap_optimization::asap_aware_cse::{ - pane_candidates, rollup_candidates, topk_reuse_candidates, window_composition_pass2_candidates, - with_pane_chain, + pane_candidates, rollup_candidates, rollup_pass2_candidates, topk_reuse_candidates, + window_composition_pass2_candidates, with_pane_chain, RollupCandidate, }; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; use asap_types::ir::flat::flatten; +use asap_types::ir::properties::AccuracyGuarantee; use asap_types::ir::schema::sketch::SketchKind; use asap_types::ir::schema::FieldDataType; use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot}; @@ -39,6 +43,7 @@ use asap_types::workload::{greatest_common_divisor, QueryWorkloadEntry}; use super::{ CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo, MemoGroup, + RollupRead, }; /// One threshold of a shared pane chain: its grid, the longest window it @@ -367,13 +372,24 @@ pub(crate) fn read_at_threshold( Ok(Rc::clone(node)) } -/// Rollup: each coarse exact accumulator of a group's Pass 1 (or -/// shared-build) candidate that another group's candidate can answer by -/// regrouping its finalized values (`rollup_candidates`), as a -/// [`CandidateOrigin::Rollup`] candidate that reads that group: the finer -/// value is the finer group's output, so the finer candidate must be that -/// group's whole implementation. Sketch rollups, which merge a finer -/// build's states rather than a group's output, are not offered. +/// Rollup: each coarse build of a group's Pass 1 (or shared-build) +/// candidate that another group's candidate can answer +/// (`rollup_candidates` over the two) gets a [`CandidateOrigin::Rollup`] +/// candidate of the coarse group: +/// +/// - an exact accumulator regroups the finer group's finalized values: the +/// finer value is the finer group's output, which must be exact +/// ([`ImplementationCandidate::requires_exact_children`]); +/// - a sketch merges the finer build's states by the coarse keys: it reads +/// that state ([`ImplementationCandidate::rollup_of`]), and the +/// finer build is sized for both groups' readers, a +/// [`CouplingKind::SharedBuildSize`] site ([`state_size_thresholds`]) +/// that the finer group's candidate also gets a copy with. The joint +/// rule rolls a coarse build up only from the finest build over its rows, +/// and every coarse build of one finest build together +/// ([`requirements_met`] checks both). +/// +/// Ends with the groups in topological order ([`in_topological_order`]). pub(crate) fn add_rollup(memo: &mut Memo) { let group_of = group_of_node(&memo.groups); let mut added: Vec<(usize, ImplementationCandidate)> = Vec::new(); @@ -385,10 +401,7 @@ pub(crate) fn add_rollup(memo: &mut Memo) { }; for (coarse_group, coarse_candidates) in memo.candidates.iter().enumerate() { for coarse in coarse_candidates.iter().filter(|c| base(c)) { - let reached: HashSet<*const OperatorNode> = OperatorNode::reachable(&coarse.node) - .iter() - .map(Rc::as_ptr) - .collect(); + let in_coarse = reached(&coarse.node); for (finer_group, finer_candidates) in memo.candidates.iter().enumerate() { if finer_group == coarse_group { continue; @@ -398,45 +411,36 @@ pub(crate) fn add_rollup(memo: &mut Memo) { QueryRoot::Operator(Rc::clone(&coarse.node)), QueryRoot::Operator(Rc::clone(&finer.node)), ]; + let in_finer = reached(&finer.node); for rollup in rollup_candidates(&roots) { - let exact = matches!( - rollup.finer.asap(), - Some(ASAPOp::FinalizeExactAccumulator { .. }) - ); - if !exact - || !Rc::ptr_eq(&rollup.finer, &finer.node) - || !reached.contains(&Rc::as_ptr(&rollup.coarse)) + if !in_coarse.contains(&Rc::as_ptr(&rollup.coarse)) + || in_coarse.contains(&Rc::as_ptr(&rollup.finer)) + || !in_finer.contains(&Rc::as_ptr(&rollup.finer)) { continue; } - // As `rollup_pass2_candidates` rewrites it, with the finer - // value read from the finer group: every finalized value - // of the coarse build replaced by the rollup. - let finer_output = &memo.groups[finer_group].node; - let to_group = HashMap::from([(Rc::as_ptr(&finer.node), finer_output)]); - let Ok(rolled) = replace(&rollup.rollup, &to_group, &mut HashMap::new()) - else { - continue; + let made = if is_sketch_build(&rollup.finer) { + sketch_rollup(&group_of, coarse, finer_group, finer, &roots, rollup) + } else { + exact_rollup( + &memo.groups, + &group_of, + coarse, + finer_group, + finer, + &rollup, + ) }; - let finalized: Vec> = OperatorNode::reachable( - &coarse.node, - ) - .into_iter() - .filter(|n| { - matches!(n.asap(), Some(ASAPOp::FinalizeExactAccumulator { child }) - if Rc::ptr_eq(child, &rollup.coarse)) - }) - .collect(); - let by_rollup: HashMap<*const OperatorNode, &Rc> = - finalized.iter().map(|n| (Rc::as_ptr(n), &rolled)).collect(); - let Ok(node) = replace(&coarse.node, &by_rollup, &mut HashMap::new()) - else { - continue; - }; - added.push(( - coarse_group, - rewritten_candidate(&group_of, coarse, node, CandidateOrigin::Rollup), - )); + added.extend(made.into_iter().map(|(g, c)| { + ( + if g == FINER { + finer_group + } else { + coarse_group + }, + c, + ) + })); } } } @@ -445,33 +449,432 @@ pub(crate) fn add_rollup(memo: &mut Memo) { for (g, candidate) in added { push_new(memo, g, candidate); } + in_topological_order(memo); +} + +/// Marks a candidate [`sketch_rollup`] adds to the finer group. +const FINER: usize = usize::MAX; + +/// The exact rollup of `rollup` (over `coarse`'s and `finer`'s nodes): the +/// coarse candidate with every finalized value of the coarse build +/// replaced by the regrouping of the finer group's output, as +/// `rollup_pass2_candidates` rewrites it. `finer` must be the finer +/// group's whole implementation (its node the finalized finer value). +fn exact_rollup( + groups: &[MemoGroup], + group_of: &HashMap<*const OperatorNode, usize>, + coarse: &ImplementationCandidate, + finer_group: usize, + finer: &ImplementationCandidate, + rollup: &RollupCandidate, +) -> Vec<(usize, ImplementationCandidate)> { + if !Rc::ptr_eq(&rollup.finer, &finer.node) { + return Vec::new(); + } + let finer_output = &groups[finer_group].node; + let to_group = HashMap::from([(Rc::as_ptr(&finer.node), finer_output)]); + let Ok(rolled) = replace(&rollup.rollup, &to_group, &mut HashMap::new()) else { + return Vec::new(); + }; + let by_rollup: HashMap<*const OperatorNode, &Rc> = + OperatorNode::reachable(&coarse.node) + .into_iter() + .filter(|n| { + matches!(n.asap(), Some(ASAPOp::FinalizeExactAccumulator { child }) + if Rc::ptr_eq(child, &rollup.coarse)) + }) + .map(|n| (Rc::as_ptr(&n), &rolled)) + .collect(); + let Ok(node) = replace(&coarse.node, &by_rollup, &mut HashMap::new()) else { + return Vec::new(); + }; + let mut candidate = rewritten_candidate(group_of, coarse, node, CandidateOrigin::Rollup); + let finer_id = GroupId(finer_group as u32); + if !candidate.children.contains(&finer_id) { + candidate.children.push(finer_id); + } + candidate.requires_exact_children = vec![finer_id]; + candidate.rollup_of = Some(RollupRead { + finer_group: finer_id, + finer_state: None, + coarse_candidate: Rc::clone(&coarse.node), + coarse_build: Rc::clone(&rollup.coarse), + }); + vec![(0, candidate)] +} + +/// The sketch rollup of `rollup` (over `roots`: `coarse`'s and `finer`'s +/// nodes): the coarse candidate as `rollup_pass2_candidates` rewrites it +/// (the coarse build replaced by the regrouping merge of the finer build, +/// sized for both), reading the finer build's state, and the finer +/// candidate's copy that exposes that state; both with the build's +/// [`CouplingKind::SharedBuildSize`] site. +fn sketch_rollup( + group_of: &HashMap<*const OperatorNode, usize>, + coarse: &ImplementationCandidate, + finer_group: usize, + finer: &ImplementationCandidate, + roots: &[QueryRoot], + rollup: RollupCandidate, +) -> Vec<(usize, ImplementationCandidate)> { + let site = CouplingSite { + key: state_key(&rollup.finer), + kind: CouplingKind::SharedBuildSize, + }; + let coarse_build = Rc::clone(&rollup.coarse); + let Ok(mut pass2) = rollup_pass2_candidates(roots, &[rollup]) else { + return Vec::new(); + }; + let Some(Ok(rolled)) = pass2.pop().map(|p| p.roots) else { + return Vec::new(); + }; + let QueryRoot::Operator(node) = &rolled[0] else { + return Vec::new(); + }; + // The resized finer build the regrouping merge reads. + let Some(finer_build) = + OperatorNode::reachable(node) + .into_iter() + .find_map(|n| match n.asap() { + Some(ASAPOp::SummaryMerge { + children, + group_by: Some(_), + .. + }) if is_sketch_build(&children[0]) => Some(Rc::clone(&children[0])), + _ => None, + }) + else { + return Vec::new(); + }; + let mut candidate = + rewritten_candidate(group_of, coarse, Rc::clone(node), CandidateOrigin::Rollup); + // Its private nodes stop at the finer build, which the finer group + // builds: the walk below it found the finer group's children. + candidate.children = coarse.children.clone(); + let finer_id = GroupId(finer_group as u32); + candidate.children.push(finer_id); + candidate.coupling_sites.push(site); + candidate.rollup_of = Some(RollupRead { + finer_group: finer_id, + finer_state: Some(finer_build), + coarse_candidate: Rc::clone(&coarse.node), + coarse_build, + }); + let mut exposed = finer.clone(); + exposed.coupling_sites.push(site); + vec![(0, candidate), (FINER, exposed)] +} + +/// Whether `node` is a sketch `SummaryAgg`. +fn is_sketch_build(node: &OperatorNode) -> bool { + matches!( + node.asap(), + Some(ASAPOp::SummaryAgg { + family: FieldDataType::Sketch(..), + .. + }) + ) +} + +/// The key of a sketch build's state, its size masked: the build with its +/// family's algorithm and grouping only. +fn state_key(build: &OperatorNode) -> u64 { + let mut value = serde_json::to_value(&build.operator).unwrap_or_default(); + if let (Some(ASAPOp::SummaryAgg { family, .. }), Some(agg)) = ( + build.asap(), + value + .pointer_mut("/ASAP/SummaryAgg") + .and_then(|v| v.as_object_mut()), + ) { + agg.insert( + "family".to_string(), + serde_json::Value::String(family_without_size(family)), + ); + } + hash_of(&format!("rollup state {value}")) +} + +/// The sizes the finer build of a sketch rollup's `site` may take: the +/// shared size ([`size_for_all_readers`]) of every set of its readers' +/// sizes (the finer group's own build and each rolled-up coarse build), +/// without one larger than the shared size of the readers it admits. +pub(crate) fn state_size_thresholds(memo: &Memo, site: &CouplingSite) -> Vec { + let readers: Vec = memo + .candidates + .iter() + .flatten() + .filter(|c| c.coupling_sites.contains(site)) + .filter_map(|c| state_reader_size(c, site)) + .collect(); + let mut sizes: Vec = Vec::new(); + for reader in &readers { + let mut grown: Vec = sizes + .iter() + .filter_map(|size| shared_size(size, reader)) + .collect(); + grown.push(reader.clone()); + for size in grown { + if !sizes.iter().any(|known| same_size(known, &size)) { + sizes.push(size); + } + } + } + sizes + .into_iter() + .filter(|threshold| { + let joined = readers + .iter() + .filter(|r| no_stricter(r, threshold)) + .try_fold(None::, |acc, r| { + Some(Some(match acc { + None => r.clone(), + Some(acc) => shared_size(&acc, r)?, + })) + }); + matches!(joined, Some(Some(joined)) if same_size(&joined, threshold)) + }) + .collect() +} + +/// Whether `candidate`, a reader of the sketch rollup state `site`, may +/// read it at size `threshold`: its size is no stricter. +pub(crate) fn may_read_state( + threshold: &FieldDataType, + candidate: &ImplementationCandidate, + site: &CouplingSite, +) -> bool { + state_reader_size(candidate, site).is_some_and(|size| no_stricter(&size, threshold)) +} + +/// The size a reader of the rollup state `site` needs: a rollup's coarse +/// build's, or the finer group's own build's. +fn state_reader_size( + candidate: &ImplementationCandidate, + site: &CouplingSite, +) -> Option { + let build = match &candidate.rollup_of { + Some(read) => Rc::clone(&read.coarse_build), + None => OperatorNode::reachable(&candidate.node) + .into_iter() + .find(|n| is_sketch_build(n) && state_key(n) == site.key)?, + }; + match build.asap() { + Some(ASAPOp::SummaryAgg { family, .. }) => Some(family.clone()), + _ => None, + } +} + +/// Whether a build of size `size` serves a reader needing `required`. +fn no_stricter(required: &FieldDataType, size: &FieldDataType) -> bool { + shared_size(size, required).is_some_and(|joined| same_size(&joined, size)) +} + +/// Equal families and sizes, the recorded failure probabilities and errors +/// included (which `SketchParams` equality ignores). +fn same_size(a: &FieldDataType, b: &FieldDataType) -> bool { + a == b && format!("{a:?}") == format!("{b:?}") +} + +/// Whether the candidates `chosen` (one per group, with every group they +/// name and every group with a summary build over the same rows as a +/// rolled-up build, in workload order) meet each other's requirements: +/// +/// - never an exact top-k reuse with panes ([`combinable`]); +/// - every group in a candidate's `requires_exact_children` chose an exact +/// implementation ([`is_exact`]); +/// - every rollup reads the build the joint rule would roll its coarse +/// build up from, over the chosen plan with each rollup undone: an exact +/// accumulator's coarsest finer build, a sketch's finest, in its finer +/// group's chosen implementation; +/// - a sketch rollup's finer group chose a candidate exposing the state at +/// the same site, and every chosen coarse build whose finest build that is +/// rolls up too (the joint rule's choice of one finest build). +pub(crate) fn requirements_met(chosen: &[(GroupId, &ImplementationCandidate)]) -> bool { + if !combinable(chosen.iter().map(|(_, c)| *c)) { + return false; + } + let of = |g: GroupId| chosen.iter().find(|(h, _)| *h == g).map(|(_, c)| *c); + for (_, candidate) in chosen { + for &g in &candidate.requires_exact_children { + if !of(g).is_some_and(is_exact) { + return false; + } + } + } + let rolled: Vec<&RollupRead> = chosen + .iter() + .filter_map(|(_, c)| c.rollup_of.as_ref()) + .collect(); + if rolled.is_empty() { + return true; + } + // The joint rule's rollups over the chosen plan, each rollup undone. + let roots: Vec = chosen + .iter() + .map(|(_, c)| match &c.rollup_of { + Some(read) => QueryRoot::Operator(Rc::clone(&read.coarse_candidate)), + None => QueryRoot::Operator(Rc::clone(&c.node)), + }) + .collect(); + let source: Vec<(*const OperatorNode, Rc)> = rollup_candidates(&roots) + .into_iter() + .map(|r| (Rc::as_ptr(&r.coarse), r.finer)) + .collect(); + let mut used_state: HashSet<*const OperatorNode> = HashSet::new(); + let mut rolled_coarse: HashSet<*const OperatorNode> = HashSet::new(); + for read in &rolled { + let Some(finer) = of(read.finer_group) else { + return false; + }; + let Some((_, build)) = source + .iter() + .find(|(coarse, _)| *coarse == Rc::as_ptr(&read.coarse_build)) + else { + return false; + }; + if !reached(&finer.node).contains(&Rc::as_ptr(build)) { + return false; + } + rolled_coarse.insert(Rc::as_ptr(&read.coarse_build)); + let Some(state) = &read.finer_state else { + continue; + }; + let site = CouplingSite { + key: state_key(state), + kind: CouplingKind::SharedBuildSize, + }; + if !finer.coupling_sites.contains(&site) || state_key(build) != site.key { + return false; + } + used_state.insert(Rc::as_ptr(build)); + } + source + .iter() + .filter(|(_, build)| used_state.contains(&Rc::as_ptr(build))) + .all(|(coarse, _)| rolled_coarse.contains(coarse)) +} + +/// Whether `candidate` computes exactly: no node it reaches builds an +/// approximate summary (every summary an exact accumulator). +pub(crate) fn is_exact(candidate: &ImplementationCandidate) -> bool { + OperatorNode::reachable(&candidate.node).iter().all(|n| { + !matches!( + n.asap(), + Some( + ASAPOp::SummaryAgg { family, .. } | ASAPOp::PaneChain { family, .. } + ) if !matches!(family, FieldDataType::ExactAggregate(..)) + ) + }) +} + +/// The addresses of every node `node` reaches. +fn reached(node: &Rc) -> HashSet<*const OperatorNode> { + OperatorNode::reachable(node) + .iter() + .map(Rc::as_ptr) + .collect() +} + +/// Renumber `memo`'s groups so that every group comes after the groups it +/// reads, through its node or any candidate (children and required +/// groups), keeping the current order where it allows. Unchanged when +/// those reads form a cycle. +pub(crate) fn in_topological_order(memo: &mut Memo) { + let n = memo.groups.len(); + let mut reads: Vec> = vec![BTreeSet::new(); n]; + for (g, group) in memo.groups.iter().enumerate() { + reads[g].extend(group.children.iter().map(|c| c.0 as usize)); + for candidate in &memo.candidates[g] { + reads[g].extend(candidate.children.iter().map(|c| c.0 as usize)); + reads[g].extend( + candidate + .requires_exact_children + .iter() + .map(|c| c.0 as usize), + ); + reads[g].extend(candidate.rollup_of.iter().map(|r| r.finer_group.0 as usize)); + } + reads[g].remove(&g); + } + // Repeatedly take the first group whose reads are all placed. + let mut placed = vec![false; n]; + let mut order: Vec = Vec::with_capacity(n); + while order.len() < n { + let Some(next) = (0..n).find(|&g| !placed[g] && reads[g].iter().all(|&r| placed[r])) else { + return; + }; + placed[next] = true; + order.push(next); + } + if order.iter().enumerate().all(|(new, &old)| new == old) { + return; + } + let mut new_of = vec![0u32; n]; + for (new, &old) in order.iter().enumerate() { + new_of[old] = new as u32; + } + let map = |id: &mut GroupId| id.0 = new_of[id.0 as usize]; + let mut groups: Vec> = std::mem::take(&mut memo.groups) + .into_iter() + .map(Some) + .collect(); + let mut candidates: Vec>> = + std::mem::take(&mut memo.candidates) + .into_iter() + .map(Some) + .collect(); + for &old in &order { + let mut group = groups[old].take().expect("each group once"); + group.children.iter_mut().for_each(map); + group.parents.iter_mut().for_each(map); + memo.groups.push(group); + let mut own = candidates[old].take().expect("each group once"); + for candidate in &mut own { + candidate.implements.iter_mut().for_each(map); + candidate.children.iter_mut().for_each(map); + candidate.requires_exact_children.iter_mut().for_each(map); + if let Some(read) = &mut candidate.rollup_of { + map(&mut read.finer_group); + } + } + memo.candidates.push(own); + } } -/// Exact top-k reuse: the groups whose Pass 1 candidates are exact top-k -/// computations of one input and partition with at least two different -/// `k` (`topk_reuse_candidates` over the candidates' nodes, which carry -/// their guarantees) each get a [`CandidateOrigin::TopKReuse`] candidate +/// Exact top-k reuse: the groups that are top-k computations of one input +/// group and partition with at least two different `k` +/// (`topk_reuse_candidates` over the groups' nodes) each get a [`CandidateOrigin::TopKReuse`] candidate /// reading the ranking group, `Limit(k_max)` over a descending `Sort` of /// the input group (the top-k group's child). The ranking's two groups /// are added when the memo has no equal ones, each with its one candidate /// (also of origin `TopKReuse`). The rule reads an exact input: the input -/// group's chosen implementation must be exact. +/// group's chosen implementation must be exact +/// ([`ImplementationCandidate::requires_exact_children`]). Ends with the +/// groups in topological order ([`in_topological_order`]). pub(crate) fn add_topk_reuse(memo: &mut Memo) { - let mut group_of_candidate: HashMap<*const OperatorNode, usize> = HashMap::new(); + // The groups' nodes, each input marked exact: the rule reads an exact + // input, which `requires_exact_children` asks of the input group. + let mut group_of_copy: HashMap<*const OperatorNode, usize> = HashMap::new(); let mut roots: Vec = Vec::new(); - for (g, candidates) in memo.candidates.iter().enumerate() { - for candidate in candidates { - if candidate.origin == CandidateOrigin::Pass1 { - group_of_candidate.insert(Rc::as_ptr(&candidate.node), g); - roots.push(QueryRoot::Operator(Rc::clone(&candidate.node))); - } - } + for (g, group) in memo.groups.iter().enumerate() { + let [input] = &group.node.children()[..] else { + continue; + }; + let exact_input = Rc::new(OperatorNode::clone(input).with_guarantee(Some( + AccuracyGuarantee::exact("the input group chosen exact"), + ))); + let Ok(copy) = group.node.with_new_children(|_| Rc::clone(&exact_input)) else { + continue; + }; + let copy = Rc::new(copy); + group_of_copy.insert(Rc::as_ptr(©), g); + roots.push(QueryRoot::Operator(copy)); } for candidate in topk_reuse_candidates(&roots) { let mut members: Vec = candidate .members .iter() - .filter_map(|m| group_of_candidate.get(&Rc::as_ptr(m)).copied()) + .filter_map(|m| group_of_copy.get(&Rc::as_ptr(m)).copied()) .collect(); members.sort_unstable(); members.dedup(); @@ -485,17 +888,16 @@ pub(crate) fn add_topk_reuse(memo: &mut Memo) { .into_iter() .collect(); // The ranking over the input group's node. - let Some(&input) = memo.groups[members[0]].node.children().first() else { + let Some(&input_group) = memo.groups[members[0]].children.first() else { continue; }; - let input = Rc::clone(input); + let input = Rc::clone(&memo.groups[input_group.0 as usize].node); let sort = candidate.ranking.children()[0]; let Ok(sort) = sort.with_new_children(|_| Rc::clone(&input)) else { continue; }; let sort = Rc::new(sort); - let input_group = find_group(&memo.groups, &input); - let sort_group = group_for(memo, &sort, input_group.into_iter().collect(), &queries); + let sort_group = group_for(memo, &sort, vec![input_group], &queries); let sort = Rc::clone(&memo.groups[sort_group.0 as usize].node); let ranking = match candidate.ranking.with_new_children(|_| Rc::clone(&sort)) { Ok(ranking) => Rc::new(ranking), @@ -516,10 +918,13 @@ pub(crate) fn add_topk_reuse(memo: &mut Memo) { children: vec![ranking_group], origin: CandidateOrigin::TopKReuse, coupling_sites: Vec::new(), + requires_exact_children: vec![input_group], + rollup_of: None, }; push_new(memo, g, reuse); } } + in_topological_order(memo); } /// The group whose node equals `node`, or a new one over `children`, read @@ -550,6 +955,8 @@ fn group_for( children, origin: CandidateOrigin::TopKReuse, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, }]); id } @@ -594,6 +1001,8 @@ fn rewritten_candidate( children, origin, coupling_sites: from.coupling_sites.clone(), + requires_exact_children: Vec::new(), + rollup_of: None, } } @@ -742,11 +1151,23 @@ mod tests { /// `intervals_ms[q]`), its node the lowered query, its candidates every /// Pass 1 candidate of the query alone, each wholly private. fn memo(queries: &[&str], intervals_ms: &[u32]) -> (Memo, Vec) { + let exact: Vec<(&str, AccuracyTarget)> = queries + .iter() + .map(|&q| (q, AccuracyTarget::exact())) + .collect(); + memo_with(&exact, intervals_ms) + } + + /// [`memo`] with each query's own accuracy target. + fn memo_with( + queries: &[(&str, AccuracyTarget)], + intervals_ms: &[u32], + ) -> (Memo, Vec) { let entries: Vec = intervals_ms.iter().map(|&ms| repeating(ms)).collect(); let mut memo = Memo::default(); - for (q, query) in queries.iter().enumerate() { - let node = lower_promql(query, AccuracyTarget::exact()); + for (q, (query, target)) in queries.iter().enumerate() { + let node = lower_promql(query, target.clone()); let node = with_promql_series_identity(&node).unwrap(); let inventory = generate_per_computation_candidates(vec![( q, @@ -766,6 +1187,8 @@ mod tests { children: Vec::new(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + requires_exact_children: Vec::new(), + rollup_of: None, }) .collect(); memo.groups.push(MemoGroup { @@ -898,18 +1321,25 @@ mod tests { .all(|c| c.origin != CandidateOrigin::Rollup)); } - /// Exact top-5 and top-10 of one per-job sum: the ranking (`Limit(10)` - /// over a `Sort`) is added as a group, and each top-k gets a candidate - /// reading it; it cannot be chosen with a pane candidate. + /// Exact top-5 and top-10 of one per-job sum (the memo of the + /// workload): the ranking (`Limit(10)` over a `Sort` of the sum group) + /// is added as a group before both top-k groups, and each top-k gets a + /// candidate reading it that requires the sum group exact; it cannot be + /// chosen with a pane candidate. #[test] fn a_topk_reuse_candidate_reads_the_ranking() { - let (mut memo, entries) = memo( - &[ - "topk(5, sum by (job) (sum_over_time(m[5m])))", - "topk(10, sum by (job) (sum_over_time(m[5m])))", - ], - &[60_000, 60_000], - ); + let queries = [ + "topk(5, sum by (job) (sum_over_time(m[5m])))", + "topk(10, sum by (job) (sum_over_time(m[5m])))", + ]; + let roots = queries + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, AccuracyTarget::exact()))) + .enumerate() + .collect(); + let mut memo = + super::super::groups::memo_of(&generate_per_computation_candidates(roots).unwrap()); + let entries = vec![repeating(60_000); 2]; add_pass2_candidates(&mut memo, &entries); let ranking = memo .groups @@ -921,22 +1351,235 @@ mod tests { ) }) .expect("a ranking group"); - for g in 0..2 { + let topk: Vec = (0..memo.groups.len()) + .filter(|&g| { + memo.candidates[g] + .iter() + .any(|c| c.origin == CandidateOrigin::TopKReuse && g != ranking) + && !memo.groups[g].roots_of.is_empty() + }) + .collect(); + assert_eq!(topk.len(), 2); + for g in topk { + assert!(ranking < g, "children before parents"); let reuse = memo.candidates[g] .iter() .find(|c| c.origin == CandidateOrigin::TopKReuse) - .expect("a top-k reuse candidate"); + .unwrap(); assert_eq!(reuse.children, vec![GroupId(ranking as u32)]); assert!(Rc::ptr_eq( reuse.node.children()[0], &memo.groups[ranking].node )); + let input = memo.groups[g].children[0]; + assert_eq!(reuse.requires_exact_children, vec![input]); + assert!(input.0 < ranking as u32); let panes = ImplementationCandidate { origin: CandidateOrigin::Panes, ..reuse.clone() }; assert!(!combinable([reuse, &panes])); - assert!(combinable([reuse])); + let exact_input = memo.candidates[input.0 as usize] + .iter() + .find(|c| is_exact(c)) + .unwrap(); + assert!(requirements_met(&[ + (GroupId(g as u32), reuse), + (input, exact_input) + ])); + assert!(!requirements_met(&[(GroupId(g as u32), reuse)])); } } + + /// An exact rollup requires the finer group exact and implemented by + /// the build it rolls up from: met with the finer group's grouped + /// exact accumulator, not with its per-series one (no rollup source), + /// nor with a sketch implementation chosen. + #[test] + fn an_exact_rollup_requires_the_finer_group_exact() { + let (mut memo, _) = memo( + &[ + "sum by (job, instance) (sum_over_time(m[5m]))", + "sum by (job) (sum_over_time(m[5m]))", + ], + &[60_000, 60_000], + ); + add_rollup(&mut memo); + let rollup = memo.candidates[1] + .iter() + .find(|c| c.origin == CandidateOrigin::Rollup) + .unwrap(); + assert_eq!(rollup.requires_exact_children, vec![GroupId(0)]); + let met: Vec = memo.candidates[0] + .iter() + .map(|finer| requirements_met(&[(GroupId(0), finer), (GroupId(1), rollup)])) + .collect(); + assert!(met.contains(&true) && met.contains(&false), "{met:?}"); + let source = &memo.candidates[0][met.iter().position(|&m| m).unwrap()]; + assert!(is_exact(source)); + // A KLL quantile stands in for an approximate implementation. + let (sketched, _) = memo_with( + &[( + "quantile_over_time(0.9, m[5m])", + AccuracyTarget::max_error(0.01), + )], + &[60_000], + ); + let approximate = sketched.candidates[0] + .iter() + .find(|c| !is_exact(c)) + .unwrap(); + assert!(!requirements_met(&[ + (GroupId(0), approximate), + (GroupId(1), rollup) + ])); + } + + /// KLL quantiles by `(host, job)` at ε 0.01 and by `job` at 0.001: the + /// per-job group gets a rollup reading the finer group's KLL state, the + /// finer group a copy exposing it, both at one size site whose largest + /// threshold is the coarse reader's (the larger `k`); the requirement + /// holds only with the exposing copy chosen. + #[test] + fn a_sketch_rollup_reads_the_finer_groups_state() { + let (mut memo, _) = memo_with( + &[ + ( + "quantile by (host, job) (0.5, data)", + AccuracyTarget::max_error(0.01), + ), + ( + "quantile by (job) (0.5, data)", + AccuracyTarget::max_error(0.001), + ), + ], + &[60_000, 60_000], + ); + add_rollup(&mut memo); + let rollup = memo.candidates[1] + .iter() + .find(|c| { + c.rollup_of + .as_ref() + .is_some_and(|r| r.finer_state.is_some()) + }) + .expect("a sketch rollup"); + let read = rollup.rollup_of.as_ref().unwrap(); + assert_eq!(read.finer_group, GroupId(0)); + assert!(rollup.children.contains(&GroupId(0))); + let site = *rollup + .coupling_sites + .iter() + .find(|s| s.kind == CouplingKind::SharedBuildSize) + .unwrap(); + let exposed = memo.candidates[0] + .iter() + .find(|c| c.coupling_sites.contains(&site)) + .expect("the finer group exposes the state"); + let plain = memo.candidates[0] + .iter() + .find(|c| c.coupling_sites.is_empty() && Rc::ptr_eq(&c.node, &exposed.node)) + .unwrap(); + assert!(requirements_met(&[ + (GroupId(0), exposed), + (GroupId(1), rollup) + ])); + assert!(!requirements_met(&[ + (GroupId(0), plain), + (GroupId(1), rollup) + ])); + let k = |family: &FieldDataType| match family { + FieldDataType::Sketch(kind, _) => format!("{:?}", kind.params()), + _ => String::new(), + }; + let thresholds = state_size_thresholds(&memo, &site); + let coarse_size = state_reader_size(rollup, &site).unwrap(); + assert!(thresholds.iter().any(|t| k(t) == k(&coarse_size))); + assert!(thresholds.iter().all(|t| may_read_state(t, exposed, &site) + || !no_stricter(&state_reader_size(exposed, &site).unwrap(), t))); + let largest = thresholds + .iter() + .find(|t| may_read_state(t, rollup, &site) && may_read_state(t, exposed, &site)) + .expect("one size serves both"); + assert_eq!(k(largest), k(&coarse_size)); + } + + /// Three KLL groupings `job` ⊂ `(job, instance)` ⊂ `(job, instance, + /// zone)`: the joint rule rolls both coarse builds up from the finest + /// together. Rolling one alone, or rolling from the middle build while + /// the finest is chosen, fails the requirement. + #[test] + fn sketch_rollups_follow_the_finest_build_all_or_nothing() { + let target = AccuracyTarget::max_error(0.01); + let (mut memo, _) = memo_with( + &[ + ("quantile by (job) (0.9, m)", target.clone()), + ("quantile by (job, instance) (0.9, m)", target.clone()), + ("quantile by (job, instance, zone) (0.9, m)", target), + ], + &[60_000, 60_000, 60_000], + ); + add_rollup(&mut memo); + // Rollups reorder the groups finer first: query `q`'s group. + let of = |q: usize| memo.groups.iter().position(|g| g.roots_of == [q]).unwrap(); + let (job, instance, zone) = (of(0), of(1), of(2)); + assert!(zone < instance && instance < job, "children before parents"); + let rolled_from = |g: usize, from: usize| { + memo.candidates[g] + .iter() + .find(|c| { + c.rollup_of.as_ref().is_some_and(|r| { + r.finer_state.is_some() && r.finer_group == GroupId(from as u32) + }) + }) + .unwrap_or_else(|| panic!("group {g} rolls up from {from}")) + }; + let exposing = |g: usize, rollup: &ImplementationCandidate| { + let site = rollup + .coupling_sites + .iter() + .find(|s| s.kind == CouplingKind::SharedBuildSize) + .unwrap(); + memo.candidates[g] + .iter() + .find(|c| c.rollup_of.is_none() && c.coupling_sites.contains(site)) + .unwrap() + }; + let id = |g: usize| GroupId(g as u32); + let job_from_zone = rolled_from(job, zone); + let instance_from_zone = rolled_from(instance, zone); + let zone_exposed = exposing(zone, job_from_zone); + assert!(requirements_met(&[ + (id(job), job_from_zone), + (id(instance), instance_from_zone), + (id(zone), zone_exposed), + ])); + // The `(job, instance)` build itself, sketched, not rolled up. + let instance_own = memo.candidates[instance] + .iter() + .find(|c| c.origin == CandidateOrigin::Pass1 && !is_exact(c)) + .unwrap(); + assert!(!requirements_met(&[ + (id(job), job_from_zone), + (id(instance), instance_own), + (id(zone), zone_exposed), + ])); + // From the middle build while the finest is chosen. + let job_from_instance = rolled_from(job, instance); + let instance_exposed = exposing(instance, job_from_instance); + let zone_own = memo.candidates[zone] + .iter() + .find(|c| c.origin == CandidateOrigin::Pass1 && !is_exact(c)) + .unwrap(); + assert!(!requirements_met(&[ + (id(job), job_from_instance), + (id(instance), instance_exposed), + (id(zone), zone_own), + ])); + // Without the finest build, the middle one is the finest. + assert!(requirements_met(&[ + (id(job), job_from_instance), + (id(instance), instance_exposed), + ])); + } } From 1152ad1449ca148cd4423f370a20fc581f7f7f06 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:16:29 +0000 Subject: [PATCH 05/25] Merge W3 (per-metric scan widening, raw horizon sites); memo stats hook (temporary) --- .../src/pass/stage1_to_stage3.rs | 10 ++++ .../plan-selection/src/pass/subexpr/scans.rs | 3 ++ .../plan-selection/src/pass/subexpr/search.rs | 48 ++++++++++++++++++- 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index 388358aa7..6ad05932f 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -228,6 +228,16 @@ pub fn run_stage1_to_stage3_from_pass1( models: PlanningModels<'_>, search: Stage1ToStage3Search, ) -> Stage1ToStage3Run { + if std::env::var("JP_MEMO_STATS").is_ok() { + let started = std::time::Instant::now(); + let memo = super::subexpr::search::memo(&per_computation, entries); + eprintln!( + "JPMEMO {} in {:?}", + super::subexpr::search::stats(&memo), + started.elapsed() + ); + std::process::exit(0); + } let components = super::components::sharing_components(&per_computation, entries, &models); let searched_as_one = match super::components::searched_as_one(&components, &models, search) { Some(reason) => reason, diff --git a/crates/plan-selection/src/pass/subexpr/scans.rs b/crates/plan-selection/src/pass/subexpr/scans.rs index 3684e950d..669f565a9 100644 --- a/crates/plan-selection/src/pass/subexpr/scans.rs +++ b/crates/plan-selection/src/pass/subexpr/scans.rs @@ -214,6 +214,9 @@ pub(crate) fn widened_at_threshold( children: child_groups(&mut Boundary::of(memo, candidate), node), origin: CandidateOrigin::WidenedScans, coupling_sites: candidate.coupling_sites.clone(), + budget_share: candidate.budget_share, + requires_exact_children: candidate.requires_exact_children.clone(), + rollup_of: candidate.rollup_of.clone(), })) } diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index e6f41fc1e..aefca92f6 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -1 +1,47 @@ -//! See the module doc of [`super`]. +//! Clusters, coupling thresholds with branch and bound, the dynamic +//! program, build-late (design "Memo-group search", steps 3–7). + +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; +use asap_types::workload::QueryWorkloadEntry; + +use super::{groups, Memo}; + +/// The workload's memo: groups, Pass 1 candidates and the Pass 2 +/// transformations' candidates. +pub(crate) fn memo( + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], +) -> Memo { + let mut memo = groups::memo_of(per_computation); + super::capability::add_shared_builds(&mut memo); + super::windows::add_pass2_candidates(&mut memo, entries); + memo +} + +/// A one-line summary of `memo` for logs: groups, shared groups (two or +/// more parents or queries), candidates, distinct coupling sites. +pub(crate) fn stats(memo: &Memo) -> String { + let shared = memo + .groups + .iter() + .filter(|g| g.parents.len() > 1 || (g.parents.is_empty() && g.roots_of.len() > 1)) + .count(); + let candidates: usize = memo.candidates.iter().map(Vec::len).sum(); + let mut sites: Vec = memo + .candidates + .iter() + .flatten() + .flat_map(|c| c.coupling_sites.iter().map(|s| s.key)) + .collect(); + sites.sort_unstable(); + sites.dedup(); + let most = memo.candidates.iter().map(Vec::len).max().unwrap_or(0); + format!( + "groups {} shared {} candidates {} (most in one group {}) coupling sites {}", + memo.groups.len(), + shared, + candidates, + most, + sites.len() + ) +} From 9292ed3ea68c9f484d142b4847dbb6e97293eaf4 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:08:15 +0000 Subject: [PATCH 06/25] plan-selection: Stage 3 prices one node at a time (price_node) price_at_scale's per-node loop body is now price_node, which takes the node's producers' outputs, its computation rate, a query-time scan's look-back, a chain's retained panes and an ingestion-time node's window, and returns the node's output statistics, usage, cost and the charges known only once the whole plan is (held memory, partial copies, raw data). price_at_scale calls it per node; prices are unchanged. --- crates/plan-selection/src/lib.rs | 1296 ++++++++++++++++-------------- 1 file changed, 711 insertions(+), 585 deletions(-) diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 28146e320..76958156d 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -3291,8 +3291,8 @@ fn price_at_scale( CostModel::UsageOverTime(prices) => Some(prices), _ => None, }; - // The sources each node reads, filled children first. - let mut sources_read_by: HashMap> = HashMap::new(); + // Each node priced, filled children first. + let mut priced: HashMap> = HashMap::new(); let multipliers = computations( dag, recurrence, @@ -3335,586 +3335,68 @@ fn price_at_scale( let sketch_size = |params: &SketchParams| sketch_shape(params, StateInputRows::UNKNOWN).1 as f64; for node in &dag.nodes { - let ingestion = node.output_state.timing == ExecutionTiming::IngestionTime; - let producers = producers_of.get(&node.id).map_or(&[][..], Vec::as_slice); - let inputs: Vec<_> = producers.iter().map(|p| output[p]).collect(); - let sources: BTreeSet<&str> = match &node.payload { - Payload::NonASAP(NonASAPOp::Scan { source, .. }) => [source_name(source)].into(), - _ => producers - .iter() - .flat_map(|p| sources_read_by[p].iter().copied()) - .collect(), - }; - let series = priced_series(&sources, data); - // The rows per ms its sources receive, which its scans and time - // ranges read. - let rows_per_ms = source_rows_per_ms(&sources, data); - sources_read_by.insert(node.id, sources); - let family = match &node.payload { - Payload::ASAP( - ASAPOp::SummaryAgg { - family, grouping, .. - } - | ASAPOp::PaneChain { - family, grouping, .. - }, - ) => match family { - FieldDataType::Sketch(kind, _) - if *grouping == GroupingStrategy::PerSubpopulationInstance => - { - StateFamily::Sketch(kind.params()) - } - FieldDataType::ExactAggregate(kind, _) => StateFamily::Exact(kind), - _ => StateFamily::None, - }, - Payload::ASAP( - ASAPOp::SummaryMerge { .. } - | ASAPOp::SummaryEstimate { .. } - | ASAPOp::FinalizeExactAccumulator { .. }, - ) => producers - .first() - .and_then(|p| family_of.get(p).copied()) - .unwrap_or(StateFamily::None), - _ => StateFamily::None, - }; - if family != StateFamily::None { - family_of.insert(node.id, family); - } - let operator = OperatorKind::of(&node.payload); - let unit_price = match &node.payload { - Payload::NonASAP(NonASAPOp::Aggregate { measures, .. }) => { - aggregate_lookup(table, node.output_state.timing, measures) - } - _ => lookup( - table, - operator, - node.output_state.timing, - family, - sketch_size, - ), - }; - // The kernel's measured bytes per state, when the table has them. - let measured_state_bytes = unit_price - .and_then(|p| p.memory_bytes_per_state) - .map(|bytes| bytes.round() as u64); - let mut price_source = NodePriceSource::Analytical; - let input = inputs - .first() - .copied() - .unwrap_or(EdgeStatistics { rows: 1, bytes: 1 }); - let width = row_bytes(&node.output_schema); - let edge = |rows: u64| EdgeStatistics { - rows, - bytes: rows * width, - }; - // Whether a summary built here sees raw data samples (its producer a - // scan or a time range) of sources all declared non-negative. - let raw_data_samples_non_negative = || { - let raw = producers.first().is_some_and(|producer| { - matches!( - nodes[producer].payload, - Payload::NonASAP(NonASAPOp::Scan { .. } | NonASAPOp::TimeRange { .. }) - ) - }); - raw && sources_read_by[&node.id] - .iter() - .all(|source| declared_non_negative(data, source)) - }; - let unary = |output| UnaryEdgeStatistics { - input, - output, - promql: None, - }; - // The schema the node's group and partition keys index. - let input_schema = producers.first().map(|p| &nodes[p].output_schema); - // The rows a grouping reads per evaluation: its input at query - // time; at ingestion time, the arriving rows of the window it - // maintains (one second of them does not bound its groups). - let rows_read = match ingestion { - false => Some(input.rows.max(1)), - true => maintained_window_ms(node.id, nodes, &producers_of) - .map(|window_ms| ((rows_per_ms * window_ms as f64).round() as u64).max(1)), - }; - // Query time always has its input rows; an ingestion-time grouping - // with no window below it is an error (Stage 2 builds none). - - let estimated_groups = |keys: &GroupKeys| { - let rows_read = - rows_read.ok_or((node.id, AnalyticalCostError::IngestionGroupingWithoutWindow))?; - Ok(group_count_bound( - keys, - input_schema, - &sources_read_by[&node.id], - data, - rows_read, - )) - }; - let groups = |reduction: &Reduction| match reduction { - Reduction::Reduce(keys) if keys.keys().is_empty() && !keys.is_without() => Ok(1), - Reduction::Reduce(keys) => estimated_groups(keys), - // One group per series. - Reduction::PerEntity => Ok(match ingestion { - true => series, - false => series.min(input.rows.max(1)), - }), - }; - let (out, estimate, detail) = match &node.payload { - Payload::NonASAP(operator) => match operator { - NonASAPOp::Scan { .. } => { - // At ingestion time, one second of arriving rows; at - // query time, the longest window read above it. - let lookback = match ingestion { - true => 1_000, - false => query_time_lookback_ms(dag, nodes, node.id, kept_chains), - }; - let out = edge(((rows_per_ms * lookback as f64).round() as u64).max(1)); - let estimate = estimate_operator( - PhysicalOperator::Scan, - OperatorStatistics::Scan { - edges: UnaryEdgeStatistics { - input: out, - output: out, - promql: None, - }, - source_read_bytes: out.bytes, - }, - ); - (out, estimate, format!("scan {} samples", out.rows)) - } - NonASAPOp::Aggregate { - reduction, - measures, - .. - } => { - let group_count = groups(reduction)?; - let keys = match reduction { - Reduction::Reduce(keys) => keys.keys().len() as u64, - _ => 1, - }; - let out = edge(group_count); - let held = input.rows.div_ceil(group_count.max(1)); - let estimate = estimate_operator( - PhysicalOperator::HashAggregate { - grouping_key_count: keys, - accumulator_count: measures.len().max(1) as u64, - }, - OperatorStatistics::HashAggregate { - edges: unary(out), - group_count, - key_bytes: 16 * keys, - accumulator_bytes_per_group: measures - .iter() - .map(|m| if keeps_every_value(m) { 8 * held } else { 8 }) - .sum::() - .max(8), - }, - ); - ( - out, - estimate, - format!( - "hash aggregate {} rows into {group_count} groups", - input.rows - ), - ) - } - NonASAPOp::Sort { - keys, partition_by, .. - } => { - let partitions = if partition_by.keys().is_empty() { - 1 - } else { - estimated_groups(partition_by)? - }; - let estimate = estimate_operator( - PhysicalOperator::InMemoryComparisonSort { - ordering_key_count: keys.len() as u64, - partitioned: partitions > 1, - }, - OperatorStatistics::InMemoryComparisonSort { - edges: unary(input), - input_partitioning: split(input, partitions), - }, - ); - ( - input, - estimate, - format!("sort {} rows in {partitions} partitions", input.rows), - ) - } - NonASAPOp::Limit { - n, - offset, - partition_by, - .. - } => { - let partitions = if partition_by.keys().is_empty() { - 1 - } else { - estimated_groups(partition_by)? - }; - let limit = n.map_or(u64::MAX, |n| (n as u64).saturating_mul(partitions)); - let offset = (*offset as u64).saturating_mul(partitions); - let out = edge(input.rows.saturating_sub(offset).min(limit)); - let estimate = estimate_operator( - PhysicalOperator::Limit { limit, offset }, - OperatorStatistics::Limit { edges: unary(out) }, - ); - (out, estimate, format!("limit to {} rows", out.rows)) - } - NonASAPOp::TimeShift { .. } => ( - input, - estimate_operator( - PhysicalOperator::TimeShift, - OperatorStatistics::PassThrough { - edges: unary(input), - }, - ), - format!("time shift: pass {} rows, free", input.rows), - ), - NonASAPOp::TimeRange { range, kind, .. } => { - // A range keeps its window's samples; an instant - // selector the latest sample of each series. - let kept = match kind { - TimeRangeKind::Range => { - (rows_per_ms * range.as_millis() as f64).round() as u64 - } - TimeRangeKind::Instant => series, - }; - let out = edge(input.rows.min(kept.max(1))); - let out = EdgeStatistics { - rows: out.rows, - bytes: out.bytes.min(input.bytes), - }; - let estimate = estimate_operator( - PhysicalOperator::TimeRange, - OperatorStatistics::PassThrough { edges: unary(out) }, - ); - ( - out, - estimate, - format!( - "time range {range:?}: keep {} of {} rows", - out.rows, input.rows - ), - ) - } - _ => { - let estimate = estimate_operator( - PhysicalOperator::PassThrough, - OperatorStatistics::PassThrough { - edges: unary(input), - }, - ); - ( - input, - estimate, - format!("operator: pass {} rows", input.rows), - ) - } - }, - Payload::ASAP(ASAPOp::SummaryAgg { - family, reduction, .. - }) => { - let group_count = groups(reduction)?; - // At query time, each state sees its group's input rows; at - // ingestion time, the rows of a window one second of - // arriving rows does not bound. - let input_rows = StateInputRows { - rows_per_state: (!ingestion).then(|| input.rows.div_ceil(group_count.max(1))), - non_negative: raw_data_samples_non_negative(), - }; - let (depth, state_bytes) = summary_shape(family, input_rows); - let state_bytes = measured_state_bytes.unwrap_or(state_bytes); - let out = EdgeStatistics { - rows: group_count, - bytes: group_count * state_bytes, - }; - ( - out, - Ok(ResourceEstimate::new( - input.rows as f64 * depth, - out.bytes, - 0, - )), - format!( - "build {} into {group_count} states: {} rows x depth {depth}", - family_name(family), - input.rows - ), - ) - } - Payload::ASAP(ASAPOp::PaneChain { - family, - reduction, - pane_width_ms, - pane_count, - .. - }) => { - // Each input row is kept by the one pane its time falls in, - // as a time range keeps it (one operation), and updates - // that pane. - let group_count = groups(reduction)?; - let states = group_count * u64::from(*pane_count); - // Each pane's state sees its group's rows in one pane: of - // one second's rows at ingestion time, of the rows read at - // query time (all of them for a chain kept across - // evaluations, whose scan Stage 2 bounds to its newest pane, - // the only one it builds). - let rows_per_pane = match (ingestion, node.materialization) { - (true, _) => (input.rows * *pane_width_ms).div_ceil(1_000), - (false, Materialization::MatAtQueryKeptAcrossEvaluations(_)) => input.rows, - (false, _) => input.rows.div_ceil(u64::from(*pane_count).max(1)), - }; - let input_rows = StateInputRows { - rows_per_state: Some(rows_per_pane.div_ceil(group_count.max(1))), - non_negative: raw_data_samples_non_negative(), - }; - let (depth, state_bytes) = summary_shape(family, input_rows); - let state_bytes = measured_state_bytes.unwrap_or(state_bytes); - let out = EdgeStatistics { - rows: states, - bytes: states * state_bytes, - }; - ( - out, - Ok(ResourceEstimate::new( - input.rows as f64 * (depth + 1.0), - out.bytes, - 0, - )), - format!( - "build {} in {pane_count} panes of {pane_width_ms} ms for \ - {group_count} groups: {} rows x (1 + depth {depth})", - family_name(family), - input.rows - ), - ) - } - Payload::ASAP(ASAPOp::SummaryMerge { - newest_panes, - group_by, - .. - }) => { - // One operation per input state read; one state per group - // out, each pane of a chain being a state of its group. A - // merge of a chain's newest `n` panes reads `n` per group. A - // rollup's groups are those of its keys, at most its input - // groups. - let mut group_count = 0; - for producer in producers { - let panes = match &nodes[producer].payload { - Payload::ASAP(ASAPOp::PaneChain { pane_count, .. }) => { - u64::from(*pane_count) - } - _ => 1, - }; - group_count = group_count.max(output[producer].rows / panes.max(1)); - } - let states = merged_states(*newest_panes, producers, &output, nodes); - if let Some(keys) = group_by { - group_count = group_count.min(groups(&Reduction::by(keys.clone()))?); - } - let out = EdgeStatistics { - rows: group_count, - bytes: input.bytes.div_ceil(input.rows.max(1)) * group_count, - }; - ( - out, - Ok(ResourceEstimate::new(states as f64, out.bytes, 0)), - format!("merge {states} states into {group_count}"), - ) - } - Payload::ASAP(ASAPOp::SummaryEstimate { query, .. }) => { - let per_group = match query { - SketchStatistic::TopK { k } => *k as u64, - _ => 1, - }; - let out = edge(input.rows * per_group); - ( - out, - Ok(ResourceEstimate::new(out.rows as f64, 0, 0)), - format!("estimate {} groups x {per_group}", input.rows), - ) - } - Payload::ASAP(ASAPOp::FinalizeExactAccumulator { .. }) => ( - edge(input.rows), - Ok(ResourceEstimate::new(input.rows as f64, 0, 0)), - format!("finalize {} accumulators", input.rows), - ), - _ => ( - edge(input.rows), - Ok(ResourceEstimate::new(input.rows as f64, 0, 0)), - format!("{} rows", input.rows), - ), - }; - let estimate = estimate.map_err(|error| (node.id, error))?; - let (multiplier, per_second) = multipliers[&node.id]; - let formula_ms = estimate.cpu_ops() * speed.ms_per_cpu_op - + estimate.scan_bytes() as f64 * speed.ms_per_scanned_byte; - // What one computation processes, the unit its entry is priced per - // ([`OperatorKind`]): the samples a scan reads, the states a merge - // reads, otherwise the rows of its (first) input. - let units = match &node.payload { - Payload::NonASAP(NonASAPOp::Scan { .. }) => out.rows, - Payload::ASAP(ASAPOp::SummaryMerge { newest_panes, .. }) => { - merged_states(*newest_panes, producers, &output, nodes) - } - _ => input.rows, - }; - let mut detail = detail; - // A node with an entry takes its measured time, scanning included; - // any other node its formulas' CPU operations times the calibration - // factor, its scanned bytes unscaled (scaling only the CPU keeps it - // non-negative whatever the factor). - let estimate = match unit_price { - Some(unit) => { - price_source = NodePriceSource::Measured; - let measured_ms = (units as f64 * unit.ns_per_unit + unit.ns_per_invocation) * 1e-6; - fit.measured_ms_per_second += measured_ms * per_second; - fit.formula_ms_per_second += formula_ms * per_second; - detail += &format!( - "; measured {units} x {:.3} ns + {:.3} ns", - unit.ns_per_unit, unit.ns_per_invocation - ); - ResourceEstimate::new( - measured_ms / speed.ms_per_cpu_op, - estimate.peak_memory_bytes(), - 0, - ) - } - None if formula_scale != 1.0 => { - detail += &format!("; formula x {formula_scale:.3} (calibration)"); - ResourceEstimate::new( - estimate.cpu_ops() * formula_scale, - estimate.peak_memory_bytes(), - estimate.scan_bytes(), - ) - } - None => estimate, - }; - let node_ms = estimate.cpu_ops() * speed.ms_per_cpu_op - + estimate.scan_bytes() as f64 * speed.ms_per_scanned_byte; - work_ms.insert(node.id, node_ms); - usage.cpu_ops_per_second += estimate.cpu_ops() * per_second; - usage.scanned_bytes_per_second += estimate.scan_bytes() as f64 * per_second; - // What this node keeps: memory averaged over time (byte-seconds per - // second), its peak memory and its disk bytes, added to the - // candidate's usage below. - let mut average_memory_bytes = 0.0; - let mut peak_memory_bytes = 0u64; - let mut disk_bytes = 0u64; - let mut detail = format!("{detail}; x {per_second:.4}/s"); - // State kept between computations: at ingestion time, the window - // being built and the completed ones queries read (a chain of N - // panes keeps N + 1), every second. At query time the output is - // reused only within one evaluation, so it is held only while that - // evaluation runs (counted below, once latencies are known). The - // state is what the node holds: its output, or the values an exact - // aggregate keeps per group when that is more. A computation that - // keeps nothing holds its working memory while it runs, once per - // consumer that recomputes it; at query time, that is held for the - // evaluation too (counted below, as a MatAtQuery state is). - let state = out.bytes.max(estimate.peak_memory_bytes()); - let (kept, medium) = match (node.materialization, &node.payload) { - (Materialization::NotMaterialized, _) => { - let working_memory_bytes = estimate.peak_memory_bytes(); - peak_memory_bytes = working_memory_bytes.saturating_mul(multiplier); - if !ingestion && working_memory_bytes > 0 { - held.insert(node.id, (working_memory_bytes, per_second, true)); - } - (0, Medium::Memory) - } - // A pane chain kept across evaluations holds the panes its - // retention covers ([`pane_retention_ms`]: N, plus one per pane - // width of the readers' range-evaluation span), all but the - // newest between evaluations and all while one runs. - (Materialization::MatAtQueryKeptAcrossEvaluations(medium), _) => { - let panes = panes_of(node); - let retained = retained_panes(node, &retention_ms); - let between = state / panes * (retained - 1); - match medium { - Medium::Memory => { - average_memory_bytes = between as f64; - peak_memory_bytes = match retained == panes { - true => state, - false => state / panes * retained, - }; - } - Medium::Disk => disk_bytes = between, - } - detail += &format!("; keeps {between} bytes across evaluations"); - (0, medium) - } - (Materialization::MatAtQuery(medium), _) => { - match medium { - Medium::Memory => { - held.insert(node.id, (state, per_second, false)); - peak_memory_bytes = state; - } - Medium::Disk => disk_bytes = state, + let views: Vec> = producers_of + .get(&node.id) + .map_or(&[][..], Vec::as_slice) + .iter() + .map(|p| { + let produced: &PricedNode<'_> = &priced[p]; + ProducerOutput { + payload: &nodes[p].payload, + output_schema: &nodes[p].output_schema, + output: produced.output, + sources: produced.sources.clone(), + family: produced.family, } - (0, medium) - } - // The retained panes and the one being built. - (Materialization::MatAtIngestion(medium), Payload::ASAP(ASAPOp::PaneChain { .. })) => { - let panes = panes_of(node); - ( - state / panes * (retained_panes(node, &retention_ms) + 1), - medium, - ) + }) + .collect(); + let lookback_ms = match (&node.payload, node.output_state.timing) { + (Payload::NonASAP(NonASAPOp::Scan { .. }), ExecutionTiming::QueryTime) => { + query_time_lookback_ms(dag, nodes, node.id, kept_chains) } - (Materialization::MatAtIngestion(medium), _) => (2 * state, medium), + _ => 0, }; - if node.materialization == Materialization::MatAtIngestion(Medium::Memory) { - maintained.push((node.id, state / panes_of(node))); - } - if kept > 0 { - match medium { - Medium::Memory => { - average_memory_bytes = kept as f64; - peak_memory_bytes = kept; - } - Medium::Disk => disk_bytes = kept, - } - detail += &format!("; keeps {kept} bytes"); - } - if let Payload::NonASAP(NonASAPOp::Scan { source, .. }) = &node.payload { - if !ingestion { - let largest = raw_data.entry(source).or_insert((out.bytes, node.id)); - if out.bytes > largest.0 { - *largest = (out.bytes, node.id); - } + let node_priced = price_node( + node, + &views, + multipliers[&node.id], + lookback_ms, + retained_panes(node, &retention_ms), + (node.output_state.timing == ExecutionTiming::IngestionTime) + .then(|| maintained_window_ms(node.id, nodes, &producers_of)) + .flatten(), + data, + settings, + formula_scale, + ) + .map_err(|error| (node.id, error))?; + usage.cpu_ops_per_second += node_priced.cpu_ops_per_second; + usage.scanned_bytes_per_second += node_priced.scanned_bytes_per_second; + usage.average_memory_bytes += node_priced.average_memory_bytes; + usage.peak_memory_bytes = usage + .peak_memory_bytes + .saturating_add(node_priced.peak_memory_bytes); + usage.disk_bytes = usage.disk_bytes.saturating_add(node_priced.disk_bytes); + fit.add(node_priced.fit); + work_ms.insert(node.id, node_priced.work_ms); + output.insert(node.id, node_priced.output); + if node_priced.family != StateFamily::None { + family_of.insert(node.id, node_priced.family); + } + if let Some(entry) = node_priced.held { + held.insert(node.id, entry); + } + if let Some(window) = node_priced.maintained { + maintained.push((node.id, window)); + } + if let Some((source, bytes)) = node_priced.raw_scan { + let largest = raw_data.entry(source).or_insert((bytes, node.id)); + if bytes > largest.0 { + *largest = (bytes, node.id); } } - // One computation's CPU and scan price times its rate, then the - // kept state's. - let cost = usage_prices.map(|prices| { - (estimate.cpu_ops() * prices.cost_per_cpu_op - + estimate.scan_bytes() as f64 * prices.cost_per_scanned_byte) - * per_second - + prices.cost_per_memory_byte_second * average_memory_bytes - + prices.cost_per_disk_byte_second * disk_bytes as f64 - }); - usage.average_memory_bytes += average_memory_bytes; - usage.peak_memory_bytes = usage.peak_memory_bytes.saturating_add(peak_memory_bytes); - usage.disk_bytes = usage.disk_bytes.saturating_add(disk_bytes); - if cost.is_some_and(|cost| !cost.is_finite()) { - return Err((node.id, AnalyticalCostError::Overflow)); - } - output.insert(node.id, out); - per_node.insert( - node.id, - NodeCost { - operator, - timing: node.output_state.timing, - cost_per_second: cost, - cpu_core_seconds_per_second: node_ms * per_second / 1_000.0, - multiplier, - detail, - price_source, - }, - ); + per_node.insert(node.id, node_priced.cost.clone()); + priced.insert(node.id, node_priced); } // The runtime's fixed cost per query request, charged to the node the // query reads first and added to its latency below. @@ -4066,6 +3548,655 @@ fn price_at_scale( )) } +/// One producer of a node being priced ([`price_node`]): its payload and +/// schema, its output statistics, the sources it reads and its state +/// family. +#[derive(Debug, Clone)] +pub(crate) struct ProducerOutput<'a> { + pub(crate) payload: &'a Payload, + pub(crate) output_schema: &'a Schema, + pub(crate) output: EdgeStatistics, + pub(crate) sources: BTreeSet<&'a str>, + pub(crate) family: StateFamily<'a>, +} + +/// What one node costs and outputs ([`price_node`]): its output statistics, +/// sources and state family for its consumers; its work for one +/// computation in ms on one core; its usage per second; what is charged +/// once the whole plan is known (memory held for a query's evaluation, an +/// in-memory ingestion window's partial copies, the raw data a query-time +/// scan keeps on disk); its calibration fit and its [`NodeCost`]. +#[derive(Debug, Clone)] +pub(crate) struct PricedNode<'a> { + pub(crate) output: EdgeStatistics, + pub(crate) sources: BTreeSet<&'a str>, + pub(crate) family: StateFamily<'a>, + pub(crate) work_ms: f64, + pub(crate) cpu_ops_per_second: f64, + pub(crate) scanned_bytes_per_second: f64, + pub(crate) average_memory_bytes: f64, + pub(crate) peak_memory_bytes: u64, + pub(crate) disk_bytes: u64, + /// Query-time memory held while an evaluation runs: (bytes, computations + /// per second, whether each reading query computes its own copy). + pub(crate) held: Option<(u64, f64, bool)>, + /// An in-memory ingestion-time state's bytes per window. + pub(crate) maintained: Option, + /// A query-time scan's source and the raw bytes it reads. + pub(crate) raw_scan: Option<(&'a Source, u64)>, + pub(crate) fit: CalibrationFit, + pub(crate) cost: NodeCost, +} + +/// Measure and price one node of a physical DAG (one step of +/// [`price_at_scale`]): `producers` are its producers in edge order, +/// `rate` its `(multiplier, per_second)` ([`computations`]), `lookback_ms` +/// the raw data a query-time scan reads ([`query_time_lookback_ms`]; read +/// only for scans), `retained_pane_count` the panes a pane chain keeps +/// completed ([`retained_panes`]), and `ingestion_window_ms` the window an +/// ingestion-time node maintains ([`maintained_window_ms`]). +#[allow(clippy::too_many_arguments)] +pub(crate) fn price_node<'a>( + node: &'a asap_types::ir::physical_export::PhysicalASAPDAGNode, + producers: &[ProducerOutput<'a>], + rate: (u64, f64), + lookback_ms: u64, + retained_pane_count: u64, + ingestion_window_ms: Option, + data: &DataWorkload, + settings: &'a PlanSelectionSettings, + formula_scale: f64, +) -> Result, AnalyticalCostError> { + let table = &settings.measured_operation_costs; + let speed = settings.resource_speed; + let usage_prices = match &settings.cost_model { + CostModel::UsageOverTime(prices) => Some(prices), + _ => None, + }; + let sketch_size = + |params: &SketchParams| sketch_shape(params, StateInputRows::UNKNOWN).1 as f64; + let mut fit = CalibrationFit::default(); + let mut held = None; + let mut maintained = None; + let mut raw_scan = None; + let ingestion = node.output_state.timing == ExecutionTiming::IngestionTime; + let inputs: Vec = producers.iter().map(|p| p.output).collect(); + let sources: BTreeSet<&str> = match &node.payload { + Payload::NonASAP(NonASAPOp::Scan { source, .. }) => [source_name(source)].into(), + _ => producers + .iter() + .flat_map(|p| p.sources.iter().copied()) + .collect(), + }; + let series = priced_series(&sources, data); + // The rows per ms its sources receive, which its scans and time + // ranges read. + let rows_per_ms = source_rows_per_ms(&sources, data); + let family = match &node.payload { + Payload::ASAP( + ASAPOp::SummaryAgg { + family, grouping, .. + } + | ASAPOp::PaneChain { + family, grouping, .. + }, + ) => match family { + FieldDataType::Sketch(kind, _) + if *grouping == GroupingStrategy::PerSubpopulationInstance => + { + StateFamily::Sketch(kind.params()) + } + FieldDataType::ExactAggregate(kind, _) => StateFamily::Exact(kind), + _ => StateFamily::None, + }, + Payload::ASAP( + ASAPOp::SummaryMerge { .. } + | ASAPOp::SummaryEstimate { .. } + | ASAPOp::FinalizeExactAccumulator { .. }, + ) => producers + .first() + .map_or(StateFamily::None, |p| p.family), + _ => StateFamily::None, + }; + let operator = OperatorKind::of(&node.payload); + let unit_price = match &node.payload { + Payload::NonASAP(NonASAPOp::Aggregate { measures, .. }) => { + aggregate_lookup(table, node.output_state.timing, measures) + } + _ => lookup( + table, + operator, + node.output_state.timing, + family, + sketch_size, + ), + }; + // The kernel's measured bytes per state, when the table has them. + let measured_state_bytes = unit_price + .and_then(|p| p.memory_bytes_per_state) + .map(|bytes| bytes.round() as u64); + let mut price_source = NodePriceSource::Analytical; + let input = inputs + .first() + .copied() + .unwrap_or(EdgeStatistics { rows: 1, bytes: 1 }); + let width = row_bytes(&node.output_schema); + let edge = |rows: u64| EdgeStatistics { + rows, + bytes: rows * width, + }; + // Whether a summary built here sees raw data samples (its producer a + // scan or a time range) of sources all declared non-negative. + let raw_data_samples_non_negative = || { + let raw = producers.first().is_some_and(|producer| { + matches!( + producer.payload, + Payload::NonASAP(NonASAPOp::Scan { .. } | NonASAPOp::TimeRange { .. }) + ) + }); + raw && sources + .iter() + .all(|source| declared_non_negative(data, source)) + }; + let unary = |output| UnaryEdgeStatistics { + input, + output, + promql: None, + }; + // The schema the node's group and partition keys index. + let input_schema = producers.first().map(|p| p.output_schema); + // The rows a grouping reads per evaluation: its input at query + // time; at ingestion time, the arriving rows of the window it + // maintains (one second of them does not bound its groups). + let rows_read = match ingestion { + false => Some(input.rows.max(1)), + true => ingestion_window_ms.map(|window_ms| ((rows_per_ms * window_ms as f64).round() as u64).max(1)), + }; + // Query time always has its input rows; an ingestion-time grouping + // with no window below it is an error (Stage 2 builds none). + + let estimated_groups = |keys: &GroupKeys| { + let rows_read = + rows_read.ok_or(AnalyticalCostError::IngestionGroupingWithoutWindow)?; + Ok(group_count_bound( + keys, + input_schema, + &sources, + data, + rows_read, + )) + }; + let groups = |reduction: &Reduction| match reduction { + Reduction::Reduce(keys) if keys.keys().is_empty() && !keys.is_without() => Ok(1), + Reduction::Reduce(keys) => estimated_groups(keys), + // One group per series. + Reduction::PerEntity => Ok(match ingestion { + true => series, + false => series.min(input.rows.max(1)), + }), + }; + let (out, estimate, detail) = match &node.payload { + Payload::NonASAP(operator) => match operator { + NonASAPOp::Scan { .. } => { + // At ingestion time, one second of arriving rows; at + // query time, the longest window read above it. + let lookback = match ingestion { + true => 1_000, + false => lookback_ms, + }; + let out = edge(((rows_per_ms * lookback as f64).round() as u64).max(1)); + let estimate = estimate_operator( + PhysicalOperator::Scan, + OperatorStatistics::Scan { + edges: UnaryEdgeStatistics { + input: out, + output: out, + promql: None, + }, + source_read_bytes: out.bytes, + }, + ); + (out, estimate, format!("scan {} samples", out.rows)) + } + NonASAPOp::Aggregate { + reduction, + measures, + .. + } => { + let group_count = groups(reduction)?; + let keys = match reduction { + Reduction::Reduce(keys) => keys.keys().len() as u64, + _ => 1, + }; + let out = edge(group_count); + let held = input.rows.div_ceil(group_count.max(1)); + let estimate = estimate_operator( + PhysicalOperator::HashAggregate { + grouping_key_count: keys, + accumulator_count: measures.len().max(1) as u64, + }, + OperatorStatistics::HashAggregate { + edges: unary(out), + group_count, + key_bytes: 16 * keys, + accumulator_bytes_per_group: measures + .iter() + .map(|m| if keeps_every_value(m) { 8 * held } else { 8 }) + .sum::() + .max(8), + }, + ); + ( + out, + estimate, + format!( + "hash aggregate {} rows into {group_count} groups", + input.rows + ), + ) + } + NonASAPOp::Sort { + keys, partition_by, .. + } => { + let partitions = if partition_by.keys().is_empty() { + 1 + } else { + estimated_groups(partition_by)? + }; + let estimate = estimate_operator( + PhysicalOperator::InMemoryComparisonSort { + ordering_key_count: keys.len() as u64, + partitioned: partitions > 1, + }, + OperatorStatistics::InMemoryComparisonSort { + edges: unary(input), + input_partitioning: split(input, partitions), + }, + ); + ( + input, + estimate, + format!("sort {} rows in {partitions} partitions", input.rows), + ) + } + NonASAPOp::Limit { + n, + offset, + partition_by, + .. + } => { + let partitions = if partition_by.keys().is_empty() { + 1 + } else { + estimated_groups(partition_by)? + }; + let limit = n.map_or(u64::MAX, |n| (n as u64).saturating_mul(partitions)); + let offset = (*offset as u64).saturating_mul(partitions); + let out = edge(input.rows.saturating_sub(offset).min(limit)); + let estimate = estimate_operator( + PhysicalOperator::Limit { limit, offset }, + OperatorStatistics::Limit { edges: unary(out) }, + ); + (out, estimate, format!("limit to {} rows", out.rows)) + } + NonASAPOp::TimeShift { .. } => ( + input, + estimate_operator( + PhysicalOperator::TimeShift, + OperatorStatistics::PassThrough { + edges: unary(input), + }, + ), + format!("time shift: pass {} rows, free", input.rows), + ), + NonASAPOp::TimeRange { range, kind, .. } => { + // A range keeps its window's samples; an instant + // selector the latest sample of each series. + let kept = match kind { + TimeRangeKind::Range => { + (rows_per_ms * range.as_millis() as f64).round() as u64 + } + TimeRangeKind::Instant => series, + }; + let out = edge(input.rows.min(kept.max(1))); + let out = EdgeStatistics { + rows: out.rows, + bytes: out.bytes.min(input.bytes), + }; + let estimate = estimate_operator( + PhysicalOperator::TimeRange, + OperatorStatistics::PassThrough { edges: unary(out) }, + ); + ( + out, + estimate, + format!( + "time range {range:?}: keep {} of {} rows", + out.rows, input.rows + ), + ) + } + _ => { + let estimate = estimate_operator( + PhysicalOperator::PassThrough, + OperatorStatistics::PassThrough { + edges: unary(input), + }, + ); + ( + input, + estimate, + format!("operator: pass {} rows", input.rows), + ) + } + }, + Payload::ASAP(ASAPOp::SummaryAgg { + family, reduction, .. + }) => { + let group_count = groups(reduction)?; + // At query time, each state sees its group's input rows; at + // ingestion time, the rows of a window one second of + // arriving rows does not bound. + let input_rows = StateInputRows { + rows_per_state: (!ingestion).then(|| input.rows.div_ceil(group_count.max(1))), + non_negative: raw_data_samples_non_negative(), + }; + let (depth, state_bytes) = summary_shape(family, input_rows); + let state_bytes = measured_state_bytes.unwrap_or(state_bytes); + let out = EdgeStatistics { + rows: group_count, + bytes: group_count * state_bytes, + }; + ( + out, + Ok(ResourceEstimate::new( + input.rows as f64 * depth, + out.bytes, + 0, + )), + format!( + "build {} into {group_count} states: {} rows x depth {depth}", + family_name(family), + input.rows + ), + ) + } + Payload::ASAP(ASAPOp::PaneChain { + family, + reduction, + pane_width_ms, + pane_count, + .. + }) => { + // Each input row is kept by the one pane its time falls in, + // as a time range keeps it (one operation), and updates + // that pane. + let group_count = groups(reduction)?; + let states = group_count * u64::from(*pane_count); + // Each pane's state sees its group's rows in one pane: of + // one second's rows at ingestion time, of the rows read at + // query time (all of them for a chain kept across + // evaluations, whose scan Stage 2 bounds to its newest pane, + // the only one it builds). + let rows_per_pane = match (ingestion, node.materialization) { + (true, _) => (input.rows * *pane_width_ms).div_ceil(1_000), + (false, Materialization::MatAtQueryKeptAcrossEvaluations(_)) => input.rows, + (false, _) => input.rows.div_ceil(u64::from(*pane_count).max(1)), + }; + let input_rows = StateInputRows { + rows_per_state: Some(rows_per_pane.div_ceil(group_count.max(1))), + non_negative: raw_data_samples_non_negative(), + }; + let (depth, state_bytes) = summary_shape(family, input_rows); + let state_bytes = measured_state_bytes.unwrap_or(state_bytes); + let out = EdgeStatistics { + rows: states, + bytes: states * state_bytes, + }; + ( + out, + Ok(ResourceEstimate::new( + input.rows as f64 * (depth + 1.0), + out.bytes, + 0, + )), + format!( + "build {} in {pane_count} panes of {pane_width_ms} ms for \ + {group_count} groups: {} rows x (1 + depth {depth})", + family_name(family), + input.rows + ), + ) + } + Payload::ASAP(ASAPOp::SummaryMerge { + newest_panes, + group_by, + .. + }) => { + // One operation per input state read; one state per group + // out, each pane of a chain being a state of its group. A + // merge of a chain's newest `n` panes reads `n` per group. A + // rollup's groups are those of its keys, at most its input + // groups. + let mut group_count = 0; + for producer in producers { + let panes = match producer.payload { + Payload::ASAP(ASAPOp::PaneChain { pane_count, .. }) => { + u64::from(*pane_count) + } + _ => 1, + }; + group_count = group_count.max(producer.output.rows / panes.max(1)); + } + let states = merged_states(*newest_panes, producers); + if let Some(keys) = group_by { + group_count = group_count.min(groups(&Reduction::by(keys.clone()))?); + } + let out = EdgeStatistics { + rows: group_count, + bytes: input.bytes.div_ceil(input.rows.max(1)) * group_count, + }; + ( + out, + Ok(ResourceEstimate::new(states as f64, out.bytes, 0)), + format!("merge {states} states into {group_count}"), + ) + } + Payload::ASAP(ASAPOp::SummaryEstimate { query, .. }) => { + let per_group = match query { + SketchStatistic::TopK { k } => *k as u64, + _ => 1, + }; + let out = edge(input.rows * per_group); + ( + out, + Ok(ResourceEstimate::new(out.rows as f64, 0, 0)), + format!("estimate {} groups x {per_group}", input.rows), + ) + } + Payload::ASAP(ASAPOp::FinalizeExactAccumulator { .. }) => ( + edge(input.rows), + Ok(ResourceEstimate::new(input.rows as f64, 0, 0)), + format!("finalize {} accumulators", input.rows), + ), + _ => ( + edge(input.rows), + Ok(ResourceEstimate::new(input.rows as f64, 0, 0)), + format!("{} rows", input.rows), + ), + }; + let estimate = estimate?; + let (multiplier, per_second) = rate; + let formula_ms = estimate.cpu_ops() * speed.ms_per_cpu_op + + estimate.scan_bytes() as f64 * speed.ms_per_scanned_byte; + // What one computation processes, the unit its entry is priced per + // ([`OperatorKind`]): the samples a scan reads, the states a merge + // reads, otherwise the rows of its (first) input. + let units = match &node.payload { + Payload::NonASAP(NonASAPOp::Scan { .. }) => out.rows, + Payload::ASAP(ASAPOp::SummaryMerge { newest_panes, .. }) => { + merged_states(*newest_panes, producers) + } + _ => input.rows, + }; + let mut detail = detail; + // A node with an entry takes its measured time, scanning included; + // any other node its formulas' CPU operations times the calibration + // factor, its scanned bytes unscaled (scaling only the CPU keeps it + // non-negative whatever the factor). + let estimate = match unit_price { + Some(unit) => { + price_source = NodePriceSource::Measured; + let measured_ms = (units as f64 * unit.ns_per_unit + unit.ns_per_invocation) * 1e-6; + fit.measured_ms_per_second += measured_ms * per_second; + fit.formula_ms_per_second += formula_ms * per_second; + detail += &format!( + "; measured {units} x {:.3} ns + {:.3} ns", + unit.ns_per_unit, unit.ns_per_invocation + ); + ResourceEstimate::new( + measured_ms / speed.ms_per_cpu_op, + estimate.peak_memory_bytes(), + 0, + ) + } + None if formula_scale != 1.0 => { + detail += &format!("; formula x {formula_scale:.3} (calibration)"); + ResourceEstimate::new( + estimate.cpu_ops() * formula_scale, + estimate.peak_memory_bytes(), + estimate.scan_bytes(), + ) + } + None => estimate, + }; + let node_ms = estimate.cpu_ops() * speed.ms_per_cpu_op + + estimate.scan_bytes() as f64 * speed.ms_per_scanned_byte; + + // What this node keeps: memory averaged over time (byte-seconds per + // second), its peak memory and its disk bytes, added to the + // candidate's usage below. + let mut average_memory_bytes = 0.0; + let mut peak_memory_bytes = 0u64; + let mut disk_bytes = 0u64; + let mut detail = format!("{detail}; x {per_second:.4}/s"); + // State kept between computations: at ingestion time, the window + // being built and the completed ones queries read (a chain of N + // panes keeps N + 1), every second. At query time the output is + // reused only within one evaluation, so it is held only while that + // evaluation runs (counted below, once latencies are known). The + // state is what the node holds: its output, or the values an exact + // aggregate keeps per group when that is more. A computation that + // keeps nothing holds its working memory while it runs, once per + // consumer that recomputes it; at query time, that is held for the + // evaluation too (counted below, as a MatAtQuery state is). + let state = out.bytes.max(estimate.peak_memory_bytes()); + let (kept, medium) = match (node.materialization, &node.payload) { + (Materialization::NotMaterialized, _) => { + let working_memory_bytes = estimate.peak_memory_bytes(); + peak_memory_bytes = working_memory_bytes.saturating_mul(multiplier); + if !ingestion && working_memory_bytes > 0 { + held = Some((working_memory_bytes, per_second, true)); + } + (0, Medium::Memory) + } + // A pane chain kept across evaluations holds the panes its + // retention covers ([`pane_retention_ms`]: N, plus one per pane + // width of the readers' range-evaluation span), all but the + // newest between evaluations and all while one runs. + (Materialization::MatAtQueryKeptAcrossEvaluations(medium), _) => { + let panes = panes_of(node); + let retained = retained_pane_count; + let between = state / panes * (retained - 1); + match medium { + Medium::Memory => { + average_memory_bytes = between as f64; + peak_memory_bytes = match retained == panes { + true => state, + false => state / panes * retained, + }; + } + Medium::Disk => disk_bytes = between, + } + detail += &format!("; keeps {between} bytes across evaluations"); + (0, medium) + } + (Materialization::MatAtQuery(medium), _) => { + match medium { + Medium::Memory => { + held = Some((state, per_second, false)); + peak_memory_bytes = state; + } + Medium::Disk => disk_bytes = state, + } + (0, medium) + } + // The retained panes and the one being built. + (Materialization::MatAtIngestion(medium), Payload::ASAP(ASAPOp::PaneChain { .. })) => { + let panes = panes_of(node); + ( + state / panes * (retained_pane_count + 1), + medium, + ) + } + (Materialization::MatAtIngestion(medium), _) => (2 * state, medium), + }; + if node.materialization == Materialization::MatAtIngestion(Medium::Memory) { + maintained = Some(state / panes_of(node)); + } + if kept > 0 { + match medium { + Medium::Memory => { + average_memory_bytes = kept as f64; + peak_memory_bytes = kept; + } + Medium::Disk => disk_bytes = kept, + } + detail += &format!("; keeps {kept} bytes"); + } + if let Payload::NonASAP(NonASAPOp::Scan { source, .. }) = &node.payload { + if !ingestion { + raw_scan = Some((source, out.bytes)); + } + } + // One computation's CPU and scan price times its rate, then the + // kept state's. + let cost = usage_prices.map(|prices| { + (estimate.cpu_ops() * prices.cost_per_cpu_op + + estimate.scan_bytes() as f64 * prices.cost_per_scanned_byte) + * per_second + + prices.cost_per_memory_byte_second * average_memory_bytes + + prices.cost_per_disk_byte_second * disk_bytes as f64 + }); + if cost.is_some_and(|cost| !cost.is_finite()) { + return Err(AnalyticalCostError::Overflow); + } + Ok(PricedNode { + output: out, + sources, + family, + work_ms: node_ms, + cpu_ops_per_second: estimate.cpu_ops() * per_second, + scanned_bytes_per_second: estimate.scan_bytes() as f64 * per_second, + average_memory_bytes, + peak_memory_bytes, + disk_bytes, + held, + maintained, + raw_scan, + fit, + cost: NodeCost { + operator, + timing: node.output_state.timing, + cost_per_second: cost, + cpu_core_seconds_per_second: node_ms * per_second / 1_000.0, + multiplier, + detail, + price_source, + }, + }) +} + /// [`price_at_scale`] of one candidate alone, its calibration factor fitted /// over its own measured nodes (as [`select_plan`] fits it over all). #[cfg(test)] @@ -4309,17 +4440,12 @@ fn chain_input_offset_ms( /// The input states a `SummaryMerge` over `producers` reads: every row of /// each, or of the one pane chain whose newest `newest_panes` it merges, /// that many panes' states per group. -fn merged_states( - newest_panes: Option, - producers: &[PhysicalASAPNodeId], - output: &HashMap, - nodes: &HashMap, -) -> u64 { +fn merged_states(newest_panes: Option, producers: &[ProducerOutput<'_>]) -> u64 { producers .iter() .map(|p| { - let rows = output[p].rows; - match (newest_panes, &nodes[p].payload) { + let rows = p.output.rows; + match (newest_panes, p.payload) { (Some(n), Payload::ASAP(ASAPOp::PaneChain { pane_count, .. })) => { rows / u64::from(*pane_count).max(1) * u64::from(n) } From 9a70f643821eb9358b82d34de014785197ac6ea8 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:25:24 +0000 Subject: [PATCH 07/25] plan-selection: sub-expression memo: shared builds over equal input groups, scan materializations A shared-build site keys its input by structure, so equal windows of two queries (one group each without a unique key) meet. Candidates list the materializations Stage 2 may give them; a scan group read through one window is only built once, one read through different windows may also be read per reader. Equal SQL table scans are one group. --- .../src/pass/subexpr/capability.rs | 70 +++++++++-- .../plan-selection/src/pass/subexpr/groups.rs | 2 + crates/plan-selection/src/pass/subexpr/mod.rs | 7 ++ .../src/pass/subexpr/pricing.rs | 1 + .../plan-selection/src/pass/subexpr/scans.rs | 116 ++++++++++++++++++ .../src/pass/subexpr/windows.rs | 4 + 6 files changed, 193 insertions(+), 7 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index 3ebe2546a..f876425ef 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -32,6 +32,7 @@ use asap_types::ir::schema::sketch::{GroupingStrategy, SketchParams}; use asap_types::ir::schema::FieldDataType; use asap_types::ir::{ASAPOp, OperatorNode}; +use super::groups::NodeKeys; use super::{CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo}; /// One size a shared build may take: the per-parameter maximum @@ -159,7 +160,8 @@ fn same_size(a: &SketchParams, b: &SketchParams) -> bool { /// The site key of `candidate`'s summary build, when it is a shareable /// one over a group of `memo`: a hash of the family (algorithm and -/// grouping, or the exact accumulator), the input group, the reduction, +/// grouping, or the exact accumulator), the input group's structure (equal +/// input groups are one input), the reduction, /// the grouping, the filter, and the item and weight. fn site_key(memo: &Memo, candidate: &ImplementationCandidate) -> Option { let (build, _) = reader_build(candidate)?; @@ -174,10 +176,14 @@ fn site_key(memo: &Memo, candidate: &ImplementationCandidate) -> Option { else { return None; }; - let input_group = candidate.children.iter().find(|&&GroupId(g)| { - let node = &memo.groups[g as usize].node; - Rc::ptr_eq(node, child) || **node == **child - })?; + // The build reads one of its child groups; equal input groups (two + // queries' identical windows, which are not one group without a + // unique key) are one input, as the rule's `same_input_data` decides. + candidate + .children + .iter() + .find(|&&GroupId(g)| Rc::ptr_eq(&memo.groups[g as usize].node, child))?; + let input = (NodeKeys::default().key(child), input); let family_key = match family { FieldDataType::Sketch(kind, strategy) => format!("{:?}/{strategy:?}", kind.algorithm()), exact => format!("{exact:?}"), @@ -185,8 +191,8 @@ fn site_key(memo: &Memo, candidate: &ImplementationCandidate) -> Option { let mut hasher = DefaultHasher::new(); ( family_key, - input_group, - format!("{reduction:?}/{grouping:?}/{filter:?}/{input:?}"), + input.0, + format!("{reduction:?}/{grouping:?}/{filter:?}/{:?}", input.1), ) .hash(&mut hasher); Some(hasher.finish()) @@ -315,6 +321,7 @@ mod tests { children: vec![GroupId(1)], origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, @@ -442,4 +449,53 @@ mod tests { .collect(); assert_eq!(parents_roots, vec![vec![0], vec![1]]); } + + /// On a memo built by `memo_of`: the 0.9 quantile of one 5-minute + /// window within 1% and within 2% error. The scan is one group; the two + /// windows (no unique key) are a group each but equal, so each query's + /// KLL candidate gets a shared-build twin with one site, whose + /// thresholds are the two sizes. + #[test] + fn kll_readers_of_equal_windows_share_a_site_on_a_real_memo() { + use crate::pass::subexpr::groups::memo_of; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::ir::QueryRoot; + let roots = [0.01, 0.02] + .iter() + .map(|&e| { + QueryRoot::Operator(lower_promql( + "quantile_over_time(0.9, m[5m])", + AccuracyTarget::max_error(e), + )) + }) + .enumerate() + .collect(); + let mut memo = memo_of(&generate_per_computation_candidates(roots).unwrap()); + add_shared_builds(&mut memo); + let (readers, sites) = shared(&memo); + let kll_sites: Vec<&CouplingSite> = sites + .iter() + .filter(|site| { + size_thresholds(&memo, site) + .iter() + .all(|t| matches!(t.params, Some(SketchParams::Kll { .. }))) + }) + .collect(); + assert_eq!(kll_sites.len(), 1); + let kll_readers: Vec<_> = readers + .iter() + .filter(|r| r.coupling_sites.contains(kll_sites[0])) + .collect(); + assert_eq!(kll_readers.len(), 2); + assert_ne!(kll_readers[0].implements, kll_readers[1].implements); + let mut ks: Vec = size_thresholds(&memo, kll_sites[0]) + .iter() + .map(k_of) + .collect(); + ks.sort_unstable(); + assert_eq!(ks.len(), 2); + assert!(ks[0] < ks[1]); + } } diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index deaf4e3ea..4d216df2b 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -194,6 +194,7 @@ fn pass1_candidates( children: groups[g].children.clone(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, @@ -334,6 +335,7 @@ fn target_candidates( children, origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share, requires_exact_children: Vec::new(), rollup_of: None, diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 6057cde38..00162a388 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -90,7 +90,14 @@ pub(crate) struct ImplementationCandidate { pub requires_exact_children: Vec, /// For a rollup, what it rolls up and from where; `None` otherwise. pub rollup_of: Option, + /// The materializations Stage 2 may give it; empty for every one + /// Stage 2 allows. A scan group whose readers all read one window is + /// only built once (`MatAtQuery`); one read through different windows + /// may also be read per reader (`NotMaterialized`) + /// ([`scans::scan_materializations`]). + pub materializations: Vec, } + /// What a [`CandidateOrigin::Rollup`] candidate reads, and what it /// rewrites: the search checks that the finer group's chosen /// implementation is the one the joint rule would roll up from. diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index fd12a0c8e..5f3db3817 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -559,6 +559,7 @@ mod tests { children, origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, diff --git a/crates/plan-selection/src/pass/subexpr/scans.rs b/crates/plan-selection/src/pass/subexpr/scans.rs index 669f565a9..8f887215d 100644 --- a/crates/plan-selection/src/pass/subexpr/scans.rs +++ b/crates/plan-selection/src/pass/subexpr/scans.rs @@ -31,6 +31,7 @@ use asap_logical_asap_optimization::asap_aware_cse::with_union_scan_labels; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; use asap_types::ir::operator::Source; use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; +use asap_types::physical::{Materialization, Medium}; use super::groups::NodeKeys; use super::{CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo}; @@ -214,12 +215,56 @@ pub(crate) fn widened_at_threshold( children: child_groups(&mut Boundary::of(memo, candidate), node), origin: CandidateOrigin::WidenedScans, coupling_sites: candidate.coupling_sites.clone(), + materializations: candidate.materializations.clone(), budget_share: candidate.budget_share, requires_exact_children: candidate.requires_exact_children.clone(), rollup_of: candidate.rollup_of.clone(), })) } +/// The materializations of each scan group's candidates (a user +/// decision): a scan whose readers (its parent groups) all read one window +/// is built once, `MatAtQuery`, with no per-reader choice; a scan read +/// through different windows may be built once or read per reader +/// (`NotMaterialized`). Equal scans are one group (any scan is shareable, +/// a SQL table's keyless rows included), so this is where a scan read by +/// several queries is decided. +pub(crate) fn scan_materializations(memo: &mut Memo) { + for (g, group) in memo.groups.iter().enumerate() { + if !matches!(group.node.non_asap(), Some(NonASAPOp::Scan { .. })) { + continue; + } + let windows: HashSet> = group + .parents + .iter() + .map(|p| window_of(&memo.groups[p.0 as usize].node)) + .collect(); + let allowed = if windows.len() <= 1 { + vec![Materialization::MatAtQuery(Medium::Memory)] + } else { + vec![ + Materialization::NotMaterialized, + Materialization::MatAtQuery(Medium::Memory), + ] + }; + for candidate in &mut memo.candidates[g] { + candidate.materializations = allowed.clone(); + } + } +} + +/// The window a scan's reader reads: a range, lookback or shift with its +/// length; `None` for a reader of every row the scan outputs. +fn window_of(reader: &OperatorNode) -> Option { + match reader.non_asap() { + Some(NonASAPOp::TimeRange { range, kind, .. }) => Some(format!("{kind:?}/{range:?}")), + Some(NonASAPOp::TimeShift { .. }) => { + Some(format!("{:?}", reader.operator.map_children(|_| ()))) + } + _ => None, + } +} + /// Raw horizon: for each source whose raw windows candidates of two or /// more groups read, the source's [`CouplingKind::RawHorizon`] site added /// to each of those candidates. @@ -596,4 +641,75 @@ mod tests { add_raw_horizons(&mut memo); assert!(sites(&memo, CouplingKind::RawHorizon).is_empty()); } + + /// The scan groups of `memo` and their candidates' materializations, + /// after [`scan_materializations`]. + fn scan_options(mut memo: Memo) -> Vec<(Vec, Vec)> { + scan_materializations(&mut memo); + memo.groups + .iter() + .zip(&memo.candidates) + .filter(|(g, _)| matches!(g.node.non_asap(), Some(NonASAPOp::Scan { .. }))) + .map(|(g, c)| (g.queries.clone(), c[0].materializations.clone())) + .collect() + } + + /// Two queries over one 5-minute window: one scan group read through + /// equal windows, built once with no per-reader choice. Windows of 1 + /// and 5 minutes keep both choices. + #[test] + fn a_scan_read_through_one_window_is_built_once() { + let at_query = Materialization::MatAtQuery(Medium::Memory); + assert_eq!( + scan_options(memo(&["sum_over_time(m[5m])", "max_over_time(m[5m])"])), + vec![(vec![0, 1], vec![at_query])] + ); + assert_eq!( + scan_options(memo(&["sum_over_time(m[1m])", "sum_over_time(m[5m])"])), + vec![(vec![0, 1], vec![Materialization::NotMaterialized, at_query])] + ); + } + + /// Two SQL sums over one keyless table grouped differently: the equal + /// table scans are one group (any scan is shareable), read by both + /// aggregates over every row, so it is built once. + #[test] + fn equal_sql_table_scans_are_one_group_built_once() { + use asap_logical_asap_optimization::test_support::{agg, scan_from}; + use asap_types::ir::operator::AggIntent; + use asap_types::ir::schema::{DataType, Field, Schema}; + let table = || { + scan_from( + Source::Table { + table_ref: "t".into(), + }, + Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("value", DataType::Float64, false), + Field::plain("a", DataType::Utf8, true), + Field::plain("b", DataType::Utf8, true), + ], + 0, + vec![], + ), + ) + }; + let sum = || AggIntent::Sum { col: Some(1) }; + let roots = vec![ + QueryRoot::Operator(agg(vec![2], sum(), table())), + QueryRoot::Operator(agg(vec![3], sum(), table())), + ] + .into_iter() + .enumerate() + .collect(); + let memo = memo_of(&generate_per_computation_candidates(roots).unwrap()); + assert_eq!( + scan_options(memo), + vec![( + vec![0, 1], + vec![Materialization::MatAtQuery(Medium::Memory)] + )] + ); + } } diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index c8a5b7d92..97f25bd96 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -918,6 +918,7 @@ pub(crate) fn add_topk_reuse(memo: &mut Memo) { children: vec![ranking_group], origin: CandidateOrigin::TopKReuse, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share: 1, requires_exact_children: vec![input_group], rollup_of: None, @@ -956,6 +957,7 @@ fn group_for( children, origin: CandidateOrigin::TopKReuse, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, @@ -1003,6 +1005,7 @@ fn rewritten_candidate( children, origin, coupling_sites: from.coupling_sites.clone(), + materializations: from.materializations.clone(), budget_share: from.budget_share, requires_exact_children: Vec::new(), rollup_of: None, @@ -1190,6 +1193,7 @@ mod tests { children: Vec::new(), origin: CandidateOrigin::Pass1, coupling_sites: Vec::new(), + materializations: Vec::new(), budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, From 115ea2f4034688314a358129e98c412b52ca8869 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:25:26 +0000 Subject: [PATCH 08/25] plan-selection: sub-expression memo: a coarser grouping reads a finer grouping's pane chain; sketch rollup state resized to a threshold; required sites; one topological order --- .../src/asap_aware_cse.rs | 68 ++- .../src/pass/subexpr/capability.rs | 1 + .../plan-selection/src/pass/subexpr/groups.rs | 2 + crates/plan-selection/src/pass/subexpr/mod.rs | 4 + .../src/pass/subexpr/pricing.rs | 1 + .../plan-selection/src/pass/subexpr/scans.rs | 1 + .../src/pass/subexpr/windows.rs | 462 ++++++++++++++---- 7 files changed, 448 insertions(+), 91 deletions(-) diff --git a/crates/logical-asap-optimization/src/asap_aware_cse.rs b/crates/logical-asap-optimization/src/asap_aware_cse.rs index cf97221b1..417421154 100644 --- a/crates/logical-asap-optimization/src/asap_aware_cse.rs +++ b/crates/logical-asap-optimization/src/asap_aware_cse.rs @@ -1708,8 +1708,9 @@ fn in_panes( /// of `roots`) reading the newest panes of another chain over the same /// input: `pane_count` panes of `pane_width_ms` of family `family` (a size /// that serves every reader of the chain). The merge reads the newest -/// panes covering its own chain's window, which `pane_width_ms` divides -/// and the new chain covers; every guarantee is composed again. +/// panes covering its own window (its newest panes, or its whole chain), +/// which `pane_width_ms` divides and the new chain covers; every guarantee +/// is composed again. pub fn with_pane_chain( roots: &[QueryRoot], merge: &Rc, @@ -1717,7 +1718,12 @@ pub fn with_pane_chain( pane_width_ms: u64, pane_count: u32, ) -> Result, PerComputationCandidateError> { - let Some(ASAPOp::SummaryMerge { children, .. }) = merge.asap() else { + let Some(ASAPOp::SummaryMerge { + children, + newest_panes: own_newest, + .. + }) = merge.asap() + else { unreachable!("a merge over a pane chain"); }; let own = &children[0]; @@ -1734,7 +1740,8 @@ pub fn with_pane_chain( else { unreachable!("a merge over a pane chain"); }; - let window_ms = own_width * u64::from(*own_count); + // The merge's window: its newest panes, or the whole chain. + let window_ms = own_width * u64::from(own_newest.unwrap_or(*own_count)); let newest = u32::try_from(window_ms / pane_width_ms).expect("a reader's window is at most the chain's"); let mut chain = OperatorNode::new(Operator::ASAP(ASAPOp::PaneChain { @@ -1778,6 +1785,59 @@ pub fn with_pane_chain( recompose_guarantees(&out) } +/// `roots` with the summary build `build` (a `SummaryAgg` of `roots`) +/// built at family `family` instead (the same algorithm at another size), +/// its output names kept; every guarantee is composed again. +pub fn with_resized_build( + roots: &[QueryRoot], + build: &Rc, + family: FieldDataType, +) -> Result, PerComputationCandidateError> { + let Some(ASAPOp::SummaryAgg { + child, + input, + reduction, + grouping, + filter, + .. + }) = build.asap() + else { + unreachable!("a summary build"); + }; + let mut resized = OperatorNode::new(Operator::ASAP(ASAPOp::SummaryAgg { + child: Rc::clone(child), + family, + input: input.clone(), + reduction: reduction.clone(), + grouping: grouping.clone(), + filter: filter.clone(), + }))?; + for (field, kept) in resized.schema.fields.iter_mut().zip(&build.schema.fields) { + field.name = kept.name.clone(); + field.table = kept.table.clone(); + } + let replaced = HashMap::from([(Rc::as_ptr(build), Rc::new(resized))]); + let mut rebuilt = HashMap::new(); + let mut out = Vec::new(); + for root in roots { + out.push(match root { + QueryRoot::Operator(node) => { + QueryRoot::Operator(rebuild(node, &replaced, &mut rebuilt)?) + } + QueryRoot::Scalar(expr) => { + let mut nodes = HashMap::new(); + for node in expr.operator_refs() { + nodes.insert(Rc::as_ptr(node), rebuild(node, &replaced, &mut rebuilt)?); + } + QueryRoot::Scalar( + expr.map_operator_refs(&mut |node| nodes[&Rc::as_ptr(node)].clone()), + ) + } + }); + } + recompose_guarantees(&out) +} + /// The [`ASAPOp::PaneChain`] of `candidate`'s panes, with the build's /// output names. fn pane_chain(candidate: &PaneCandidate) -> Result, PerComputationCandidateError> { diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index 3ebe2546a..1ec7c5443 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -318,6 +318,7 @@ mod tests { budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), }]); } add_shared_builds(&mut memo); diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index deaf4e3ea..e3f30bb18 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -197,6 +197,7 @@ fn pass1_candidates( budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), }; // Each query's most approximate inputs: its targets that offer a // summary, each one sketch state at most. @@ -337,6 +338,7 @@ fn target_candidates( budget_share, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), }); } // Mixed-radix over this group's targets only. diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 6057cde38..07ced263f 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -90,6 +90,10 @@ pub(crate) struct ImplementationCandidate { pub requires_exact_children: Vec, /// For a rollup, what it rolls up and from where; `None` otherwise. pub rollup_of: Option, + /// Sites other groups' chosen implementations must carry for this one + /// to be valid: a sketch rollup's finer group exposing its state, a + /// coarser grouping's host group building the pane chain it reads. + pub requires_sites: Vec<(GroupId, CouplingSite)>, } /// What a [`CandidateOrigin::Rollup`] candidate reads, and what it /// rewrites: the search checks that the finer group's chosen diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index fd12a0c8e..8154c193d 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -562,6 +562,7 @@ mod tests { budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), }; let memo = Memo { groups: vec![group(&below, vec![]), group(top, vec![GroupId(0)])], diff --git a/crates/plan-selection/src/pass/subexpr/scans.rs b/crates/plan-selection/src/pass/subexpr/scans.rs index 669f565a9..d3b3c1cf3 100644 --- a/crates/plan-selection/src/pass/subexpr/scans.rs +++ b/crates/plan-selection/src/pass/subexpr/scans.rs @@ -217,6 +217,7 @@ pub(crate) fn widened_at_threshold( budget_share: candidate.budget_share, requires_exact_children: candidate.requires_exact_children.clone(), rollup_of: candidate.rollup_of.clone(), + requires_sites: candidate.requires_sites.clone(), })) } diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index c8a5b7d92..5bfe9c8a8 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -31,7 +31,7 @@ use std::rc::Rc; use asap_logical_asap_optimization::accuracy::estimators::size_for_all_readers; use asap_logical_asap_optimization::asap_aware_cse::{ pane_candidates, rollup_candidates, rollup_pass2_candidates, topk_reuse_candidates, - window_composition_pass2_candidates, with_pane_chain, RollupCandidate, + window_composition_pass2_candidates, with_pane_chain, with_resized_build, RollupCandidate, }; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; use asap_types::ir::flat::flatten; @@ -77,10 +77,12 @@ pub(crate) struct ChainReader { pub merge: Rc, } -/// Every transformation of this module, in order: window composition, -/// rollup, exact top-k reuse. `entries[q]` is workload query `q`'s entry. +/// Every transformation of this module, in order: window composition, a +/// coarser grouping reading a finer grouping's chain, rollup, exact top-k +/// reuse. `entries[q]` is workload query `q`'s entry. pub(crate) fn add_pass2_candidates(memo: &mut Memo, entries: &[QueryWorkloadEntry]) { add_window_composition(memo, entries); + add_rollup_windows(memo, entries); add_rollup(memo); add_topk_reuse(memo); } @@ -187,7 +189,12 @@ pub(crate) fn chain_readers(node: &Rc) -> Vec { if !seen.insert(Rc::as_ptr(&merge)) { continue; } - let Some(ASAPOp::SummaryMerge { children, .. }) = merge.asap() else { + let Some(ASAPOp::SummaryMerge { + children, + newest_panes, + .. + }) = merge.asap() + else { continue; }; let [chain] = &children[..] else { @@ -229,7 +236,7 @@ pub(crate) fn chain_readers(node: &Rc) -> Vec { readers.push(ChainReader { key, pane_width_ms: *pane_width_ms, - window_ms: pane_width_ms * u64::from(*pane_count), + window_ms: pane_width_ms * u64::from(newest_panes.unwrap_or(*pane_count)), family: family.clone(), merge: Rc::clone(&merge), }); @@ -372,6 +379,184 @@ pub(crate) fn read_at_threshold( Ok(Rc::clone(node)) } +/// `candidate`'s node (a sketch rollup reading the state of `site`, or +/// the finer group's candidate exposing it) with that build resized to +/// `threshold` ([`state_size_thresholds`], which the candidate may read: +/// [`may_read_state`]); every guarantee composed again. +pub(crate) fn resize_state_at( + candidate: &ImplementationCandidate, + site: &CouplingSite, + threshold: &FieldDataType, +) -> Result, PerComputationCandidateError> { + let build = match candidate + .rollup_of + .as_ref() + .and_then(|r| r.finer_state.as_ref()) + { + Some(state) => Some(Rc::clone(state)), + None => OperatorNode::reachable(&candidate.node) + .into_iter() + .find(|n| is_sketch_build(n) && state_key(n) == site.key), + }; + let Some(build) = build else { + return Ok(Rc::clone(&candidate.node)); + }; + let roots = [QueryRoot::Operator(Rc::clone(&candidate.node))]; + let resized = with_resized_build(&roots, &build, threshold.clone())?; + let [QueryRoot::Operator(node)] = &resized[..] else { + unreachable!("an operator root"); + }; + Ok(Rc::clone(node)) +} + +/// A coarser grouping reading a finer grouping's pane chain (the +/// window-composition rule's rollup groups): for each Pass 1 (or +/// shared-build) candidate of a group with a sketch build over a SQL +/// table's rows whose readers keep a merged guarantee, and each such +/// candidate of another group building the same family, input, grouping +/// and filter over the same rows grouped by more keys over a window at +/// least as long, a [`CandidateOrigin::Panes`] candidate of the coarser +/// group that regroups the newest panes of the finer build's chain +/// (`window_composition_pass2_candidates` on the two, the coarser group's +/// pane width the finer build's grid). It carries the finer chain's +/// [`CouplingKind::PaneChain`] site, whose threshold sizes the chain for +/// every reader, and reads the finer group, which must build that chain: +/// the finer group gets a copy of its own pane candidate with the site +/// ([`ImplementationCandidate::requires_sites`]). +pub(crate) fn add_rollup_windows(memo: &mut Memo, entries: &[QueryWorkloadEntry]) { + let group_of = group_of_node(&memo.groups); + let grids: Vec> = memo + .groups + .iter() + .map(|group| group_grid_ms(group, entries)) + .collect(); + let base = |c: &ImplementationCandidate| { + matches!( + c.origin, + CandidateOrigin::Pass1 | CandidateOrigin::SharedBuild + ) + }; + let mut added: Vec<(usize, ImplementationCandidate)> = Vec::new(); + // Per finer group, the chain keys its pane candidates must expose. + let mut hosts: Vec<(usize, u64)> = Vec::new(); + for (coarse_group, coarse_candidates) in memo.candidates.iter().enumerate() { + let Some(coarse_grid) = grids[coarse_group] else { + continue; + }; + for coarse in coarse_candidates.iter().filter(|c| base(c)) { + let coarse_root = [QueryRoot::Operator(Rc::clone(&coarse.node))]; + // The coarser build's own pane width. + let own = pane_candidates(&coarse_root, &[Some(coarse_grid)]); + let [own] = &own[..] else { + continue; + }; + for (finer_group, finer_candidates) in memo.candidates.iter().enumerate() { + if finer_group == coarse_group || grids[finer_group].is_none() { + continue; + } + for finer in finer_candidates.iter().filter(|c| base(c)) { + let roots = [ + QueryRoot::Operator(Rc::clone(&coarse.node)), + QueryRoot::Operator(Rc::clone(&finer.node)), + ]; + let grids_of_pair = [Some(coarse_grid), Some(own.pane_width_ms)]; + let panes = pane_candidates(&roots, &grids_of_pair); + if panes.len() != 2 { + continue; + } + let Ok(pass2_candidates) = window_composition_pass2_candidates(&roots, &panes) + else { + continue; + }; + for pass2_candidate in pass2_candidates { + if pass2_candidate.pane_candidates.len() != 2 + || pass2_candidate.shared_chains.is_empty() + { + continue; + } + let Ok(rewritten) = pass2_candidate.roots else { + continue; + }; + let [QueryRoot::Operator(node), QueryRoot::Operator(finer_node)] = + &rewritten[..] + else { + continue; + }; + // The finer build's chain, which the coarser one + // regroups. + let [finer_chain] = &chain_readers(finer_node)[..] else { + continue; + }; + let [coarse_read] = &chain_readers(node)[..] else { + continue; + }; + if coarse_read.key != finer_chain.key || !regroups(node, &coarse_read.merge) + { + continue; + } + let site = CouplingSite { + key: finer_chain.key, + kind: CouplingKind::PaneChain, + }; + let mut candidate = rewritten_candidate( + &group_of, + coarse, + Rc::clone(node), + CandidateOrigin::Panes, + ); + let finer_id = GroupId(finer_group as u32); + if !candidate.children.contains(&finer_id) { + candidate.children.push(finer_id); + } + candidate.coupling_sites.push(site); + candidate.requires_sites.push((finer_id, site)); + added.push((coarse_group, candidate)); + hosts.push((finer_group, site.key)); + } + } + } + } + } + for (g, candidate) in added { + push_new(memo, g, candidate); + } + // Each finer group's own pane candidates building a host chain, with + // its site. + hosts.sort_unstable(); + hosts.dedup(); + for (g, key) in hosts { + let site = CouplingSite { + key, + kind: CouplingKind::PaneChain, + }; + let exposing: Vec = memo.candidates[g] + .iter() + .filter(|c| c.origin == CandidateOrigin::Panes && !c.coupling_sites.contains(&site)) + .filter(|c| c.requires_sites.is_empty()) + .filter(|c| chain_readers(&c.node).iter().any(|r| r.key == key)) + .cloned() + .map(|mut c| { + c.coupling_sites + .retain(|s| s.kind != CouplingKind::PaneChain); + c.coupling_sites.push(site); + c + }) + .collect(); + for candidate in exposing { + push_new(memo, g, candidate); + } + } +} + +/// Whether `merge` (over a pane chain) is read by a regrouping merge in +/// `node`: a coarser grouping's read of a finer chain. +fn regroups(node: &Rc, merge: &Rc) -> bool { + OperatorNode::reachable(node).iter().any(|n| { + matches!(n.asap(), Some(ASAPOp::SummaryMerge { children, group_by: Some(_), .. }) + if Rc::ptr_eq(&children[0], merge)) + }) +} + /// Rollup: each coarse build of a group's Pass 1 (or shared-build) /// candidate that another group's candidate can answer /// (`rollup_candidates` over the two) gets a [`CandidateOrigin::Rollup`] @@ -388,8 +573,6 @@ pub(crate) fn read_at_threshold( /// rule rolls a coarse build up only from the finest build over its rows, /// and every coarse build of one finest build together /// ([`requirements_met`] checks both). -/// -/// Ends with the groups in topological order ([`in_topological_order`]). pub(crate) fn add_rollup(memo: &mut Memo) { let group_of = group_of_node(&memo.groups); let mut added: Vec<(usize, ImplementationCandidate)> = Vec::new(); @@ -449,7 +632,6 @@ pub(crate) fn add_rollup(memo: &mut Memo) { for (g, candidate) in added { push_new(memo, g, candidate); } - in_topological_order(memo); } /// Marks a candidate [`sketch_rollup`] adds to the finer group. @@ -560,6 +742,7 @@ fn sketch_rollup( coarse_candidate: Rc::clone(&coarse.node), coarse_build, }); + candidate.requires_sites.push((finer_id, site)); let mut exposed = finer.clone(); exposed.coupling_sites.push(site); vec![(0, candidate), (FINER, exposed)] @@ -681,7 +864,8 @@ fn same_size(a: &FieldDataType, b: &FieldDataType) -> bool { /// /// - never an exact top-k reuse with panes ([`combinable`]); /// - every group in a candidate's `requires_exact_children` chose an exact -/// implementation ([`is_exact`]); +/// implementation ([`is_exact`]), and every group in its +/// `requires_sites` one carrying that site; /// - every rollup reads the build the joint rule would roll its coarse /// build up from, over the chosen plan with each rollup undone: an exact /// accumulator's coarsest finer build, a sketch's finest, in its finer @@ -700,6 +884,11 @@ pub(crate) fn requirements_met(chosen: &[(GroupId, &ImplementationCandidate)]) - return false; } } + for (g, site) in &candidate.requires_sites { + if !of(*g).is_some_and(|c| c.coupling_sites.contains(site)) { + return false; + } + } } let rolled: Vec<&RollupRead> = chosen .iter() @@ -775,72 +964,6 @@ fn reached(node: &Rc) -> HashSet<*const OperatorNode> { .collect() } -/// Renumber `memo`'s groups so that every group comes after the groups it -/// reads, through its node or any candidate (children and required -/// groups), keeping the current order where it allows. Unchanged when -/// those reads form a cycle. -pub(crate) fn in_topological_order(memo: &mut Memo) { - let n = memo.groups.len(); - let mut reads: Vec> = vec![BTreeSet::new(); n]; - for (g, group) in memo.groups.iter().enumerate() { - reads[g].extend(group.children.iter().map(|c| c.0 as usize)); - for candidate in &memo.candidates[g] { - reads[g].extend(candidate.children.iter().map(|c| c.0 as usize)); - reads[g].extend( - candidate - .requires_exact_children - .iter() - .map(|c| c.0 as usize), - ); - reads[g].extend(candidate.rollup_of.iter().map(|r| r.finer_group.0 as usize)); - } - reads[g].remove(&g); - } - // Repeatedly take the first group whose reads are all placed. - let mut placed = vec![false; n]; - let mut order: Vec = Vec::with_capacity(n); - while order.len() < n { - let Some(next) = (0..n).find(|&g| !placed[g] && reads[g].iter().all(|&r| placed[r])) else { - return; - }; - placed[next] = true; - order.push(next); - } - if order.iter().enumerate().all(|(new, &old)| new == old) { - return; - } - let mut new_of = vec![0u32; n]; - for (new, &old) in order.iter().enumerate() { - new_of[old] = new as u32; - } - let map = |id: &mut GroupId| id.0 = new_of[id.0 as usize]; - let mut groups: Vec> = std::mem::take(&mut memo.groups) - .into_iter() - .map(Some) - .collect(); - let mut candidates: Vec>> = - std::mem::take(&mut memo.candidates) - .into_iter() - .map(Some) - .collect(); - for &old in &order { - let mut group = groups[old].take().expect("each group once"); - group.children.iter_mut().for_each(map); - group.parents.iter_mut().for_each(map); - memo.groups.push(group); - let mut own = candidates[old].take().expect("each group once"); - for candidate in &mut own { - candidate.implements.iter_mut().for_each(map); - candidate.children.iter_mut().for_each(map); - candidate.requires_exact_children.iter_mut().for_each(map); - if let Some(read) = &mut candidate.rollup_of { - map(&mut read.finer_group); - } - } - memo.candidates.push(own); - } -} - /// Exact top-k reuse: the groups that are top-k computations of one input /// group and partition with at least two different `k` /// (`topk_reuse_candidates` over the groups' nodes) each get a [`CandidateOrigin::TopKReuse`] candidate @@ -849,8 +972,9 @@ pub(crate) fn in_topological_order(memo: &mut Memo) { /// are added when the memo has no equal ones, each with its one candidate /// (also of origin `TopKReuse`). The rule reads an exact input: the input /// group's chosen implementation must be exact -/// ([`ImplementationCandidate::requires_exact_children`]). Ends with the -/// groups in topological order ([`in_topological_order`]). +/// ([`ImplementationCandidate::requires_exact_children`]).. +/// Appended groups come after their readers; [`Memo::topological_order`] +/// orders them. pub(crate) fn add_topk_reuse(memo: &mut Memo) { // The groups' nodes, each input marked exact: the rule reads an exact // input, which `requires_exact_children` asks of the input group. @@ -921,11 +1045,11 @@ pub(crate) fn add_topk_reuse(memo: &mut Memo) { budget_share: 1, requires_exact_children: vec![input_group], rollup_of: None, + requires_sites: Vec::new(), }; push_new(memo, g, reuse); } } - in_topological_order(memo); } /// The group whose node equals `node`, or a new one over `children`, read @@ -959,6 +1083,7 @@ fn group_for( budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), }]); id } @@ -1006,6 +1131,7 @@ fn rewritten_candidate( budget_share: from.budget_share, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), } } @@ -1165,13 +1291,25 @@ mod tests { fn memo_with( queries: &[(&str, AccuracyTarget)], intervals_ms: &[u32], + ) -> (Memo, Vec) { + let nodes = queries + .iter() + .map(|(query, target)| { + with_promql_series_identity(&lower_promql(query, target.clone())).unwrap() + }) + .collect(); + memo_of_nodes(nodes, intervals_ms) + } + + /// [`memo`] of query roots `nodes`, already lowered. + fn memo_of_nodes( + nodes: Vec>, + intervals_ms: &[u32], ) -> (Memo, Vec) { let entries: Vec = intervals_ms.iter().map(|&ms| repeating(ms)).collect(); let mut memo = Memo::default(); - for (q, (query, target)) in queries.iter().enumerate() { - let node = lower_promql(query, target.clone()); - let node = with_promql_series_identity(&node).unwrap(); + for (q, node) in nodes.into_iter().enumerate() { let inventory = generate_per_computation_candidates(vec![( q, QueryRoot::Operator(Rc::clone(&node)), @@ -1193,6 +1331,7 @@ mod tests { budget_share: 1, requires_exact_children: Vec::new(), rollup_of: None, + requires_sites: Vec::new(), }) .collect(); memo.groups.push(MemoGroup { @@ -1365,7 +1504,9 @@ mod tests { .collect(); assert_eq!(topk.len(), 2); for g in topk { - assert!(ranking < g, "children before parents"); + let order = memo.topological_order().unwrap(); + let at = |g: usize| order.iter().position(|x| x.0 as usize == g).unwrap(); + assert!(at(ranking) < at(g), "children before parents"); let reuse = memo.candidates[g] .iter() .find(|c| c.origin == CandidateOrigin::TopKReuse) @@ -1377,7 +1518,7 @@ mod tests { )); let input = memo.groups[g].children[0]; assert_eq!(reuse.requires_exact_children, vec![input]); - assert!(input.0 < ranking as u32); + assert!(at(input.0 as usize) < at(ranking)); let panes = ImplementationCandidate { origin: CandidateOrigin::Panes, ..reuse.clone() @@ -1524,10 +1665,13 @@ mod tests { &[60_000, 60_000, 60_000], ); add_rollup(&mut memo); - // Rollups reorder the groups finer first: query `q`'s group. - let of = |q: usize| memo.groups.iter().position(|g| g.roots_of == [q]).unwrap(); - let (job, instance, zone) = (of(0), of(1), of(2)); - assert!(zone < instance && instance < job, "children before parents"); + let (job, instance, zone) = (0, 1, 2); + let order = memo.topological_order().unwrap(); + let at = |g: usize| order.iter().position(|x| x.0 as usize == g).unwrap(); + assert!( + at(zone) < at(instance) && at(instance) < at(job), + "finer first" + ); let rolled_from = |g: usize, from: usize| { memo.candidates[g] .iter() @@ -1586,4 +1730,148 @@ mod tests { (id(instance), instance_exposed), ])); } + + /// The workload of #878 over a SQL table `metrics(ts, city, value)`, + /// every 10 s at ε = 0.01: the global median over 5 minutes and the + /// per-city p90 over 20 minutes. The global group gets a candidate that + /// regroups the newest 30 of the per-city chain's 120 panes of 10 s, + /// at the per-city chain's site, which the per-city group must build; + /// one threshold (10 s, 20 minutes) serves both. + #[test] + fn a_coarser_grouping_reads_the_finer_groupings_chain() { + use asap_logical_asap_optimization::test_support::{agg, scan_from, time_range}; + use asap_types::ir::operator::{AggIntent, Source}; + use asap_types::ir::schema::{DataType, Field, Schema}; + use std::time::Duration; + let scan = || { + scan_from( + Source::Table { + table_ref: "metrics".into(), + }, + Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("city", DataType::Utf8, true), + Field::plain("value", DataType::Float64, true), + ], + 0, + vec![], + ), + ) + }; + let quantile = |q: f64| AggIntent::Quantile { + col: Some(2), + q, + accuracy: AccuracyTarget::max_error(0.01), + }; + let minutes = |m: u64| Duration::from_secs(60 * m); + let global = agg(vec![], quantile(0.5), time_range(minutes(5), scan())); + let per_city = agg(vec![1], quantile(0.9), time_range(minutes(20), scan())); + let (mut memo, entries) = memo_of_nodes(vec![global, per_city], &[10_000, 10_000]); + add_window_composition(&mut memo, &entries); + add_rollup_windows(&mut memo, &entries); + let reads: Vec<&ImplementationCandidate> = memo.candidates[0] + .iter() + .filter(|c| !c.requires_sites.is_empty()) + .collect(); + assert!(!reads.is_empty()); + for read in &reads { + assert_eq!(read.origin, CandidateOrigin::Panes); + let (host, site) = read.requires_sites[0]; + assert_eq!(host, GroupId(1)); + assert!(read.children.contains(&GroupId(1))); + assert!(read.coupling_sites.contains(&site)); + let [reader] = &chain_readers(&read.node)[..] else { + panic!("one chain read"); + }; + assert!(regroups(&read.node, &reader.merge)); + assert_eq!((reader.pane_width_ms, reader.window_ms), (10_000, 300_000)); + let host_candidate = memo.candidates[1] + .iter() + .find(|c| c.coupling_sites.contains(&site) && chain_readers(&c.node).len() == 1) + .expect("the per-city group builds the chain"); + assert!(requirements_met(&[ + (GroupId(0), read), + (GroupId(1), host_candidate) + ])); + let own = memo.candidates[1] + .iter() + .find(|c| c.origin == CandidateOrigin::Pass1) + .unwrap(); + assert!(!requirements_met(&[(GroupId(0), read), (GroupId(1), own)])); + let threshold = chain_thresholds(&memo, &site) + .into_iter() + .find(|t| { + t.pane_width_ms == 10_000 + && t.longest_window_ms == 1_200_000 + && may_read(t, reader) + && chain_readers(&host_candidate.node) + .iter() + .all(|r| may_read(t, r)) + }) + .expect("one threshold serves both"); + let at = read_at_threshold(read, &site, &threshold).unwrap(); + let [at_reader] = &chain_readers(&at)[..] else { + panic!("one chain read"); + }; + assert_eq!(at_reader.window_ms, 300_000); + assert!(regroups(&at, &at_reader.merge)); + } + } + + /// A sketch rollup and its finer group's exposing candidate, resized to + /// a size threshold, both build the finer KLL at that size. + #[test] + fn resizing_the_rolled_up_state_to_a_threshold() { + let (mut memo, _) = memo_with( + &[ + ( + "quantile by (host, job) (0.5, data)", + AccuracyTarget::max_error(0.01), + ), + ( + "quantile by (job) (0.5, data)", + AccuracyTarget::max_error(0.001), + ), + ], + &[60_000, 60_000], + ); + add_rollup(&mut memo); + let rollup = memo.candidates[1] + .iter() + .find(|c| { + c.rollup_of + .as_ref() + .is_some_and(|r| r.finer_state.is_some()) + }) + .unwrap(); + let site = rollup.requires_sites[0].1; + let exposed = memo.candidates[0] + .iter() + .find(|c| c.coupling_sites.contains(&site)) + .unwrap(); + let sizes = |node: &Rc| -> Vec { + OperatorNode::reachable(node) + .iter() + .filter(|n| is_sketch_build(n) && state_key(n) == site.key) + .map(|n| match n.asap() { + Some(ASAPOp::SummaryAgg { family, .. }) => format!("{family:?}"), + _ => unreachable!(), + }) + .collect() + }; + let mut resized = 0; + for threshold in state_size_thresholds(&memo, &site) { + let expected = vec![format!("{threshold:?}")]; + for candidate in [exposed, rollup] { + if may_read_state(&threshold, candidate, &site) { + let node = resize_state_at(candidate, &site, &threshold).unwrap(); + assert_eq!(sizes(&node), expected); + resized += 1; + } + } + } + // The larger threshold serves both, the smaller only the finer group. + assert_eq!(resized, 3); + } } From 6bf7c2699df365e7b19013e5d4e4034e6bba803b Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:27:17 +0000 Subject: [PATCH 09/25] plan-selection: sub-expression memo Pass 1 parity test; top-k roots keep their own columns A lib test composes every whole-workload Pass 1 choice of the random workloads of joint_search_equivalence and checks each query's plan is a combination of the memo's per-group candidates. It found the fused top-k heap candidates of an ungrouped top-k root dropped because their output columns differ from the logical node's; a candidate keeps its own schema when its columns differ, and every candidate of a group a parent reads outputs the group's schema. Equal SQL table scans are one group. --- .../plan-selection/src/pass/subexpr/groups.rs | 301 +++++++++++++++++- 1 file changed, 291 insertions(+), 10 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index deaf4e3ea..1bd93640a 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -15,7 +15,8 @@ //! equal to one of the group's descendant groups, which are its children. //! The private 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. +//! 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 query whose value combines `n` approximate inputs splits its budget //! into `n` equal shares, so a summary's size depends on its readers: each @@ -309,16 +310,15 @@ fn target_candidates( &mut HashMap::new(), ); // A parent reads the candidate in the group's place, so it - // outputs the group's schema: the same columns, with the - // group's metadata (composition may derive a closed label - // set where the frontend's node leaves it open). + // 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].node; - if node.schema.fields != group_node.schema.fields - || node.schema.time_index != group_node.schema.time_index - { - continue; - } - let node = if node.schema == group_node.schema { + 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 { node } else { Rc::new(OperatorNode { @@ -854,4 +854,285 @@ mod tests { .push(GroupId(quantile as u32)); assert_eq!(memo.topological_order(), None); } + + /// The workload generator of `joint_search_equivalence`: a seeded + /// xorshift generator. + struct Rng(u64); + + impl Rng { + fn below(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 % n as u64) as usize + } + + fn pick<'a, T>(&mut self, items: &'a [T]) -> &'a T { + &items[self.below(items.len())] + } + } + + /// One random PromQL query of `joint_search_equivalence`. + fn random_query(rng: &mut Rng, gauges: &[&str]) -> String { + let m = rng.pick(gauges); + let group = rng.pick(&["job", "job, instance", "instance"]); + let window = rng.pick(&["1m", "5m", "10m"]); + let phi = rng.pick(&["0.5", "0.9", "0.99"]); + let k = rng.pick(&["3", "5"]); + match rng.below(10) { + 0 => format!("sum by ({group}) (sum_over_time({m}[{window}]))"), + 1 => format!("avg by ({group}) (quantile_over_time({phi}, {m}[{window}]))"), + 2 => format!("quantile_over_time({phi}, {m}[{window}])"), + 3 => format!("topk by ({group}) ({k}, sum_over_time({m}[{window}]))"), + 4 => format!("sum by ({group}) (rate(c_total[{window}]))"), + 5 => format!("max by ({group}) (max_over_time({m}[{window}]))"), + 6 => format!("count_over_time({m}[{window}])"), + 7 => format!("sum by ({group}) (increase(c_total[{window}]))"), + 8 => format!("stddev_over_time({m}[{window}])"), + _ => format!("topk({k}, sum by ({group}) (rate(c_total[{window}])))"), + } + } + + /// The roots of `joint_search_equivalence`'s random workload `seed` + /// over `gauges`, lowered as the planner lowers PromQL (with series + /// identity). + fn random_roots(seed: u64, gauges: &[&str]) -> Vec { + use asap_types::ir::properties::AccuracyMetric; + use asap_types::workload::{ + DataArrival, DataWorkload, DurationMs, MetricType, PlanningWorkload, Predictability, + Query, QueryLanguage, QueryRequirements, QueryWorkload, RepeatedDemand, RepeatingEntry, + RepetitionInterval, SourceFacts, + }; + let mut rng = Rng(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1); + let queries = 2 + rng.below(3); + let entries = (0..queries) + .map(|_| RepeatingEntry { + query: Query(random_query(&mut rng, gauges)), + demand: RepeatedDemand::FixedInterval(RepetitionInterval( + *rng.pick(&[60_000u32, 300_000]), + )), + requirements: QueryRequirements { + accuracy: rng + .pick(&[ + AccuracyTarget::exact(), + AccuracyTarget::max_error(0.05), + AccuracyTarget::max_error(0.01).failure_probability(0.01), + AccuracyTarget::max_error(0.05).accuracy_metric(AccuracyMetric::Rank), + ]) + .clone(), + max_latency_ms: None, + }, + predictability: Predictability::Predictable { known_at: None }, + range_evaluation: None, + }) + .collect(); + let source = |metric_type| SourceFacts { + series_count: Some(1_000), + metric_type: Some(metric_type), + sample_interval: Some(DurationMs(15_000)), + ..SourceFacts::default() + }; + let workload = PlanningWorkload { + query_workload: QueryWorkload { + language: QueryLanguage::PromQL, + query_batch: None, + repeating_queries: Some(entries), + }, + data_workload: Some(DataWorkload { + arrival: DataArrival::Ingesting { + rows_per_second: 2_000.0 / 15.0, + }, + sources: gauges + .iter() + .map(|g| (g.to_string(), source(MetricType::Gauge))) + .chain([("c_total".to_string(), source(MetricType::Counter))]) + .collect(), + }), + }; + asap_frontend_promql::lower_promql_query_workload(&workload) + .unwrap() + .into_iter() + .map(|root| match root { + QueryRoot::Operator(node) => QueryRoot::Operator( + asap_types::ir::schema::schema_support::with_promql_series_identity(&node) + .unwrap(), + ), + scalar => scalar, + }) + .collect() + } + + /// Whether a composed sub-DAG is some combination of the memo's + /// per-group candidates: a candidate's private nodes equal the composed + /// nodes operator by operator (schemas and guarantees aside, which + /// composition derives), and each child group's node is matched by a + /// candidate of that group in turn. + struct Combination<'a> { + memo: &'a Memo, + matched: HashMap<(usize, *const OperatorNode), bool>, + own: HashMap<*const OperatorNode, (Rc, serde_json::Value)>, + } + + impl Combination<'_> { + fn own_fields(&mut self, node: &Rc) -> serde_json::Value { + self.own + .entry(Rc::as_ptr(node)) + .or_insert_with(|| { + let operator = node.operator.map_children(|_| ()); + (Rc::clone(node), serde_json::to_value(&operator).unwrap()) + }) + .1 + .clone() + } + + fn group(&mut self, g: usize, composed: &Rc) -> bool { + let key = (g, Rc::as_ptr(composed)); + if let Some(&done) = self.matched.get(&key) { + return done; + } + self.matched.insert(key, false); + let memo = self.memo; + let found = memo.candidates[g] + .iter() + .any(|c| self.private(&c.node, composed, &c.children)); + self.matched.insert(key, found); + found + } + + fn private( + &mut self, + node: &Rc, + composed: &Rc, + children: &[GroupId], + ) -> bool { + let memo = self.memo; + if let Some(child) = children + .iter() + .find(|c| Rc::ptr_eq(&memo.groups[c.0 as usize].node, node)) + { + return self.group(child.0 as usize, composed); + } + let (ours, theirs) = (node.children(), composed.children()); + ours.len() == theirs.len() + && self.own_fields(node) == self.own_fields(composed) + && ours + .iter() + .zip(&theirs) + .all(|(a, b)| self.private(a, b, children)) + } + } + + /// For the random workloads of `joint_search_equivalence` (over one + /// gauge and over two), every whole-workload Pass 1 choice composes, + /// for each query, a plan that is a combination of the memo's per-group + /// candidates: the memo misses no Pass 1 candidate. Every candidate of + /// a group a parent reads outputs the group's schema. + #[test] + fn every_composed_pass1_candidate_is_a_combination_of_group_candidates() { + let mut checked = 0; + let mut missing = Vec::new(); + let mut other_schema = Vec::new(); + for gauges in [&["m"][..], &["m", "n"]] { + for seed in 1..=60u64 { + let roots = random_roots(seed, gauges); + let inventory = generate_per_computation_candidates( + roots.iter().cloned().enumerate().collect(), + ) + .unwrap(); + if inventory.combinations() > 400 { + continue; + } + let memo = memo_of(&inventory); + for (g, group) in memo.groups.iter().enumerate() { + for candidate in &memo.candidates[g] { + if !group.parents.is_empty() && candidate.node.schema != group.node.schema { + other_schema.push(format!("{gauges:?} seed {seed} g{g}")); + } + } + } + let root_group = |q: usize| { + memo.groups + .iter() + .position(|group| group.roots_of.contains(&q)) + .unwrap() + }; + let mut combination = Combination { + memo: &memo, + matched: HashMap::new(), + own: HashMap::new(), + }; + let mut choice = vec![0; inventory.targets.len()]; + loop { + if !inventory.absorbed_choice_varies(&choice) { + if let Ok(composed) = compose_logical_candidate(&inventory, &choice) { + for (q, (_, root)) in composed.iter().enumerate() { + let QueryRoot::Operator(node) = root else { + continue; + }; + checked += 1; + if !combination.group(root_group(q), node) { + missing.push(format!("{gauges:?} seed {seed} q{q} {choice:?}")); + } + } + } + } + let mut advanced = false; + for t in (0..choice.len()).rev() { + choice[t] += 1; + if choice[t] < inventory.targets[t].choice_count() { + advanced = true; + break; + } + choice[t] = 0; + } + if !advanced { + break; + } + } + } + } + assert!(checked > 100, "{checked} plans checked"); + assert!( + missing.is_empty(), + "{} of {checked} missing: {missing:#?}", + missing.len() + ); + } + + /// Two SQL queries over one table, each lowered with its own scan, read + /// one scan group: equal scans are shared, a table's keyless rows + /// included. Their aggregates, grouped differently, stay two groups. + #[test] + fn equal_sql_table_scans_are_one_group() { + use asap_logical_asap_optimization::test_support::{agg, scan_from}; + use asap_types::ir::operator::{AggIntent, Source}; + use asap_types::ir::schema::{DataType, Field, Schema}; + let table = || { + scan_from( + Source::Table { + table_ref: "t".into(), + }, + Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("value", DataType::Float64, false), + Field::plain("a", DataType::Utf8, true), + Field::plain("b", DataType::Utf8, true), + ], + 0, + vec![], + ), + ) + }; + let sum = || AggIntent::Sum { col: Some(1) }; + let roots = vec![ + (0, QueryRoot::Operator(agg(vec![2], sum(), table()))), + (1, QueryRoot::Operator(agg(vec![3], sum(), table()))), + ]; + let memo = memo_of(&generate_per_computation_candidates(roots).unwrap()); + assert_eq!(memo.groups.len(), 3); + assert_eq!(memo.groups[0].node.operator.kind_name(), "Scan"); + assert_eq!(memo.groups[0].queries, vec![0, 1]); + assert_eq!(memo.groups[0].parents, vec![GroupId(1), GroupId(2)]); + } } From 10112d3963a9ba71be7d31b6e26bc20c3db45fe8 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:27:26 +0000 Subject: [PATCH 10/25] memo: reference brute-force search over the memo, behind memo_groups; memo equivalence test --- .../tests/joint_search_equivalence.rs | 84 +++++ .../src/asap_aware_cse.rs | 10 + .../src/pass/stage1_to_stage3.rs | 36 ++ .../plan-selection/src/pass/subexpr/search.rs | 326 +++++++++++++++++- 4 files changed, 454 insertions(+), 2 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 74be1785a..2a8764ac8 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -503,6 +503,7 @@ fn a_known_bound_that_misses_is_not_discarded_before_composition() { Stage1ToStage3Search { keep_every_candidate: true, exhaustive: false, + memo_groups: false, }, ) .unwrap(); @@ -1157,3 +1158,86 @@ fn queries_over_other_windows_or_with_a_latency_bound_are_in_no_class() { let search = &run.stage3_plan_selection.as_ref().unwrap().search; assert!(search.interchangeable_queries.is_empty()); } + +/// The memo-group search (not the default yet) on seeded random workloads +/// over one or two gauges, with raw data kept or not, under cost model 1: +/// it selects a plan of the exhaustive search's price, or both find none. +#[test] +fn memo_group_search_selects_the_exhaustive_price() { + let mut searched = 0; + let mut mismatches = Vec::new(); + for seed in 1..=60u64 { + let gauges: &[&str] = match seed % 3 { + 0 => &["m", "n"], + _ => &["m"], + }; + let workload = random_workload_over(seed, gauges); + let roots = lower_promql_roots(&workload).unwrap(); + let inventory = + asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates( + roots.iter().cloned().enumerate().collect(), + ) + .unwrap(); + if inventory.combinations() > MAX_EXHAUSTIVE_COMBINATIONS { + continue; + } + searched += 1; + let entries: Vec<_> = workload.query_workload.entries().collect(); + let data = workload.data_workload.clone().unwrap(); + let deployment = match seed % 2 { + 0 => no_raw_data(), + _ => asap_logical_asap_optimization::DeploymentCapabilities::default(), + }; + let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); + let indexed = || roots.iter().cloned().enumerate().collect::>(); + let exhaustive = run_stage1_to_stage3_with( + indexed(), + &entries, + &data, + models, + Stage1ToStage3Search::exhaustive(), + ) + .unwrap(); + let memo = run_stage1_to_stage3_with( + indexed(), + &entries, + &data, + models, + Stage1ToStage3Search { + memo_groups: true, + ..Stage1ToStage3Search::default() + }, + ) + .unwrap(); + let price = |run: &Stage1ToStage3Run| { + run.stage3_plan_selection + .as_ref() + .ok() + .map(|s| s.costs[&s.selected].cost_per_second.unwrap()) + }; + let (all, found) = (price(&exhaustive), price(&memo)); + let same = match (all, found) { + (Some(a), Some(f)) => (a - f).abs() <= 1e-9 * a.abs(), + (None, None) => true, + _ => false, + }; + if !same { + mismatches.push(format!( + "seed {seed}: exhaustive {all:?}, memo {found:?}: {:?}", + workload + .query_workload + .repeating_queries + .as_ref() + .unwrap() + .iter() + .map(|e| e.query.0.clone()) + .collect::>() + )); + } + } + assert!( + searched >= 20, + "only {searched} workloads were small enough" + ); + assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); +} diff --git a/crates/logical-asap-optimization/src/asap_aware_cse.rs b/crates/logical-asap-optimization/src/asap_aware_cse.rs index cf97221b1..e5bf67424 100644 --- a/crates/logical-asap-optimization/src/asap_aware_cse.rs +++ b/crates/logical-asap-optimization/src/asap_aware_cse.rs @@ -502,6 +502,16 @@ pub fn identical_expression_pass2_candidates( .collect() } +/// `roots` with every one of `candidates` (sharing candidates of `roots`, +/// or some of them) made one node: the identical-expression Pass 2 +/// candidate that shares them all. +pub fn share_identical_expressions( + roots: &[QueryRoot], + candidates: &[SharingCandidate], +) -> Vec { + apply(roots, candidates, &vec![true; candidates.len()]).0 +} + /// `roots` with the sharing candidates in `subset` shared, and the ones /// that leaves shared: those whose every node still read became one node (a /// node below a shared node other than the one kept is no longer read). diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index 6ad05932f..bd4859718 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -159,6 +159,11 @@ pub struct Stage1ToStage3Search { /// own candidates first: the exhaustive search the pruned one is tested /// against; default `false`. pub exhaustive: bool, + /// Search the memo over logical sub-expressions (joint-planning-search + /// design, "Memo-group search") instead of whole-workload candidates; + /// ignored with `exhaustive`. Default `false` until it selects the + /// exhaustive search's price on every equivalence test. + pub memo_groups: bool, } impl Stage1ToStage3Search { @@ -168,6 +173,7 @@ impl Stage1ToStage3Search { Self { keep_every_candidate: true, exhaustive: true, + memo_groups: false, } } } @@ -238,6 +244,36 @@ pub fn run_stage1_to_stage3_from_pass1( ); std::process::exit(0); } + if search.memo_groups && !search.exhaustive { + let pruned = per_computation + .clone() + .without_unsupported_summaries(models.deployment_capabilities); + let combinations = pruned.combinations(); + let found = super::subexpr::search::brute_force(&pruned, entries, data, models); + let (selection, selected, logical) = match found { + Ok((selection, selected, logical)) => (Ok(selection), Some(selected), Some(logical)), + Err(error) => (Err(error), None, None), + }; + let selected_logical = logical + .as_ref() + .and_then(|roots| selected_logical(roots, entries, data)); + return Stage1ToStage3Run { + stage1_pass1_per_computation: pruned, + combinations, + stage1_logical_asap_optimization: Vec::new(), + stage2_physical_asap_optimization: Vec::new(), + stage2_pruned: Vec::new(), + stage3_plan_selection: selection, + stage3_selected: selected, + stage3_selected_logical: selected_logical, + stage1_selected: logical.map(|roots| LogicalCandidate { + id: "L1".to_string(), + choice: Vec::new(), + roots: Ok(roots), + pass2: None, + }), + }; + } let components = super::components::sharing_components(&per_computation, entries, &models); let searched_as_one = match super::components::searched_as_one(&components, &models, search) { Some(reason) => reason, diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index aefca92f6..b9c7337c4 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -1,10 +1,24 @@ //! Clusters, coupling thresholds with branch and bound, the dynamic //! program, build-late (design "Memo-group search", steps 3–7). +//! +//! [`brute_force`] is the reference the decomposed search is checked +//! against: every consistent choice of one candidate per needed group and +//! one threshold per coupling site, composed into a plan and priced by +//! Stage 2 and Stage 3 in every materialization. +use std::collections::HashMap; +use std::rc::Rc; + +use asap_logical_asap_optimization::asap_aware_cse::{ + share_identical_expressions, sharing_candidates, +}; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; -use asap_types::workload::QueryWorkloadEntry; +use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate}; +use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; -use super::{groups, Memo}; +use super::{groups, CouplingKind, GroupId, ImplementationCandidate, Memo}; +use crate::{PlanSelector, PlanningModels, Selection, SelectionError}; /// The workload's memo: groups, Pass 1 candidates and the Pass 2 /// transformations' candidates. @@ -15,6 +29,8 @@ pub(crate) fn memo( let mut memo = groups::memo_of(per_computation); super::capability::add_shared_builds(&mut memo); super::windows::add_pass2_candidates(&mut memo, entries); + super::scans::add_scan_widening(&mut memo); + super::scans::add_raw_horizons(&mut memo); memo } @@ -45,3 +61,309 @@ pub(crate) fn stats(memo: &Memo) -> String { sites.len() ) } + +/// One plan of the memo: a candidate per needed group (`None` for a group +/// no chosen candidate needs, or one a fused candidate implements). +type GroupChoice = Vec>; + +/// Every consistent choice of one candidate per needed group, parents +/// before children: a group is needed when it is a query's root or a child +/// of a chosen candidate, and not implemented by a chosen fused candidate. +fn group_choices(memo: &Memo) -> Option> { + let order = memo.topological_order()?; + let parents_first: Vec = order.into_iter().rev().collect(); + let n = memo.groups.len(); + let mut needed = vec![0u32; n]; + let mut covered = vec![0u32; n]; + for (g, group) in memo.groups.iter().enumerate() { + if !group.roots_of.is_empty() { + needed[g] += 1; + } + } + let mut choice: GroupChoice = vec![None; n]; + let mut out = Vec::new(); + fn walk( + memo: &Memo, + order: &[GroupId], + at: usize, + needed: &mut [u32], + covered: &mut [u32], + choice: &mut GroupChoice, + out: &mut Vec, + ) { + let Some(&GroupId(g)) = order.get(at) else { + out.push(choice.clone()); + return; + }; + let g = g as usize; + if needed[g] == 0 || covered[g] > 0 { + walk(memo, order, at + 1, needed, covered, choice, out); + return; + } + for (index, candidate) in memo.candidates[g].iter().enumerate() { + for absorbed in &candidate.implements[1..] { + covered[absorbed.0 as usize] += 1; + } + for child in &candidate.children { + needed[child.0 as usize] += 1; + } + choice[g] = Some(index); + walk(memo, order, at + 1, needed, covered, choice, out); + choice[g] = None; + for child in &candidate.children { + needed[child.0 as usize] -= 1; + } + for absorbed in &candidate.implements[1..] { + covered[absorbed.0 as usize] -= 1; + } + } + } + walk( + memo, + &parents_first, + 0, + &mut needed, + &mut covered, + &mut choice, + &mut out, + ); + Some(out) +} + +/// The chosen candidates of `choice`, each rewritten at a threshold of each +/// of its coupling sites: every combination of thresholds every reader may +/// read. A site without a rewrite (raw horizon, a shared build until it can +/// be resized) leaves its readers as they are; Stage 3 prices it. +fn at_thresholds(memo: &Memo, choice: &GroupChoice) -> Vec>> { + let chosen: Vec> = choice + .iter() + .enumerate() + .map(|(g, c)| c.map(|c| memo.candidates[g][c].clone())) + .collect(); + // The sites the chosen candidates read, each once. + let mut sites = Vec::new(); + for candidate in chosen.iter().flatten() { + for site in &candidate.coupling_sites { + if !sites.contains(site) { + sites.push(*site); + } + } + } + let mut plans = vec![chosen]; + for site in sites { + let mut next = Vec::new(); + for plan in plans { + match site.kind { + CouplingKind::PaneChain => { + for threshold in super::windows::chain_thresholds(memo, &site) { + let mut rewritten = plan.clone(); + let mut valid = true; + for candidate in rewritten.iter_mut().flatten() { + if !candidate.coupling_sites.contains(&site) { + continue; + } + let readers: Vec<_> = super::windows::chain_readers(&candidate.node) + .into_iter() + .filter(|r| r.key == site.key) + .collect(); + if !readers + .iter() + .all(|r| super::windows::may_read(&threshold, r)) + { + valid = false; + break; + } + match super::windows::read_at_threshold(candidate, &site, &threshold) { + Ok(node) => candidate.node = node, + Err(_) => { + valid = false; + break; + } + } + } + if valid { + next.push(rewritten); + } + } + } + CouplingKind::ScanLabels => { + for threshold in super::scans::label_thresholds(memo, &site) { + let mut rewritten = plan.clone(); + let mut valid = true; + for candidate in rewritten.iter_mut().flatten() { + if !candidate.coupling_sites.contains(&site) { + continue; + } + match super::scans::widened_at_threshold( + memo, candidate, &site, &threshold, + ) { + Ok(Some(widened)) => *candidate = widened, + _ => { + valid = false; + break; + } + } + } + if valid { + next.push(rewritten); + } + } + } + CouplingKind::SharedBuildSize | CouplingKind::RawHorizon => next.push(plan), + } + } + plans = next; + } + plans +} + +/// The query roots of a plan: each chosen candidate's node with every +/// child group's node replaced by that group's built node, children first, +/// so a group read by several parents is one node. +fn composed_roots( + memo: &Memo, + per_computation: &PerComputationCandidates, + chosen: &[Option], +) -> Option> { + let order = memo.topological_order()?; + let mut built: HashMap<*const OperatorNode, Rc> = HashMap::new(); + let mut built_of: Vec>> = vec![None; memo.groups.len()]; + for GroupId(g) in order { + let g = g as usize; + let Some(candidate) = &chosen[g] else { + continue; + }; + let mut rebuilt = HashMap::new(); + let node = substituted(&candidate.node, &built, &mut rebuilt); + for implemented in &candidate.implements { + built.insert( + Rc::as_ptr(&memo.groups[implemented.0 as usize].node), + Rc::clone(&node), + ); + } + built_of[g] = Some(node); + } + per_computation + .roots + .iter() + .enumerate() + .map(|(q, (_, root))| { + let group_of = |node: &Rc| -> Option> { + memo.groups + .iter() + .enumerate() + .find(|(_, group)| group.roots_of.contains(&q) && *group.node == **node) + .and_then(|(g, _)| built_of[g].clone()) + }; + Some(match root { + QueryRoot::Operator(node) => QueryRoot::Operator(group_of(node)?), + QueryRoot::Scalar(expr) => { + let mut missing = false; + let mapped = expr.map_operator_refs(&mut |node| { + group_of(node).unwrap_or_else(|| { + missing = true; + Rc::clone(node) + }) + }); + if missing { + return None; + } + QueryRoot::Scalar(mapped) + } + }) + }) + .collect() +} + +/// `node` with every node `built` maps (by address) replaced by its built +/// node, rebuilt above them; `rebuilt` memoizes by address. +fn substituted( + node: &Rc, + built: &HashMap<*const OperatorNode, Rc>, + rebuilt: &mut HashMap<*const OperatorNode, Rc>, +) -> Rc { + if let Some(replacement) = built.get(&Rc::as_ptr(node)) { + return Rc::clone(replacement); + } + if let Some(done) = rebuilt.get(&Rc::as_ptr(node)) { + return Rc::clone(done); + } + let mut changed = false; + let operator = node.operator.map_children(|child| { + let new = substituted(child, built, rebuilt); + changed |= !Rc::ptr_eq(&new, child); + new + }); + let new = match changed { + false => Rc::clone(node), + true => Rc::new( + OperatorNode::with_schema(operator, node.schema.clone()) + .with_guarantee(node.guarantee.clone()) + .with_timing(node.timing) + .with_materialization(node.materialization), + ), + }; + rebuilt.insert(Rc::as_ptr(node), Rc::clone(&new)); + new +} + +/// The reference search: every plan of the memo ([`group_choices`] × +/// [`at_thresholds`]), its identical nodes made one node, implemented in +/// every Stage 2 materialization and priced by Stage 3; the cheapest +/// valid one, with its logical roots and the number of physical candidates +/// checked. +pub(crate) fn brute_force( + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result<(Selection, PhysicalCandidate, Vec), SelectionError> { + let memo = memo(per_computation, entries); + let choices = group_choices(&memo).unwrap_or_default(); + let mut selector = PlanSelector::new(entries, data, models, false)?; + let mut logical: HashMap> = HashMap::new(); + let mut number = 0usize; + for choice in &choices { + let chosen: Vec<(GroupId, &ImplementationCandidate)> = choice + .iter() + .enumerate() + .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) + .collect(); + if !super::windows::requirements_met(&chosen) { + continue; + } + for plan in at_thresholds(&memo, choice) { + let Some(roots) = composed_roots(&memo, per_computation, &plan) else { + continue; + }; + let shareable = sharing_candidates(&roots); + let roots = match shareable.is_empty() { + true => roots, + false => share_identical_expressions(&roots, &shareable), + }; + number += 1; + let id = format!("L{number}"); + let Ok(implemented) = + optimize_physical_asap(&id, &roots, entries, data, models.deployment_capabilities) + else { + continue; + }; + for (k, mut physical) in implemented.into_iter().enumerate() { + physical.id = match k { + 0 => format!("P{number}"), + k => format!("P{number}-m{k}"), + }; + selector.offer(&physical)?; + } + if selector.keeps_from_logical(&id) { + logical.insert(id, roots); + } + logical.retain(|id, _| selector.keeps_from_logical(id)); + } + } + let (selection, selected) = selector.finish()?; + let roots = logical + .remove(&selected.from_logical) + .expect("the selection's logical roots are kept"); + Ok((selection, selected, roots)) +} From 0c4f61efeacbbf3e30f38176936dbba1c418dcd0 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:28:05 +0000 Subject: [PATCH 11/25] Merge W1 (fused heap candidates keep their schema, Pass 1 parity test); scan materializations in the memo --- crates/plan-selection/src/pass/subexpr/search.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index b9c7337c4..b191b5b49 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -31,6 +31,7 @@ pub(crate) fn memo( super::windows::add_pass2_candidates(&mut memo, entries); super::scans::add_scan_widening(&mut memo); super::scans::add_raw_horizons(&mut memo); + super::scans::scan_materializations(&mut memo); memo } From ecd42bf5eb5fb406357171da0319d1d69616acec Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:31:48 +0000 Subject: [PATCH 12/25] memo: debugging test for a top-k memo --- .../plan-selection/src/pass/subexpr/search.rs | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index b191b5b49..693367129 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -368,3 +368,94 @@ pub(crate) fn brute_force( .expect("the selection's logical roots are kept"); Ok((selection, selected, roots)) } + +#[cfg(test)] +mod tests { + use super::*; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; + use asap_logical_asap_optimization::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::ir::schema::schema_support::with_promql_series_identity; + use asap_types::workload::{ + DataArrival, DurationMs, MetricType, Predictability, QueryRecurrence, QueryRequirements, + RepeatedDemand, RepetitionInterval, SourceFacts, + }; + + #[test] + #[ignore = "debugging aid"] + fn debug_topk_memo() { + let queries = ["topk by (job) (3, sum_over_time(m[1m]))"]; + let roots: Vec<(usize, QueryRoot)> = queries + .iter() + .map(|q| { + let root = lower_promql(q, AccuracyTarget::max_error(0.05)); + QueryRoot::Operator(with_promql_series_identity(&root).unwrap()) + }) + .enumerate() + .collect(); + let pc = generate_per_computation_candidates(roots).unwrap(); + let entries: Vec = queries + .iter() + .map(|q| QueryWorkloadEntry { + query: asap_types::workload::Query((*q).into()), + requirements: QueryRequirements { + accuracy: AccuracyTarget::max_error(0.05), + max_latency_ms: None, + }, + predictability: Predictability::Predictable { known_at: None }, + recurrence: QueryRecurrence::Repeated(RepeatedDemand::FixedInterval( + RepetitionInterval(60_000), + )), + range_evaluation: None, + }) + .collect(); + let data = DataWorkload { + arrival: DataArrival::Ingesting { + rows_per_second: 2_000.0 / 15.0, + }, + sources: [( + "m".to_string(), + SourceFacts { + series_count: Some(1_000), + metric_type: Some(MetricType::Gauge), + sample_interval: Some(DurationMs(15_000)), + ..SourceFacts::default() + }, + )] + .into(), + }; + let memo = memo(&pc, &entries); + eprintln!("{}", stats(&memo)); + for (g, group) in memo.groups.iter().enumerate() { + eprintln!( + "G{g} {:?} children {:?} roots_of {:?} candidates {}", + group.node.operator.kind_name(), + group.children, + group.roots_of, + memo.candidates[g].len() + ); + for c in &memo.candidates[g] { + eprintln!( + " {:?} implements {:?} children {:?}", + c.origin, c.implements, c.children + ); + } + } + let choices = group_choices(&memo).unwrap(); + eprintln!("choices {}", choices.len()); + for choice in &choices { + for plan in at_thresholds(&memo, choice) { + let roots = composed_roots(&memo, &pc, &plan); + eprintln!("choice {choice:?} composed {}", roots.is_some()); + if let Some(roots) = roots { + let r = + optimize_physical_asap("L", &roots, &entries, &data, &Default::default()); + eprintln!( + " stage2 {:?}", + r.as_ref().map(|v| v.len()).map_err(|e| e.to_string()) + ); + } + } + } + } +} From d220e45f16035f67b975234cd04b01b4c673ef44 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:35:42 +0000 Subject: [PATCH 13/25] physical: materialization_slots, the options a search chooses per materialization candidate --- crates/physical-asap-optimization/src/lib.rs | 42 ++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/crates/physical-asap-optimization/src/lib.rs b/crates/physical-asap-optimization/src/lib.rs index 0c34f3b12..12ff0a64e 100644 --- a/crates/physical-asap-optimization/src/lib.rs +++ b/crates/physical-asap-optimization/src/lib.rs @@ -526,6 +526,48 @@ fn assigned_candidate( }) } +/// One materialization candidate of a plan, as Stage 2 finds it: a summary +/// build, or the highest node of a shared non-summary sub-DAG. +#[derive(Debug, Clone)] +pub struct MaterializationSlot { + /// The node, in the implemented plan. + pub node: Rc, + /// The queries (indices into the roots) that read it. + pub readers: Vec, + /// Its options, `NotMaterialized` first. + pub options: Vec, + /// The slots below it (by index): one at ingestion time needs every one + /// below it at ingestion time too. + pub below: Vec, +} + +/// The materialization candidates of `roots`, in the order +/// [`optimize_physical_asap_assigned`]'s `choose` is called for them, with +/// their options: what a search that chooses each one itself needs. +pub fn materialization_slots( + roots: &[QueryRoot], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + deployment: &DeploymentCapabilities, +) -> Result, PhysicalASAPOptimizationError> { + let mut memo = HashMap::new(); + let implemented = map_roots(roots, &mut |node| implement(node, &mut memo))?; + let mat_candidates = materialization_candidates(&implemented, entries, data, deployment); + check_panes_on_range_steps(&mat_candidates, entries)?; + let nodes: Vec<_> = mat_candidates.iter().map(|c| Rc::clone(&c.node)).collect(); + let below = candidates_below(&nodes); + Ok(mat_candidates + .iter() + .zip(below) + .map(|(c, below)| MaterializationSlot { + node: Rc::clone(&c.node), + readers: c.readers.clone(), + options: c.options(), + below, + }) + .collect()) +} + /// The number of physical candidates [`optimize_physical_asap`] returns for /// `roots`, without building them; `Err` when a candidate cannot be /// implemented. From f8bcdba3c3bbd9a5672200bdd7c6950c234894a4 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:36:52 +0000 Subject: [PATCH 14/25] plan-selection: price a memo candidate's private nodes in its consumer context price_at_scale's node loop is split into price_each_node, and lib.rs gains price_private_nodes: Stage 3's per-node pricing of a candidate's private nodes for given consumers (queries, or materialized ancestors with their rate), reading its children's given outputs, with computations()' counting restricted to them. Memory held for an evaluation and raw data on disk stay plan-level terms (HeldMemory, raw scans), priced by plan_cost once each query's latency is known. subexpr/pricing.rs: private_usage, output_stats, plan_cost and frozen_calibration, cached by (candidate key, consumers, materialization, child outputs); frontiers dropped (cost model 1 only). Tests: the groups' local prices add up to price_at_scale of the composed plan for several one-query plans, a shared group not materialized and materialized at query time (equal and different intervals), and a pane chain at ingestion. --- crates/plan-selection/src/lib.rs | 792 ++++++-- .../src/measured_operation_costs.rs | 2 +- crates/plan-selection/src/pass/subexpr/mod.rs | 22 +- .../src/pass/subexpr/pricing.rs | 1693 +++++++++++------ 4 files changed, 1795 insertions(+), 714 deletions(-) diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 76958156d..aae16cc0c 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1185,52 +1185,510 @@ pub struct PlanSelector<'a> { runner_up: Option, } -/// Stage 3's per-node pricing of `candidate`, a part of a workload whose -/// roots are read by the queries of `entries` (one entry per root), sized -/// for their failure probabilities when that resizes, with the formulas' -/// CPU time scaled by `formula_scale`: its cost and calibration fit. -/// `None` when it fails a check no summary size repairs: a per-series -/// aggregate the executor cannot run, or raw data the deployment does not -/// keep. The memo over sub-expressions prices each group's sub-DAG this -/// way; accuracy and latency are checked on the whole plan. -pub(crate) fn price_part( - candidate: &PhysicalCandidate, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: &PlanningModels<'_>, - formula_scale: f64, -) -> Result, SelectionError> { - let targets = accuracy_targets(entries); - let sized = sized_for_failure_probability(candidate, &targets, data) - .unwrap_or_else(|_| candidate.clone()); - if per_entity_aggregate_violation(&sized.dag).is_some() { +/// Evaluations per second of a node computed once for the queries of +/// `entries` together, as a node they share at query time is +/// ([`computations`]). +pub(crate) fn shared_evaluation_rate( + entries: &[&QueryWorkloadEntry], + settings: &PlanSelectionSettings, +) -> Result { + evaluation_rate( + entries.iter().map(|e| (&e.recurrence, e.range_evaluation)), + settings.non_repeating_query_horizon_seconds, + ) +} + +/// Memory a query-time node holds while an evaluation runs, charged once +/// its readers' latencies are known (Stage 3's "held" memory): +/// `byte_seconds_per_ms` times the readers' rate-weighted mean latency in +/// ms (`weighted`), or their longest latency (otherwise), is the byte-seconds +/// per second it adds to the average memory. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct HeldMemory { + pub(crate) byte_seconds_per_ms: f64, + /// The queries reading it, each with its evaluations per second, by + /// position among the part's roots (or, once mapped, in the workload). + pub(crate) readers: Vec<(usize, f64)>, + /// Whether each reader computes its own copy (rate-weighted mean + /// latency) rather than one copy held until the longest finishes. + pub(crate) weighted: bool, +} + +impl HeldMemory { + /// Its byte-seconds per second given each reader's latency in ms. + pub(crate) fn byte_seconds(&self, latency_ms: impl Fn(usize) -> f64) -> f64 { + let longest = || { + self.readers + .iter() + .map(|&(q, _)| latency_ms(q)) + .fold(0.0, f64::max) + }; + let rate: f64 = self.readers.iter().map(|&(_, r)| r).sum(); + let evaluation_ms = match self.weighted && rate > 0.0 { + true => { + self.readers + .iter() + .map(|&(q, r)| r * latency_ms(q)) + .sum::() + / rate + } + false => longest(), + }; + self.byte_seconds_per_ms * evaluation_ms + } +} + +/// Who computes a set of nodes: a reader at `rate_per_second` +/// computations per second, through which the `queries` (positions among +/// [`PrivatePricing::queries`]) reach them at query time; at ingestion time +/// (`at_ingestion`, a build above them materialized at ingestion) no query +/// reaches them. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct NodeConsumer { + pub(crate) rate_per_second: f64, + pub(crate) queries: Vec, + pub(crate) at_ingestion: bool, +} + +/// Stage 3's pricing of some nodes of a physical DAG on their own +/// ([`price_private_nodes`]): a candidate's private nodes, the others being +/// its children's, whose outputs are given. +pub(crate) struct PrivatePricing<'a> { + /// The DAG (Stage 2's, the candidate over its children), its first + /// root the candidate's node. + pub(crate) dag: &'a PhysicalASAPDAG, + /// Which nodes are priced, by position in `dag.nodes`. + pub(crate) private: &'a [bool], + /// The output of each other node a private node reads (a child's root). + pub(crate) child_outputs: &'a HashMap, + /// Who computes the candidate's node; all at ingestion time or all not. + pub(crate) consumers: &'a [NodeConsumer], + /// The queries the consumers name. + pub(crate) queries: &'a [QueryWorkloadEntry], + /// Per query, whether the candidate's node is its root (its request + /// overhead is charged here). + pub(crate) root_of: &'a [bool], + /// The look-back of a query-time scan no window among the private + /// nodes bounds ([`DEFAULT_LOOKBACK_MS`] when `None`). + pub(crate) scan_lookback_ms: Option, + pub(crate) data: &'a DataWorkload, + pub(crate) settings: &'a PlanSelectionSettings, + pub(crate) formula_scale: f64, +} + +/// The usage of a candidate's private nodes ([`price_private_nodes`]). +/// Charges that need the whole plan are kept apart: memory held for an +/// evaluation ([`HeldMemory`]) and each query-time raw scan (Stage 3 keeps +/// each source's longest on disk). Per query (position): the query-time +/// work it reaches among the private nodes, in ms, and the longest +/// compaction among the ingestion-time ones it reads first. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct PrivateNodesUsage { + /// Under cost model 1, the nodes' prices without held memory and raw + /// data; `None` otherwise. + pub(crate) cost_per_second: Option, + pub(crate) cpu_ops_per_second: f64, + pub(crate) scanned_bytes_per_second: f64, + pub(crate) average_memory_bytes: f64, + pub(crate) peak_memory_bytes: u64, + pub(crate) disk_bytes: u64, + pub(crate) query_ms: Vec, + pub(crate) compaction_ms: Vec, + pub(crate) held: Vec, + pub(crate) raw_scans: Vec<(String, u64)>, + /// The candidate's node's output per computation. + pub(crate) output: EdgeStatistics, + pub(crate) fit: CalibrationFit, +} + +/// Each node of `dag` below and including `id`: the sources it reads, +/// without pricing (for a child's root, whose output is given). +fn sources_below<'a>( + dag: &'a PhysicalASAPDAG, + nodes: &NodeMap<'a>, + id: PhysicalASAPNodeId, +) -> BTreeSet<&'a str> { + let mut sources = BTreeSet::new(); + let mut stack = vec![id]; + let mut seen = HashSet::new(); + while let Some(id) = stack.pop() { + if !seen.insert(id) { + continue; + } + if let Payload::NonASAP(NonASAPOp::Scan { source, .. }) = &nodes[&id].payload { + sources.insert(source_name(source)); + } + stack.extend( + dag.edges + .iter() + .filter(|e| e.consumer == id) + .map(|e| e.producer), + ); + } + sources +} + +/// The state family of node `id` of `dag`, without pricing. +fn family_below<'a>( + dag: &'a PhysicalASAPDAG, + nodes: &NodeMap<'a>, + id: PhysicalASAPNodeId, +) -> StateFamily<'a> { + let first = dag + .edges + .iter() + .find(|e| e.consumer == id) + .map(|e| e.producer); + let producer_family = match &nodes[&id].payload { + Payload::ASAP( + ASAPOp::SummaryMerge { .. } + | ASAPOp::SummaryEstimate { .. } + | ASAPOp::FinalizeExactAccumulator { .. }, + ) => first.map(|p| family_below(dag, nodes, p)), + _ => None, + }; + state_family(&nodes[&id].payload, producer_family) +} + +/// Price `input`'s private nodes as Stage 3 prices them in the whole plan +/// ([`price_node`]): each computed for its consumers as [`computations`] +/// counts them (a materialized private node once, at its own rate: +/// 1 at ingestion time, the queries reaching it together at query time, +/// once per request when kept across evaluations), reading its children's +/// given outputs. With consumers at ingestion time every private node runs +/// at ingestion time. `None` when a private node fails a check no summary +/// size repairs: a per-series aggregate the executor cannot run, or a +/// query-time scan of raw data the deployment does not keep; also when a +/// node materialized at query time runs under an ingestion-time consumer. +pub(crate) fn price_private_nodes( + input: &PrivatePricing<'_>, + raw_data_retained: bool, +) -> Result, (PhysicalASAPNodeId, AnalyticalCostError)> { + let at_ingestion = input.consumers.iter().any(|c| c.at_ingestion); + // The DAG with the private nodes placed for the consumers. + let mut placed = input.dag.clone(); + if at_ingestion { + for (node, private) in placed.nodes.iter_mut().zip(input.private) { + if !*private { + continue; + } + if node.materialization.timing() == ExecutionTiming::QueryTime + && node.materialization != Materialization::NotMaterialized + { + return Ok(None); + } + node.output_state.timing = ExecutionTiming::IngestionTime; + } + } + let dag = &placed; + let private_ids: HashSet = dag + .nodes + .iter() + .zip(input.private) + .filter(|(_, p)| **p) + .map(|(n, _)| n.id) + .collect(); + let private_dag = PhysicalASAPDAG { + nodes: dag + .nodes + .iter() + .filter(|n| private_ids.contains(&n.id)) + .cloned() + .collect(), + edges: Vec::new(), + roots: Vec::new(), + scalar_roots: Vec::new(), + }; + if per_entity_aggregate_violation(&private_dag).is_some() { + return Ok(None); + } + let nodes = node_map(dag); + let kept_chains = kept_pane_chains(dag); + let first = dag.nodes.first().map(|n| n.id).unwrap_or_default(); + let Some(&root) = dag.roots.first() else { return Ok(None); + }; + let settings = input.settings; + let horizon = settings.non_repeating_query_horizon_seconds; + let mut producers_of: HashMap> = HashMap::new(); + for e in &dag.edges { + producers_of.entry(e.consumer).or_default().push(e.producer); + } + let private_inputs = |id: PhysicalASAPNodeId| { + producers_of + .get(&id) + .into_iter() + .flatten() + .copied() + .filter(|p| private_ids.contains(p)) + .collect::>() + }; + // Each query's reach among the private nodes at query time, stopping + // at (and including) ingestion-time nodes, as Stage 3's. + let query_count = input.queries.len(); + let mut reach: Vec> = vec![HashSet::new(); query_count]; + if private_ids.contains(&root) { + for consumer in input.consumers.iter().filter(|c| !c.at_ingestion) { + for &query in &consumer.queries { + let mut stack = vec![root]; + while let Some(id) = stack.pop() { + if reach[query].insert(id) + && nodes[&id].output_state.timing != ExecutionTiming::IngestionTime + { + stack.extend(private_inputs(id)); + } + } + } + } } - let nodes = node_map(&sized.dag); - let kept_chains = kept_pane_chains(&sized.dag); - if !models.deployment_capabilities.raw_data_retained - && raw_data_violation(&sized.dag, &nodes, &kept_chains).is_some() + let rate_of = |queries: &mut dyn Iterator, per_request: bool| { + evaluation_rate( + queries.filter_map(|q| { + input.queries.get(q).map(|e| { + let range = if per_request { + None + } else { + e.range_evaluation + }; + (&e.recurrence, range) + }) + }), + horizon, + ) + }; + // How often each private node is computed ([`computations`]). + let mut rate_of_materialized: HashMap = HashMap::new(); + for node in dag.nodes.iter().filter(|n| private_ids.contains(&n.id)) { + let rate = match node.materialization { + Materialization::NotMaterialized => continue, + Materialization::MatAtIngestion(_) => 1.0, + m => { + let per_request = matches!(m, Materialization::MatAtQueryKeptAcrossEvaluations(_)); + let reaching: Vec = (0..query_count) + .filter(|&q| reach[q].contains(&node.id)) + .collect(); + rate_of(&mut reaching.into_iter(), per_request).map_err(|error| (node.id, error))? + } + }; + rate_of_materialized.insert(node.id, rate); + } + let mut walks: Vec<(Vec, f64)> = input + .consumers + .iter() + .map(|c| (vec![root], c.rate_per_second)) + .collect(); + for (&id, &rate) in &rate_of_materialized { + walks.push((private_inputs(id), rate)); + } + let mut multipliers: HashMap = HashMap::new(); + for (start, rate) in walks { + let mut seen = HashSet::new(); + let mut stack = start; + while let Some(id) = stack.pop() { + if !private_ids.contains(&id) + || rate_of_materialized.contains_key(&id) + || !seen.insert(id) + { + continue; + } + let entry = multipliers.entry(id).or_default(); + entry.0 += 1; + entry.1 += rate; + stack.extend(private_inputs(id)); + } + } + for (&id, &rate) in &rate_of_materialized { + multipliers.insert(id, (1, rate)); + } + let retention_ms = pane_retention_ms( + dag, + &reach, + &input + .queries + .iter() + .map(|e| e.range_evaluation) + .collect::>(), + ); + // Price each private node, children first. + let mut priced: HashMap> = HashMap::new(); + for node in dag.nodes.iter().filter(|n| private_ids.contains(&n.id)) { + let views: Vec> = producers_of + .get(&node.id) + .map_or(&[][..], Vec::as_slice) + .iter() + .map(|p| match priced.get(p) { + Some(produced) => ProducerOutput { + payload: &nodes[p].payload, + output_schema: &nodes[p].output_schema, + output: produced.output, + sources: produced.sources.clone(), + family: produced.family, + }, + None => ProducerOutput { + payload: &nodes[p].payload, + output_schema: &nodes[p].output_schema, + output: input + .child_outputs + .get(p) + .copied() + .unwrap_or(EdgeStatistics { rows: 1, bytes: 1 }), + sources: sources_below(dag, &nodes, *p), + family: family_below(dag, &nodes, *p), + }, + }) + .collect(); + let lookback_ms = match (&node.payload, node.output_state.timing) { + (Payload::NonASAP(NonASAPOp::Scan { .. }), ExecutionTiming::QueryTime) => { + windows_read(dag, &nodes, node.id, &kept_chains) + .into_iter() + .max() + .or(input.scan_lookback_ms) + .unwrap_or(DEFAULT_LOOKBACK_MS) + } + _ => 0, + }; + let node_priced = price_node( + node, + &views, + multipliers.get(&node.id).copied().unwrap_or((0, 0.0)), + lookback_ms, + retained_panes(node, &retention_ms), + (node.output_state.timing == ExecutionTiming::IngestionTime) + .then(|| maintained_window_ms(node.id, &nodes, &producers_of)) + .flatten(), + input.data, + settings, + input.formula_scale, + ) + .map_err(|error| (node.id, error))?; + priced.insert(node.id, node_priced); + } + if !raw_data_retained + && priced.values().any(|p| p.raw_scan.is_some()) + && raw_data_violation(dag, &nodes, &kept_chains).is_some() { return Ok(None); } - let recurrence: Vec = entries.iter().map(|e| e.recurrence.clone()).collect(); - let range_evaluations: Vec<_> = entries.iter().map(|e| e.range_evaluation).collect(); - price_at_scale( - &sized.dag, - &nodes, - &kept_chains, - &recurrence, - &range_evaluations, - data, - models.plan_selection, - formula_scale, - ) - .map(Some) - .map_err(|(node, error)| SelectionError::Cost { - candidate: candidate.id.clone(), - node, - error, - }) + let work_ms: HashMap = dag + .nodes + .iter() + .map(|n| (n.id, priced.get(&n.id).map_or(0.0, |p| p.work_ms))) + .collect(); + let mut usage = PrivateNodesUsage { + cost_per_second: match settings.cost_model { + CostModel::UsageOverTime(_) => Some(0.0), + _ => None, + }, + cpu_ops_per_second: 0.0, + scanned_bytes_per_second: 0.0, + average_memory_bytes: 0.0, + peak_memory_bytes: 0, + disk_bytes: 0, + query_ms: vec![0.0; query_count], + compaction_ms: vec![0.0; query_count], + held: Vec::new(), + raw_scans: Vec::new(), + output: priced + .get(&root) + .map_or(EdgeStatistics { rows: 0, bytes: 0 }, |p| p.output), + fit: CalibrationFit::default(), + }; + let table = &settings.measured_operation_costs; + let sketch_size = + |params: &SketchParams| sketch_shape(params, StateInputRows::UNKNOWN).1 as f64; + let formula_merge_ms = settings.resource_speed.ms_per_cpu_op * input.formula_scale; + for node in dag.nodes.iter().filter(|n| private_ids.contains(&n.id)) { + let node_priced = &priced[&node.id]; + if let (Some(total), Some(cost)) = + (&mut usage.cost_per_second, node_priced.cost.cost_per_second) + { + *total += cost; + } + usage.cpu_ops_per_second += node_priced.cpu_ops_per_second; + usage.scanned_bytes_per_second += node_priced.scanned_bytes_per_second; + usage.average_memory_bytes += node_priced.average_memory_bytes; + usage.peak_memory_bytes = usage + .peak_memory_bytes + .saturating_add(node_priced.peak_memory_bytes); + usage.disk_bytes = usage.disk_bytes.saturating_add(node_priced.disk_bytes); + usage.fit.add(node_priced.fit); + // Partial copies: the workers this state's ingestion work below it + // takes, counting the private nodes only. + if let Some(window) = node_priced.maintained { + let copies = ingest_workers(dag, &work_ms, node.id) - 1; + usage.peak_memory_bytes = usage + .peak_memory_bytes + .saturating_add(copies.saturating_mul(window)); + } + if let Some((source, bytes)) = node_priced.raw_scan { + usage + .raw_scans + .push((source_name(source).to_string(), bytes)); + } + if let Some((state, per_second, weighted)) = node_priced.held { + let mut readers = Vec::new(); + for (query, reached) in reach.iter().enumerate() { + if reached.contains(&node.id) { + let rate = match weighted { + true => rate_of(&mut std::iter::once(query), false) + .map_err(|error| (node.id, error))?, + false => 0.0, + }; + readers.push((query, rate)); + } + } + usage.held.push(HeldMemory { + byte_seconds_per_ms: state as f64 / 1_000.0 * per_second, + readers, + weighted, + }); + } + for (query, reached) in reach.iter().enumerate() { + if !reached.contains(&node.id) { + continue; + } + if node.output_state.timing == ExecutionTiming::IngestionTime { + let workers = ingest_workers(dag, &work_ms, node.id) as f64; + let states = (node_priced.output.rows / panes_of(node)) as f64; + let per_merge_ms = lookup( + table, + OperatorKind::SummaryMerge, + ExecutionTiming::QueryTime, + node_priced.family, + sketch_size, + ) + .map_or(formula_merge_ms, |unit| { + (unit.ns_per_unit + unit.ns_per_invocation) * 1e-6 + }); + let compaction = (workers - 1.0) * states * per_merge_ms; + usage.compaction_ms[query] = usage.compaction_ms[query].max(compaction); + } else { + usage.query_ms[query] += node_priced.work_ms; + } + } + } + // A query's fixed request cost, where the candidate is its root. + let evaluation_ms = table.ns_per_evaluation * 1e-6; + for query in 0..query_count { + if evaluation_ms <= 0.0 || !input.root_of.get(query).copied().unwrap_or(false) { + continue; + } + let requests_per_second = evaluation_rate( + input.queries.get(query).map(|e| (&e.recurrence, None)), + horizon, + ) + .map_err(|error| (first, error))?; + let ops_per_second = + evaluation_ms / settings.resource_speed.ms_per_cpu_op * requests_per_second; + usage.cpu_ops_per_second += ops_per_second; + usage.query_ms[query] += evaluation_ms; + if let (Some(cost), CostModel::UsageOverTime(prices)) = + (&mut usage.cost_per_second, &settings.cost_model) + { + *cost += prices.cost_per_cpu_op * ops_per_second; + } + } + Ok(Some(usage)) } /// Prices within this fraction of each other are the same sum of charges @@ -2206,7 +2664,7 @@ fn at_most_one_row(node: &OperatorNode) -> bool { /// stays unmet. Stage 1 sizes for ε and reads no data workload; only this /// stage binds `G`. `Err` is the reason the candidate is rejected when its /// resized summaries do not rebuild into a plan. -fn sized_for_failure_probability( +pub(crate) fn sized_for_failure_probability( candidate: &PhysicalCandidate, targets: &[AccuracyTarget], data: &DataWorkload, @@ -3173,6 +3631,7 @@ fn computations( recurrence: &[QueryRecurrence], range_evaluations: &[Option], non_repeating_query_horizon_seconds: f64, + root_rates: Option<&[f64]>, ) -> Result, AnalyticalCostError> { let range_of = |query: usize| range_evaluations.get(query).copied().flatten(); let materialized: HashMap<_, _> = dag @@ -3223,10 +3682,13 @@ fn computations( // Each consumer's starting nodes and its rate. let mut consumers = Vec::new(); for (query, start) in starts.into_iter().enumerate() { - let rate = evaluation_rate( - recurrence.get(query).map(|r| (r, range_of(query))), - non_repeating_query_horizon_seconds, - )?; + let rate = match root_rates.and_then(|rates| rates.get(query)) { + Some(&rate) => rate, + None => evaluation_rate( + recurrence.get(query).map(|r| (r, range_of(query))), + non_repeating_query_horizon_seconds, + )?, + }; consumers.push((start, rate)); } for (&id, &rate) in &rate_of_materialized { @@ -3291,17 +3753,22 @@ fn price_at_scale( CostModel::UsageOverTime(prices) => Some(prices), _ => None, }; - // Each node priced, filled children first. - let mut priced: HashMap> = HashMap::new(); - let multipliers = computations( + let EachNodePriced { + priced, + multipliers, + reach, + } = price_each_node( dag, + nodes, + kept_chains, recurrence, range_evaluations, - settings.non_repeating_query_horizon_seconds, - ) - .map_err(|error| (first, error))?; - let reach = query_reach(dag, nodes); - let retention_ms = pane_retention_ms(dag, &reach, range_evaluations); + data, + settings, + formula_scale, + None, + None, + )?; let mut output: HashMap = HashMap::new(); let mut per_node = BTreeMap::new(); // The candidate's usage, accumulated node by node; peak cores and @@ -3322,11 +3789,6 @@ fn price_at_scale( // Each in-memory ingestion-time state's bytes per window (what it keeps // for one window, its output or an exact aggregate's values). let mut maintained: Vec<(PhysicalASAPNodeId, u64)> = Vec::new(); - // Each node's producers, in edge order. - let mut producers_of: HashMap> = HashMap::new(); - for e in &dag.edges { - producers_of.entry(e.consumer).or_default().push(e.producer); - } // Each summary node's state family, to look up in the measured table: // a build's own (one sketch instance per group; a shared Hydra // structure is not that sketch), inherited by the merges, estimates and @@ -3335,41 +3797,7 @@ fn price_at_scale( let sketch_size = |params: &SketchParams| sketch_shape(params, StateInputRows::UNKNOWN).1 as f64; for node in &dag.nodes { - let views: Vec> = producers_of - .get(&node.id) - .map_or(&[][..], Vec::as_slice) - .iter() - .map(|p| { - let produced: &PricedNode<'_> = &priced[p]; - ProducerOutput { - payload: &nodes[p].payload, - output_schema: &nodes[p].output_schema, - output: produced.output, - sources: produced.sources.clone(), - family: produced.family, - } - }) - .collect(); - let lookback_ms = match (&node.payload, node.output_state.timing) { - (Payload::NonASAP(NonASAPOp::Scan { .. }), ExecutionTiming::QueryTime) => { - query_time_lookback_ms(dag, nodes, node.id, kept_chains) - } - _ => 0, - }; - let node_priced = price_node( - node, - &views, - multipliers[&node.id], - lookback_ms, - retained_panes(node, &retention_ms), - (node.output_state.timing == ExecutionTiming::IngestionTime) - .then(|| maintained_window_ms(node.id, nodes, &producers_of)) - .flatten(), - data, - settings, - formula_scale, - ) - .map_err(|error| (node.id, error))?; + let node_priced = &priced[&node.id]; usage.cpu_ops_per_second += node_priced.cpu_ops_per_second; usage.scanned_bytes_per_second += node_priced.scanned_bytes_per_second; usage.average_memory_bytes += node_priced.average_memory_bytes; @@ -3396,7 +3824,6 @@ fn price_at_scale( } } per_node.insert(node.id, node_priced.cost.clone()); - priced.insert(node.id, node_priced); } // The runtime's fixed cost per query request, charged to the node the // query reads first and added to its latency below. @@ -3548,6 +3975,135 @@ fn price_at_scale( )) } +/// Every node of a physical DAG priced ([`price_each_node`]), with each +/// node's `(multiplier, per_second)` ([`computations`]) and the nodes each +/// query reaches at query time ([`query_reach`]). +pub(crate) struct EachNodePriced<'a> { + pub(crate) priced: HashMap>, + pub(crate) multipliers: HashMap, + pub(crate) reach: Vec>, +} + +/// [`price_node`] over every node of `dag`, children first, with the +/// arguments of [`price_at_scale`]. `root_scan_lookback_ms` is the +/// look-back of a query-time scan no window above it in `dag` bounds (a +/// scan priced without its consumers); `None` reads +/// [`DEFAULT_LOOKBACK_MS`] as Stage 3 does. `root_rates`, when given, are +/// the computations per second of each root's consumer in place of its +/// query's own rate ([`computations`]). +#[allow(clippy::too_many_arguments)] +pub(crate) fn price_each_node<'a>( + dag: &'a PhysicalASAPDAG, + nodes: &NodeMap<'a>, + kept_chains: &HashSet, + recurrence: &[QueryRecurrence], + range_evaluations: &[Option], + data: &DataWorkload, + settings: &'a PlanSelectionSettings, + formula_scale: f64, + root_scan_lookback_ms: Option, + root_rates: Option<&[f64]>, +) -> Result, (PhysicalASAPNodeId, AnalyticalCostError)> { + let first = dag.nodes.first().map(|n| n.id).unwrap_or_default(); + let multipliers = computations( + dag, + recurrence, + range_evaluations, + settings.non_repeating_query_horizon_seconds, + root_rates, + ) + .map_err(|error| (first, error))?; + let reach = query_reach(dag, nodes); + let retention_ms = pane_retention_ms(dag, &reach, range_evaluations); + let mut producers_of: HashMap> = HashMap::new(); + for e in &dag.edges { + producers_of.entry(e.consumer).or_default().push(e.producer); + } + let mut priced: HashMap> = HashMap::new(); + for node in &dag.nodes { + let views: Vec> = producers_of + .get(&node.id) + .map_or(&[][..], Vec::as_slice) + .iter() + .map(|p| { + let produced: &PricedNode<'a> = &priced[p]; + ProducerOutput { + payload: &nodes[p].payload, + output_schema: &nodes[p].output_schema, + output: produced.output, + sources: produced.sources.clone(), + family: produced.family, + } + }) + .collect(); + let lookback_ms = match (&node.payload, node.output_state.timing) { + (Payload::NonASAP(NonASAPOp::Scan { .. }), ExecutionTiming::QueryTime) => { + let windows = windows_read(dag, nodes, node.id, kept_chains); + match (windows.into_iter().max(), root_scan_lookback_ms) { + (Some(window), _) => window, + (None, Some(lookback)) => lookback, + (None, None) => DEFAULT_LOOKBACK_MS, + } + } + _ => 0, + }; + let node_priced = price_node( + node, + &views, + multipliers[&node.id], + lookback_ms, + retained_panes(node, &retention_ms), + (node.output_state.timing == ExecutionTiming::IngestionTime) + .then(|| maintained_window_ms(node.id, nodes, &producers_of)) + .flatten(), + data, + settings, + formula_scale, + ) + .map_err(|error| (node.id, error))?; + priced.insert(node.id, node_priced); + } + Ok(EachNodePriced { + priced, + multipliers, + reach, + }) +} + +/// The state family of a node with `payload` whose first producer has +/// `producer_family`: a build's own (one sketch instance per group; a +/// shared Hydra structure is not that sketch), inherited by the merges, +/// estimates and finalizations over it, none otherwise. +fn state_family<'a>( + payload: &'a Payload, + producer_family: Option>, +) -> StateFamily<'a> { + match payload { + Payload::ASAP( + ASAPOp::SummaryAgg { + family, grouping, .. + } + | ASAPOp::PaneChain { + family, grouping, .. + }, + ) => match family { + FieldDataType::Sketch(kind, _) + if *grouping == GroupingStrategy::PerSubpopulationInstance => + { + StateFamily::Sketch(kind.params()) + } + FieldDataType::ExactAggregate(kind, _) => StateFamily::Exact(kind), + _ => StateFamily::None, + }, + Payload::ASAP( + ASAPOp::SummaryMerge { .. } + | ASAPOp::SummaryEstimate { .. } + | ASAPOp::FinalizeExactAccumulator { .. }, + ) => producer_family.unwrap_or(StateFamily::None), + _ => StateFamily::None, + } +} + /// One producer of a node being priced ([`price_node`]): its payload and /// schema, its output statistics, the sources it reads and its state /// family. @@ -3632,32 +4188,7 @@ pub(crate) fn price_node<'a>( // The rows per ms its sources receive, which its scans and time // ranges read. let rows_per_ms = source_rows_per_ms(&sources, data); - let family = match &node.payload { - Payload::ASAP( - ASAPOp::SummaryAgg { - family, grouping, .. - } - | ASAPOp::PaneChain { - family, grouping, .. - }, - ) => match family { - FieldDataType::Sketch(kind, _) - if *grouping == GroupingStrategy::PerSubpopulationInstance => - { - StateFamily::Sketch(kind.params()) - } - FieldDataType::ExactAggregate(kind, _) => StateFamily::Exact(kind), - _ => StateFamily::None, - }, - Payload::ASAP( - ASAPOp::SummaryMerge { .. } - | ASAPOp::SummaryEstimate { .. } - | ASAPOp::FinalizeExactAccumulator { .. }, - ) => producers - .first() - .map_or(StateFamily::None, |p| p.family), - _ => StateFamily::None, - }; + let family = state_family(&node.payload, producers.first().map(|p| p.family)); let operator = OperatorKind::of(&node.payload); let unit_price = match &node.payload { Payload::NonASAP(NonASAPOp::Aggregate { measures, .. }) => { @@ -3710,14 +4241,14 @@ pub(crate) fn price_node<'a>( // maintains (one second of them does not bound its groups). let rows_read = match ingestion { false => Some(input.rows.max(1)), - true => ingestion_window_ms.map(|window_ms| ((rows_per_ms * window_ms as f64).round() as u64).max(1)), + true => ingestion_window_ms + .map(|window_ms| ((rows_per_ms * window_ms as f64).round() as u64).max(1)), }; // Query time always has its input rows; an ingestion-time grouping // with no window below it is an error (Stage 2 builds none). let estimated_groups = |keys: &GroupKeys| { - let rows_read = - rows_read.ok_or(AnalyticalCostError::IngestionGroupingWithoutWindow)?; + let rows_read = rows_read.ok_or(AnalyticalCostError::IngestionGroupingWithoutWindow)?; Ok(group_count_bound( keys, input_schema, @@ -3853,9 +4384,7 @@ pub(crate) fn price_node<'a>( // A range keeps its window's samples; an instant // selector the latest sample of each series. let kept = match kind { - TimeRangeKind::Range => { - (rows_per_ms * range.as_millis() as f64).round() as u64 - } + TimeRangeKind::Range => (rows_per_ms * range.as_millis() as f64).round() as u64, TimeRangeKind::Instant => series, }; let out = edge(input.rows.min(kept.max(1))); @@ -3981,9 +4510,7 @@ pub(crate) fn price_node<'a>( let mut group_count = 0; for producer in producers { let panes = match producer.payload { - Payload::ASAP(ASAPOp::PaneChain { pane_count, .. }) => { - u64::from(*pane_count) - } + Payload::ASAP(ASAPOp::PaneChain { pane_count, .. }) => u64::from(*pane_count), _ => 1, }; group_count = group_count.max(producer.output.rows / panes.max(1)); @@ -4134,10 +4661,7 @@ pub(crate) fn price_node<'a>( // The retained panes and the one being built. (Materialization::MatAtIngestion(medium), Payload::ASAP(ASAPOp::PaneChain { .. })) => { let panes = panes_of(node); - ( - state / panes * (retained_pane_count + 1), - medium, - ) + (state / panes * (retained_pane_count + 1), medium) } (Materialization::MatAtIngestion(medium), _) => (2 * state, medium), }; diff --git a/crates/plan-selection/src/measured_operation_costs.rs b/crates/plan-selection/src/measured_operation_costs.rs index 40b8fe8f0..dbbf2da2d 100644 --- a/crates/plan-selection/src/measured_operation_costs.rs +++ b/crates/plan-selection/src/measured_operation_costs.rs @@ -206,7 +206,7 @@ impl OperatorKind { /// The state family a summary node works on, which selects its entries. #[derive(Debug, Clone, Copy, PartialEq)] -pub(super) enum StateFamily<'a> { +pub(crate) enum StateFamily<'a> { /// No family: a non-summary operator, or a shared (Hydra) structure. None, Sketch(&'a SketchParams), diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 6057cde38..efe023af3 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -159,16 +159,32 @@ pub(crate) struct Context { pub materialization: Vec<(GroupId, Materialization)>, } -/// One candidate's usage priced in a context (Stage 3's per-node pricing). +/// One candidate's usage priced in a context: Stage 3's per-node pricing of +/// its private nodes ([`pricing::local_price`]). The charges Stage 3 makes +/// only once the whole plan is known are kept apart and added by +/// [`pricing::plan_cost`]: memory held for an evaluation (which needs the +/// readers' whole latency) and the raw data query-time scans keep on disk +/// (each source's longest scan, once). #[derive(Debug, Clone, PartialEq)] pub(crate) struct LocalUsage { + /// Under cost model 1, its private nodes' price per second, without + /// held memory and raw data: a lower bound of what it adds. pub cost_per_second: Option, pub cpu_ops_per_second: f64, pub scanned_bytes_per_second: f64, pub average_memory_bytes: f64, pub peak_memory_bytes: u64, - pub peak_cpu_cores: u32, pub disk_bytes: u64, - /// Each query's latency the candidate adds per evaluation, in ms. + /// Per reading query, the query-time work in ms its private nodes add + /// to the query's latency per evaluation (a query's latency is the sum + /// over its groups plus its longest compaction). pub latency_ms: Vec<(usize, f64)>, + /// Per reading query, the longest compaction in ms of an + /// ingestion-time state among its private nodes that the query reads. + pub compaction_ms: Vec<(usize, f64)>, + /// Memory its private nodes hold while an evaluation runs, readers by + /// query index in the workload. + pub held_memory: Vec, + /// Each query-time raw scan among its private nodes: source and bytes. + pub raw_scans: Vec<(String, u64)>, } diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index fd12a0c8e..7900f1919 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -1,55 +1,142 @@ -//! Local prices per (group, candidate, context), cached, the frozen -//! calibration factor, and frontiers (design "Memo-group search": Local -//! price, Frontier, Coupled charges). +//! Local prices of implementation candidates, cached, and the frozen +//! calibration factor (design "Memo-group search": Local price, Coupled +//! charges). Cost model 1 only: cost models 2 and 3 keep the current +//! search. //! -//! A candidate's local price is the Stage 3 usage of its private nodes: -//! from its node down to, but excluding, its child groups' nodes. It is -//! priced as the difference of two Stage 3 prices of whole sub-DAGs, each -//! built by Stage 2 and priced by [`crate::price_part`] with the readers' -//! entries: +//! A candidate's local price is Stage 3's per-node price +//! ([`crate::price_node`]) of its private nodes, from its node down to, but +//! excluding, its child groups' nodes, computed for the consumers the +//! search fixed top-down ([`ConsumerContext`]) and reading its children's +//! given outputs. Stage 2 builds the candidate's sub-DAG once per +//! materialization; a child's nodes in it are recognized by structure +//! (payload over its producers, and schema) from the child group's own +//! sub-DAG. //! -//! - the candidate's node with everything below it, its children being the -//! child groups' nodes, so its private nodes see their children's output -//! statistics and the readers' rates; -//! - minus each child group's node with everything below it, priced alike. +//! Each of Stage 3's charges is either per node, priced exactly here, or a +//! plan-level term kept apart in [`LocalUsage`] and added by [`plan_cost`]: //! -//! Additive usage (price under cost model 1, CPU, scans, average memory, -//! peak memory, disk) is the difference; peak cores are the candidate's -//! whole sub-DAG's (frontiers combine them by maximum); each reader's -//! latency is what the candidate adds to its longest child's. Summed over a -//! plan's groups, the differences telescope to Stage 3's price of the whole -//! plan, coupled charges included (memory held for a query's whole latency, -//! a source's longest raw scan), which the root's sub-DAG carries. +//! - CPU, scans, kept state (memory, disk), pane retention: per node. A +//! chain's retention is read from the consumers' queries (their range +//! spans), as Stage 3 reads it from the queries reaching it. +//! - Memory held while an evaluation runs (a query-time state, or a +//! recomputed node's working memory): plan level ([`crate::HeldMemory`]), +//! since it lasts the readers' whole latency; [`plan_cost`] prices it with +//! each query's latency, the sum of its groups' work plus its longest +//! compaction. +//! - Raw data on disk: plan level, each source's longest query-time scan +//! ([`LocalUsage::raw_scans`]). +//! - A query's request overhead: per node, on the group that is its root. +//! - Compaction and partial copies of an ingestion-time state: per node, +//! with the ingest workers counted from the private nodes' ingestion +//! work only (exact while a state's whole ingestion work below it fits +//! on one core, as in every current test; otherwise a lower bound). //! -//! Materialization: the context's materializations, in order, are the -//! options of the sub-DAG's builds in Stage 2's discovery order; builds -//! past the list are `NotMaterialized`. A child's builds are the trailing -//! ones of its parent's sub-DAG. +//! Known limits: summaries are sized for the failure probability only in +//! a group that is the root of every query computing it directly; a +//! query-time scan group priced on its own reads the longest window its +//! parent groups' logical nodes put on it. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; use asap_physical_asap_optimization::{optimize_physical_asap_assigned, PhysicalCandidate}; -use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::ir::physical_export::{PhysicalASAPDAG, PhysicalASAPNodeId}; +use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; use asap_types::physical::Materialization; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; -use std::rc::Rc; -use super::{Context, GroupId, ImplementationCandidate, LocalUsage, Memo}; +use super::{GroupId, ImplementationCandidate, LocalUsage, Memo}; +use crate::cost::physical_operator_statistics::EdgeStatistics; use crate::{ - price_part, CalibrationFit, CandidateCost, ClusterLimits, CostModel, PlanningModels, - SelectionError, + accuracy_targets, price_each_node, price_private_nodes, shared_evaluation_rate, + sized_for_failure_probability, CalibrationFit, CostModel, NodeConsumer, PlanningModels, + PrivatePricing, SelectionError, }; -/// The cached prices of one search: valid for one workload, data workload -/// and set of planning models, priced with one frozen calibration factor. -#[derive(Debug)] +/// Who computes a group's private nodes: its consumers, as Stage 3's +/// `computations` derives them from the whole plan. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ConsumerContext { + pub consumers: Vec, +} + +/// One consumer of a group. +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum Consumer { + /// A query reaching the group through unmaterialized nodes only: at + /// the query's rate, range steps included. + Query { query: usize }, + /// A materialized ancestor: at its rate (1 at ingestion time, the + /// queries reaching it together at query time, once per request when + /// kept across evaluations). `queries` are the queries reaching it at + /// query time (none at ingestion time), whose latency and range spans + /// the group's held memory and pane retention read. + Materialized { + rate_per_second: f64, + queries: Vec, + at_ingestion: bool, + }, +} + +/// A consumer as a cache key (its rate by bits). +type ConsumerKey = (Option, u64, Vec, bool); + +impl Consumer { + fn key(&self) -> ConsumerKey { + match self { + Consumer::Query { query } => (Some(*query), 0, Vec::new(), false), + Consumer::Materialized { + rate_per_second, + queries, + at_ingestion, + } => ( + None, + rate_per_second.to_bits(), + queries.clone(), + *at_ingestion, + ), + } + } + + fn queries(&self) -> Vec { + match self { + Consumer::Query { query } => vec![*query], + Consumer::Materialized { queries, .. } => queries.clone(), + } + } + + fn at_ingestion(&self) -> bool { + matches!( + self, + Consumer::Materialized { + at_ingestion: true, + .. + } + ) + } +} + +type PartKey = (usize, Materialization, Vec); +type PriceKey = (u32, Vec, Materialization, Vec<(u64, u64)>); + +/// The cached prices of one search over one memo: valid while the memo's +/// nodes live (parts are keyed by address), for one workload, data +/// workload and set of planning models, priced with one frozen +/// calibration factor. +#[derive(Debug, Default)] pub(crate) struct PriceCache { /// The frozen calibration factor ([`frozen_calibration`]). - formula_scale: f64, - /// Local prices by (group, candidate's node, context, readers). - local: HashMap>, - /// Whole sub-DAG prices by (node, readers, materializations). - trees: HashMap>, + formula_scale: Option, + /// Structural node keys, interned. + keys: HashMap, + /// Each candidate's sub-DAG by (node address, materialization, the + /// queries it is built for), with its nodes' keys. + parts: HashMap>>, + /// Each child node's sub-DAG's node keys and root key, by its address. + child_keys: HashMap>, + /// Local prices by (candidate's structural key, consumers, + /// materialization, children's outputs). + prices: HashMap>, /// Local prices found in the cache. pub hits: u64, /// Local prices computed. @@ -60,435 +147,636 @@ impl PriceCache { /// An empty cache pricing with the calibration factor `formula_scale`. pub(crate) fn new(formula_scale: f64) -> Self { Self { - formula_scale, - local: HashMap::new(), - trees: HashMap::new(), - hits: 0, - misses: 0, + formula_scale: Some(formula_scale), + ..Self::default() } } } -/// A node by its structure: its serialized tree (a hash alone does not -/// prove two nodes equal). -type NodeKey = String; -type LocalKey = (GroupId, NodeKey, Context, Vec); -type TreeKey = (NodeKey, Vec, Vec); - -/// One whole sub-DAG's Stage 3 price. -#[derive(Debug, Clone, PartialEq)] -struct TreePrice { - usage: LocalUsage, - /// The builds Stage 2 found in it, in discovery order. - builds: usize, +/// A candidate's sub-DAG as Stage 2 builds it, with its nodes' keys. +#[derive(Debug)] +struct Part { + physical: PhysicalCandidate, + keys: Vec, } -fn node_key(node: &OperatorNode) -> NodeKey { - serde_json::to_string(node).expect("an operator node serializes") +/// A child group's sub-DAG's node keys, and its root's. +#[derive(Debug, Default)] +struct ChildKeys { + nodes: HashSet, + root: Option, } -/// `candidate` of `group` priced in `context` for the queries `readers` (by -/// index in the workload, whose `entries` they are): the Stage 3 usage of -/// its private nodes (module doc). `None` when it fails a check sharing -/// cannot repair, or Stage 2 does not build it with the context's -/// materializations. +/// The usage of `candidate`'s private nodes (from its node down to, +/// excluding, its child groups' nodes) computed for `consumers`, +/// materialized as `materialization` (which makes it the consumer of its +/// children), reading its children's outputs `child_stats` (in the order +/// of `candidate.children`). Cost model 1, the frozen calibration factor. +/// `None` when a check sharing cannot repair fails (a per-series aggregate +/// the executor lacks, a query-time scan of raw data the deployment does +/// not keep) or Stage 2 does not build it so. #[allow(clippy::too_many_arguments)] -pub(crate) fn local_price( +pub(crate) fn private_usage( memo: &Memo, - group: GroupId, candidate: &ImplementationCandidate, - context: &Context, - readers: &[usize], + consumers: &ConsumerContext, + materialization: Materialization, + child_stats: &[EdgeStatistics], entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, cache: &mut PriceCache, ) -> Result, SelectionError> { - let key = ( - group, - node_key(&candidate.node), - context.clone(), - readers.to_vec(), + Ok(priced( + memo, + candidate, + consumers, + materialization, + child_stats, + entries, + data, + models, + cache, + )? + .map(|(usage, _)| usage)) +} + +/// `candidate`'s output per computation given its children's outputs +/// `child_stats`, computed at query time for its group's queries with +/// nothing materialized; `None` when it is not priced. +#[allow(clippy::too_many_arguments)] +pub(crate) fn output_stats( + memo: &Memo, + candidate: &ImplementationCandidate, + child_stats: &[EdgeStatistics], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + cache: &mut PriceCache, +) -> Result, SelectionError> { + let group = &memo.groups[candidate.implements[0].0 as usize]; + let consumers = ConsumerContext { + consumers: group + .queries + .iter() + .map(|&query| Consumer::Query { query }) + .collect(), + }; + Ok(priced( + memo, + candidate, + &consumers, + Materialization::NotMaterialized, + child_stats, + entries, + data, + models, + cache, + )? + .map(|(_, output)| output)) +} + +#[allow(clippy::too_many_arguments)] +fn priced( + memo: &Memo, + candidate: &ImplementationCandidate, + consumers: &ConsumerContext, + materialization: Materialization, + child_stats: &[EdgeStatistics], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + cache: &mut PriceCache, +) -> Result, SelectionError> { + let group = candidate.implements[0]; + let roots_of = &memo.groups[group.0 as usize].roots_of; + // The queries computing it directly; the part is built for them (one + // root each) and sized for them when it is all their roots. + let direct: Vec = consumers + .consumers + .iter() + .filter_map(|c| match c { + Consumer::Query { query } => Some(*query), + _ => None, + }) + .collect(); + let built_for = match direct.is_empty() { + false => direct.clone(), + true => memo.groups[group.0 as usize] + .queries + .first() + .copied() + .into_iter() + .collect(), + }; + let sized = !direct.is_empty() + && direct.len() == consumers.consumers.len() + && direct.iter().all(|q| roots_of.contains(q)); + let Some(part) = part( + candidate, + materialization, + &built_for, + sized, + entries, + data, + &models, + cache, + ) else { + return Ok(None); + }; + let key: PriceKey = ( + root_key(&part.physical.dag, &part.keys).unwrap_or(u32::MAX), + consumers.consumers.iter().map(Consumer::key).collect(), + materialization, + child_stats.iter().map(|s| (s.rows, s.bytes)).collect(), ); - if let Some(price) = cache.local.get(&key) { + if let Some(price) = cache.prices.get(&key) { cache.hits += 1; return Ok(price.clone()); } cache.misses += 1; - let reader_entries: Vec = - readers.iter().map(|&q| entries[q].clone()).collect(); - let materializations: Vec = - context.materialization.iter().map(|(_, m)| *m).collect(); - let pricing = TreePricing { - readers, - entries: &reader_entries, + let price = price_part( + memo, + candidate, + &part, + consumers, + child_stats, + entries, data, - models: &models, - }; - let Some(whole) = pricing.price(&candidate.node, &materializations, cache)? else { - cache.local.insert(key, None); - return Ok(None); - }; - let mut children = Vec::with_capacity(candidate.children.len()); - // Each child's builds trail its parent's, in child order. - let mut end = whole.builds; - let child_builds: Vec = candidate - .children - .iter() - .map(|c| pricing.builds(&memo.groups[c.0 as usize].node)) - .collect(); - let mut offsets = vec![0; child_builds.len()]; - for (i, builds) in child_builds.iter().enumerate().rev() { - end = end.saturating_sub(*builds); - offsets[i] = end; - } - for (child, (&offset, &builds)) in candidate - .children - .iter() - .zip(offsets.iter().zip(&child_builds)) - { - let node = &memo.groups[child.0 as usize].node; - let own = materializations - .get(offset..(offset + builds).min(materializations.len())) - .unwrap_or(&[]); - // A child that does not price alone (e.g. a raw scan only its - // parent's ingestion build makes valid) is charged inside its - // parent: nothing is subtracted for it. - if let Some(price) = pricing.price(node, own, cache)? { - children.push(price.usage); - } - } - let local = difference(&whole.usage, &children, &models.plan_selection.cost_model); - cache.local.insert(key, Some(local.clone())); - Ok(Some(local)) + &models, + cache, + )?; + cache.prices.insert(key, price.clone()); + Ok(price) } -/// Pricing whole sub-DAGs for one set of readers. -struct TreePricing<'a> { - readers: &'a [usize], - entries: &'a [QueryWorkloadEntry], - data: &'a DataWorkload, - models: &'a PlanningModels<'a>, +/// The key of `dag`'s first root, `keys` being its nodes' by position. +fn root_key(dag: &PhysicalASAPDAG, keys: &[u32]) -> Option { + let root = dag.roots.first()?; + let i = dag.nodes.iter().position(|n| n.id == *root)?; + Some(keys[i]) } -impl TreePricing<'_> { - /// Stage 2's sub-DAG of `node`, one root per reader, its builds given - /// `materializations` in order (`NotMaterialized` past them), and the - /// number of builds; `None` when Stage 2 does not build it so. - fn built( - &self, - node: &Rc, - materializations: &[Materialization], - ) -> Option<(PhysicalCandidate, usize)> { - let roots = vec![QueryRoot::Operator(Rc::clone(node)); self.readers.len()]; - let mut next = materializations.iter().copied(); - let mut builds = 0; +/// `candidate`'s sub-DAG, built by Stage 2 for the queries `built_for` +/// (one root each) with `materialization` on its first materialization +/// candidate (its topmost build, or its node when several read it) and +/// nothing on the others, sized for the queries' failure probabilities +/// when `sized`; cached. +#[allow(clippy::too_many_arguments)] +fn part( + candidate: &ImplementationCandidate, + materialization: Materialization, + built_for: &[usize], + sized: bool, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, + cache: &mut PriceCache, +) -> Option> { + let mut queries = built_for.to_vec(); + if sized { + queries.push(usize::MAX); + } + let key = ( + Rc::as_ptr(&candidate.node) as usize, + materialization, + queries, + ); + if let Some(part) = cache.parts.get(&key) { + return part.clone(); + } + let part_entries: Vec = + built_for.iter().map(|&q| entries[q].clone()).collect(); + let physical = built( + &candidate.node, + &part_entries, + data, + models, + materialization, + ) + .map(|physical| match sized { + true => sized_for_failure_probability(&physical, &accuracy_targets(&part_entries), data) + .unwrap_or(physical), + false => physical, + }); + let part = physical.map(|physical| { + let keys = keys_of(&physical.dag, &mut cache.keys); + Rc::new(Part { physical, keys }) + }); + cache.parts.insert(key, part.clone()); + part +} + +/// Stage 2's sub-DAG of `node`, one root per entry (two copies of a +/// single entry for `MatAtQuery`, which needs two readers), with +/// `materialization` on one materialization candidate and nothing on the +/// others: `MatAtQuery` on its node (the last candidate, the highest +/// shared node), any other on its topmost build (the first). `None` when +/// Stage 2 does not build that. +fn built( + node: &Rc, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, + materialization: Materialization, +) -> Option { + let at_query = matches!(materialization, Materialization::MatAtQuery(_)); + let entries: Vec = match (at_query, entries) { + (true, [one]) => vec![one.clone(), one.clone()], + _ => entries.to_vec(), + }; + let roots = vec![QueryRoot::Operator(Rc::clone(node)); entries.len()]; + let build = |target: Option| { + let mut index = 0; let mut choose = |_: &[usize]| { - builds += 1; - Some(next.next().unwrap_or(Materialization::NotMaterialized)) + let chosen = match Some(index) == target { + true => materialization, + false => Materialization::NotMaterialized, + }; + index += 1; + Some(chosen) }; let built = optimize_physical_asap_assigned( "S", &roots, - self.entries, - self.data, - self.models.deployment_capabilities, + &entries, + data, + models.deployment_capabilities, &mut choose, - ); - let physical = built.ok().flatten()?; - Some((physical, builds)) + ) + .ok() + .flatten(); + (built, index) + }; + let (plain, candidates) = build(None); + if materialization == Materialization::NotMaterialized { + return plain; } + let target = match at_query { + true => candidates.checked_sub(1)?, + false => 0, + }; + let built = build(Some(target)).0?; + built + .dag + .nodes + .iter() + .any(|n| n.materialization == materialization) + .then_some(built) +} - /// The builds Stage 2 finds in `node`'s sub-DAG. - fn builds(&self, node: &Rc) -> usize { - self.built(node, &[]).map_or(0, |(_, builds)| builds) +/// The node keys of `child`'s own sub-DAG (one reader, nothing +/// materialized) and its root's key, cached by the child's node. +fn child_keys( + memo: &Memo, + child: GroupId, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, + cache: &mut PriceCache, +) -> Rc { + let node = &memo.groups[child.0 as usize].node; + let address = Rc::as_ptr(node) as usize; + if let Some(keys) = cache.child_keys.get(&address) { + return Rc::clone(keys); } + let reader_entries: Vec = memo.groups[child.0 as usize] + .queries + .first() + .map(|&q| entries[q].clone()) + .into_iter() + .collect(); + let keys = built( + node, + &reader_entries, + data, + models, + Materialization::NotMaterialized, + ) + .map(|physical| { + let keys = keys_of(&physical.dag, &mut cache.keys); + ChildKeys { + root: root_key(&physical.dag, &keys), + nodes: keys.into_iter().collect(), + } + }) + .unwrap_or_default(); + let keys = Rc::new(keys); + cache.child_keys.insert(address, Rc::clone(&keys)); + keys +} - /// `node`'s whole sub-DAG priced by Stage 3, cached. - fn price( - &self, - node: &Rc, - materializations: &[Materialization], - cache: &mut PriceCache, - ) -> Result, SelectionError> { - let key = ( - node_key(node), - self.readers.to_vec(), - materializations.to_vec(), - ); - if let Some(price) = cache.trees.get(&key) { - return Ok(price.clone()); +/// Each node's structural key, by position: its payload over its +/// producers' keys, and its output schema; placement (timing, +/// materialization) left out, so a child's nodes match whatever is chosen +/// for them. +fn keys_of(dag: &PhysicalASAPDAG, interned: &mut HashMap) -> Vec { + let mut key_of: HashMap = HashMap::new(); + dag.nodes + .iter() + .map(|node| { + let payload = node.payload.map_children(|id| key_of.get(id).copied()); + let text = serde_json::to_string(&(&payload, &node.output_schema)) + .expect("a physical node serializes"); + let next = interned.len() as u32; + let key = *interned.entry(text).or_insert(next); + key_of.insert(node.id, key); + key + }) + .collect() +} + +#[allow(clippy::too_many_arguments)] +fn price_part( + memo: &Memo, + candidate: &ImplementationCandidate, + part: &Part, + consumers: &ConsumerContext, + child_stats: &[EdgeStatistics], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, + cache: &mut PriceCache, +) -> Result, SelectionError> { + let dag = &part.physical.dag; + // The children's nodes, and each child root's output. + let mut child_nodes: HashSet = HashSet::new(); + let mut child_roots: Vec> = Vec::new(); + for &child in &candidate.children { + let keys = child_keys(memo, child, entries, data, models, cache); + child_nodes.extend(keys.nodes.iter().copied()); + child_roots.push(keys.root); + } + let private: Vec = part.keys.iter().map(|k| !child_nodes.contains(k)).collect(); + let mut child_outputs = HashMap::new(); + for (node, key) in dag.nodes.iter().zip(&part.keys) { + if let Some(i) = child_roots.iter().position(|k| *k == Some(*key)) { + if let Some(stats) = child_stats.get(i) { + child_outputs.insert(node.id, *stats); + } + } + } + // The queries the consumers name, by position. + let mut queries: Vec = Vec::new(); + for consumer in &consumers.consumers { + for q in consumer.queries() { + if !queries.contains(&q) { + queries.push(q); + } + } + } + let position = |q: usize| queries.iter().position(|&p| p == q).expect("named"); + let query_entries: Vec = + queries.iter().map(|&q| entries[q].clone()).collect(); + let roots_of = &memo.groups[candidate.implements[0].0 as usize].roots_of; + let root_of: Vec = queries.iter().map(|q| roots_of.contains(q)).collect(); + let settings = models.plan_selection; + let cost_error = |(node, error)| SelectionError::Cost { + candidate: format!("group {}", candidate.implements[0].0), + node, + error, + }; + let node_consumer = |consumer: &Consumer| -> Result { + Ok(match consumer { + Consumer::Query { query } => NodeConsumer { + rate_per_second: shared_evaluation_rate(&[&entries[*query]], settings) + .map_err(|error| cost_error((Default::default(), error)))?, + queries: vec![position(*query)], + at_ingestion: false, + }, + Consumer::Materialized { + rate_per_second, + queries, + at_ingestion, + } => NodeConsumer { + rate_per_second: *rate_per_second, + queries: queries.iter().map(|&q| position(q)).collect(), + at_ingestion: *at_ingestion, + }, + }) + }; + // A copy at query time for the query-time consumers, one at ingestion + // time for the others (Stage 2 builds one per timing). + let mut usage = LocalUsage { + cost_per_second: Some(0.0), + cpu_ops_per_second: 0.0, + scanned_bytes_per_second: 0.0, + average_memory_bytes: 0.0, + peak_memory_bytes: 0, + disk_bytes: 0, + latency_ms: queries.iter().map(|&q| (q, 0.0)).collect(), + compaction_ms: queries.iter().map(|&q| (q, 0.0)).collect(), + held_memory: Vec::new(), + raw_scans: Vec::new(), + }; + let mut output = None; + for at_ingestion in [false, true] { + let copy: Vec = consumers + .consumers + .iter() + .filter(|c| c.at_ingestion() == at_ingestion) + .map(node_consumer) + .collect::>()?; + if copy.is_empty() { + continue; } - let price = match self.built(node, materializations) { - None => None, - Some((physical, builds)) => price_part( - &physical, - self.entries, - self.data, - self.models, - cache.formula_scale, - )? - .map(|(cost, _)| TreePrice { - usage: usage_of(&cost, self.readers), - builds, - }), + let input = PrivatePricing { + dag, + private: &private, + child_outputs: &child_outputs, + consumers: ©, + queries: &query_entries, + root_of: &root_of, + scan_lookback_ms: scan_lookback_ms(memo, candidate.implements[0]), + data, + settings, + formula_scale: cache.formula_scale.unwrap_or(1.0), }; - cache.trees.insert(key, price.clone()); - Ok(price) + let Some(nodes) = + price_private_nodes(&input, models.deployment_capabilities.raw_data_retained) + .map_err(cost_error)? + else { + return Ok(None); + }; + usage.cost_per_second = match (usage.cost_per_second, nodes.cost_per_second) { + (Some(a), Some(b)) => Some(a + b), + _ => None, + }; + usage.cpu_ops_per_second += nodes.cpu_ops_per_second; + usage.scanned_bytes_per_second += nodes.scanned_bytes_per_second; + usage.average_memory_bytes += nodes.average_memory_bytes; + usage.peak_memory_bytes = usage + .peak_memory_bytes + .saturating_add(nodes.peak_memory_bytes); + usage.disk_bytes = usage.disk_bytes.saturating_add(nodes.disk_bytes); + for (i, ms) in nodes.query_ms.iter().enumerate() { + usage.latency_ms[i].1 += ms; + } + for (i, ms) in nodes.compaction_ms.iter().enumerate() { + usage.compaction_ms[i].1 = usage.compaction_ms[i].1.max(*ms); + } + usage + .held_memory + .extend(nodes.held.into_iter().map(|mut held| { + for (reader, _) in &mut held.readers { + *reader = queries[*reader]; + } + held + })); + usage.raw_scans.extend(nodes.raw_scans); + output.get_or_insert(nodes.output); } + Ok(output.map(|output| (usage, output))) } -/// `cost` as a usage, each root's latency that of the reader at its index. -fn usage_of(cost: &CandidateCost, readers: &[usize]) -> LocalUsage { - let usage = &cost.usage; - LocalUsage { - cost_per_second: cost.cost_per_second, - cpu_ops_per_second: usage.cpu_ops_per_second, - scanned_bytes_per_second: usage.scanned_bytes_per_second, - average_memory_bytes: usage.average_memory_bytes, - peak_memory_bytes: usage.peak_memory_bytes, - peak_cpu_cores: usage.peak_cpu_cores, - disk_bytes: usage.disk_bytes, - latency_ms: readers - .iter() - .zip(&usage.latency_per_query) - .map(|(&q, latency)| (q, latency.total_ms())) - .collect(), +/// The raw data a query-time scan group priced on its own reads: the +/// longest window its parent groups' nodes put on it (through time +/// shifts); `None` when it is not a scan or no window bounds it. +fn scan_lookback_ms(memo: &Memo, group: GroupId) -> Option { + let scan = &memo.groups[group.0 as usize]; + if !matches!(scan.node.non_asap(), Some(NonASAPOp::Scan { .. })) { + return None; } + let mut windows = Vec::new(); + let mut stack = scan.parents.clone(); + let mut seen = HashSet::new(); + while let Some(g) = stack.pop() { + if !seen.insert(g) { + continue; + } + let parent = &memo.groups[g.0 as usize]; + match parent.node.non_asap() { + Some(NonASAPOp::TimeRange { range, .. }) => windows.push(range.as_millis() as u64), + Some(NonASAPOp::TimeShift { .. }) => stack.extend(parent.parents.iter().copied()), + _ => {} + } + } + windows.into_iter().max() } -/// `whole` minus `children`: additive components subtracted, peak cores -/// `whole`'s, each reader's latency over its longest child's. Under cost -/// model 2 the price is recomputed from the difference's peaks. -fn difference(whole: &LocalUsage, children: &[LocalUsage], cost_model: &CostModel) -> LocalUsage { - let mut local = whole.clone(); - for child in children { - local.cost_per_second = match (local.cost_per_second, child.cost_per_second) { - (Some(a), Some(b)) => Some(a - b), - (a, _) => a, - }; - local.cpu_ops_per_second -= child.cpu_ops_per_second; - local.scanned_bytes_per_second -= child.scanned_bytes_per_second; - local.average_memory_bytes -= child.average_memory_bytes; - local.peak_memory_bytes = local - .peak_memory_bytes - .saturating_sub(child.peak_memory_bytes); - local.disk_bytes = local.disk_bytes.saturating_sub(child.disk_bytes); - } - for (query, latency) in &mut local.latency_ms { - let longest_child = children - .iter() - .flat_map(|c| &c.latency_ms) - .filter(|(q, _)| q == query) - .map(|(_, l)| *l) - .fold(0.0, f64::max); - *latency -= longest_child; +/// The whole plan's price per second under cost model 1 from its groups' +/// local usages: their prices, plus the memory each holds while an +/// evaluation runs (each query's latency being the sum of its groups' +/// work plus its longest compaction), plus each source's longest +/// query-time raw scan kept on disk. `None` under cost models 2 and 3. +pub(crate) fn plan_cost<'u>( + parts: impl IntoIterator + Clone, + cost_model: &CostModel, +) -> Option { + let CostModel::UsageOverTime(prices) = cost_model else { + return None; + }; + let mut cost = 0.0; + let mut work: HashMap = HashMap::new(); + let mut compaction: HashMap = HashMap::new(); + let mut raw: HashMap<&str, u64> = HashMap::new(); + for part in parts.clone() { + cost += part.cost_per_second?; + for &(q, ms) in &part.latency_ms { + *work.entry(q).or_default() += ms; + } + for &(q, ms) in &part.compaction_ms { + let longest = compaction.entry(q).or_default(); + *longest = longest.max(ms); + } + for (source, bytes) in &part.raw_scans { + let longest = raw.entry(source.as_str()).or_default(); + *longest = (*longest).max(*bytes); + } } - if let CostModel::ProvisionedCapacity(_) = cost_model { - local.cost_per_second = provisioned_cost(&local, cost_model); + let latency_ms = |q: usize| { + work.get(&q).copied().unwrap_or(0.0) + compaction.get(&q).copied().unwrap_or(0.0) + }; + for part in parts { + for held in &part.held_memory { + cost += prices.cost_per_memory_byte_second * held.byte_seconds(latency_ms); + } } - local + cost += raw + .values() + .map(|&bytes| prices.cost_per_disk_byte_second * bytes as f64) + .sum::(); + Some(cost) } /// The calibration factor of the formulas' CPU time to a measured cost -/// table, fitted once over every candidate of every group, each with all -/// its group's queries as readers, before the search, and frozen (design -/// "Coupled charges"); 1 without a measured table. +/// table, fitted once before the search over every candidate's private +/// nodes (each candidate's sub-DAG built for its group's queries with +/// nothing materialized) and frozen (design "Coupled charges"); 1 without +/// a measured table. pub(crate) fn frozen_calibration( memo: &Memo, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, ) -> Result { - if models - .plan_selection - .measured_operation_costs - .operators - .is_empty() - { + let settings = models.plan_selection; + if settings.measured_operation_costs.operators.is_empty() { return Ok(1.0); } let mut fit = CalibrationFit::default(); - for (group, candidates) in memo.groups.iter().zip(&memo.candidates) { - let reader_entries: Vec = - group.queries.iter().map(|&q| entries[q].clone()).collect(); - let pricing = TreePricing { - readers: &group.queries, - entries: &reader_entries, - data, - models: &models, - }; + let mut cache = PriceCache::new(1.0); + for (g, candidates) in memo.candidates.iter().enumerate() { + let queries = memo.groups[g].queries.clone(); + let part_entries: Vec = + queries.iter().map(|&q| entries[q].clone()).collect(); + let recurrence: Vec<_> = part_entries.iter().map(|e| e.recurrence.clone()).collect(); + let ranges: Vec<_> = part_entries.iter().map(|e| e.range_evaluation).collect(); for candidate in candidates { - if let Some((physical, _)) = pricing.built(&candidate.node, &[]) { - if let Some((_, candidate_fit)) = - price_part(&physical, &reader_entries, data, &models, 1.0)? - { - fit.add(candidate_fit); - } - } - } - } - Ok(fit.factor()) -} - -/// Under cost model 2, `usage`'s provisioned price; otherwise its own. -fn provisioned_cost(usage: &LocalUsage, cost_model: &CostModel) -> Option { - match cost_model { - CostModel::ProvisionedCapacity(prices) => Some( - prices.cost_per_provisioned_memory_byte_second * usage.peak_memory_bytes as f64 - + prices.cost_per_provisioned_cpu_core_second * f64::from(usage.peak_cpu_cores) - + prices.cost_per_disk_byte_second * usage.disk_bytes as f64, - ), - _ => usage.cost_per_second, - } -} - -/// The candidates of a subproblem no other dominates in usage: under cost -/// model 1 the cheapest (Stage 3's tie tolerance; ties keep the earlier), -/// under cost models 2 and 3 the Pareto front over (price, peak memory, -/// peak cores, disk, worst latency), exact ties kept. Under cost model 3, -/// usage over the cluster's limits is never kept. -#[derive(Debug, Clone)] -pub(crate) struct Frontier { - cost_model: CostModel, - /// The kept usages with what each stands for, in insertion order. - pub entries: Vec<(LocalUsage, T)>, -} - -impl Frontier { - pub(crate) fn new(cost_model: &CostModel) -> Self { - Self { - cost_model: cost_model.clone(), - entries: Vec::new(), - } - } - - /// Add `usage` standing for `item` unless a kept entry dominates it, - /// dropping the kept entries it dominates; whether it was kept. - pub(crate) fn insert(&mut self, usage: LocalUsage, item: T) -> bool { - if let CostModel::ClusterCapacity(limits) = &self.cost_model { - if over_limits(&usage, limits) { - return false; + let Some(part) = part( + candidate, + Materialization::NotMaterialized, + &queries, + false, + entries, + data, + &models, + &mut cache, + ) else { + continue; + }; + let mut children = HashSet::new(); + for &child in &candidate.children { + let keys = child_keys(memo, child, entries, data, &models, &mut cache); + children.extend(keys.nodes.iter().copied()); } - } - if let CostModel::UsageOverTime(_) = self.cost_model { - let price = usage.cost_per_second.unwrap_or(f64::INFINITY); - match self.entries.first() { - Some((best, _)) - if !cheaper(price, best.cost_per_second.unwrap_or(f64::INFINITY)) => - { - return false - } - _ => { - self.entries.clear(); - self.entries.push((usage, item)); - return true; + let dag = &part.physical.dag; + let nodes = dag.nodes.iter().map(|n| (n.id, n)).collect(); + let each = price_each_node( + dag, + &nodes, + &HashSet::new(), + &recurrence, + &ranges, + data, + settings, + 1.0, + None, + None, + ) + .map_err(|(node, error)| SelectionError::Cost { + candidate: format!("group {g}"), + node, + error, + })?; + for (node, key) in dag.nodes.iter().zip(&part.keys) { + if !children.contains(key) { + fit.add(each.priced[&node.id].fit); } } } - let point = front_point(&usage); - if self - .entries - .iter() - .any(|(kept, _)| dominates(&front_point(kept), &point)) - { - return false; - } - self.entries - .retain(|(kept, _)| !dominates(&point, &front_point(kept))); - self.entries.push((usage, item)); - true - } - - /// Every pair of `self`'s and `other`'s entries combined, two disjoint - /// parts of a plan: additive usage summed, peak memory summed (Stage 3 - /// adds every node's peak), peak cores by maximum, each query's latency - /// by maximum, the price summed (recomputed from the peaks under cost - /// model 2); `join` combines the items. Kept as [`Self::insert`] keeps. - pub(crate) fn combine( - &self, - other: &Frontier, - join: impl Fn(&T, &U) -> V, - ) -> Frontier { - let mut combined = Frontier::new(&self.cost_model); - for (a, x) in &self.entries { - for (b, y) in &other.entries { - combined.insert(sum(a, b, &self.cost_model), join(x, y)); - } - } - combined } -} - -/// `a` and `b` of two disjoint parts of a plan, combined ([`Frontier::combine`]). -fn sum(a: &LocalUsage, b: &LocalUsage, cost_model: &CostModel) -> LocalUsage { - let mut latency_ms = a.latency_ms.clone(); - for &(query, latency) in &b.latency_ms { - match latency_ms.iter_mut().find(|(q, _)| *q == query) { - Some((_, l)) => *l = l.max(latency), - None => latency_ms.push((query, latency)), - } - } - latency_ms.sort_by_key(|(q, _)| *q); - let mut usage = LocalUsage { - cost_per_second: match (a.cost_per_second, b.cost_per_second) { - (Some(x), Some(y)) => Some(x + y), - _ => None, - }, - cpu_ops_per_second: a.cpu_ops_per_second + b.cpu_ops_per_second, - scanned_bytes_per_second: a.scanned_bytes_per_second + b.scanned_bytes_per_second, - average_memory_bytes: a.average_memory_bytes + b.average_memory_bytes, - peak_memory_bytes: a.peak_memory_bytes.saturating_add(b.peak_memory_bytes), - peak_cpu_cores: a.peak_cpu_cores.max(b.peak_cpu_cores), - disk_bytes: a.disk_bytes.saturating_add(b.disk_bytes), - latency_ms, - }; - usage.cost_per_second = provisioned_cost(&usage, cost_model); - usage -} - -fn over_limits(usage: &LocalUsage, limits: &ClusterLimits) -> bool { - usage.peak_memory_bytes > limits.memory_bytes - || usage.disk_bytes > limits.disk_bytes - || limits - .cpu_cores - .is_some_and(|cores| usage.peak_cpu_cores > cores) -} - -/// Stage 3's price comparison: `cost` is cheaper than `best` beyond its -/// tie tolerance (1e-9). -fn cheaper(cost: f64, best: f64) -> bool { - cost < best - best.abs() * 1e-9 -} - -/// The Pareto front's coordinates: price (when there is one), peak memory, -/// peak cores, disk, worst latency. -fn front_point(usage: &LocalUsage) -> Vec { - let mut point: Vec = usage.cost_per_second.into_iter().collect(); - point.extend([ - usage.peak_memory_bytes as f64, - f64::from(usage.peak_cpu_cores), - usage.disk_bytes as f64, - usage.latency_ms.iter().map(|(_, l)| *l).fold(0.0, f64::max), - ]); - point -} - -/// Whether `a` is no greater than `b` everywhere and smaller somewhere -/// (points of different shapes are incomparable). -fn dominates(a: &[f64], b: &[f64]) -> bool { - a.len() == b.len() - && a.iter().zip(b).all(|(x, y)| x <= y) - && a.iter().zip(b).any(|(x, y)| x < y) + Ok(fit.factor()) } #[cfg(test)] mod tests { use super::*; - use crate::pass::subexpr::{CandidateOrigin, MemoGroup}; - use crate::{ClusterObjective, ResourceUnitCost}; + use crate::pass::subexpr::groups::memo_of; + use asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates; use asap_logical_asap_optimization::test_support::{entry_with, lower_promql}; use asap_types::ir::properties::AccuracyTarget; + use asap_types::physical::Medium; use asap_types::workload::{ DataArrival, QueryRecurrence, RepeatedDemand, RepetitionInterval, SourceFacts, }; @@ -510,210 +798,463 @@ mod tests { } } - fn entries() -> Vec { - let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); - vec![entry_with( - target, - QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval(60_000))), - )] - } - - /// A one-query workload's selected plan split into two groups, its root - /// node over the rest, and the context of the selected - /// materializations. - fn split_plan(models: PlanningModels<'_>) -> (Memo, Context, Option) { - let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); - let root = lower_promql("quantile_over_time(0.99, m[5m])", target); - let run = crate::pass::run_stage1_to_stage3( - vec![(0, QueryRoot::Operator(root))], - &entries(), - &data(), - models, + /// Queries with exact targets repeating every + /// `intervals_ms`. + fn entries(intervals_ms: &[u32]) -> Vec { + intervals_ms + .iter() + .map(|&interval| { + entry_with( + AccuracyTarget::exact(), + QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval( + interval, + ))), + ) + }) + .collect() + } + + /// The memo of PromQL `queries` with exact targets. + fn memo(queries: &[&str]) -> Memo { + let roots: Vec<(usize, QueryRoot)> = queries + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, AccuracyTarget::exact()))) + .enumerate() + .collect(); + memo_of(&generate_per_computation_candidates(roots).unwrap()) + } + + /// The candidate chosen for group `g`: its first. + fn own(memo: &Memo, g: usize) -> &ImplementationCandidate { + &memo.candidates[g][0] + } + + /// Group `g`'s chosen candidate over its children's, composed into one + /// node (each group composed once, so a shared group is one node). + fn composed( + memo: &Memo, + g: usize, + done: &mut HashMap>, + ) -> Rc { + if let Some(node) = done.get(&g) { + return Rc::clone(node); + } + let candidate = own(memo, g); + let children: Vec<(Rc, Rc)> = candidate + .children + .iter() + .map(|c| { + let c = c.0 as usize; + (Rc::clone(&memo.groups[c].node), composed(memo, c, done)) + }) + .collect(); + fn rebuild( + node: &Rc, + children: &[(Rc, Rc)], + ) -> Rc { + if let Some((_, child)) = children.iter().find(|(group, _)| **group == **node) { + return Rc::clone(child); + } + if node.children().is_empty() { + return Rc::clone(node); + } + Rc::new( + node.with_new_children(|child| rebuild(child, children)) + .expect("rebuilds") + .with_timing(node.timing) + .with_guarantee(node.guarantee.clone()), + ) + } + let node = rebuild(&candidate.node, &children); + done.insert(g, Rc::clone(&node)); + node + } + + /// Stage 3's price ([`crate::price_at_scale`]) of the whole plan of `memo`'s roots, each group of + /// `materialized` materialized as given (as Stage 2's materialization + /// candidate its readers name), nothing else. + fn stage3_price( + memo: &Memo, + entries: &[QueryWorkloadEntry], + materialized: &[(usize, Materialization)], + ) -> f64 { + let data = data(); + let models = PlanningModels::builtin(); + let mut done = HashMap::new(); + let roots: Vec = (0..entries.len()) + .map(|q| { + let g = memo + .groups + .iter() + .position(|group| group.roots_of.contains(&q)) + .unwrap(); + QueryRoot::Operator(composed(memo, g, &mut done)) + }) + .collect(); + // Stage 2's materialization candidates, in discovery order: builds + // first, then each highest shared node, which a group materialized + // at query time is (the last one its readers name). + let mut readers_of: Vec> = Vec::new(); + let mut record = |readers: &[usize]| { + readers_of.push(readers.to_vec()); + Some(Materialization::NotMaterialized) + }; + optimize_physical_asap_assigned( + "L1", + &roots, + entries, + &data, + models.deployment_capabilities, + &mut record, + ) + .unwrap(); + let chosen: Vec = (0..readers_of.len()) + .map(|i| { + materialized + .iter() + .find(|(g, _)| { + let queries = &memo.groups[*g].queries; + readers_of.iter().rposition(|r| r == queries) == Some(i) + }) + .map_or(Materialization::NotMaterialized, |(_, m)| *m) + }) + .collect(); + let mut next = chosen.into_iter(); + let mut choose = |_: &[usize]| next.next(); + let candidate = optimize_physical_asap_assigned( + "L1", + &roots, + entries, + &data, + models.deployment_capabilities, + &mut choose, + ) + .unwrap() + .unwrap(); + let dag = &candidate.dag; + let recurrence: Vec<_> = entries.iter().map(|e| e.recurrence.clone()).collect(); + let ranges: Vec<_> = entries.iter().map(|e| e.range_evaluation).collect(); + let (cost, _) = crate::price_at_scale( + dag, + &crate::node_map(dag), + &crate::kept_pane_chains(dag), + &recurrence, + &ranges, + &data, + models.plan_selection, + 1.0, ) .unwrap(); - let selection = run.stage3_plan_selection.as_ref().unwrap(); - let price = selection.costs[&selection.selected].cost_per_second; - let assignment = run.stage3_selected.as_ref().unwrap().assignment.clone(); - let roots = run - .stage1_selected - .as_ref() - .unwrap() - .roots - .as_ref() - .unwrap(); - let QueryRoot::Operator(top) = &roots[0] else { - panic!("an operator root") + cost.cost_per_second.unwrap() + } + + /// Every group's own candidate priced in its true context, as the + /// search fixes it: outputs bottom-up, consumers top-down (each query + /// at its root; a parent passes its consumers on, or is the one + /// consumer when materialized), the groups of `materialized` + /// materialized as given. + fn local_prices( + memo: &Memo, + entries: &[QueryWorkloadEntry], + materialized: &[(usize, Materialization)], + cache: &mut PriceCache, + ) -> Vec { + let data = data(); + let models = PlanningModels::builtin(); + let settings = models.plan_selection; + let materialization_of = |g: usize| { + materialized + .iter() + .find(|(m, _)| *m == g) + .map_or(Materialization::NotMaterialized, |(_, m)| *m) }; - let below = Rc::clone(top.children()[0]); - let group = |node: &Rc, children: Vec| MemoGroup { - node: Rc::clone(node), - children, - parents: Vec::new(), - queries: vec![0], - roots_of: Vec::new(), + let child_stats = |outputs: &[EdgeStatistics], g: usize| -> Vec { + own(memo, g) + .children + .iter() + .map(|c| outputs[c.0 as usize]) + .collect() }; - let candidate = - |id: u32, node: &Rc, children: Vec| ImplementationCandidate { - implements: vec![GroupId(id)], - node: Rc::clone(node), - children, - origin: CandidateOrigin::Pass1, - coupling_sites: Vec::new(), - budget_share: 1, - requires_exact_children: Vec::new(), - rollup_of: None, + let mut outputs: Vec = Vec::new(); + for g in 0..memo.groups.len() { + let stats = child_stats(&outputs, g); + let output = output_stats(memo, own(memo, g), &stats, entries, &data, models, cache) + .unwrap() + .expect("an output"); + outputs.push(output); + } + let mut consumers: Vec> = vec![Vec::new(); memo.groups.len()]; + for g in (0..memo.groups.len()).rev() { + let group = &memo.groups[g]; + for &q in &group.roots_of { + consumers[g].push(Consumer::Query { query: q }); + } + let passed: Vec = match materialization_of(g) { + Materialization::NotMaterialized => consumers[g].clone(), + m => { + let at_ingestion = matches!(m, Materialization::MatAtIngestion(_)); + let reaching: Vec<&QueryWorkloadEntry> = + group.queries.iter().map(|&q| &entries[q]).collect(); + vec![Consumer::Materialized { + rate_per_second: match at_ingestion { + true => 1.0, + false => shared_evaluation_rate(&reaching, settings).unwrap(), + }, + queries: match at_ingestion { + true => Vec::new(), + false => group.queries.clone(), + }, + at_ingestion, + }] + } }; - let memo = Memo { - groups: vec![group(&below, vec![]), group(top, vec![GroupId(0)])], - candidates: vec![ - vec![candidate(0, &below, vec![])], - vec![candidate(1, top, vec![GroupId(0)])], - ], - }; - let context = Context { - thresholds: Vec::new(), - materialization: assignment.into_iter().map(|m| (GroupId(0), m)).collect(), - }; - (memo, context, price) + for &child in &own(memo, g).children { + consumers[child.0 as usize].extend(passed.iter().cloned()); + } + } + (0..memo.groups.len()) + .map(|g| { + private_usage( + memo, + own(memo, g), + &ConsumerContext { + consumers: consumers[g].clone(), + }, + materialization_of(g), + &child_stats(&outputs, g), + entries, + &data, + models, + cache, + ) + .unwrap() + .unwrap_or_else(|| panic!("group {g} is priced")) + }) + .collect() } - #[test] - fn local_prices_add_up_to_the_whole_plans_price() { - let models = PlanningModels::builtin(); - let (memo, context, expected) = split_plan(models); - let (entries, data) = (entries(), data()); - let mut cache = - PriceCache::new(frozen_calibration(&memo, &entries, &data, models).unwrap()); - let mut total = 0.0; - for (g, candidates) in memo.candidates.iter().enumerate() { - let local = local_price( - &memo, - GroupId(g as u32), - &candidates[0], - &context, - &[0], - &entries, - &data, - models, - &mut cache, - ) - .unwrap() - .unwrap(); - // Each group carries a part of the price, not the whole. - let price = local.cost_per_second.unwrap(); - assert!( - price > 0.0 && price < expected.unwrap(), - "group {g}: {price}" - ); - total += price; - } - let expected = expected.unwrap(); + fn assert_close(actual: f64, expected: f64) { assert!( - (total - expected).abs() <= 1e-9 * expected.abs(), - "{total} vs {expected}" + (actual - expected).abs() <= 1e-9 * expected.abs(), + "{actual} vs {expected}" ); } - #[test] - fn the_cache_hits_on_a_repeated_context() { - let models = PlanningModels::builtin(); - let (memo, context, _) = split_plan(models); - let (entries, data) = (entries(), data()); + /// The groups' local prices of `queries`' plan, `materialized` as + /// given, add up to Stage 3's price of the whole plan. + fn check( + queries: &[&str], + intervals_ms: &[u32], + materialized: impl Fn(&Memo) -> Vec<(usize, Materialization)>, + ) { + let memo = memo(queries); + let entries = entries(intervals_ms); + let materialized = materialized(&memo); let mut cache = PriceCache::new(1.0); - let mut price = || { - local_price( - &memo, - GroupId(1), - &memo.candidates[1][0], - &context, - &[0], - &entries, - &data, - models, - &mut cache, - ) - .unwrap() - }; - let first = price(); - let second = price(); - assert!(first.is_some()); - assert_eq!(first, second); - assert_eq!((cache.misses, cache.hits), (1, 1)); - } - - fn usage(cost: f64, peak_memory_bytes: u64, peak_cpu_cores: u32, latency: f64) -> LocalUsage { - LocalUsage { - cost_per_second: Some(cost), - cpu_ops_per_second: 1.0, - scanned_bytes_per_second: 1.0, - average_memory_bytes: 1.0, - peak_memory_bytes, - peak_cpu_cores, - disk_bytes: 1, - latency_ms: vec![(0, latency)], - } + let parts = local_prices(&memo, &entries, &materialized, &mut cache); + let model = &PlanningModels::builtin().plan_selection.cost_model; + let expected = stage3_price(&memo, &entries, &materialized); + assert_close(plan_cost(&parts, model).unwrap(), expected); } - fn items(frontier: &Frontier) -> Vec { - frontier.entries.iter().map(|(_, t)| t.clone()).collect() + /// The highest group both queries read. + fn shared(memo: &Memo) -> usize { + memo.groups + .iter() + .rposition(|g| g.queries.len() == 2) + .expect("a shared group") } #[test] - fn cost_model_1_keeps_the_cheapest_and_the_earlier_tie() { - let mut frontier = - Frontier::new(&CostModel::UsageOverTime(ResourceUnitCost::illustrative())); - assert!(frontier.insert(usage(2.0, 1, 1, 1.0), "b")); - assert!(frontier.insert(usage(1.0, 9, 9, 9.0), "a")); - assert!(!frontier.insert(usage(1.0, 1, 1, 1.0), "tie")); - assert!(!frontier.insert(usage(3.0, 0, 0, 0.0), "c")); - assert_eq!(items(&frontier), vec!["a"]); + fn local_prices_add_up_to_a_one_query_plans_price() { + for query in [ + "sum by (job) (rate(m[5m]))", + "max(avg_over_time(m[10m]))", + "topk(3, sum by (job) (rate(m[5m])))", + "sum(m offset 5m)", + "count by (job) (max_over_time(m[1h]))", + ] { + check(&[query], &[30_000], |_| Vec::new()); + } } + /// The per-job rate sum both queries read is one group: computed per + /// reader when not materialized, once for both when materialized at + /// query time, its inputs with it. #[test] - fn cost_model_2_keeps_the_pareto_front() { - let mut frontier = Frontier::new(&CostModel::ProvisionedCapacity( - ResourceUnitCost::illustrative(), - )); - assert!(frontier.insert(usage(2.0, 10, 1, 5.0), "cheap memory")); - assert!(frontier.insert(usage(1.0, 20, 1, 5.0), "cheap price")); - // Dominated by "cheap price". - assert!(!frontier.insert(usage(1.0, 20, 1, 6.0), "slower")); - // An exact tie is kept. - assert!(frontier.insert(usage(1.0, 20, 1, 5.0), "tie")); - // Dominates "cheap memory". - assert!(frontier.insert(usage(2.0, 10, 1, 4.0), "faster")); - assert_eq!(items(&frontier), vec!["cheap price", "tie", "faster"]); + fn local_prices_add_up_with_a_shared_group() { + let queries = [ + "sum by (job) (rate(m[5m]))", + "sum by (job) (rate(m[5m])) * 2", + ]; + let memo = memo(&queries); + let entries = entries(&[60_000, 60_000]); + let at_query = [(shared(&memo), Materialization::MatAtQuery(Medium::Memory))]; + assert!(stage3_price(&memo, &entries, &[]) != stage3_price(&memo, &entries, &at_query)); + check(&queries, &[60_000, 60_000], |_| Vec::new()); + check(&queries, &[60_000, 60_000], |memo| { + vec![(shared(memo), Materialization::MatAtQuery(Medium::Memory))] + }); + check(&queries, &[60_000, 15_000], |memo| { + vec![(shared(memo), Materialization::MatAtQuery(Medium::Memory))] + }); + // A shared group under each query's own aggregate, materialized or not. + let queries = ["max(rate(m[5m]))", "min(rate(m[5m]))"]; + check(&queries, &[60_000, 30_000], |_| Vec::new()); + check(&queries, &[60_000, 30_000], |memo| { + vec![(shared(memo), Materialization::MatAtQuery(Medium::Memory))] + }); } + /// Materialized, the shared group's input runs half as often for two + /// readers evaluating together. #[test] - fn combining_sums_usage_and_takes_the_largest_peak_cores_and_latency() { - let cluster = CostModel::ClusterCapacity(ClusterLimits { - memory_bytes: 25, - disk_bytes: u64::MAX, - cpu_cores: Some(4), - objective: ClusterObjective::ParetoFront, - }); - let mut a = Frontier::new(&cluster); - let none = |mut u: LocalUsage| { - u.cost_per_second = None; - u + fn a_materialized_shared_group_computes_its_input_once() { + let memo = memo(&[ + "sum by (job) (rate(m[5m]))", + "sum by (job) (rate(m[5m])) * 2", + ]); + let entries = entries(&[60_000, 60_000]); + let s = shared(&memo); + let below = memo.groups[s].children[0].0 as usize; + let mut cache = PriceCache::new(1.0); + let apart = local_prices(&memo, &entries, &[], &mut cache); + let together = local_prices( + &memo, + &entries, + &[(s, Materialization::MatAtQuery(Medium::Memory))], + &mut cache, + ); + assert_close( + together[below].cpu_ops_per_second * 2.0, + apart[below].cpu_ops_per_second, + ); + } + + /// A tumbling-pane chain built at ingestion time for a 5-minute sum + /// every minute: its local price is Stage 3's price of the plan. + #[test] + fn a_pane_chain_at_ingestion_is_priced_as_stage_3_prices_it() { + use crate::pass::subexpr::windows::add_window_composition; + use crate::pass::subexpr::{CandidateOrigin, MemoGroup}; + use asap_logical_asap_optimization::workload_candidates::stage1_logical_asap_optimization; + use asap_types::ir::schema::schema_support::with_promql_series_identity; + use asap_types::workload::Predictability; + + let mut entries = entries(&[60_000]); + entries[0].predictability = Predictability::Predictable { known_at: None }; + let node = with_promql_series_identity(&lower_promql( + "sum_over_time(m[5m])", + AccuracyTarget::exact(), + )) + .unwrap(); + let inventory = + generate_per_computation_candidates(vec![(0, QueryRoot::Operator(Rc::clone(&node)))]) + .unwrap(); + let candidates = stage1_logical_asap_optimization(&inventory, &entries) + .filter(|c| c.pass2.is_none()) + .filter_map(|c| match c.roots.ok()?.pop()? { + QueryRoot::Operator(node) => Some(node), + QueryRoot::Scalar(_) => None, + }) + .map(|node| ImplementationCandidate { + implements: vec![GroupId(0)], + node, + children: Vec::new(), + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + budget_share: 1, + requires_exact_children: Vec::new(), + rollup_of: None, + }) + .collect(); + let mut memo = Memo { + groups: vec![MemoGroup { + node, + children: Vec::new(), + parents: Vec::new(), + queries: vec![0], + roots_of: vec![0], + }], + candidates: vec![candidates], }; - a.insert(none(usage(0.0, 10, 2, 5.0)), 1); - a.insert(none(usage(0.0, 5, 3, 7.0)), 2); - let mut b = Frontier::new(&cluster); - b.insert(none(usage(0.0, 12, 4, 6.0)), 10); - let combined = a.combine(&b, |x, y| x + y); - // 10 + 12 and 5 + 12 bytes both fit under 25. - assert_eq!(items(&combined), vec![11, 12]); - let (first, _) = &combined.entries[0]; - assert_eq!(first.peak_memory_bytes, 22); - assert_eq!(first.peak_cpu_cores, 4); - assert_eq!(first.latency_ms, vec![(0, 6.0)]); - assert_eq!(first.cpu_ops_per_second, 2.0); - assert_eq!(first.disk_bytes, 2); - // Over the memory limit: dropped. - let mut c = Frontier::new(&cluster); - c.insert(none(usage(0.0, 21, 1, 1.0)), 100); - assert!(a.combine(&c, |x, y| x + y).entries.is_empty()); + add_window_composition(&mut memo, &entries); + let panes = memo.candidates[0] + .iter() + .find(|c| c.origin == CandidateOrigin::Panes && c.coupling_sites.is_empty()) + .expect("a pane candidate"); + let at_ingestion = Materialization::MatAtIngestion(Medium::Memory); + let data = data(); + let models = PlanningModels::builtin(); + let mut cache = PriceCache::new(1.0); + let local = private_usage( + &memo, + panes, + &ConsumerContext { + consumers: vec![Consumer::Query { query: 0 }], + }, + at_ingestion, + &[], + &entries, + &data, + models, + &mut cache, + ) + .unwrap() + .expect("priced"); + let mut first = true; + let mut choose = |_: &[usize]| { + let chosen = match first { + true => at_ingestion, + false => Materialization::NotMaterialized, + }; + first = false; + Some(chosen) + }; + let plan = optimize_physical_asap_assigned( + "L1", + &[QueryRoot::Operator(Rc::clone(&panes.node))], + &entries, + &data, + models.deployment_capabilities, + &mut choose, + ) + .unwrap() + .unwrap(); + assert!(plan + .dag + .nodes + .iter() + .any(|n| n.materialization == at_ingestion)); + let dag = &plan.dag; + let (cost, _) = crate::price_at_scale( + dag, + &crate::node_map(dag), + &crate::kept_pane_chains(dag), + &[entries[0].recurrence.clone()], + &[entries[0].range_evaluation], + &data, + models.plan_selection, + 1.0, + ) + .unwrap(); + let model = &models.plan_selection.cost_model; + assert_close( + plan_cost([&local], model).unwrap(), + cost.cost_per_second.unwrap(), + ); + } + + #[test] + fn the_cache_hits_on_a_repeated_context() { + let memo = memo(&["sum by (job) (rate(m[5m]))"]); + let entries = entries(&[60_000]); + let mut cache = PriceCache::new(1.0); + let first = local_prices(&memo, &entries, &[], &mut cache); + let (misses, hits) = (cache.misses, cache.hits); + let second = local_prices(&memo, &entries, &[], &mut cache); + assert_eq!(first, second); + assert_eq!(cache.misses, misses); + assert_eq!(cache.hits, hits + 2 * memo.groups.len() as u64); } } From 323c49f7bbba2029104fd455b4e2219b55e5058e Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:42:01 +0000 Subject: [PATCH 15/25] plan-selection: sub-expression memo: shared builds resized at a size threshold resized_at_threshold resizes a shared-build reader's build to a threshold (with_resized_build, guarantees composed again, child groups' nodes kept). The brute-force reference rewrites SharedBuildSize sites with it and makes resized builds over equal windows one build. --- .../src/pass/subexpr/capability.rs | 215 +++++++++++++++++- .../plan-selection/src/pass/subexpr/search.rs | 135 ++++++++++- 2 files changed, 340 insertions(+), 10 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index f7c91ee52..1f2f87a25 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -22,14 +22,17 @@ //! iff [`may_read`]. use std::collections::hash_map::DefaultHasher; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::hash::{Hash, Hasher}; use std::rc::Rc; use asap_logical_asap_optimization::accuracy::estimators::size_for_all_readers; +use asap_logical_asap_optimization::asap_aware_cse::with_resized_build; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; use asap_types::ir::operator::Reduction; -use asap_types::ir::schema::sketch::{GroupingStrategy, SketchParams}; +use asap_types::ir::schema::sketch::{GroupingStrategy, SketchKind, SketchParams}; use asap_types::ir::schema::FieldDataType; +use asap_types::ir::QueryRoot; use asap_types::ir::{ASAPOp, OperatorNode}; use super::groups::NodeKeys; @@ -147,6 +150,214 @@ pub(crate) fn may_read(threshold: &SizeThreshold, reader: &ImplementationCandida } } +/// `candidate`'s node with its shared build (of `site`) built at +/// `threshold`, every guarantee composed again: the node itself when the +/// build already has that size (or is an exact accumulator, which has +/// none); `None` when `candidate` does not read `site` or may not read the +/// threshold ([`may_read`]). +pub(crate) fn resized_at_threshold( + memo: &Memo, + candidate: &ImplementationCandidate, + site: &CouplingSite, + threshold: &SizeThreshold, +) -> Result>, PerComputationCandidateError> { + if !candidate.coupling_sites.contains(site) + || site_key(memo, candidate) != Some(site.key) + || !may_read(threshold, candidate) + { + return Ok(None); + } + let (Some((build, Some(own))), Some(size)) = (reader_build(candidate), &threshold.params) + else { + return Ok(Some(Rc::clone(&candidate.node))); + }; + if same_size(&own, size) { + return Ok(Some(Rc::clone(&candidate.node))); + } + let Some(ASAPOp::SummaryAgg { + family: FieldDataType::Sketch(kind, strategy), + .. + }) = build.asap() + else { + unreachable!("a reader's build is a sketch when it has a size"); + }; + let family = FieldDataType::Sketch( + SketchKind::new(kind.algorithm().clone(), size.clone()), + strategy.clone(), + ); + let roots = with_resized_build( + &[QueryRoot::Operator(Rc::clone(&candidate.node))], + &build, + family, + )?; + let [QueryRoot::Operator(node)] = &roots[..] else { + unreachable!("an operator root"); + }; + // Composing the guarantees again rebuilds the child groups' nodes too; + // the candidate reads them as the groups' own nodes. + let child_nodes: HashSet<*const OperatorNode> = candidate + .children + .iter() + .map(|c| Rc::as_ptr(&memo.groups[c.0 as usize].node)) + .collect(); + Ok(Some(with_child_nodes( + &candidate.node, + node, + &child_nodes, + &mut HashMap::new(), + ))) +} + +/// `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 +/// rebuilt over them, keeping their own schema and guarantee. +fn with_child_nodes( + original: &Rc, + resized: &Rc, + child_nodes: &HashSet<*const OperatorNode>, + done: &mut HashMap<*const OperatorNode, Rc>, +) -> Rc { + if child_nodes.contains(&Rc::as_ptr(original)) { + return Rc::clone(original); + } + if let Some(found) = done.get(&Rc::as_ptr(resized)) { + return Rc::clone(found); + } + let pairs: HashMap<*const OperatorNode, Rc> = original + .children() + .into_iter() + .zip(resized.children()) + .map(|(o, r)| (Rc::as_ptr(r), with_child_nodes(o, r, child_nodes, done))) + .collect(); + let mut changed = false; + let operator = resized.operator.map_children(|child| { + let new = Rc::clone(&pairs[&Rc::as_ptr(child)]); + changed |= !Rc::ptr_eq(&new, child); + new + }); + let out = match changed { + false => Rc::clone(resized), + true => Rc::new( + OperatorNode::with_schema(operator, resized.schema.clone()) + .with_guarantee(resized.guarantee.clone()) + .with_timing(resized.timing) + .with_materialization(resized.materialization), + ), + }; + done.insert(Rc::as_ptr(resized), Rc::clone(&out)); + out +} + +/// The shared build `candidate` reads, if any ([`add_shared_builds`]). +pub(crate) fn shared_build_of(candidate: &ImplementationCandidate) -> Option> { + reader_build(candidate).map(|(build, _)| build) +} + +/// `roots` with each set of equal summary builds that one of `resized` +/// equals made one build: the summary-capability rule's builds resized to +/// one size, over equal inputs (two queries' equal windows are two groups, +/// since a window has no unique key, so identical-expression sharing +/// leaves their builds apart). As the rule shares only builds whose sizes +/// are not all equal, equal builds none of which was resized stay apart. +/// Each build is replaced by the first equal one a children-first walk +/// meets, and the nodes above rebuilt. +pub(crate) fn share_equal_builds( + roots: &[QueryRoot], + resized: &[Rc], +) -> Vec { + let mut kept: Vec> = Vec::new(); + let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); + let mut seen = HashSet::new(); + for node in roots.iter().flat_map(root_nodes) { + if !seen.insert(Rc::as_ptr(&node)) + || !matches!(node.asap(), Some(ASAPOp::SummaryAgg { .. })) + || !resized.iter().any(|r| **r == *node) + { + continue; + } + match kept.iter().find(|k| ***k == *node) { + Some(first) => { + replaced.insert(Rc::as_ptr(&node), Rc::clone(first)); + } + None => kept.push(node), + } + } + if replaced.is_empty() { + return roots.to_vec(); + } + let mut done = HashMap::new(); + roots + .iter() + .map(|root| match root { + QueryRoot::Operator(node) => { + QueryRoot::Operator(replaced_in(node, &replaced, &mut done)) + } + QueryRoot::Scalar(expr) => QueryRoot::Scalar( + expr.map_operator_refs(&mut |node| replaced_in(node, &replaced, &mut done)), + ), + }) + .collect() +} + +/// Every node of `root`'s operators, each once, children first. +fn root_nodes(root: &QueryRoot) -> Vec> { + let starts: Vec> = match root { + QueryRoot::Operator(node) => vec![Rc::clone(node)], + QueryRoot::Scalar(expr) => expr.operator_refs().into_iter().cloned().collect(), + }; + let mut out = Vec::new(); + let mut seen = HashSet::new(); + fn walk( + node: &Rc, + seen: &mut HashSet<*const OperatorNode>, + out: &mut Vec>, + ) { + if !seen.insert(Rc::as_ptr(node)) { + return; + } + for child in node.children() { + walk(child, seen, out); + } + out.push(Rc::clone(node)); + } + for start in &starts { + walk(start, &mut seen, &mut out); + } + out +} + +/// `node` with each node of `replaced` (by address) swapped, the nodes +/// above rebuilt over the swapped ones. +fn replaced_in( + node: &Rc, + replaced: &HashMap<*const OperatorNode, Rc>, + done: &mut HashMap<*const OperatorNode, Rc>, +) -> Rc { + if let Some(found) = replaced.get(&Rc::as_ptr(node)) { + return Rc::clone(found); + } + if let Some(found) = done.get(&Rc::as_ptr(node)) { + return Rc::clone(found); + } + let mut changed = false; + let operator = node.operator.map_children(|child| { + let new = replaced_in(child, replaced, done); + changed |= !Rc::ptr_eq(&new, child); + new + }); + let out = match changed { + false => Rc::clone(node), + true => Rc::new( + OperatorNode::with_schema(operator, node.schema.clone()) + .with_guarantee(node.guarantee.clone()) + .with_timing(node.timing) + .with_materialization(node.materialization), + ), + }; + done.insert(Rc::as_ptr(node), Rc::clone(&out)); + out +} + /// Whether a build at `threshold` serves a reader that requires `required`. fn admits(threshold: &SketchParams, required: &SketchParams) -> bool { size_for_all_readers(&[threshold, required]).is_some_and(|join| same_size(&join, threshold)) diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 693367129..122ab5f01 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -133,8 +133,8 @@ fn group_choices(memo: &Memo) -> Option> { /// The chosen candidates of `choice`, each rewritten at a threshold of each /// of its coupling sites: every combination of thresholds every reader may -/// read. A site without a rewrite (raw horizon, a shared build until it can -/// be resized) leaves its readers as they are; Stage 3 prices it. +/// read. A raw-horizon site has no rewrite and leaves its readers as they +/// are; Stage 3 prices it. fn at_thresholds(memo: &Memo, choice: &GroupChoice) -> Vec>> { let chosen: Vec> = choice .iter() @@ -210,7 +210,30 @@ fn at_thresholds(memo: &Memo, choice: &GroupChoice) -> Vec next.push(plan), + CouplingKind::SharedBuildSize => { + for threshold in super::capability::size_thresholds(memo, &site) { + let mut rewritten = plan.clone(); + let mut valid = true; + for candidate in rewritten.iter_mut().flatten() { + if !candidate.coupling_sites.contains(&site) { + continue; + } + match super::capability::resized_at_threshold( + memo, candidate, &site, &threshold, + ) { + Ok(Some(node)) => candidate.node = node, + _ => { + valid = false; + break; + } + } + } + if valid { + next.push(rewritten); + } + } + } + CouplingKind::RawHorizon => next.push(plan), } } plans = next; @@ -308,6 +331,46 @@ fn substituted( new } +/// `roots` with its identical nodes made one node, then its summary builds +/// equal to one of `resized` (over equal windows of two groups) made one +/// build and the identical nodes above them shared in turn. +fn shared(roots: &[QueryRoot], resized: &[Rc]) -> Vec { + let identical = |roots: Vec| { + let shareable = sharing_candidates(&roots); + match shareable.is_empty() { + true => roots, + false => share_identical_expressions(&roots, &shareable), + } + }; + identical(super::capability::share_equal_builds( + &identical(roots.to_vec()), + resized, + )) +} + +/// The shared builds `plan` resized: each shared-build reader's build that +/// differs from the build of the memo's candidate it was chosen as +/// (`choice`). +fn resized_builds( + memo: &Memo, + choice: &GroupChoice, + plan: &[Option], +) -> Vec> { + choice + .iter() + .zip(plan) + .enumerate() + .filter_map(|(g, (c, candidate))| { + let candidate = candidate.as_ref()?; + let original = &memo.candidates[g][(*c)?]; + let build = super::capability::shared_build_of(candidate)?; + let own = super::capability::shared_build_of(original)?; + (candidate.origin == super::CandidateOrigin::SharedBuild && *build != *own) + .then_some(build) + }) + .collect() +} + /// The reference search: every plan of the memo ([`group_choices`] × /// [`at_thresholds`]), its identical nodes made one node, implemented in /// every Stage 2 materialization and priced by Stage 3; the cheapest @@ -337,11 +400,7 @@ pub(crate) fn brute_force( let Some(roots) = composed_roots(&memo, per_computation, &plan) else { continue; }; - let shareable = sharing_candidates(&roots); - let roots = match shareable.is_empty() { - true => roots, - false => share_identical_expressions(&roots, &shareable), - }; + let roots = shared(&roots, &resized_builds(&memo, choice, &plan)); number += 1; let id = format!("L{number}"); let Ok(implemented) = @@ -381,6 +440,66 @@ mod tests { RepeatedDemand, RepetitionInterval, SourceFacts, }; + /// The 0.9 quantile of `m[5m]` within 1% and within 2% error, each + /// query reading its shared-build twin: at the larger threshold both + /// builds are resized to one `k`, and the composed plan, its identical + /// nodes shared, has one KLL build, although the two queries' windows + /// are two groups. At the smaller threshold the stricter reader may not + /// read the build, so it is the only plan. + #[test] + fn two_kll_readers_at_a_threshold_compose_to_one_build() { + use super::super::CandidateOrigin; + let roots: Vec<(usize, QueryRoot)> = [0.01, 0.02] + .iter() + .map(|&e| { + QueryRoot::Operator(lower_promql( + "quantile_over_time(0.9, m[5m])", + AccuracyTarget::max_error(e), + )) + }) + .enumerate() + .collect(); + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let mut memo = groups::memo_of(&per_computation); + super::super::capability::add_shared_builds(&mut memo); + let is_kll = |c: &ImplementationCandidate| { + OperatorNode::reachable(&c.node).iter().any(|n| { + matches!(n.asap(), Some(asap_types::ir::ASAPOp::SummaryAgg { family, .. }) + if format!("{family:?}").contains("Kll")) + }) + }; + // Each query's root group reads its KLL twin; every other group + // its one candidate. + let choice: GroupChoice = memo + .groups + .iter() + .enumerate() + .map(|(g, group)| { + if group.roots_of.is_empty() { + return Some(0); + } + memo.candidates[g] + .iter() + .position(|c| c.origin == CandidateOrigin::SharedBuild && is_kll(c)) + }) + .collect(); + assert!(choice.iter().all(Option::is_some)); + let plans = at_thresholds(&memo, &choice); + assert_eq!(plans.len(), 1); + let roots = composed_roots(&memo, &per_computation, &plans[0]).unwrap(); + let roots = shared(&roots, &resized_builds(&memo, &choice, &plans[0])); + let builds: std::collections::HashSet<*const OperatorNode> = roots + .iter() + .flat_map(|root| match root { + QueryRoot::Operator(node) => OperatorNode::reachable(node), + QueryRoot::Scalar(_) => Vec::new(), + }) + .filter(|n| matches!(n.asap(), Some(asap_types::ir::ASAPOp::SummaryAgg { .. }))) + .map(|n| Rc::as_ptr(&n)) + .collect(); + assert_eq!(builds.len(), 1); + } + #[test] #[ignore = "debugging aid"] fn debug_topk_memo() { From 71e03d5a4c5037d1574e69fefe684e3563b78cf6 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:42:28 +0000 Subject: [PATCH 16/25] plan-selection: memo plans compose every node's guarantee again; Pass 1 parity test checks every target's candidates A memo plan built from per-group candidates kept the frontend's empty guarantee on a group's own node (an exact top-k over the per-series sums' exact accumulator), so Stage 3 rejected it for having no guarantee. composed_roots now recomposes every guarantee, as Pass 1 composition does (recompose_guarantees made public). The parity test also checks that each target's group, absorbed or not, has the target's own and fused candidates. --- .../src/per_computation_candidates.rs | 2 +- .../plan-selection/src/pass/subexpr/groups.rs | 33 +++++++++++++++++-- .../plan-selection/src/pass/subexpr/search.rs | 10 ++++-- 3 files changed, 40 insertions(+), 5 deletions(-) diff --git a/crates/logical-asap-optimization/src/per_computation_candidates.rs b/crates/logical-asap-optimization/src/per_computation_candidates.rs index e9c2f4a1e..7b5363c16 100644 --- a/crates/logical-asap-optimization/src/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/per_computation_candidates.rs @@ -1519,7 +1519,7 @@ fn rewrite( /// ([`composed_guarantee`]), children first; a sub-DAG shared between roots /// stays shared. Pass 2 calls it after re-sizing a shared summary, so each /// consumer's guarantee is the shared state's. -pub(crate) fn recompose_guarantees( +pub fn recompose_guarantees( roots: &[QueryRoot], ) -> Result, PerComputationCandidateError> { let mut memo = HashMap::new(); diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index fed7dc3c2..032d3e64f 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -1029,8 +1029,10 @@ mod tests { /// For the random workloads of `joint_search_equivalence` (over one /// gauge and over two), every whole-workload Pass 1 choice composes, /// for each query, a plan that is a combination of the memo's per-group - /// candidates: the memo misses no Pass 1 candidate. Every candidate of - /// a group a parent reads outputs the group's schema. + /// candidates: the memo misses no Pass 1 candidate. Every target's + /// group has the target's own and fused candidates, also a target a + /// fused computation absorbs; every candidate of a group a parent reads + /// outputs the group's schema. #[test] fn every_composed_pass1_candidate_is_a_combination_of_group_candidates() { let mut checked = 0; @@ -1047,6 +1049,33 @@ mod tests { continue; } let memo = memo_of(&inventory); + // Each target's own and fused candidates, absorbed or not. + let group_of = |t: usize| { + memo.groups + .iter() + .position(|group| *group.node == *inventory.targets[t].aggregate) + .unwrap() + }; + for (t, target) in inventory.targets.iter().enumerate() { + let g = group_of(t); + if inventory.targets.iter().enumerate().any(|(u, other)| { + u != t && group_of(u) == g && other.measure != target.measure + }) { + continue; + } + let whole_budget = |fused: bool| { + memo.candidates[g] + .iter() + .filter(|c| c.budget_share == 1 && (c.implements.len() > 1) == fused) + .count() + }; + let fused = target.fused.as_ref().map_or(0, |f| f.candidates.len()); + if whole_budget(false) != target.choices.candidates.len() + || whole_budget(true) != fused + { + missing.push(format!("{gauges:?} seed {seed} target {t} candidates")); + } + } for (g, group) in memo.groups.iter().enumerate() { for candidate in &memo.candidates[g] { if !group.parents.is_empty() && candidate.node.schema != group.node.schema { diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 693367129..399d87b40 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -17,6 +17,8 @@ use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate} use asap_types::ir::{OperatorNode, QueryRoot}; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; +use asap_logical_asap_optimization::per_computation_candidates::recompose_guarantees; + use super::{groups, CouplingKind, GroupId, ImplementationCandidate, Memo}; use crate::{PlanSelector, PlanningModels, Selection, SelectionError}; @@ -244,7 +246,7 @@ fn composed_roots( } built_of[g] = Some(node); } - per_computation + let roots: Vec = per_computation .roots .iter() .enumerate() @@ -273,7 +275,11 @@ fn composed_roots( } }) }) - .collect() + .collect::>()?; + // A group's own node carries the frontend's guarantee (none); what it + // guarantees depends on the implementations chosen below it, so every + // node's guarantee is composed again, as Pass 1 composition does. + recompose_guarantees(&roots).ok() } /// `node` with every node `built` maps (by address) replaced by its built From 436cad4060d9d1b5fb8aa70324b7f6828344830f Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:43:50 +0000 Subject: [PATCH 17/25] plan-selection: memo reference search composes guarantees again over the chosen implementations --- .../tests/joint_search_equivalence.rs | 114 ++++++++++++++++++ .../src/per_computation_candidates.rs | 2 +- .../plan-selection/src/pass/subexpr/search.rs | 10 +- 3 files changed, 124 insertions(+), 2 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 2a8764ac8..b9e88ca5d 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1241,3 +1241,117 @@ fn memo_group_search_selects_the_exhaustive_price() { ); assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); } + +/// Debugging aid: the memo's and the exhaustive search's selected roots on +/// one seed. +#[test] +#[ignore = "debugging aid"] +fn debug_memo_seed() { + use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode}; + fn show(node: &std::rc::Rc, depth: usize, out: &mut String) { + let detail = match (node.asap(), node.non_asap()) { + ( + Some(ASAPOp::SummaryAgg { + family, reduction, .. + }), + _, + ) => { + format!("{family:?} {reduction:?}") + } + ( + Some(ASAPOp::SummaryMerge { + newest_panes, + group_by, + .. + }), + _, + ) => { + format!("newest {newest_panes:?} group_by {group_by:?}") + } + ( + Some(ASAPOp::PaneChain { + family, + reduction, + pane_width_ms, + pane_count, + .. + }), + _, + ) => { + format!("{family:?} {reduction:?} {pane_width_ms}x{pane_count}") + } + (_, Some(NonASAPOp::TimeRange { range, .. })) => format!("{range:?}"), + ( + _, + Some(NonASAPOp::Aggregate { + reduction, + measures, + .. + }), + ) => { + format!("{reduction:?} {measures:?}") + } + _ => String::new(), + }; + let detail: String = detail.chars().take(160).collect(); + out.push_str(&format!( + "{}{} {:p} {detail}\n", + " ".repeat(depth), + node.operator.kind_name(), + std::rc::Rc::as_ptr(node) + )); + for child in node.children() { + show(child, depth + 1, out); + } + } + let seed: u64 = std::env::var("SEED") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(59); + let gauges: &[&str] = match seed % 3 { + 0 => &["m", "n"], + _ => &["m"], + }; + let workload = random_workload_over(seed, gauges); + let roots = lower_promql_roots(&workload).unwrap(); + let entries: Vec<_> = workload.query_workload.entries().collect(); + for e in &entries { + eprintln!("{} {:?}", e.query.0, e.recurrence); + } + let data = workload.data_workload.clone().unwrap(); + let deployment = match seed % 2 { + 0 => no_raw_data(), + _ => asap_logical_asap_optimization::DeploymentCapabilities::default(), + }; + let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); + let indexed = || roots.iter().cloned().enumerate().collect::>(); + for (name, search) in [ + ("exhaustive", Stage1ToStage3Search::exhaustive()), + ( + "memo", + Stage1ToStage3Search { + memo_groups: true, + ..Stage1ToStage3Search::default() + }, + ), + ] { + let run = run_stage1_to_stage3_with(indexed(), &entries, &data, models, search).unwrap(); + let s = run.stage3_plan_selection.as_ref().unwrap(); + eprintln!("== {name}: {:?}", s.costs[&s.selected].cost_per_second); + let mut out = String::new(); + for root in run + .stage1_selected + .as_ref() + .unwrap() + .roots + .as_ref() + .unwrap() + { + if let QueryRoot::Operator(node) = root { + show(node, 0, &mut out); + } + out.push_str("--\n"); + } + eprintln!("{out}"); + } +} diff --git a/crates/logical-asap-optimization/src/per_computation_candidates.rs b/crates/logical-asap-optimization/src/per_computation_candidates.rs index e9c2f4a1e..7b5363c16 100644 --- a/crates/logical-asap-optimization/src/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/per_computation_candidates.rs @@ -1519,7 +1519,7 @@ fn rewrite( /// ([`composed_guarantee`]), children first; a sub-DAG shared between roots /// stays shared. Pass 2 calls it after re-sizing a shared summary, so each /// consumer's guarantee is the shared state's. -pub(crate) fn recompose_guarantees( +pub fn recompose_guarantees( roots: &[QueryRoot], ) -> Result, PerComputationCandidateError> { let mut memo = HashMap::new(); diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 693367129..eda6e8def 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -12,7 +12,9 @@ use std::rc::Rc; use asap_logical_asap_optimization::asap_aware_cse::{ share_identical_expressions, sharing_candidates, }; -use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; +use asap_logical_asap_optimization::per_computation_candidates::{ + recompose_guarantees, PerComputationCandidates, +}; use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate}; use asap_types::ir::{OperatorNode, QueryRoot}; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; @@ -337,6 +339,12 @@ pub(crate) fn brute_force( let Some(roots) = composed_roots(&memo, per_computation, &plan) else { continue; }; + // Each candidate's guarantees were composed over its own + // children's Pass 1 choice: compose them again over the chosen + // implementations. + let Ok(roots) = recompose_guarantees(&roots) else { + continue; + }; let shareable = sharing_candidates(&roots); let roots = match shareable.is_empty() { true => roots, From 25098ff96a949c44a6ed2d75e287d7a426909006 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:21:31 +0000 Subject: [PATCH 18/25] plan-selection: memo reference search composes each query's own nodes and lists the identical-expression and scan-widening variants as the exhaustive search does --- .../tests/joint_search_equivalence.rs | 78 +++++++ .../src/pass/subexpr/capability.rs | 15 +- .../plan-selection/src/pass/subexpr/search.rs | 218 +++++++++++++----- 3 files changed, 246 insertions(+), 65 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index b9e88ca5d..0c2120b9c 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1325,6 +1325,25 @@ fn debug_memo_seed() { }; let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); let indexed = || roots.iter().cloned().enumerate().collect::>(); + /// A plan's shape: its flattened DAG without what composition derives. + fn shape(roots: &[QueryRoot]) -> String { + fn strip(value: &mut serde_json::Value) { + match value { + serde_json::Value::Object(map) => { + for field in ["guarantee", "schema", "timing", "materialization"] { + map.remove(field); + } + map.values_mut().for_each(strip); + } + serde_json::Value::Array(items) => items.iter_mut().for_each(strip), + _ => {} + } + } + let mut value = serde_json::to_value(flatten(roots).0).unwrap(); + strip(&mut value); + value.to_string() + } + let mut runs = Vec::new(); for (name, search) in [ ("exhaustive", Stage1ToStage3Search::exhaustive()), ( @@ -1353,5 +1372,64 @@ fn debug_memo_seed() { out.push_str("--\n"); } eprintln!("{out}"); + let selected = run.selected().unwrap(); + eprintln!( + "materialization: {:?}", + selected.materialization.values().collect::>() + ); + runs.push(run); + } + let memo_shape = shape( + runs[1] + .stage1_selected + .as_ref() + .unwrap() + .roots + .as_ref() + .unwrap(), + ); + let found: Vec<&str> = runs[0] + .stage1_logical_asap_optimization + .iter() + .filter(|c| c.roots.as_ref().is_ok_and(|r| shape(r) == memo_shape)) + .map(|c| c.id.as_str()) + .collect(); + eprintln!( + "memo's plan among the exhaustive candidates ({} listed): {found:?}", + runs[0].stage1_logical_asap_optimization.len() + ); + // The exhaustive search's price of each of them. + if let Ok(selection) = &runs[0].stage3_plan_selection { + for id in &found { + for (physical, cost) in &selection.costs { + if runs[0] + .stage2_physical_asap_optimization + .iter() + .any(|p| &p.id == physical && p.from_logical == *id) + { + eprintln!(" {physical}: {:?}", cost.cost_per_second); + } + } + } + for rejection in selection_rejections(&runs[0], &found) { + eprintln!(" rejected: {rejection}"); + } } } + +/// The Stage 3 rejections of the physical candidates of `logical`. +fn selection_rejections(run: &Stage1ToStage3Run, logical: &[&str]) -> Vec { + let Ok(selection) = &run.stage3_plan_selection else { + return Vec::new(); + }; + selection + .rejected + .iter() + .filter(|r| { + run.stage2_physical_asap_optimization + .iter() + .any(|p| p.id == r.id && logical.contains(&p.from_logical.as_str())) + }) + .map(|r| format!("{} {:?} {}", r.id, r.cause, r.reason)) + .collect() +} diff --git a/crates/plan-selection/src/pass/subexpr/capability.rs b/crates/plan-selection/src/pass/subexpr/capability.rs index 1f2f87a25..8b8f070c6 100644 --- a/crates/plan-selection/src/pass/subexpr/capability.rs +++ b/crates/plan-selection/src/pass/subexpr/capability.rs @@ -265,21 +265,28 @@ pub(crate) fn share_equal_builds( roots: &[QueryRoot], resized: &[Rc], ) -> Vec { - let mut kept: Vec> = Vec::new(); + // Equal by node key: guarantees, composed again over a plan, may differ + // from the resized candidate's. + let mut keys = super::groups::NodeKeys::default(); + let resized: Vec = resized.iter().map(|r| keys.key(r)).collect(); + let mut kept: Vec<(u64, Rc)> = Vec::new(); let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); let mut seen = HashSet::new(); for node in roots.iter().flat_map(root_nodes) { if !seen.insert(Rc::as_ptr(&node)) || !matches!(node.asap(), Some(ASAPOp::SummaryAgg { .. })) - || !resized.iter().any(|r| **r == *node) { continue; } - match kept.iter().find(|k| ***k == *node) { + let key = keys.key(&node); + if !resized.contains(&key) { + continue; + } + match kept.iter().find(|(k, _)| *k == key).map(|(_, n)| n) { Some(first) => { replaced.insert(Rc::as_ptr(&node), Rc::clone(first)); } - None => kept.push(node), + None => kept.push((key, node)), } } if replaced.is_empty() { diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 39a37a693..9f5944cfb 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -10,13 +10,13 @@ use std::collections::HashMap; use std::rc::Rc; use asap_logical_asap_optimization::asap_aware_cse::{ - share_identical_expressions, sharing_candidates, + identical_expression_pass2_candidates, sharing_candidates, with_union_scan_labels, }; use asap_logical_asap_optimization::per_computation_candidates::{ recompose_guarantees, PerComputationCandidates, }; use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate}; -use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; use super::{groups, CouplingKind, GroupId, ImplementationCandidate, Memo}; @@ -243,43 +243,87 @@ fn at_thresholds(memo: &Memo, choice: &GroupChoice) -> Vec) -> Vec> { + let mut variants = Vec::new(); + for shared in identical_expression_pass2_candidates(&roots, &sharing_candidates(&roots)) { + variants.push(shared.roots); + } + if let Ok(Some(widened)) = with_union_scan_labels(&roots) { + for shared in identical_expression_pass2_candidates(&widened, &sharing_candidates(&widened)) + { + let shares_a_scan = shared + .sharing_candidates + .iter() + .any(|c| matches!(c.nodes[0].non_asap(), Some(NonASAPOp::Scan { .. }))); + if shares_a_scan { + variants.push(shared.roots); + } + } + } + variants.insert(0, roots); + variants +} + +/// The query roots of a plan, each query's own nodes as Pass 1 composes +/// them: each chosen candidate's node with every child group's node +/// replaced by that group's built node, children first. A group below +/// several queries' roots is built once per query (the identical-expression +/// rule shares them: [`identical_expression_variants`]); a group a Pass 2 +/// candidate reads from outside its query (a rollup's finer group, a +/// ranking, a host chain) is built once, by the first query below whose +/// root it is, or once for all. fn composed_roots( memo: &Memo, per_computation: &PerComputationCandidates, chosen: &[Option], ) -> Option> { - let order = memo.topological_order()?; - let mut built: HashMap<*const OperatorNode, Rc> = HashMap::new(); - let mut built_of: Vec>> = vec![None; memo.groups.len()]; - for GroupId(g) in order { - let g = g as usize; - let Some(candidate) = &chosen[g] else { - continue; - }; - let mut rebuilt = HashMap::new(); - let node = substituted(&candidate.node, &built, &mut rebuilt); - for implemented in &candidate.implements { - built.insert( - Rc::as_ptr(&memo.groups[implemented.0 as usize].node), - Rc::clone(&node), - ); + memo.topological_order()?; + let queries = per_computation.roots.len(); + // Per query, the groups below its root through the groups' children. + let mut below: Vec> = vec![vec![false; memo.groups.len()]; queries]; + for (g, group) in memo.groups.iter().enumerate() { + for &q in &group.roots_of { + mark_below(memo, g, &mut below[q]); } - built_of[g] = Some(node); } + // A child below the reading query's root is built in that query; any + // other in the first query it is below, or once for all (`queries`). + let scope_of = |scope: usize, c: usize| { + if scope < queries && below[scope][c] { + scope + } else { + (0..queries).find(|&q| below[q][c]).unwrap_or(queries) + } + }; + // The candidate building each group: its own, or the fused candidate + // implementing it. + let building: Vec> = (0..memo.groups.len()) + .map(|g| { + chosen[g].as_ref().or_else(|| { + chosen + .iter() + .flatten() + .find(|c| c.implements.contains(&GroupId(g as u32))) + }) + }) + .collect(); + let mut built: Vec>> = vec![HashMap::new(); queries + 1]; let roots: Vec = per_computation .roots .iter() .enumerate() .map(|(q, (_, root))| { - let group_of = |node: &Rc| -> Option> { - memo.groups + let mut group_of = |node: &Rc| -> Option> { + let g = memo + .groups .iter() - .enumerate() - .find(|(_, group)| group.roots_of.contains(&q) && *group.node == **node) - .and_then(|(g, _)| built_of[g].clone()) + .position(|group| group.roots_of.contains(&q) && *group.node == **node)?; + build(memo, &building, q, g, &scope_of, &mut built) }; Some(match root { QueryRoot::Operator(node) => QueryRoot::Operator(group_of(node)?), @@ -305,6 +349,52 @@ fn composed_roots( recompose_guarantees(&roots).ok() } +/// Group `g` built in `scope` (a query, or the scope of groups no query +/// root reaches): its building candidate's node over its children's built +/// nodes, each child built in `scope_of(scope, child)`; memoized in +/// `built`. +fn build( + memo: &Memo, + building: &[Option<&ImplementationCandidate>], + scope: usize, + g: usize, + scope_of: &dyn Fn(usize, usize) -> usize, + built: &mut Vec>>, +) -> Option> { + if let Some(node) = built[scope].get(&g) { + return Some(Rc::clone(node)); + } + let candidate = building[g]?; + let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); + for child in &candidate.children { + let c = child.0 as usize; + let node = build(memo, building, scope_of(scope, c), c, scope_of, built)?; + replaced.insert(Rc::as_ptr(&memo.groups[c].node), node); + } + // A node of its own in each scope, even a leaf's (a scan), so that two + // queries share it only through the identical-expression rule. + let node = Rc::new(OperatorNode::clone(&substituted( + &candidate.node, + &replaced, + &mut HashMap::new(), + ))); + for implemented in &candidate.implements { + built[scope].insert(implemented.0 as usize, Rc::clone(&node)); + } + Some(node) +} + +/// Mark `g` and every group below it through the groups' children. +fn mark_below(memo: &Memo, g: usize, below: &mut [bool]) { + if below[g] { + return; + } + below[g] = true; + for child in &memo.groups[g].children { + mark_below(memo, child.0 as usize, below); + } +} + /// `node` with every node `built` maps (by address) replaced by its built /// node, rebuilt above them; `rebuilt` memoizes by address. fn substituted( @@ -337,23 +427,6 @@ fn substituted( new } -/// `roots` with its identical nodes made one node, then its summary builds -/// equal to one of `resized` (over equal windows of two groups) made one -/// build and the identical nodes above them shared in turn. -fn shared(roots: &[QueryRoot], resized: &[Rc]) -> Vec { - let identical = |roots: Vec| { - let shareable = sharing_candidates(&roots); - match shareable.is_empty() { - true => roots, - false => share_identical_expressions(&roots, &shareable), - } - }; - identical(super::capability::share_equal_builds( - &identical(roots.to_vec()), - resized, - )) -} - /// The shared builds `plan` resized: each shared-build reader's build that /// differs from the build of the memo's candidate it was chosen as /// (`choice`). @@ -402,6 +475,14 @@ pub(crate) fn brute_force( if !super::windows::requirements_met(&chosen) { continue; } + // Scans are widened as the identical-expression rule widens them, + // every metric's at once ([`identical_expression_variants`]). + if chosen + .iter() + .any(|(_, c)| c.origin == super::CandidateOrigin::WidenedScans) + { + continue; + } for plan in at_thresholds(&memo, choice) { let Some(roots) = composed_roots(&memo, per_computation, &plan) else { continue; @@ -412,25 +493,37 @@ pub(crate) fn brute_force( let Ok(roots) = recompose_guarantees(&roots) else { continue; }; - let roots = shared(&roots, &resized_builds(&memo, choice, &plan)); - number += 1; - let id = format!("L{number}"); - let Ok(implemented) = - optimize_physical_asap(&id, &roots, entries, data, models.deployment_capabilities) - else { - continue; - }; - for (k, mut physical) in implemented.into_iter().enumerate() { - physical.id = match k { - 0 => format!("P{number}"), - k => format!("P{number}-m{k}"), + // Shared builds are one build (the summary-capability rule), + // then identical nodes shared or not, as the exhaustive search + // lists them. + let roots = super::capability::share_equal_builds( + &roots, + &resized_builds(&memo, choice, &plan), + ); + for roots in identical_expression_variants(roots) { + number += 1; + let id = format!("L{number}"); + let Ok(implemented) = optimize_physical_asap( + &id, + &roots, + entries, + data, + models.deployment_capabilities, + ) else { + continue; }; - selector.offer(&physical)?; - } - if selector.keeps_from_logical(&id) { - logical.insert(id, roots); + for (k, mut physical) in implemented.into_iter().enumerate() { + physical.id = match k { + 0 => format!("P{number}"), + k => format!("P{number}-m{k}"), + }; + selector.offer(&physical)?; + } + if selector.keeps_from_logical(&id) { + logical.insert(id, roots); + } + logical.retain(|id, _| selector.keeps_from_logical(id)); } - logical.retain(|id, _| selector.keeps_from_logical(id)); } } let (selection, selected) = selector.finish()?; @@ -499,7 +592,10 @@ mod tests { let plans = at_thresholds(&memo, &choice); assert_eq!(plans.len(), 1); let roots = composed_roots(&memo, &per_computation, &plans[0]).unwrap(); - let roots = shared(&roots, &resized_builds(&memo, &choice, &plans[0])); + let roots = super::super::capability::share_equal_builds( + &roots, + &resized_builds(&memo, &choice, &plans[0]), + ); let builds: std::collections::HashSet<*const OperatorNode> = roots .iter() .flat_map(|root| match root { From c364c225be7a50c1369c4c8d109d5dfe1470a44e Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:19:18 +0000 Subject: [PATCH 19/25] plan-selection: the selection report carries the memo's size (SearchReport.memo) MemoStats counts the memo's groups, shared groups, candidates by origin, distinct coupling sites by kind and its build time; the reference search fills it whenever it builds the memo. Replaces the string stats() and the temporary JP_MEMO_STATS hook. --- crates/plan-selection/src/lib.rs | 24 +++ .../src/pass/stage1_to_stage3.rs | 10 -- .../plan-selection/src/pass/subexpr/search.rs | 162 +++++++++++++++--- 3 files changed, 160 insertions(+), 36 deletions(-) diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index aae16cc0c..e15543682 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1028,6 +1028,30 @@ pub struct SearchReport { /// candidate costs the same, under cost model 3, or for a combination /// of components. pub runner_up_cost_per_second: Option, + /// The size of the sub-expression memo the search built; `None` when + /// the search built none. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub memo: Option, +} + +/// The size of a sub-expression memo (joint-planning-search design, +/// "Memo-group search"), for the selection report. +#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct MemoStats { + /// Its groups: the workload's distinct logical sub-expressions. + pub groups: usize, + /// The groups two or more parents, or two or more queries' roots, read. + pub shared_groups: usize, + /// Its implementation candidates, over every group. + pub candidates: usize, + /// The candidates by how they were made (`Pass1`, `SharedBuild`, + /// `Panes`, `Rollup`, `TopKReuse`). + pub candidates_by_origin: BTreeMap, + /// The distinct coupling sites by kind (`SharedBuildSize`, `PaneChain`, + /// `ScanLabels`, `RawHorizon`). + pub coupling_sites_by_kind: BTreeMap, + /// The time building it took, in milliseconds. + pub build_ms: f64, } /// One independent component of the workload: queries that can read no diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index bd4859718..419739b86 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -234,16 +234,6 @@ pub fn run_stage1_to_stage3_from_pass1( models: PlanningModels<'_>, search: Stage1ToStage3Search, ) -> Stage1ToStage3Run { - if std::env::var("JP_MEMO_STATS").is_ok() { - let started = std::time::Instant::now(); - let memo = super::subexpr::search::memo(&per_computation, entries); - eprintln!( - "JPMEMO {} in {:?}", - super::subexpr::search::stats(&memo), - started.elapsed() - ); - std::process::exit(0); - } if search.memo_groups && !search.exhaustive { let pruned = per_computation .clone() diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 9f5944cfb..f45dbb8c7 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -6,7 +6,7 @@ //! one threshold per coupling site, composed into a plan and priced by //! Stage 2 and Stage 3 in every materialization. -use std::collections::HashMap; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::rc::Rc; use asap_logical_asap_optimization::asap_aware_cse::{ @@ -20,7 +20,7 @@ use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; use super::{groups, CouplingKind, GroupId, ImplementationCandidate, Memo}; -use crate::{PlanSelector, PlanningModels, Selection, SelectionError}; +use crate::{MemoStats, PlanSelector, PlanningModels, Selection, SelectionError}; /// The workload's memo: groups, Pass 1 candidates and the Pass 2 /// transformations' candidates. @@ -37,32 +37,38 @@ pub(crate) fn memo( memo } -/// A one-line summary of `memo` for logs: groups, shared groups (two or -/// more parents or queries), candidates, distinct coupling sites. -pub(crate) fn stats(memo: &Memo) -> String { - let shared = memo +/// The size of `memo`, built in `build_ms` milliseconds, for the +/// selection report: groups, shared groups (read by two or more parents, +/// or the root of two or more queries), candidates by origin and distinct +/// coupling sites by kind. +pub(crate) fn stats(memo: &Memo, build_ms: f64) -> MemoStats { + let shared_groups = memo .groups .iter() .filter(|g| g.parents.len() > 1 || (g.parents.is_empty() && g.roots_of.len() > 1)) .count(); - let candidates: usize = memo.candidates.iter().map(Vec::len).sum(); - let mut sites: Vec = memo - .candidates - .iter() - .flatten() - .flat_map(|c| c.coupling_sites.iter().map(|s| s.key)) - .collect(); - sites.sort_unstable(); - sites.dedup(); - let most = memo.candidates.iter().map(Vec::len).max().unwrap_or(0); - format!( - "groups {} shared {} candidates {} (most in one group {}) coupling sites {}", - memo.groups.len(), - shared, - candidates, - most, - sites.len() - ) + let mut candidates_by_origin: BTreeMap = BTreeMap::new(); + let mut sites = BTreeSet::new(); + for candidate in memo.candidates.iter().flatten() { + *candidates_by_origin + .entry(format!("{:?}", candidate.origin)) + .or_default() += 1; + sites.extend(candidate.coupling_sites.iter().copied()); + } + let mut coupling_sites_by_kind: BTreeMap = BTreeMap::new(); + for site in sites { + *coupling_sites_by_kind + .entry(format!("{:?}", site.kind)) + .or_default() += 1; + } + MemoStats { + groups: memo.groups.len(), + shared_groups, + candidates: memo.candidates.iter().map(Vec::len).sum(), + candidates_by_origin, + coupling_sites_by_kind, + build_ms, + } } /// One plan of the memo: a candidate per needed group (`None` for a group @@ -461,7 +467,9 @@ pub(crate) fn brute_force( data: &DataWorkload, models: PlanningModels<'_>, ) -> Result<(Selection, PhysicalCandidate, Vec), SelectionError> { + let started = std::time::Instant::now(); let memo = memo(per_computation, entries); + let memo_stats = stats(&memo, started.elapsed().as_secs_f64() * 1e3); let choices = group_choices(&memo).unwrap_or_default(); let mut selector = PlanSelector::new(entries, data, models, false)?; let mut logical: HashMap> = HashMap::new(); @@ -526,7 +534,8 @@ pub(crate) fn brute_force( } } } - let (selection, selected) = selector.finish()?; + let (mut selection, selected) = selector.finish()?; + selection.search.memo = Some(memo_stats); let roots = logical .remove(&selected.from_logical) .expect("the selection's logical roots are kept"); @@ -608,6 +617,107 @@ mod tests { assert_eq!(builds.len(), 1); } + /// A quantile's memo, Pass 1 only, has three groups (scan, window, + /// quantile) and five candidates (the quantile's exact, KLL and + /// DDSketch); with one candidate made a pane candidate and coupling + /// sites added, the stats count candidates by origin and distinct sites + /// by kind, a site two candidates share once. + #[test] + fn stats_count_candidates_by_origin_and_sites_by_kind() { + use super::super::{CandidateOrigin, CouplingSite}; + let roots = vec![( + 0usize, + QueryRoot::Operator(lower_promql( + "quantile_over_time(0.9, m[5m])", + AccuracyTarget::max_error(0.01), + )), + )]; + let mut memo = groups::memo_of(&generate_per_computation_candidates(roots).unwrap()); + let shared = CouplingSite { + key: 1, + kind: CouplingKind::SharedBuildSize, + }; + let chain = CouplingSite { + key: 2, + kind: CouplingKind::PaneChain, + }; + memo.candidates[2][1].coupling_sites = vec![shared]; + memo.candidates[2][2].coupling_sites = vec![shared, chain]; + memo.candidates[2][2].origin = CandidateOrigin::Panes; + let stats = stats(&memo, 1.5); + assert_eq!( + stats, + MemoStats { + groups: 3, + shared_groups: 0, + candidates: 5, + candidates_by_origin: [("Pass1".to_string(), 4), ("Panes".to_string(), 1)].into(), + coupling_sites_by_kind: [ + ("SharedBuildSize".to_string(), 1), + ("PaneChain".to_string(), 1) + ] + .into(), + build_ms: 1.5, + } + ); + } + + /// The reference search reports the memo it built. + #[test] + fn the_brute_force_selection_reports_its_memo() { + let query = "quantile_over_time(0.9, m[5m])"; + let target = AccuracyTarget::max_error(0.01); + let roots = vec![( + 0usize, + QueryRoot::Operator(lower_promql(query, target.clone())), + )]; + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let entries = vec![QueryWorkloadEntry { + query: asap_types::workload::Query(query.into()), + requirements: QueryRequirements { + accuracy: target, + max_latency_ms: None, + }, + predictability: Predictability::Predictable { known_at: None }, + recurrence: QueryRecurrence::Repeated(RepeatedDemand::FixedInterval( + RepetitionInterval(60_000), + )), + range_evaluation: None, + }]; + let data = DataWorkload { + arrival: DataArrival::Ingesting { + rows_per_second: 2_000.0 / 15.0, + }, + sources: [( + "m".to_string(), + SourceFacts { + series_count: Some(1_000), + metric_type: Some(MetricType::Gauge), + sample_interval: Some(DurationMs(15_000)), + ..SourceFacts::default() + }, + )] + .into(), + }; + let (selection, _, _) = brute_force( + &per_computation, + &entries, + &data, + crate::PlanningModels::builtin(), + ) + .unwrap(); + let reported = selection.search.memo.clone().expect("the memo's stats"); + let built = memo(&per_computation, &entries); + assert_eq!(reported.groups, built.groups.len()); + assert_eq!( + reported.candidates, + built.candidates.iter().map(Vec::len).sum::() + ); + assert!(reported.build_ms >= 0.0); + let json = serde_json::to_value(&selection.search).unwrap(); + assert_eq!(json["memo"]["groups"], reported.groups); + } + #[test] #[ignore = "debugging aid"] fn debug_topk_memo() { @@ -652,7 +762,7 @@ mod tests { .into(), }; let memo = memo(&pc, &entries); - eprintln!("{}", stats(&memo)); + eprintln!("{:?}", stats(&memo, 0.0)); for (g, group) in memo.groups.iter().enumerate() { eprintln!( "G{g} {:?} children {:?} roots_of {:?} candidates {}", From 243b64ced9db09591845290039a838c789d29474 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:16:08 +0000 Subject: [PATCH 20/25] docs: the planning stages over a memo of sub-expressions; memo module overview --- crates/plan-selection/src/pass/subexpr/mod.rs | 59 ++++++++++++------- .../design_docs/proposals/planner-layering.md | 42 +++++++++++++ 2 files changed, 81 insertions(+), 20 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 61dcea355..e48f0a8f4 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -1,25 +1,44 @@ -//! Memo over logical sub-expressions (joint-planning-search design, -//! "Memo-group search"): the workload's logical DAG hash-consed into -//! groups, each group's implementation candidates generated by Pass 1 and -//! the Pass 2 transformations, priced per (group, candidate, context) by -//! Stage 3, and searched by enumerating coupling thresholds with bounds and -//! a dynamic program over the rest. +//! Memo over logical sub-expressions (planner-layering design, "The +//! stages over a memo of sub-expressions"): the stages run on the +//! workload's distinct sub-expressions instead of whole-workload +//! candidates. //! -//! Behind [`crate::pass::Stage1ToStage3Search::memo_groups`] until it -//! selects the exhaustive search's price on every equivalence test. +//! [`search::memo`] builds the [`Memo`] in this order: //! -//! - [`groups`] (W1): the hash-consed group DAG and Pass 1's -//! implementation candidates per group; -//! - [`capability`] (W3): identical-expression and summary-capability -//! transformations; -//! - [`scans`] (W3): scan widening to a metric's shared label set, and -//! each source's raw-data horizon; -//! - [`windows`] (W4): window-composition, rollup and top-k reuse -//! transformations; -//! - [`pricing`] (W2): local prices per (group, candidate, context), -//! cached, and frontiers; -//! - [`search`] (integrator): clusters, coupling thresholds with branch and -//! bound, the dynamic program, build-late. +//! 1. [`groups::memo_of`]: the workload's logical DAG hash-consed into +//! groups (equal scans one group, keyless nodes such as range windows one +//! per reader), children first, with each group's Pass 1 candidates. +//! 2. [`capability::add_shared_builds`]: summary capability. A reader of a +//! build shared with other groups (one family over equal inputs), with a +//! [`CouplingKind::SharedBuildSize`] site; identical expression is group +//! identity and adds nothing. +//! 3. [`windows::add_pass2_candidates`]: window composition (panes, and +//! newest-pane reads of a shared chain with a [`CouplingKind::PaneChain`] +//! site), rollup and exact top-k reuse (candidates reading another +//! group). +//! 4. [`scans::add_scan_widening`]: readers of a metric's scans widened to +//! a shared label set, with a [`CouplingKind::ScanLabels`] site. +//! 5. [`scans::add_raw_horizons`]: a [`CouplingKind::RawHorizon`] site on +//! every candidate that reads a raw window of a source read by two or +//! more groups. +//! 6. [`scans::scan_materializations`]: a scan read through one window is +//! only built once; through different windows, also per reader. +//! +//! Each site kind has its thresholds and reader check, and a candidate's +//! node at a threshold: [`capability::size_thresholds`] and +//! [`capability::resized_at_threshold`]; [`windows::chain_thresholds`] and +//! [`windows::read_at_threshold`]; [`scans::label_thresholds`] and +//! [`scans::widened_at_threshold`]; [`scans::horizon_thresholds`] (no +//! rewrite: Stage 3 prices the horizon). +//! +//! [`search::brute_force`] is the reference search, run when +//! [`crate::pass::Stage1ToStage3Search::memo_groups`] is set: every +//! consistent choice of one candidate per needed group and one threshold +//! per site its candidates read, composed into one plan (identical nodes +//! shared, equal resized builds made one), implemented by Stage 2 in every +//! materialization and priced by Stage 3. [`pricing`] holds the per-node +//! local prices, plan-level charges and frontiers the decomposed search +//! will use in its place. #![allow(dead_code)] pub(crate) mod capability; diff --git a/docs/design_docs/proposals/planner-layering.md b/docs/design_docs/proposals/planner-layering.md index 505937623..665969881 100644 --- a/docs/design_docs/proposals/planner-layering.md +++ b/docs/design_docs/proposals/planner-layering.md @@ -17,6 +17,7 @@ design is extended. - [Design Intuition](#design-intuition) - [Assumptions](#assumptions) - [ASAPPlanner System Overview: Planning Stages](#asapplanner-system-overview-planning-stages) + - [The stages over a memo of sub-expressions](#the-stages-over-a-memo-of-sub-expressions) - [ASAPPlanner Detailed Design: Planning Stages and their decisions/Strategies](#asapplanner-detailed-design-planning-stages-and-their-decisionsstrategies) - [0. Language-specific frontends](#0-language-specific-frontends) - [1. Logical ASAP-aware optimization](#1-logical-asap-aware-optimization) @@ -243,6 +244,47 @@ Stage details: - [3. Plan selection](#3-plan-selection) · - [4. Execution](#4-execution) +### The stages over a memo of sub-expressions + +The stages above exchange whole-workload candidates, and their number is the +product of every sub-DAG's choices. The memo search (being built beside the +current search, not yet the default) runs the same stages with the same +semantics and code boundaries, but on the workload's distinct logical +sub-expressions: each choice is made once per sub-expression, and only the +choices several sub-expressions depend on are made jointly. + +**Components.** + +| Component | What it is | +|---|---| +| Memo group | One distinct logical sub-expression: the workload's logical DAG hash-consed, so identical expressions of several queries are one group with several parents. Equal scans are one group; a node without a unique key (e.g. a range window) is one group per reader. | +| Implementation candidate | One way to implement a group over its child groups: its own nodes, the child groups it reads, how it was made (Pass 1 or which Pass 2 rule), the coupling sites its cost depends on, and the materializations Stage 2 may give it. A fused computation is one candidate implementing several groups. | +| Coupling site | A choice several groups' candidates depend on, keyed so that equal keys are one site, with its thresholds (the values the choice can take) and a check of whether a reader may read a threshold. | + +Each stage keeps its decision, on groups instead of whole candidates: + +| Stage | Over the memo | +|---|---| +| 1, Pass 1 | Emits each group's implementation candidates: exact, exact accumulator, each summary family and size, fused computations. | +| 1, Pass 2 | Each rule adds candidates to groups. Identical expression is group identity (nothing to add). Summary capability: a reader reading a build shared with other groups, with a shared-build size site (thresholds: the joins of the readers' sizes). Window composition: tumbling panes, and the newest panes of a shared chain with a pane-chain site (grid, longest window, size). Rollup and exact top-k reuse: a candidate that reads another group (a finer grouping, a ranking). Scan widening: a reader of a metric's scans widened to a shared label set, with a scan-labels site. Raw horizon: a source's raw-data horizon site on every candidate that reads a raw window of it. | +| 2 | Chooses each candidate's materialization, once for a group however many parents read it. A scan read through one window is built once; one read through different windows may also be read per reader. | +| 3 | Prices each candidate's own nodes, per node, given its readers and the thresholds of its sites; plan-level charges are added once per plan. | + +**Algorithm.** Hash-cons the workload into groups; generate each group's +candidates (Pass 1, then the Pass 2 transformations); choose one candidate +per needed group and one threshold per coupling site the chosen candidates +read; compose the choice into one plan and check and price it with +Stage 2 and Stage 3. The reference search enumerates every such choice +(brute force) and must select the exhaustive search's price; the decomposed +search that replaces it enumerates thresholds with bounds and chooses the +rest per group. + +**Code interface.** `crates/plan-selection/src/pass/subexpr/`: `Memo`, +`MemoGroup`, `ImplementationCandidate`, `CouplingSite` and `CouplingKind` +in `mod.rs`; `search::memo` builds the memo and `search::brute_force` is +the reference search, run by `run_stage1_to_stage3_from_pass1` when +`Stage1ToStage3Search::memo_groups` is set. + ## ASAPPlanner Detailed Design: Planning Stages and their decisions/Strategies Stages 0 to 2 (logical and physical planning stages) each output a candidate set holding every semantically equivalent and legal candidate DAG of that stage; stage 3 is the only step that chooses one candidate DAG as output currently (TODO: early selection for valid and more efficient candidates of stage 0-2 can be designed later). Candidate sets are internal to ASAPPlanner and may From ea6ff41219607acb18662f420becf60ce7086d14 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:18:30 +0000 Subject: [PATCH 21/25] plan-selection: fix a stale doc link in the memo types --- crates/plan-selection/src/pass/subexpr/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index e48f0a8f4..54c8bf4b8 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -190,7 +190,7 @@ pub(crate) struct Context { } /// One candidate's usage priced in a context: Stage 3's per-node pricing of -/// its private nodes ([`pricing::local_price`]). The charges Stage 3 makes +/// its private nodes ([`pricing::private_usage`]). The charges Stage 3 makes /// only once the whole plan is known are kept apart and added by /// [`pricing::plan_cost`]: memory held for an evaluation (which needs the /// readers' whole latency) and the raw data query-time scans keep on disk From 67a18d4eaacaeb66a7d34f65f71c16d647ab54a3 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:32:51 +0000 Subject: [PATCH 22/25] devtools: stage_pipeline --search legacy|memo-brute-force; drop the debugging tests --- crates/devtools/src/bin/stage_pipeline.rs | 33 ++- crates/devtools/src/planning_stages.rs | 26 ++- .../tests/joint_search_equivalence.rs | 192 ------------------ .../plan-selection/src/pass/subexpr/search.rs | 78 ------- 4 files changed, 47 insertions(+), 282 deletions(-) diff --git a/crates/devtools/src/bin/stage_pipeline.rs b/crates/devtools/src/bin/stage_pipeline.rs index 02961938f..a186da859 100644 --- a/crates/devtools/src/bin/stage_pipeline.rs +++ b/crates/devtools/src/bin/stage_pipeline.rs @@ -202,7 +202,7 @@ use asap_devtools::planner_layering_examples::{ example1, example2, pattern_a, pattern_b, PatternARecurrence, SHARED_ROWS_PER_SECOND, }; use asap_devtools::planning_stages::{ - lower_promql_roots, lower_sql_roots, run_default_search_over, run_planning_stages_over, + lower_promql_roots, lower_sql_roots, run_planning_stages_over, run_search_over, PlanningStagesError, PlanningStagesRun, }; use asap_devtools::SqlCatalog; @@ -214,7 +214,7 @@ use asap_logical_asap_optimization::DeploymentCapabilities; use asap_plan_selection::{ large_peak_memory_warning, ClusterLimits, ClusterObjective, CostModel, MeasuredOperationCosts, PlanSelectionSettings, PlanningModels, Rejection, ResourceUnitCost, ResourceUsage, Selection, - SelectionError, + SelectionError, Stage1ToStage3Search, }; use asap_planner::{installable_plan_of, PlannerExecutorCompileCheck}; use asap_types::ir::flat::{flatten, FlatDag}; @@ -233,7 +233,7 @@ const USAGE: &str = [--name ]... [--table ...] [--max-error --failure-probability --accuracy-metric ] [--sample-interval-ms ])) \ [--cost-model usage-over-time|aws-fargate-us-east-1|provisioned-capacity|cluster-capacity [--cluster-memory-bytes \ --cluster-disk-bytes [--cluster-cpu-cores ] --cluster-objective lowest-latency|least-memory|pareto-front]] \ -[--plan-selection-settings ] [--enforce-accuracy-guarantees] [--deployment-capabilities ] [--exhaustive [--max-combinations ]] [--max-displayed ] --out "; +[--plan-selection-settings ] [--enforce-accuracy-guarantees] [--deployment-capabilities ] [--exhaustive [--max-combinations ] | --search legacy|memo-brute-force] [--max-displayed ] --out "; /// The default `--max-combinations`: the most Stage 1 candidates (Pass 2's /// included) an exhaustive run enumerates. @@ -334,6 +334,9 @@ fn run(args: Vec) -> Result<(), Failure> { let mut plan_selection_settings_path: Option = None; let mut enforce_accuracy_guarantees = false; let mut exhaustive = false; + // `--search memo-brute-force`: the reference search over the memo of + // sub-expressions instead of the default search. + let mut memo_brute_force = false; let mut max_combinations_given = false; let mut args = args.into_iter(); while let Some(flag) = args.next() { @@ -393,6 +396,13 @@ fn run(args: Vec) -> Result<(), Failure> { "--plan-selection-settings" => plan_selection_settings_path = Some(value()?), "--enforce-accuracy-guarantees" => enforce_accuracy_guarantees = true, "--exhaustive" => exhaustive = true, + "--search" => { + memo_brute_force = match value()?.as_str() { + "legacy" => false, + "memo-brute-force" => true, + other => return Err(format!("unknown search {other}").into()), + } + } other => return Err(format!("unknown argument {other}").into()), } } @@ -614,7 +624,16 @@ fn run(args: Vec) -> Result<(), Failure> { Some(max_combinations), models, ), - false => run_default_search_over(roots.clone(), &entries, &data, models), + false => run_search_over( + roots.clone(), + &entries, + &data, + models, + Stage1ToStage3Search { + memo_groups: memo_brute_force, + ..Stage1ToStage3Search::default() + }, + ), }; // A selected plan's very large peak memory, when the deployment // declares no memory capacity to check it against. @@ -636,7 +655,11 @@ fn run(args: Vec) -> Result<(), Failure> { } document["workload"]["data_workload"] = serde_json::to_value(&data).map_err(|e| e.to_string())?; - document["search"] = json!(if exhaustive { "exhaustive" } else { "default" }); + document["search"] = json!(match (exhaustive, memo_brute_force) { + (true, _) => "exhaustive", + (false, true) => "memo-brute-force", + (false, false) => "default", + }); document["accuracy_guarantees"] = json!(if enforce_accuracy_guarantees { "enforced" } else { diff --git a/crates/devtools/src/planning_stages.rs b/crates/devtools/src/planning_stages.rs index dfc560f8e..c519a322a 100644 --- a/crates/devtools/src/planning_stages.rs +++ b/crates/devtools/src/planning_stages.rs @@ -227,17 +227,29 @@ pub fn run_default_search_over( data: &DataWorkload, models: PlanningModels<'_>, ) -> Result { - validate(&roots)?; - let per_computation = - generate_per_computation_candidates(roots.iter().cloned().enumerate().collect()) - .map_err(|e| PlanningStagesError::Pass1(e.to_string()))?; - let run = run_stage1_to_stage3_from_pass1( - per_computation, + run_search_over( + roots, entries, data, models, Stage1ToStage3Search::default(), - ); + ) +} + +/// [`run_default_search_over`] with the search `search` (e.g. the +/// reference search over the memo, `memo_groups`). +pub fn run_search_over( + roots: Vec, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + search: Stage1ToStage3Search, +) -> Result { + validate(&roots)?; + let per_computation = + generate_per_computation_candidates(roots.iter().cloned().enumerate().collect()) + .map_err(|e| PlanningStagesError::Pass1(e.to_string()))?; + let run = run_stage1_to_stage3_from_pass1(per_computation, entries, data, models, search); if let Some(Ok(roots)) = run.stage1_selected.as_ref().map(|c| &c.roots) { validate(roots)?; } diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 0c2120b9c..2a8764ac8 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1241,195 +1241,3 @@ fn memo_group_search_selects_the_exhaustive_price() { ); assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); } - -/// Debugging aid: the memo's and the exhaustive search's selected roots on -/// one seed. -#[test] -#[ignore = "debugging aid"] -fn debug_memo_seed() { - use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode}; - fn show(node: &std::rc::Rc, depth: usize, out: &mut String) { - let detail = match (node.asap(), node.non_asap()) { - ( - Some(ASAPOp::SummaryAgg { - family, reduction, .. - }), - _, - ) => { - format!("{family:?} {reduction:?}") - } - ( - Some(ASAPOp::SummaryMerge { - newest_panes, - group_by, - .. - }), - _, - ) => { - format!("newest {newest_panes:?} group_by {group_by:?}") - } - ( - Some(ASAPOp::PaneChain { - family, - reduction, - pane_width_ms, - pane_count, - .. - }), - _, - ) => { - format!("{family:?} {reduction:?} {pane_width_ms}x{pane_count}") - } - (_, Some(NonASAPOp::TimeRange { range, .. })) => format!("{range:?}"), - ( - _, - Some(NonASAPOp::Aggregate { - reduction, - measures, - .. - }), - ) => { - format!("{reduction:?} {measures:?}") - } - _ => String::new(), - }; - let detail: String = detail.chars().take(160).collect(); - out.push_str(&format!( - "{}{} {:p} {detail}\n", - " ".repeat(depth), - node.operator.kind_name(), - std::rc::Rc::as_ptr(node) - )); - for child in node.children() { - show(child, depth + 1, out); - } - } - let seed: u64 = std::env::var("SEED") - .ok() - .and_then(|s| s.parse().ok()) - .unwrap_or(59); - let gauges: &[&str] = match seed % 3 { - 0 => &["m", "n"], - _ => &["m"], - }; - let workload = random_workload_over(seed, gauges); - let roots = lower_promql_roots(&workload).unwrap(); - let entries: Vec<_> = workload.query_workload.entries().collect(); - for e in &entries { - eprintln!("{} {:?}", e.query.0, e.recurrence); - } - let data = workload.data_workload.clone().unwrap(); - let deployment = match seed % 2 { - 0 => no_raw_data(), - _ => asap_logical_asap_optimization::DeploymentCapabilities::default(), - }; - let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); - let indexed = || roots.iter().cloned().enumerate().collect::>(); - /// A plan's shape: its flattened DAG without what composition derives. - fn shape(roots: &[QueryRoot]) -> String { - fn strip(value: &mut serde_json::Value) { - match value { - serde_json::Value::Object(map) => { - for field in ["guarantee", "schema", "timing", "materialization"] { - map.remove(field); - } - map.values_mut().for_each(strip); - } - serde_json::Value::Array(items) => items.iter_mut().for_each(strip), - _ => {} - } - } - let mut value = serde_json::to_value(flatten(roots).0).unwrap(); - strip(&mut value); - value.to_string() - } - let mut runs = Vec::new(); - for (name, search) in [ - ("exhaustive", Stage1ToStage3Search::exhaustive()), - ( - "memo", - Stage1ToStage3Search { - memo_groups: true, - ..Stage1ToStage3Search::default() - }, - ), - ] { - let run = run_stage1_to_stage3_with(indexed(), &entries, &data, models, search).unwrap(); - let s = run.stage3_plan_selection.as_ref().unwrap(); - eprintln!("== {name}: {:?}", s.costs[&s.selected].cost_per_second); - let mut out = String::new(); - for root in run - .stage1_selected - .as_ref() - .unwrap() - .roots - .as_ref() - .unwrap() - { - if let QueryRoot::Operator(node) = root { - show(node, 0, &mut out); - } - out.push_str("--\n"); - } - eprintln!("{out}"); - let selected = run.selected().unwrap(); - eprintln!( - "materialization: {:?}", - selected.materialization.values().collect::>() - ); - runs.push(run); - } - let memo_shape = shape( - runs[1] - .stage1_selected - .as_ref() - .unwrap() - .roots - .as_ref() - .unwrap(), - ); - let found: Vec<&str> = runs[0] - .stage1_logical_asap_optimization - .iter() - .filter(|c| c.roots.as_ref().is_ok_and(|r| shape(r) == memo_shape)) - .map(|c| c.id.as_str()) - .collect(); - eprintln!( - "memo's plan among the exhaustive candidates ({} listed): {found:?}", - runs[0].stage1_logical_asap_optimization.len() - ); - // The exhaustive search's price of each of them. - if let Ok(selection) = &runs[0].stage3_plan_selection { - for id in &found { - for (physical, cost) in &selection.costs { - if runs[0] - .stage2_physical_asap_optimization - .iter() - .any(|p| &p.id == physical && p.from_logical == *id) - { - eprintln!(" {physical}: {:?}", cost.cost_per_second); - } - } - } - for rejection in selection_rejections(&runs[0], &found) { - eprintln!(" rejected: {rejection}"); - } - } -} - -/// The Stage 3 rejections of the physical candidates of `logical`. -fn selection_rejections(run: &Stage1ToStage3Run, logical: &[&str]) -> Vec { - let Ok(selection) = &run.stage3_plan_selection else { - return Vec::new(); - }; - selection - .rejected - .iter() - .filter(|r| { - run.stage2_physical_asap_optimization - .iter() - .any(|p| p.id == r.id && logical.contains(&p.from_logical.as_str())) - }) - .map(|r| format!("{} {:?} {}", r.id, r.cause, r.reason)) - .collect() -} diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index f45dbb8c7..15be75eb7 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -717,82 +717,4 @@ mod tests { let json = serde_json::to_value(&selection.search).unwrap(); assert_eq!(json["memo"]["groups"], reported.groups); } - - #[test] - #[ignore = "debugging aid"] - fn debug_topk_memo() { - let queries = ["topk by (job) (3, sum_over_time(m[1m]))"]; - let roots: Vec<(usize, QueryRoot)> = queries - .iter() - .map(|q| { - let root = lower_promql(q, AccuracyTarget::max_error(0.05)); - QueryRoot::Operator(with_promql_series_identity(&root).unwrap()) - }) - .enumerate() - .collect(); - let pc = generate_per_computation_candidates(roots).unwrap(); - let entries: Vec = queries - .iter() - .map(|q| QueryWorkloadEntry { - query: asap_types::workload::Query((*q).into()), - requirements: QueryRequirements { - accuracy: AccuracyTarget::max_error(0.05), - max_latency_ms: None, - }, - predictability: Predictability::Predictable { known_at: None }, - recurrence: QueryRecurrence::Repeated(RepeatedDemand::FixedInterval( - RepetitionInterval(60_000), - )), - range_evaluation: None, - }) - .collect(); - let data = DataWorkload { - arrival: DataArrival::Ingesting { - rows_per_second: 2_000.0 / 15.0, - }, - sources: [( - "m".to_string(), - SourceFacts { - series_count: Some(1_000), - metric_type: Some(MetricType::Gauge), - sample_interval: Some(DurationMs(15_000)), - ..SourceFacts::default() - }, - )] - .into(), - }; - let memo = memo(&pc, &entries); - eprintln!("{:?}", stats(&memo, 0.0)); - for (g, group) in memo.groups.iter().enumerate() { - eprintln!( - "G{g} {:?} children {:?} roots_of {:?} candidates {}", - group.node.operator.kind_name(), - group.children, - group.roots_of, - memo.candidates[g].len() - ); - for c in &memo.candidates[g] { - eprintln!( - " {:?} implements {:?} children {:?}", - c.origin, c.implements, c.children - ); - } - } - let choices = group_choices(&memo).unwrap(); - eprintln!("choices {}", choices.len()); - for choice in &choices { - for plan in at_thresholds(&memo, choice) { - let roots = composed_roots(&memo, &pc, &plan); - eprintln!("choice {choice:?} composed {}", roots.is_some()); - if let Some(roots) = roots { - let r = - optimize_physical_asap("L", &roots, &entries, &data, &Default::default()); - eprintln!( - " stage2 {:?}", - r.as_ref().map(|v| v.len()).map_err(|e| e.to_string()) - ); - } - } - } - } } From 0753935c771066cf1578570d47e3fc758e6f4fdd Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:54:36 +0000 Subject: [PATCH 23/25] plan-selection: a fused memo candidate absorbs only groups no other reader needs Pass 1 fuses an aggregate only into its single consumer within one query. The memo merges equal nodes across queries, so an absorbed group may also be another query's root: the fused candidate now keeps such a group (and what only it reads) outside its implemented groups, so a top-k heap never answers a second query's inner sum, and the plan with the top-k by the heap and the sum by its own build is listed. Doc comments: no default look-back. --- .../tests/joint_search_equivalence.rs | 50 +++++++++++++++++++ crates/plan-selection/src/lib.rs | 8 +-- .../plan-selection/src/pass/subexpr/groups.rs | 25 ++++++++++ .../plan-selection/src/pass/subexpr/search.rs | 42 ++++++++++++++++ 4 files changed, 122 insertions(+), 3 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 2a8764ac8..93ecea390 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1241,3 +1241,53 @@ fn memo_group_search_selects_the_exhaustive_price() { ); assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); } + +/// A query's top-k over a sum, beside a query of that sum alone: the top-k's +/// fused heap absorbs the sum, which is the other query's root, so the memo +/// keeps the sum its own group for the second query. The memo-group search +/// selects a plan of the exhaustive search's price, with raw data kept and +/// without. +#[test] +fn a_fused_top_k_beside_its_inner_sum_selects_the_exhaustive_price() { + let workload = workload_of( + &[ + ("topk(5, sum by (job) (sum_over_time(m[5m])))", 60_000), + ("sum by (job) (sum_over_time(m[5m]))", 60_000), + ], + AccuracyTarget::max_error(0.01), + ); + let roots = lower_promql_roots(&workload).unwrap(); + let entries: Vec<_> = workload.query_workload.entries().collect(); + let data = workload.data_workload.clone().unwrap(); + for deployment in [ + asap_logical_asap_optimization::DeploymentCapabilities::default(), + no_raw_data(), + ] { + let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); + let run = |search| { + run_stage1_to_stage3_with( + roots.iter().cloned().enumerate().collect(), + &entries, + &data, + models, + search, + ) + .unwrap() + }; + let price = |run: &Stage1ToStage3Run| { + run.stage3_plan_selection + .as_ref() + .ok() + .map(|s| s.costs[&s.selected].cost_per_second.unwrap()) + }; + let exhaustive = price(&run(Stage1ToStage3Search::exhaustive())); + let memo = price(&run(Stage1ToStage3Search { + memo_groups: true, + ..Stage1ToStage3Search::default() + })); + match (exhaustive, memo) { + (Some(a), Some(f)) => assert!((a - f).abs() <= 1e-9 * a.abs(), "{a} vs {f}"), + (a, f) => assert_eq!(a.is_some(), f.is_some(), "{a:?} vs {f:?}"), + } + } +} diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index fbbfccdce..be5133988 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1288,7 +1288,8 @@ pub(crate) struct PrivatePricing<'a> { /// overhead is charged here). pub(crate) root_of: &'a [bool], /// The look-back of a query-time scan no window among the private - /// nodes bounds ([`DEFAULT_LOOKBACK_MS`] when `None`). + /// nodes bounds; `None` when the scan has no window, so it reads the + /// source's `rows_held` or fails with `MissingSourceFact`. pub(crate) scan_lookback_ms: Option, pub(crate) data: &'a DataWorkload, pub(crate) settings: &'a PlanSelectionSettings, @@ -4003,8 +4004,9 @@ pub(crate) struct EachNodePriced<'a> { /// [`price_node`] over every node of `dag`, children first, with the /// arguments of [`price_at_scale`]. `root_scan_lookback_ms` is the /// look-back of a query-time scan no window above it in `dag` bounds (a -/// scan priced without its consumers); `None` reads -/// [`DEFAULT_LOOKBACK_MS`] as Stage 3 does. `root_rates`, when given, are +/// scan priced without its consumers); `None` when the scan has no +/// window, so it reads the source's `rows_held` or fails with +/// `MissingSourceFact`, as Stage 3 does. `root_rates`, when given, are /// the computations per second of each root's consumer in place of its /// query's own rate ([`computations`]). #[allow(clippy::too_many_arguments)] diff --git a/crates/plan-selection/src/pass/subexpr/groups.rs b/crates/plan-selection/src/pass/subexpr/groups.rs index 032d3e64f..fcb774eee 100644 --- a/crates/plan-selection/src/pass/subexpr/groups.rs +++ b/crates/plan-selection/src/pass/subexpr/groups.rs @@ -284,6 +284,7 @@ fn target_candidates( } } } + private_absorbed(groups, &mut implements); for budget_share in 1..=max_budget_share { let Ok(composed) = compose_alone(per_computation, &aggregate, targets, &choice, budget_share) @@ -360,6 +361,30 @@ fn target_candidates( (!candidates.is_empty()).then_some(candidates) } +/// Keeps in `implements` (the group, then the groups its fused computation +/// absorbs) only the absorbed groups no other reader needs: no query's root, +/// and every parent in `implements`. Pass 1 fuses an aggregate only into its +/// single consumer within one query; the memo merges equal nodes across +/// queries, so an absorbed group may also be another query's root or +/// another group's child. Such a group (and, in turn, what only it reads) +/// is computed privately inside the fused candidate and stays its own group +/// for its other readers, so e.g. `topk(5, sum by (job) (…))` by a heap +/// never answers a second query's `sum by (job) (…)`. +fn private_absorbed(groups: &[MemoGroup], implements: &mut Vec) { + loop { + let shared = implements[1..].iter().position(|a| { + let group = &groups[a.0 as usize]; + !group.roots_of.is_empty() || group.parents.iter().any(|p| !implements.contains(p)) + }); + match shared { + Some(i) => { + implements.remove(i + 1); + } + None => break, + } + } +} + /// `aggregate` composed alone with `choice[t]` for each of its measures' /// `targets`, every other target left exact: the inventory holds only /// `targets`, and the targets a chosen fused computation absorbs (not diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 1809d5a79..a7e7f93ec 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -616,6 +616,48 @@ mod tests { assert_eq!(builds.len(), 1); } + /// Q1 = `topk(5, sum by (job) (sum_over_time(m[5m])))` and Q2 = its + /// inner sum: Q1's fused heap candidate absorbs the inner sum, which is + /// Q2's root. No composed plan answers Q2 with the heap, and the plan + /// "Q1 by the heap, Q2 by its own sum" is listed. + #[test] + fn a_fused_candidate_does_not_absorb_another_querys_root() { + let target = AccuracyTarget::max_error(0.01); + let roots: Vec<(usize, QueryRoot)> = [ + "topk(5, sum by (job) (sum_over_time(m[5m])))", + "sum by (job) (sum_over_time(m[5m]))", + ] + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, target.clone()))) + .enumerate() + .collect(); + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let memo = groups::memo_of(&per_computation); + let uses_heap = |root: &QueryRoot| match root { + QueryRoot::Operator(node) => OperatorNode::reachable(node).iter().any(|n| { + matches!(n.asap(), Some(asap_types::ir::ASAPOp::SummaryAgg { family, .. }) + if format!("{family:?}").contains("Heap")) + }), + QueryRoot::Scalar(_) => false, + }; + let mut q1_heap_q2_own = false; + for choice in group_choices(&memo).unwrap() { + for plan in at_thresholds(&memo, &choice) { + let roots = composed_roots(&memo, &per_computation, &plan).unwrap(); + assert!( + !uses_heap(&roots[1]), + "Q2 answered by a heap: {:?}", + roots[1] + ); + q1_heap_q2_own |= uses_heap(&roots[0]); + } + } + assert!( + q1_heap_q2_own, + "no plan answers Q1 by the heap and Q2 by its own sum" + ); + } + /// A quantile's memo, Pass 1 only, has three groups (scan, window, /// quantile) and five candidates (the quantile's exact, KLL and /// DDSketch); with one candidate made a pane candidate and coupling From e4a685953c06a7a93582a4b4b1f9376f83f7bcf9 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:24:27 +0000 Subject: [PATCH 24/25] plan-selection: the memo search lists sketch rollups of a shared finer build A sketch rollup's state site was never sized in the brute force (its thresholds went to the shared-build sizing, which lists no rollup reader), so no plan rolled a coarse quantile up from a finer KLL build; and the requirement check keyed the site by the resized copy, whose recomposed guarantees change the key. The search now sizes rollup state sites by their readers' thresholds and keys the site by the memo's build. On the quantile workload with by (host, job) and by (job) beside a per-series quantile it selects the exhaustive search's plan (0.041256/s, was 0.047030/s). --- .../tests/joint_search_equivalence.rs | 81 +++++++++++++++---- .../plan-selection/src/pass/subexpr/search.rs | 79 ++++++++++++++++++ .../src/pass/subexpr/windows.rs | 31 +++++-- 3 files changed, 167 insertions(+), 24 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 93ecea390..8174cf5cf 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1242,21 +1242,10 @@ fn memo_group_search_selects_the_exhaustive_price() { assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); } -/// A query's top-k over a sum, beside a query of that sum alone: the top-k's -/// fused heap absorbs the sum, which is the other query's root, so the memo -/// keeps the sum its own group for the second query. The memo-group search -/// selects a plan of the exhaustive search's price, with raw data kept and -/// without. -#[test] -fn a_fused_top_k_beside_its_inner_sum_selects_the_exhaustive_price() { - let workload = workload_of( - &[ - ("topk(5, sum by (job) (sum_over_time(m[5m])))", 60_000), - ("sum by (job) (sum_over_time(m[5m]))", 60_000), - ], - AccuracyTarget::max_error(0.01), - ); - let roots = lower_promql_roots(&workload).unwrap(); +/// The memo-group search selects a plan of the exhaustive search's price +/// on `workload`, with raw data kept and without (or both find none). +fn assert_memo_selects_the_exhaustive_price(label: &str, workload: &PlanningWorkload) { + let roots = lower_promql_roots(workload).unwrap(); let entries: Vec<_> = workload.query_workload.entries().collect(); let data = workload.data_workload.clone().unwrap(); for deployment in [ @@ -1286,8 +1275,66 @@ fn a_fused_top_k_beside_its_inner_sum_selects_the_exhaustive_price() { ..Stage1ToStage3Search::default() })); match (exhaustive, memo) { - (Some(a), Some(f)) => assert!((a - f).abs() <= 1e-9 * a.abs(), "{a} vs {f}"), - (a, f) => assert_eq!(a.is_some(), f.is_some(), "{a:?} vs {f:?}"), + (Some(a), Some(f)) => assert!( + (a - f).abs() <= 1e-9 * a.abs(), + "{label}: exhaustive {a}, memo {f}" + ), + (a, f) => assert_eq!(a.is_some(), f.is_some(), "{label}: {a:?} vs {f:?}"), } } } + +/// A query's top-k over a sum, beside a query of that sum alone: the top-k's +/// fused heap absorbs the sum, which is the other query's root, so the memo +/// keeps the sum its own group for the second query. The memo-group search +/// selects a plan of the exhaustive search's price, with raw data kept and +/// without. +#[test] +fn a_fused_top_k_beside_its_inner_sum_selects_the_exhaustive_price() { + let workload = workload_of( + &[ + ("topk(5, sum by (job) (sum_over_time(m[5m])))", 60_000), + ("sum by (job) (sum_over_time(m[5m]))", 60_000), + ], + AccuracyTarget::max_error(0.01), + ); + assert_memo_selects_the_exhaustive_price("top-k beside its sum", &workload); +} + +/// A per-series quantile over `data[1m]` beside `quantile by (host, job)` +/// and `quantile by (job)` of `data` (and, in the whole workload, a 5-minute +/// quantile), every 10 s over 10,000 series: the exhaustive plan rolls the +/// `by (job)` quantile up from the `by (host, job)` KLL build, which the +/// memo lists too. +#[test] +fn a_sketch_rollup_beside_other_quantiles_selects_the_exhaustive_price() { + let workload = |queries: &[&str]| -> PlanningWorkload { + let entries: Vec = queries + .iter() + .map(|q| { + serde_json::json!({ + "query": q, + "demand": {"fixed_interval": 10000}, + "requirements": {"accuracy": {"max_error": 0.01, "failure_probability": 0.001}}, + "predictability": {"predictable": {"known_at": null}} + }) + }) + .collect(); + serde_json::from_value(serde_json::json!({ + "query_workload": {"language": "promql", "query_batch": null, "repeating_queries": entries}, + "data_workload": { + "arrival": {"ingesting": {"rows_per_second": 666.67}}, + "sources": {"data": {"series_count": 10000, "metric_type": "gauge", "sample_interval": 15000}} + } + })) + .unwrap() + }; + let (q1, q2, q3, q4) = ( + "quantile_over_time(0.5, data[1m])", + "quantile_over_time(0.9, data[5m])", + "quantile by (host, job) (0.5, data)", + "quantile by (job) (0.5, data)", + ); + assert_memo_selects_the_exhaustive_price("Q1, Q3, Q4", &workload(&[q1, q3, q4])); + assert_memo_selects_the_exhaustive_price("Q2, Q3, Q4", &workload(&[q2, q3, q4])); +} diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index a7e7f93ec..29e5348dc 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -218,6 +218,33 @@ fn at_thresholds(memo: &Memo, choice: &GroupChoice) -> Vec { + // A sketch rollup's finer build, sized for its readers: + // the rollups and the finer group's exposing candidate. + for threshold in super::windows::state_size_thresholds(memo, &site) { + let mut rewritten = plan.clone(); + let mut valid = true; + for candidate in rewritten.iter_mut().flatten() { + if !candidate.coupling_sites.contains(&site) { + continue; + } + if !super::windows::may_read_state(&threshold, candidate, &site) { + valid = false; + break; + } + match super::windows::resize_state_at(candidate, &site, &threshold) { + Ok(node) => candidate.node = node, + Err(_) => { + valid = false; + break; + } + } + } + if valid { + next.push(rewritten); + } + } + } CouplingKind::SharedBuildSize => { for threshold in super::capability::size_thresholds(memo, &site) { let mut rewritten = plan.clone(); @@ -703,6 +730,58 @@ mod tests { ); } + /// A per-series quantile over `data[1m]` beside `quantile by (host, + /// job)` and `quantile by (job)` of `data`: the `by (job)` group's + /// rollup merging the `by (host, job)` KLL build meets its + /// requirements. The rolled-up copy of the finer build carries its + /// recomposed guarantees, so its state key differs from the memo build's + /// that keys the shared-size site; the check keys the site by the memo + /// build, so the plan is listed. + #[test] + fn a_sketch_rollup_beside_another_quantile_meets_its_requirements() { + let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); + let qs = [ + "quantile_over_time(0.5, data[1m])", + "quantile by (host, job) (0.5, data)", + "quantile by (job) (0.5, data)", + ]; + let roots: Vec<(usize, QueryRoot)> = qs + .iter() + .map(|q| QueryRoot::Operator(lower_promql(q, target.clone()))) + .enumerate() + .collect(); + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let entries: Vec = qs + .iter() + .map(|q| QueryWorkloadEntry { + query: asap_types::workload::Query((*q).into()), + requirements: QueryRequirements { + accuracy: target.clone(), + max_latency_ms: None, + }, + predictability: Predictability::Predictable { known_at: None }, + recurrence: QueryRecurrence::Repeated(RepeatedDemand::FixedInterval( + RepetitionInterval(10_000), + )), + range_evaluation: None, + }) + .collect(); + let memo = memo(&per_computation, &entries); + let met = group_choices(&memo).unwrap().into_iter().any(|choice| { + let chosen: Vec<(GroupId, &ImplementationCandidate)> = choice + .iter() + .enumerate() + .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) + .collect(); + chosen.iter().any(|(_, c)| { + c.rollup_of + .as_ref() + .is_some_and(|read| read.finer_state.is_some()) + }) && super::super::windows::requirements_met(&chosen) + }); + assert!(met, "no sketch rollup meets its requirements"); + } + /// The reference search reports the memo it built. #[test] fn the_brute_force_selection_reports_its_memo() { diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index 6196df94b..8fcacfac8 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -819,6 +819,17 @@ pub(crate) fn state_size_thresholds(memo: &Memo, site: &CouplingSite) -> Vec bool { + memo.candidates.iter().flatten().any(|c| { + c.rollup_of + .as_ref() + .is_some_and(|r| r.finer_state.is_some()) + && c.coupling_sites.contains(site) + }) +} + /// Whether `candidate`, a reader of the sketch rollup state `site`, may /// read it at size `threshold`: its size is no stricter. pub(crate) fn may_read_state( @@ -890,9 +901,9 @@ pub(crate) fn requirements_met(chosen: &[(GroupId, &ImplementationCandidate)]) - } } } - let rolled: Vec<&RollupRead> = chosen + let rolled: Vec<(&ImplementationCandidate, &RollupRead)> = chosen .iter() - .filter_map(|(_, c)| c.rollup_of.as_ref()) + .filter_map(|(_, c)| Some((*c, c.rollup_of.as_ref()?))) .collect(); if rolled.is_empty() { return true; @@ -911,7 +922,7 @@ pub(crate) fn requirements_met(chosen: &[(GroupId, &ImplementationCandidate)]) - .collect(); let mut used_state: HashSet<*const OperatorNode> = HashSet::new(); let mut rolled_coarse: HashSet<*const OperatorNode> = HashSet::new(); - for read in &rolled { + for &(reader, read) in &rolled { let Some(finer) = of(read.finer_group) else { return false; }; @@ -925,14 +936,20 @@ pub(crate) fn requirements_met(chosen: &[(GroupId, &ImplementationCandidate)]) - return false; } rolled_coarse.insert(Rc::as_ptr(&read.coarse_build)); - let Some(state) = &read.finer_state else { + if read.finer_state.is_none() { continue; - }; + } + // The state the reader reads is the joint rule's finer build's: its + // site (keyed by the memo's build, as [`sketch_rollup`] keys it, not + // by the resized copy, whose recomposed guarantees change its key), + // which the finer group's chosen candidate exposes. let site = CouplingSite { - key: state_key(state), + key: state_key(build), kind: CouplingKind::SharedBuildSize, }; - if !finer.coupling_sites.contains(&site) || state_key(build) != site.key { + if !reader.requires_sites.contains(&(read.finer_group, site)) + || !finer.coupling_sites.contains(&site) + { return false; } used_state.insert(Rc::as_ptr(build)); From 3f8706a315b6876e02c708097d714a900050b108 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:39:28 +0000 Subject: [PATCH 25/25] Memo search review fixes: name the selected logical candidate, drop unused interfaces - The memo branch names its Stage 1 candidate by the selected plan's from_logical, so the document's Stage 2 candidates read a listed Stage 1 candidate (the viewer's validator rejected L1 vs L3). - Delete share_identical_expressions and materialization_slots, which nothing calls. - stage_pipeline rejects --exhaustive with --search memo-brute-force and says --search legacy is the default search. - Memoize unchanged_below by node; point doc comments at planner-layering's memo section. --- crates/devtools/src/bin/stage_pipeline.rs | 6 ++- crates/devtools/tests/stage_pipeline.rs | 51 +++++++++++++++++++ .../src/asap_aware_cse.rs | 10 ---- crates/physical-asap-optimization/src/lib.rs | 42 --------------- crates/plan-selection/src/lib.rs | 20 +++++--- .../src/pass/stage1_to_stage3.rs | 22 +++++--- .../src/pass/subexpr/pricing.rs | 4 +- .../plan-selection/src/pass/subexpr/scans.rs | 3 +- .../plan-selection/src/pass/subexpr/search.rs | 7 ++- .../src/pass/subexpr/windows.rs | 2 +- 10 files changed, 92 insertions(+), 75 deletions(-) diff --git a/crates/devtools/src/bin/stage_pipeline.rs b/crates/devtools/src/bin/stage_pipeline.rs index 0a7d11d44..698dabe2f 100644 --- a/crates/devtools/src/bin/stage_pipeline.rs +++ b/crates/devtools/src/bin/stage_pipeline.rs @@ -239,7 +239,8 @@ const USAGE: &str = [--name ]... [--table ...] [--max-error --failure-probability --accuracy-metric ] [--sample-interval-ms ] [--series-count ] [--rows-per-second ] [--rows-held ])) \ [--cost-model usage-over-time|aws-fargate-us-east-1|provisioned-capacity|cluster-capacity [--cluster-memory-bytes \ --cluster-disk-bytes [--cluster-cpu-cores ] --cluster-objective lowest-latency|least-memory|pareto-front]] \ -[--plan-selection-settings ] [--enforce-accuracy-guarantees] [--deployment-capabilities ] [--exhaustive [--max-combinations ] | --search legacy|memo-brute-force] [--max-displayed ] --out "; +[--plan-selection-settings ] [--enforce-accuracy-guarantees] [--deployment-capabilities ] [--exhaustive [--max-combinations ] | --search legacy|memo-brute-force] [--max-displayed ] --out \n\ +--search legacy is the default search (the document's \"search\" is \"default\")."; /// The default `--max-combinations`: the most Stage 1 candidates (Pass 2's /// included) an exhaustive run enumerates. @@ -437,6 +438,9 @@ fn run(args: Vec) -> Result<(), Failure> { if max_combinations_given && !exhaustive { return Err("--max-combinations limits the exhaustive search: it needs --exhaustive (the default search has no limit)".into()); } + if exhaustive && memo_brute_force { + return Err("--exhaustive and --search memo-brute-force are two searches: give one".into()); + } if plan_selection_settings_path.is_some() && (cost_model.is_some() || cluster_memory.is_some() diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index 877bdaee9..e93546d4f 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1705,3 +1705,54 @@ fn scans_projecting_different_labels_are_shared() { "{rolled_up} vs {without_rollup}" ); } + +/// The memo search's document names its selected plan's logical candidate +/// under the id its physical candidates read (`from_logical`), as the +/// viewer's validator requires: the brute force numbers its logical plans, +/// so the selected one need not be the first. +#[test] +fn the_memo_search_document_names_the_selected_logical_candidate() { + let document = document_of(&[ + "--promql", + "quantile_over_time(0.5, data[5m])", + "--promql", + "quantile_over_time(0.9, data[5m])", + "--sample-interval-ms", + "15000", + "--search", + "memo-brute-force", + ]); + assert_eq!(document["search"], "memo-brute-force"); + let logical: HashSet<&str> = document["stage1_logical_asap_optimization"]["candidates"] + .as_array() + .unwrap() + .iter() + .map(|c| c["id"].as_str().unwrap()) + .collect(); + let physical = document["stage2_physical_asap_optimization"]["candidates"] + .as_array() + .unwrap(); + assert!(!physical.is_empty()); + for candidate in physical { + let from = candidate["from_logical"].as_str().unwrap(); + assert!(logical.contains(from), "{from} is not in {logical:?}"); + } + let selected = document["stage3_plan_selection"]["selected"] + .as_str() + .unwrap(); + assert!(physical.iter().any(|c| c["id"] == selected)); +} + +/// `--exhaustive` and `--search memo-brute-force` are two searches. +#[test] +fn exhaustive_and_the_memo_search_are_not_both_given() { + let (status, stderr) = failure(&[ + "--promql", + "sum(m)", + "--exhaustive", + "--search", + "memo-brute-force", + ]); + assert_eq!(status, Some(2)); + assert!(stderr.contains("give one"), "{stderr}"); +} diff --git a/crates/logical-asap-optimization/src/asap_aware_cse.rs b/crates/logical-asap-optimization/src/asap_aware_cse.rs index 7e937586a..417421154 100644 --- a/crates/logical-asap-optimization/src/asap_aware_cse.rs +++ b/crates/logical-asap-optimization/src/asap_aware_cse.rs @@ -502,16 +502,6 @@ pub fn identical_expression_pass2_candidates( .collect() } -/// `roots` with every one of `candidates` (sharing candidates of `roots`, -/// or some of them) made one node: the identical-expression Pass 2 -/// candidate that shares them all. -pub fn share_identical_expressions( - roots: &[QueryRoot], - candidates: &[SharingCandidate], -) -> Vec { - apply(roots, candidates, &vec![true; candidates.len()]).0 -} - /// `roots` with the sharing candidates in `subset` shared, and the ones /// that leaves shared: those whose every node still read became one node (a /// node below a shared node other than the one kept is no longer read). diff --git a/crates/physical-asap-optimization/src/lib.rs b/crates/physical-asap-optimization/src/lib.rs index 929cb71cf..4f14355e0 100644 --- a/crates/physical-asap-optimization/src/lib.rs +++ b/crates/physical-asap-optimization/src/lib.rs @@ -536,48 +536,6 @@ fn assigned_candidate( }) } -/// One materialization candidate of a plan, as Stage 2 finds it: a summary -/// build, or the highest node of a shared non-summary sub-DAG. -#[derive(Debug, Clone)] -pub struct MaterializationSlot { - /// The node, in the implemented plan. - pub node: Rc, - /// The queries (indices into the roots) that read it. - pub readers: Vec, - /// Its options, `NotMaterialized` first. - pub options: Vec, - /// The slots below it (by index): one at ingestion time needs every one - /// below it at ingestion time too. - pub below: Vec, -} - -/// The materialization candidates of `roots`, in the order -/// [`optimize_physical_asap_assigned`]'s `choose` is called for them, with -/// their options: what a search that chooses each one itself needs. -pub fn materialization_slots( - roots: &[QueryRoot], - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - deployment: &DeploymentCapabilities, -) -> Result, PhysicalASAPOptimizationError> { - let mut memo = HashMap::new(); - let implemented = map_roots(roots, &mut |node| implement(node, &mut memo))?; - let mat_candidates = materialization_candidates(&implemented, entries, data, deployment); - check_panes_on_range_steps(&mat_candidates, entries)?; - let nodes: Vec<_> = mat_candidates.iter().map(|c| Rc::clone(&c.node)).collect(); - let below = candidates_below(&nodes); - Ok(mat_candidates - .iter() - .zip(below) - .map(|(c, below)| MaterializationSlot { - node: Rc::clone(&c.node), - readers: c.readers.clone(), - options: c.options(), - below, - }) - .collect()) -} - /// The number of physical candidates [`optimize_physical_asap`] returns for /// `roots`, without building them; `Err` when a candidate cannot be /// implemented. diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 2bd1ba708..8d879be79 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1029,8 +1029,8 @@ pub struct SearchReport { pub memo: Option, } -/// The size of a sub-expression memo (joint-planning-search design, -/// "Memo-group search"), for the selection report. +/// The size of a sub-expression memo (planner-layering design, "The +/// stages over a memo of sub-expressions"), for the selection report. #[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] pub struct MemoStats { /// Its groups: the workload's distinct logical sub-expressions. @@ -1371,12 +1371,17 @@ fn family_below<'a>( } /// The unchanged source columns of node `id` of `dag` -/// ([`unchanged_source_columns`]), without pricing. +/// ([`unchanged_source_columns`]), without pricing; memoized by node in +/// `known`, so a node shared below several is walked once. fn unchanged_below<'a>( dag: &'a PhysicalASAPDAG, nodes: &NodeMap<'a>, id: PhysicalASAPNodeId, + known: &mut HashMap>, ) -> BTreeSet<&'a str> { + if let Some(columns) = known.get(&id) { + return columns.clone(); + } let producers: Vec<(&'a Schema, BTreeSet<&'a str>)> = dag .edges .iter() @@ -1384,7 +1389,7 @@ fn unchanged_below<'a>( .map(|e| { ( &nodes[&e.producer].output_schema, - unchanged_below(dag, nodes, e.producer), + unchanged_below(dag, nodes, e.producer, known), ) }) .collect(); @@ -1392,7 +1397,9 @@ fn unchanged_below<'a>( .iter() .map(|(schema, columns)| (*schema, columns)) .collect(); - unchanged_source_columns(nodes[&id], &views) + let columns = unchanged_source_columns(nodes[&id], &views); + known.insert(id, columns.clone()); + columns } /// Price `input`'s private nodes as Stage 3 prices them in the whole plan @@ -1556,6 +1563,7 @@ pub(crate) fn price_private_nodes( ); // Price each private node, children first. let mut priced: HashMap> = HashMap::new(); + let mut unchanged_known = HashMap::new(); for node in dag.nodes.iter().filter(|n| private_ids.contains(&n.id)) { let views: Vec> = producers_of .get(&node.id) @@ -1580,7 +1588,7 @@ pub(crate) fn price_private_nodes( .unwrap_or(EdgeStatistics { rows: 1, bytes: 1 }), sources: sources_below(dag, &nodes, *p), family: family_below(dag, &nodes, *p), - unchanged_columns: unchanged_below(dag, &nodes, *p), + unchanged_columns: unchanged_below(dag, &nodes, *p, &mut unchanged_known), }, }) .collect(); diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index 419739b86..56bcf04c1 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -159,8 +159,9 @@ pub struct Stage1ToStage3Search { /// own candidates first: the exhaustive search the pruned one is tested /// against; default `false`. pub exhaustive: bool, - /// Search the memo over logical sub-expressions (joint-planning-search - /// design, "Memo-group search") instead of whole-workload candidates; + /// Search the memo over logical sub-expressions (planner-layering + /// design, "The stages over a memo of sub-expressions") instead of + /// whole-workload candidates; /// ignored with `exhaustive`. Default `false` until it selects the /// exhaustive search's price on every equivalence test. pub memo_groups: bool, @@ -244,6 +245,7 @@ pub fn run_stage1_to_stage3_from_pass1( Ok((selection, selected, logical)) => (Ok(selection), Some(selected), Some(logical)), Err(error) => (Err(error), None, None), }; + let from_logical = selected.as_ref().map(|s| s.from_logical.clone()); let selected_logical = logical .as_ref() .and_then(|roots| selected_logical(roots, entries, data)); @@ -256,12 +258,16 @@ pub fn run_stage1_to_stage3_from_pass1( stage3_plan_selection: selection, stage3_selected: selected, stage3_selected_logical: selected_logical, - stage1_selected: logical.map(|roots| LogicalCandidate { - id: "L1".to_string(), - choice: Vec::new(), - roots: Ok(roots), - pass2: None, - }), + // The selected plan's logical candidate, under the id its + // physical candidates name (`from_logical`). + stage1_selected: logical + .zip(from_logical) + .map(|(roots, id)| LogicalCandidate { + id, + choice: Vec::new(), + roots: Ok(roots), + pass2: None, + }), }; } let components = super::components::sharing_components(&per_computation, entries, &models); diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index 9c200c032..9ceb6180f 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -1,6 +1,6 @@ //! Local prices of implementation candidates, cached, and the frozen -//! calibration factor (design "Memo-group search": Local price, Coupled -//! charges). Cost model 1 only: cost models 2 and 3 keep the current +//! calibration factor (planner-layering design, "The stages over a memo +//! of sub-expressions"). Cost model 1 only: cost models 2 and 3 keep the current //! search. //! //! A candidate's local price is Stage 3's per-node price diff --git a/crates/plan-selection/src/pass/subexpr/scans.rs b/crates/plan-selection/src/pass/subexpr/scans.rs index 2188adaa7..ce36fd191 100644 --- a/crates/plan-selection/src/pass/subexpr/scans.rs +++ b/crates/plan-selection/src/pass/subexpr/scans.rs @@ -1,5 +1,6 @@ //! Couplings through a source's rows: scan widening and the raw-data -//! horizon (joint-planning-search design, "Memo-group search", step 4). +//! horizon (planner-layering design, "The stages over a memo of +//! sub-expressions"). //! //! - **Scan widening** ([`add_scan_widening`]): the time-series scans of a //! metric that project different labels are separate groups (an unkeyed diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 29e5348dc..a02d199b9 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -1,8 +1,7 @@ -//! Clusters, coupling thresholds with branch and bound, the dynamic -//! program, build-late (design "Memo-group search", steps 3–7). +//! The search over the memo (planner-layering design, "The stages over a +//! memo of sub-expressions"). //! -//! [`brute_force`] is the reference the decomposed search is checked -//! against: every consistent choice of one candidate per needed group and +//! [`brute_force`] is the reference the MILP search is checked against: every consistent choice of one candidate per needed group and //! one threshold per coupling site, composed into a plan and priced by //! Stage 2 and Stage 3 in every materialization. diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index 8fcacfac8..060f0ec31 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -1,5 +1,5 @@ //! The window-composition, rollup and exact top-k reuse transformations -//! (joint-planning-search design, "Memo-group search", stage interfaces): +//! (planner-layering design, "The stages over a memo of sub-expressions"): //! Pass 2 rules that add implementation candidates to the memo's groups. //! //! - **Window composition** ([`add_window_composition`]): a group whose