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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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/67] 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 57fe7ebfa2f31ad7213eb5a41c40f9b6f86607c1 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:12:04 +0000 Subject: [PATCH 18/67] plan-selection: depend on good_lp 1.15.3 with the HiGHS backend (MIT) for the memo's MILP --- Cargo.lock | 157 +++++++++++++++++++++++++++---- crates/plan-selection/Cargo.toml | 2 + 2 files changed, 141 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 461791389..0d0aba83d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -448,6 +448,7 @@ dependencies = [ "asap-physical-asap-optimization", "asap-types", "asap_sketchlib", + "good_lp", "serde", "serde_json", "thiserror 2.0.18", @@ -570,6 +571,26 @@ dependencies = [ "serde", ] +[[package]] +name = "bindgen" +version = "0.72.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools 0.13.0", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn", +] + [[package]] name = "bitflags" version = "2.11.1" @@ -683,6 +704,15 @@ dependencies = [ "shlex", ] +[[package]] +name = "cexpr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -727,6 +757,26 @@ dependencies = [ "phf 0.12.1", ] +[[package]] +name = "clang-sys" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "157a8ba7b480713b56f4c09fd13fc3e0a22a5dfab8097ba61cbc5feef950788a" +dependencies = [ + "glob", + "libc", + "libloading", +] + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "comfy-table" version = "7.2.2" @@ -919,7 +969,7 @@ dependencies = [ "flate2", "futures", "indexmap", - "itertools", + "itertools 0.14.0", "liblzma", "log", "object_store", @@ -951,7 +1001,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools", + "itertools 0.14.0", "log", "object_store", "parking_lot", @@ -976,7 +1026,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools", + "itertools 0.14.0", "log", "object_store", ] @@ -995,7 +1045,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.14.0", "libc", "log", "object_store", @@ -1042,7 +1092,7 @@ dependencies = [ "flate2", "futures", "glob", - "itertools", + "itertools 0.14.0", "liblzma", "log", "object_store", @@ -1073,7 +1123,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools", + "itertools 0.14.0", "object_store", "tokio", ] @@ -1147,7 +1197,7 @@ dependencies = [ "datafusion-pruning", "datafusion-session", "futures", - "itertools", + "itertools 0.14.0", "log", "object_store", "parking_lot", @@ -1200,7 +1250,7 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr-common", "indexmap", - "itertools", + "itertools 0.14.0", "recursive", "serde_json", "sqlparser", @@ -1215,7 +1265,7 @@ dependencies = [ "arrow", "datafusion-common", "indexmap", - "itertools", + "itertools 0.14.0", ] [[package]] @@ -1239,7 +1289,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hex", - "itertools", + "itertools 0.14.0", "log", "md-5", "memchr", @@ -1302,7 +1352,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hashbrown 0.17.1", - "itertools", + "itertools 0.14.0", "itoa", "log", "memchr", @@ -1375,7 +1425,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools", + "itertools 0.14.0", "log", "recursive", "regex", @@ -1397,7 +1447,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.14.0", "parking_lot", "petgraph", "recursive", @@ -1416,7 +1466,7 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools", + "itertools 0.14.0", ] [[package]] @@ -1431,7 +1481,7 @@ dependencies = [ "datafusion-expr-common", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.14.0", "parking_lot", "pin-project", ] @@ -1451,7 +1501,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", - "itertools", + "itertools 0.14.0", "recursive", ] @@ -1480,7 +1530,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools", + "itertools 0.14.0", "log", "num-traits", "parking_lot", @@ -1835,6 +1885,16 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +[[package]] +name = "good_lp" +version = "1.15.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4255eb1d65eea322f73a326038f2bcec19dae13ca50fa557da19b03274b4d56" +dependencies = [ + "fnv", + "highs", +] + [[package]] name = "half" version = "2.7.1" @@ -1885,6 +1945,26 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "highs" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dfc791427e78fdc1a0a37983f79c20f86a0b7b7dfdb5d6a57a6f72d77cd9074" +dependencies = [ + "highs-sys", + "log", +] + +[[package]] +name = "highs-sys" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4da0c52322648e2360a5a8896d23301fa5ed99ed7607c5faf229bd0bbb30d3f" +dependencies = [ + "bindgen", + "cmake", +] + [[package]] name = "http" version = "1.5.0" @@ -2061,6 +2141,15 @@ version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -2179,6 +2268,16 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "liblzma" version = "0.4.8" @@ -2373,6 +2472,12 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -2383,6 +2488,16 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "num-bigint" version = "0.4.6" @@ -2459,7 +2574,7 @@ dependencies = [ "futures-util", "http", "humantime", - "itertools", + "itertools 0.14.0", "parking_lot", "percent-encoding", "thiserror 2.0.18", @@ -2879,6 +2994,12 @@ dependencies = [ "serde", ] +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + [[package]] name = "rustc_version" version = "0.4.1" diff --git a/crates/plan-selection/Cargo.toml b/crates/plan-selection/Cargo.toml index abb350b64..af7854e82 100644 --- a/crates/plan-selection/Cargo.toml +++ b/crates/plan-selection/Cargo.toml @@ -15,6 +15,8 @@ thiserror = "2" serde = { version = "1", features = ["derive"] } serde_json = "1" +good_lp = { version = "=1.15.3", default-features = false, features = ["highs"] } + [dev-dependencies] asap-logical-asap-optimization = { path = "../logical-asap-optimization", features = ["test-support"] } asap-frontend-promql = { path = "../frontend-promql" } From 0ba55f1adc7e17f6fcee10614b8e5377414974d8 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:18:38 +0000 Subject: [PATCH 19/67] memo: MILP model over memo choices (reach, built-once groups, union rates, latency, cuts) --- .../plan-selection/src/pass/subexpr/milp.rs | 620 ++++++++++++++++++ crates/plan-selection/src/pass/subexpr/mod.rs | 1 + 2 files changed, 621 insertions(+) create mode 100644 crates/plan-selection/src/pass/subexpr/milp.rs diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs new file mode 100644 index 000000000..7ce9a9764 --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -0,0 +1,620 @@ +//! The memo's solver: one MILP over the memo's choices (design "The +//! solver: MILP over the memo"). +//! +//! A choice is one candidate of one group with an option for each of its +//! private materialization slots ([`Choice`]); its coefficients come from +//! Stage 3's per-node pricing ([`ChoiceCoefficients`]). The model picks one +//! choice per needed group, a materialization for each shared group's own +//! node, and propagates which consumers (queries, materialized slots) reach +//! which group, so each private node is charged once per consumer reaching +//! it, at that consumer's rate. Terms the model omits (memory held while a +//! query runs, composed accuracy) are added back by the caller's lazy cuts: +//! the model's optimum is a lower bound on the plan's Stage 3 price. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use good_lp::solvers::highs::highs; +use good_lp::{ + constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, WithMipGap, +}; + +use super::GroupId; +use asap_types::physical::Materialization; + +/// How a child group of a choice is reached: by the consumers reaching the +/// choice's own node, or by one of its materialized private slots. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Via { + Direct, + Slot(usize), +} + +/// What one part of a choice costs, from Stage 3's per-node pricing. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub(crate) struct PartCoefficients { + /// Price per second per unit of the rate of a consumer reaching it. + pub cost_per_unit_rate: f64, + /// Work in ms per computation, which a query reaching it at query time + /// adds to its latency. + pub work_ms: f64, +} + +/// One materialized private slot of a choice. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct SlotCoefficients { + pub materialization: Materialization, + /// Price per second that does not depend on a rate (ingestion-time + /// work and kept state). + pub fixed_cost: f64, + /// Price per second per unit of the slot's own rate (query-time slots: + /// the union rate of the queries reaching the group). + pub cost_per_unit_rate: f64, + /// Work in ms per computation at query time. + pub work_ms: f64, +} + +/// One choice's coefficients. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ChoiceCoefficients { + pub direct: PartCoefficients, + pub slots: Vec, + /// Per child group of the candidate, in order. + pub child_via: Vec, +} + +/// One choice of a group: a candidate and an option per private slot. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Choice { + pub group: GroupId, + pub candidate: usize, + /// The groups it implements (the first is `group`). + pub implements: Vec, + /// The child groups it reads. + pub children: Vec, + pub assignment: Vec, + pub coefficients: ChoiceCoefficients, +} + +/// A shared group's own node, built once per evaluation when chosen. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct GroupSlot { + pub group: GroupId, + /// Whether it may stay `NotMaterialized` (recomputed per consumer). + pub may_recompute: bool, +} + +/// One query of the model. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ModelQuery { + /// Its evaluations per second. + pub rate: f64, + /// The groups its root reads directly (one for an operator root). + pub roots: Vec, + /// Its latency bound in ms, if any (the tighter of its own bound and + /// its repeat interval). + pub max_latency_ms: Option, +} + +/// Everything the model reads. +pub(crate) struct ModelInput { + pub group_count: usize, + /// Per group, the queries whose plan may reach it. + pub group_queries: Vec>, + pub choices: Vec, + pub group_slots: Vec, + pub queries: Vec, + /// Per query, its schedule class: queries of one class run at the same + /// instants (one recurrence and range evaluation). + pub schedule: Vec, + /// The union rate of the classes in each set, by sorted class set: how + /// often a node built once for queries of those classes is computed. + pub union_rate: HashMap, f64>, + /// Choice sets excluded by earlier cuts (no-good cuts). + pub excluded: Vec>, + /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all + /// of a choice set is chosen (local cuts). + pub extra_bounds: Vec<(Vec, f64)>, +} + +/// The model's optimum. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ModelSolution { + /// The chosen choices, by index in [`ModelInput::choices`]. + pub chosen: Vec, + /// The shared groups built once per evaluation. + pub built_once: BTreeSet, + /// The model's objective: a lower bound on the plan's price. + pub objective: f64, +} + +/// A consumer of the model: a query, a materialized private slot of a +/// choice, or a shared group built once. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +enum Consumer { + Query(usize), + Slot { choice: usize, slot: usize }, + Group(GroupId), +} + +/// Build and solve the model. `None` when it is infeasible. +pub(crate) fn solve(input: &ModelInput) -> Result, String> { + let mut vars = ProblemVariables::new(); + let n = input.group_count; + let x: Vec = input + .choices + .iter() + .map(|_| vars.add(variable().binary())) + .collect(); + let need: Vec = (0..n).map(|_| vars.add(variable().binary())).collect(); + let once: BTreeMap = input + .group_slots + .iter() + .map(|slot| (slot.group, vars.add(variable().binary()))) + .collect(); + // Which queries' plans reach each group at all (any materialization): + // the union rate of a node built once reads them. + let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); + for (g, queries) in input.group_queries.iter().enumerate() { + for &q in queries { + reaches_at_all.insert((q, g), vars.add(variable().binary())); + } + } + // Consumers that may reach each group's own node, and through it its + // direct part. + let consumers = possible_consumers(input); + let mut reach: HashMap<(Consumer, usize), Variable> = HashMap::new(); + for (&g, ks) in &consumers { + for &k in ks { + reach.insert((k, g), vars.add(variable().min(0).max(1))); + } + } + // A consumer passes through a group into its chosen direct part when + // it reaches the group and the group is not built once (or it is the + // group's own built-once consumer). + let mut passes: HashMap<(Consumer, usize), Variable> = HashMap::new(); + for (&(k, g), _) in &reach { + passes.insert((k, g), vars.add(variable().min(0).max(1))); + } + for (&g, _) in &once { + let k = Consumer::Group(g); + passes + .entry((k, g.0 as usize)) + .or_insert_with(|| vars.add(variable().min(0).max(1))); + } + // y[k, c]: consumer k passes into choice c's direct part. + let mut y: HashMap<(Consumer, usize), Variable> = HashMap::new(); + for (c, choice) in input.choices.iter().enumerate() { + let g = choice.group.0 as usize; + for (&(k, h), _) in &passes { + if h == g { + y.insert((k, c), vars.add(variable().min(0).max(1))); + } + } + } + // Union-rate selectors for built-once consumers (shared groups and + // query-time private slots): one per set of schedule classes among the + // group's queries, exactly the classes of the queries reaching it. + let mut rate_of: HashMap, Variable)>> = HashMap::new(); + let classes_of = |g: usize| -> Vec { + let mut classes: Vec = input.group_queries[g] + .iter() + .map(|&q| input.schedule[q]) + .collect(); + classes.sort_unstable(); + classes.dedup(); + classes + }; + let class_sets = |g: usize| -> Vec> { + let classes = classes_of(g); + (1u64..(1 << classes.len())) + .map(|mask| { + (0..classes.len()) + .filter(|i| mask >> i & 1 == 1) + .map(|i| classes[i]) + .collect() + }) + .collect() + }; + let mut groups_with_rates: BTreeSet = once.keys().map(|g| g.0 as usize).collect(); + for choice in &input.choices { + if choice + .coefficients + .slots + .iter() + .any(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) + { + groups_with_rates.insert(choice.group.0 as usize); + } + } + for &g in &groups_with_rates { + let sets = class_sets(g) + .into_iter() + .map(|set| (set, vars.add(variable().binary()))) + .collect(); + rate_of.insert(g, sets); + } + + // The objective. + let mut objective = Expression::from(0.0); + // A consumer's constant rate, or `None` for one built once at query + // time, whose rate is its group's union rate. + let rate = |k: Consumer| -> Option { + match k { + Consumer::Query(q) => Some(input.queries[q].rate), + Consumer::Slot { choice, slot } => { + let m = input.choices[choice].coefficients.slots[slot].materialization; + matches!(m, Materialization::MatAtIngestion(_)).then_some(1.0) + } + Consumer::Group(_) => None, + } + }; + // Products of a variable rate (a set selector) and a reach variable. + let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); + for (&(k, c), &v) in &y { + let cost = input.choices[c].coefficients.direct.cost_per_unit_rate; + if cost == 0.0 { + continue; + } + match rate(k) { + Some(r) => objective += r * cost * v, + None => { + let g = match k { + Consumer::Group(g) => g.0 as usize, + Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, + Consumer::Query(_) => unreachable!("a query's rate is constant"), + }; + for (set, selector) in &rate_of[&g] { + rated.push((*selector, v, input.union_rate[set] * cost)); + } + } + } + } + for (c, choice) in input.choices.iter().enumerate() { + for slot in &choice.coefficients.slots { + objective += slot.fixed_cost * x[c]; + if slot.cost_per_unit_rate != 0.0 + && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) + { + for (set, selector) in &rate_of[&(choice.group.0 as usize)] { + rated.push(( + *selector, + x[c], + input.union_rate[set] * slot.cost_per_unit_rate, + )); + } + } + } + } + let mut products = Vec::new(); + for &(a, b, coefficient) in &rated { + let z = vars.add(variable().min(0).max(1)); + objective += coefficient * z; + products.push((z, a, b)); + } + // Lower bounds on the omitted terms from earlier cuts. + let extra: Vec = input + .extra_bounds + .iter() + .map(|_| vars.add(variable().min(0))) + .collect(); + for &e in &extra { + objective += e; + } + + let mut model = vars.minimise(objective.clone()).using(highs); + // Needed groups: every query's roots; children of chosen choices. + for query in &input.queries { + for root in &query.roots { + model = model.with(constraint!(need[root.0 as usize] == 1)); + } + } + for g in 0..n { + // One choice implementing it (its own, or a fused one covering + // it) when needed, none otherwise. + let mut implementing = Expression::from(0.0); + for (c, choice) in input.choices.iter().enumerate() { + if choice.implements.iter().any(|i| i.0 as usize == g) { + implementing += x[c]; + } + } + model = model.with(constraint!(implementing == need[g])); + } + for (c, choice) in input.choices.iter().enumerate() { + for child in &choice.children { + model = model.with(constraint!(need[child.0 as usize] >= x[c])); + } + } + for (&g, &v) in &once { + model = model.with(constraint!(v <= need[g.0 as usize])); + let slot = input.group_slots.iter().find(|s| s.group == g).unwrap(); + if !slot.may_recompute { + model = model.with(constraint!(v == need[g.0 as usize])); + } + } + // Reach at all: a query reaches its roots, and every child of a chosen + // choice of a group it reaches. + for (q, query) in input.queries.iter().enumerate() { + for root in &query.roots { + if let Some(&v) = reaches_at_all.get(&(q, root.0 as usize)) { + model = model.with(constraint!(v == 1)); + } + } + } + for (c, choice) in input.choices.iter().enumerate() { + let g = choice.group.0 as usize; + for child in &choice.children { + for &q in &input.group_queries[g] { + if let (Some(&above), Some(&below)) = ( + reaches_at_all.get(&(q, g)), + reaches_at_all.get(&(q, child.0 as usize)), + ) { + model = model.with(constraint!(below >= above + x[c] - 1)); + } + } + } + } + // Rate selectors: exactly the classes of the queries reaching the + // group (at most one selector; a class is in it iff one of its + // queries reaches the group). + for (&g, sets) in &rate_of { + let mut any = Expression::from(0.0); + for (_, selector) in sets { + any += *selector; + } + model = model.with(constraint!(any <= 1)); + for class in classes_of(g) { + let mut containing = Expression::from(0.0); + for (set, selector) in sets { + if set.contains(&class) { + containing += *selector; + } + } + for &q in &input.group_queries[g] { + if input.schedule[q] == class { + model = model.with(constraint!(containing.clone() >= reaches_at_all[&(q, g)])); + } + } + let mut members = Expression::from(0.0); + for &q in &input.group_queries[g] { + if input.schedule[q] == class { + members += reaches_at_all[&(q, g)]; + } + } + model = model.with(constraint!(containing <= members)); + } + } + // Consumer reach. + for (q, query) in input.queries.iter().enumerate() { + for root in &query.roots { + if let Some(&v) = reach.get(&(Consumer::Query(q), root.0 as usize)) { + model = model.with(constraint!(v >= need[root.0 as usize])); + } + } + } + for (&(k, g), &p) in &passes { + match k { + Consumer::Group(own) if own.0 as usize == g => { + model = model.with(constraint!(p >= once[&own])); + } + _ => { + let r = reach[&(k, g)]; + match once.get(&GroupId(g as u32)) { + Some(&o) => model = model.with(constraint!(p >= r - o)), + None => model = model.with(constraint!(p >= r)), + } + } + } + } + for (&(k, c), &v) in &y { + let g = input.choices[c].group.0 as usize; + model = model.with(constraint!(v >= passes[&(k, g)] + x[c] - 1)); + } + for (c, choice) in input.choices.iter().enumerate() { + for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { + let h = child.0 as usize; + match via { + Via::Direct => { + for (&(k, g), _) in &passes { + if g != choice.group.0 as usize { + continue; + } + if let Some(&r) = reach.get(&(k, h)) { + model = model.with(constraint!(r >= y[&(k, c)])); + } + } + } + Via::Slot(slot) => { + let k = Consumer::Slot { + choice: c, + slot: *slot, + }; + if let Some(&r) = reach.get(&(k, h)) { + model = model.with(constraint!(r >= x[c])); + } + } + } + } + } + for &(z, a, b) in &products { + model = model.with(constraint!(z >= a + b - 1)); + } + // Latency bounds: the work each query reaches at query time. + for (q, query) in input.queries.iter().enumerate() { + let Some(bound) = query.max_latency_ms else { + continue; + }; + let mut latency = Expression::from(0.0); + for (&(k, c), &v) in &y { + if k == Consumer::Query(q) { + latency += input.choices[c].coefficients.direct.work_ms * v; + } + } + model = model.with(constraint!(latency <= bound)); + } + // Cuts. + for set in &input.excluded { + let mut chosen = Expression::from(0.0); + for &c in set { + chosen += x[c]; + } + model = model.with(constraint!(chosen <= set.len() as f64 - 1.0)); + } + for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { + let mut chosen = Expression::from(0.0); + for &c in set { + chosen += x[c]; + } + model = model.with(constraint!( + e >= *bound * (chosen - (set.len() as f64 - 1.0)) + )); + } + let model = model + .with_mip_gap(0.0) + .map_err(|e| format!("MIP gap: {e:?}"))?; + let solution = match model.solve() { + Ok(solution) => solution, + Err(good_lp::ResolutionError::Infeasible) => return Ok(None), + Err(error) => return Err(error.to_string()), + }; + let chosen = x + .iter() + .enumerate() + .filter(|(_, v)| solution.value(**v) > 0.5) + .map(|(c, _)| c) + .collect(); + let built_once = once + .iter() + .filter(|(_, v)| solution.value(**v) > 0.5) + .map(|(g, _)| *g) + .collect(); + Ok(Some(ModelSolution { + chosen, + built_once, + objective: solution.eval(&objective), + })) +} + +/// For each group, the consumers that may reach its node: its queries +/// (that reach it without passing a materialized node), the materialized +/// private slots of choices of groups above it, and the shared groups +/// above it (or itself) that may be built once. +fn possible_consumers(input: &ModelInput) -> BTreeMap> { + let mut consumers: BTreeMap> = BTreeMap::new(); + for (g, queries) in input.group_queries.iter().enumerate() { + let entry = consumers.entry(g).or_default(); + for &q in queries { + entry.insert(Consumer::Query(q)); + } + } + // Propagate downwards until nothing changes. + let shared: BTreeSet = input + .group_slots + .iter() + .map(|s| s.group.0 as usize) + .collect(); + loop { + let mut changed = false; + for (c, choice) in input.choices.iter().enumerate() { + let g = choice.group.0 as usize; + let mut above: BTreeSet = consumers.get(&g).cloned().unwrap_or_default(); + if shared.contains(&g) { + above.insert(Consumer::Group(GroupId(g as u32))); + } + for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { + let entry = consumers.entry(child.0 as usize).or_default(); + let before = entry.len(); + match via { + Via::Direct => entry.extend(above.iter().copied()), + Via::Slot(slot) => { + entry.insert(Consumer::Slot { + choice: c, + slot: *slot, + }); + } + } + changed |= entry.len() != before; + } + } + if !changed { + break; + } + } + consumers +} + +#[cfg(test)] +mod tests { + use super::*; + + fn choice(group: u32, children: &[u32], cost: f64) -> Choice { + Choice { + group: GroupId(group), + candidate: 0, + implements: vec![GroupId(group)], + children: children.iter().map(|&c| GroupId(c)).collect(), + assignment: Vec::new(), + coefficients: ChoiceCoefficients { + direct: PartCoefficients { + cost_per_unit_rate: cost, + work_ms: 1.0, + }, + slots: Vec::new(), + child_via: children.iter().map(|_| Via::Direct).collect(), + }, + } + } + + /// A scan two queries read costs once per query recomputing it, or once + /// per evaluation when built once: the model builds it once when that + /// is cheaper, and its objective is the plan's price. + #[test] + fn a_shared_group_is_built_once_when_that_is_cheaper() { + let input = ModelInput { + group_count: 3, + group_queries: vec![vec![0, 1], vec![0], vec![1]], + choices: vec![ + choice(0, &[], 10.0), + choice(1, &[0], 1.0), + choice(2, &[0], 1.0), + ], + group_slots: vec![GroupSlot { + group: GroupId(0), + may_recompute: true, + }], + queries: vec![ + ModelQuery { + rate: 1.0, + roots: vec![GroupId(1)], + max_latency_ms: None, + }, + ModelQuery { + rate: 1.0, + roots: vec![GroupId(2)], + max_latency_ms: None, + }, + ], + schedule: vec![0, 0], + union_rate: [(vec![0], 1.0)].into(), + excluded: Vec::new(), + extra_bounds: Vec::new(), + }; + let solution = solve(&input).unwrap().unwrap(); + assert_eq!(solution.built_once, [GroupId(0)].into()); + assert!( + (solution.objective - 12.0).abs() < 1e-9, + "{}", + solution.objective + ); + // Recomputing is cheaper when the scan is cheap. + let mut cheap = input; + cheap.choices[0].coefficients.direct.cost_per_unit_rate = 0.1; + cheap.union_rate.insert(vec![0], 1.0); + // Built once still costs its union rate × 0.1 = 0.1 < 0.2. + let solution = solve(&cheap).unwrap().unwrap(); + assert!( + (solution.objective - 2.1).abs() < 1e-9, + "{}", + solution.objective + ); + } +} diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 61dcea355..6a82c2fe8 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -24,6 +24,7 @@ pub(crate) mod capability; pub(crate) mod groups; +pub(crate) mod milp; pub(crate) mod pricing; pub(crate) mod scans; pub(crate) mod search; 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 20/67] 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 21/67] 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 22/67] 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 23/67] 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 24/67] 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 5e553b8bc0c0077fcf7b1e2b7b1b9148e2ca2d03 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:42:27 +0000 Subject: [PATCH 25/67] plan-selection: memo choices and their MILP coefficients from Stage 3's per-node pricing coefficients::choices enumerates each group's choices: a candidate at each threshold of its coupling sites (chain, labels, shared size, raw horizon), with each Stage 2 assignment of its private builds. Its sub-DAG is measured per computation (price_node at rate 1) and split into the direct part and one part per materialized slot, as computations() counts consumers; PartCoefficients and SlotCoefficients come from cpu, scan, kept state and work. LazyTerms carries held memory, raw scans, the direct part priced at ingestion time, the state of the group's node built once, and the thresholds. Choice.children adds the groups a candidate reads without listing them (a pane chain over its window's scan). Tests: for every brute-force plan of one-query plans, a shared group (not materialized, built once, built at ingestion) and a shared KLL build, the coefficients times computations()' rates plus the lazy terms equal price_at_scale within 1e-9. A shared pane chain is ignored: the scan's look-back is not a raw-horizon threshold when a chain reads it. --- .../src/pass/subexpr/coefficients.rs | 1167 +++++++++++++++++ crates/plan-selection/src/pass/subexpr/mod.rs | 1 + .../src/pass/subexpr/pricing.rs | 17 +- .../plan-selection/src/pass/subexpr/scans.rs | 7 +- .../plan-selection/src/pass/subexpr/search.rs | 15 +- 5 files changed, 1194 insertions(+), 13 deletions(-) create mode 100644 crates/plan-selection/src/pass/subexpr/coefficients.rs diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs new file mode 100644 index 000000000..870d328cb --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -0,0 +1,1167 @@ +//! Each group's choices and their coefficients for the MILP ([`super::milp`]), +//! from Stage 3's per-node pricing. +//! +//! A choice is one candidate of a group with an option for each of its +//! private summary builds (its slots), as Stage 2 offers them for the +//! group's queries, down-closure applied. Its candidate's sub-DAG is built +//! by Stage 2 once per assignment, one root per query of the group, and +//! every node is measured for one computation ([`crate::price_node`] at +//! rate `(1, 1.0)`). Its private nodes (those not in a child group's own +//! sub-DAG, by structure) are split into parts by who computes them, as +//! Stage 3's `computations` counts consumers: +//! +//! - the direct part: reached from the candidate's node without passing a +//! materialized slot, computed once per consumer reaching the group; +//! - one part per materialized slot: the slot and what it reaches down to +//! the next materialized slot or a child, computed by the slot. +//! +//! Where each field comes from (all of [`crate::price_node`]'s outputs for +//! one computation, cost model 1's [`crate::ResourceUnitCost`]): +//! +//! | Field | Source | +//! | --- | --- | +//! | `cost_per_unit_rate` | `cpu_ops_per_second · cost_per_cpu_op + scanned_bytes_per_second · cost_per_scanned_byte` | +//! | `fixed_cost` | at ingestion: the per-computation price at rate 1, plus kept state; at query time: kept state only. Kept state is `average_memory_bytes · cost_per_memory_byte_second + disk_bytes · cost_per_disk_byte_second`, a pane chain's from `retained_panes` over `pane_retention_ms` with the group's queries' range spans | +//! | `work_ms` | `work_ms` of the part's query-time nodes (`query_latencies` adds each reached node once) | +//! | [`LazyTerms::held`] | `held`: bytes held while an evaluation runs, charged by Stage 3 as `bytes · per_second · E / 1000 · cost_per_memory_byte_second`, `E` the readers' rate-weighted mean latency (per reader) or longest (built once) | +//! | [`LazyTerms::raw_scans`] | `raw_scan`: Stage 3 keeps each source's longest query-time scan on disk | +//! +//! Limits: a child's output statistics are its group's own node's, as the +//! candidate's sub-DAG reads it, whatever the child's chosen +//! implementation; a query's request overhead (a measured table's +//! `ns_per_evaluation`) and an ingestion state's compaction and partial +//! copies with more than one ingest worker are not coefficients; a kept +//! chain's slot is charged per unit of its group's union rate, while +//! Stage 3 computes it once per request (range steps not counted). + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::rc::Rc; + +use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate}; +use asap_types::ir::physical_export::{PhysicalASAPDAG, PhysicalASAPNodeId}; +use asap_types::ir::properties::ExecutionTiming; +use asap_types::ir::{ASAPOp, NonASAPOp, Operator, QueryRoot}; +use asap_types::physical::{Materialization, Medium}; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; + +use super::milp::{Choice, ChoiceCoefficients, PartCoefficients, SlotCoefficients, Via}; +use super::pricing::{child_keys, keys_of, scan_lookback_ms, PriceCache}; +use super::{CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo}; +use crate::{ + kept_pane_chains, maintained_window_ms, node_map, pane_retention_ms, price_node, query_reach, + retained_panes, source_name, windows_read, CostModel, PlanningModels, PricedNode, + ProducerOutput, SelectionError, DEFAULT_LOOKBACK_MS, +}; + +/// The terms of a choice the MILP adds lazily, once its plan's latencies +/// and every source's scans are known. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct LazyTerms { + /// Memory each part holds while an evaluation runs. + pub held: Vec, + /// Each query-time raw scan of a part: source, bytes, part. + pub raw_scans: Vec<(String, u64, Via)>, + /// The direct part computed at ingestion time, as it is under a + /// parent's ingestion-time slot (scans read one second of rows); + /// `None` when it cannot run at ingestion time (a grouping with no + /// window below it). + pub direct_at_ingestion: Option, + /// The threshold of each coupling site the choice's candidate reads it + /// at (index among the site's thresholds); a raw-horizon site's is + /// Stage 3's own. + pub thresholds: Vec<(CouplingSite, usize)>, + /// When the group's own node is built once at query time (in memory): + /// the state it holds while its longest reader runs, charged in place + /// of its [`HeldTerm::own_node`] working memory. + pub built_once_state_bytes: u64, +} + +/// Bytes a part holds while an evaluation runs: Stage 3 charges +/// `bytes · rate · E / 1000` byte-seconds per second, `rate` the part's +/// consumers' rate and `E` the queries reaching it: their rate-weighted +/// mean latency when each computes its own copy (`per_reader`), else the +/// longest. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct HeldTerm { + pub via: Via, + pub bytes: u64, + pub per_reader: bool, + /// Whether it is the group's own node (replaced when it is built once). + pub own_node: bool, +} + +/// Every choice of `group` with its lazy terms, in candidate order and, +/// within a candidate, in Stage 2's assignment order. +pub(crate) fn choices( + memo: &Memo, + group: GroupId, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + cache: &mut PriceCache, +) -> Result<(Vec, Vec), SelectionError> { + let parts = choice_parts(memo, group, entries, data, models, cache)?; + Ok(parts + .into_iter() + .map(|part| (part.choice, part.lazy)) + .unzip()) +} + +/// The threshold of each coupling site a candidate is read at (index among +/// the site's thresholds). +type SiteThresholds = Vec<(CouplingSite, usize)>; + +/// A choice with the shallow keys ([`shallow_keys`]) of its nodes, by +/// which a composed plan's nodes are found (tests). +pub(super) struct ChoiceParts { + pub choice: Choice, + pub lazy: LazyTerms, + /// The shallow key of the candidate's node. + pub root_key: u32, + /// The shallow keys of the candidate node's private producers. + pub root_producer_keys: Vec, + /// Per slot of the coefficients, the shallow key of its build. + pub slot_keys: Vec, + /// The shallow keys of the private builds, in assignment order. + pub build_keys: Vec, + /// Per part (direct, then each slot), its nodes' keys. + pub part_keys: Vec>, +} + +pub(super) fn choice_parts( + memo: &Memo, + group: GroupId, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + cache: &mut PriceCache, +) -> Result, SelectionError> { + let queries = &memo.groups[group.0 as usize].queries; + let group_entries: Vec = + queries.iter().map(|&q| entries[q].clone()).collect(); + let mut out = Vec::new(); + let versions: Vec<(usize, ImplementationCandidate, SiteThresholds)> = memo.candidates + [group.0 as usize] + .iter() + .enumerate() + .flat_map(|(index, candidate)| { + at_site_thresholds(memo, candidate) + .into_iter() + .map(move |(version, thresholds)| (index, version, thresholds)) + }) + .collect(); + for (index, candidate, thresholds) in &versions { + let (index, candidate) = (*index, candidate); + let roots = vec![QueryRoot::Operator(Rc::clone(&candidate.node)); queries.len().max(1)]; + let Ok(physicals) = optimize_physical_asap( + "S", + &roots, + &group_entries, + data, + models.deployment_capabilities, + ) else { + continue; + }; + let child_groups = children_of(memo, group, candidate); + + let mut children: HashSet = HashSet::new(); + let mut child_roots: Vec> = Vec::new(); + for &child in &child_groups { + let keys = child_keys(memo, child, entries, data, &models, cache); + children.extend(keys.nodes.iter().copied()); + child_roots.push(keys.root); + } + let mut seen: Vec> = Vec::new(); + for physical in &physicals { + let keys = keys_of(&physical.dag, cache.keys_mut()); + let shallow = shallow_keys(&physical.dag, cache.keys_mut()); + let private: Vec = keys.iter().map(|k| !children.contains(k)).collect(); + let builds: Vec = (0..physical.dag.nodes.len()) + .filter(|&i| private[i] && is_build(&physical.dag.nodes[i].payload)) + .collect(); + // Only the private builds are the choice's; any other + // materialized node (the group's own node built once, a + // child's build) is another variable of the model. + let foreign = physical.dag.nodes.iter().enumerate().any(|(i, n)| { + n.materialization != Materialization::NotMaterialized && !builds.contains(&i) + }); + if foreign { + continue; + } + let assignment: Vec = builds + .iter() + .map(|&i| physical.dag.nodes[i].materialization) + .collect(); + if seen.contains(&assignment) { + continue; + } + seen.push(assignment.clone()); + let Some(parts) = coefficients_of( + memo, + group, + index, + candidate, + physical, + &keys, + &shallow, + &private, + &child_roots, + &horizons_of(memo, thresholds), + &group_entries, + data, + &models, + )? + else { + continue; + }; + out.push(ChoiceParts { + choice: Choice { + group, + candidate: index, + implements: candidate.implements.clone(), + children: child_groups.clone(), + assignment, + coefficients: parts.coefficients, + }, + lazy: LazyTerms { + thresholds: thresholds.clone(), + ..parts.lazy + }, + root_key: parts.root_key, + root_producer_keys: parts.root_producer_keys, + slot_keys: parts.slot_keys, + build_keys: builds.iter().map(|&i| shallow[i]).collect(), + part_keys: parts.part_keys, + }); + } + } + Ok(out) +} + +/// `candidate` rewritten at each combination of thresholds of its coupling +/// sites that it may read, as the brute-force search rewrites a plan's +/// candidates, with the index of each threshold. At a raw-horizon +/// threshold the candidate is unchanged and its query-time scans of the +/// site's source read the threshold's horizon ([`horizons_of`]). A +/// candidate without sites is itself. +fn at_site_thresholds( + memo: &Memo, + candidate: &ImplementationCandidate, +) -> Vec<(ImplementationCandidate, Vec<(CouplingSite, usize)>)> { + use super::{capability, scans, windows}; + let mut versions = vec![(candidate.clone(), Vec::new())]; + for site in &candidate.coupling_sites { + let mut next = Vec::new(); + for (version, chosen) in versions { + let at = |index: usize, rewritten: Option| { + rewritten.map(|rewritten| { + let mut chosen = chosen.clone(); + chosen.push((*site, index)); + (rewritten, chosen) + }) + }; + match site.kind { + CouplingKind::PaneChain => { + for (index, threshold) in + windows::chain_thresholds(memo, site).iter().enumerate() + { + let readers_ok = windows::chain_readers(&version.node) + .iter() + .filter(|r| r.key == site.key) + .all(|r| windows::may_read(threshold, r)); + let rewritten = readers_ok + .then(|| windows::read_at_threshold(&version, site, threshold).ok()) + .flatten() + .map(|node| ImplementationCandidate { + node, + ..version.clone() + }); + next.extend(at(index, rewritten)); + } + } + CouplingKind::ScanLabels => { + for (index, threshold) in scans::label_thresholds(memo, site).iter().enumerate() + { + let rewritten = + scans::widened_at_threshold(memo, &version, site, threshold) + .ok() + .flatten(); + next.extend(at(index, rewritten)); + } + } + CouplingKind::SharedBuildSize => { + for (index, threshold) in + capability::size_thresholds(memo, site).iter().enumerate() + { + let rewritten = + capability::resized_at_threshold(memo, &version, site, threshold) + .ok() + .flatten() + .map(|node| ImplementationCandidate { + node, + ..version.clone() + }); + next.extend(at(index, rewritten)); + } + } + CouplingKind::RawHorizon => { + let own: Vec = scans::raw_readers(memo, &version) + .into_iter() + .filter(|r| r.key == site.key) + .collect(); + for (index, threshold) in + scans::horizon_thresholds(memo, site).iter().enumerate() + { + if own.iter().all(|r| scans::may_read_horizon(threshold, r)) { + next.extend(at(index, Some(version.clone()))); + } + } + } + } + } + versions = next; + } + versions +} + +/// The raw-data horizon of each source key `thresholds` fixes, in ms. +fn horizons_of(memo: &Memo, thresholds: &[(CouplingSite, usize)]) -> HashMap { + thresholds + .iter() + .filter(|(site, _)| site.kind == CouplingKind::RawHorizon) + .filter_map(|(site, index)| { + super::scans::horizon_thresholds(memo, site) + .get(*index) + .map(|t| (site.key, t.horizon_ms)) + }) + .collect() +} + +/// The child groups of `candidate`: its own, and every other group whose +/// node equals a node its node reaches and that no other child already +/// reaches, which some Pass 2 candidates read without listing (a pane +/// chain over its window's scan). Their nodes are not the candidate's +/// (the composed plan makes equal nodes one), and a plan choosing it +/// needs them. +fn children_of(memo: &Memo, group: GroupId, candidate: &ImplementationCandidate) -> Vec { + use asap_types::ir::OperatorNode; + let mut children = candidate.children.clone(); + let reachable = OperatorNode::reachable(&candidate.node); + let reaches = |nodes: &[Rc], g: GroupId| { + let node = &memo.groups[g.0 as usize].node; + nodes.iter().any(|n| { + Rc::ptr_eq(n, node) || (n.operator == node.operator && n.schema == node.schema) + }) + }; + let below = |g: GroupId, children: &[GroupId]| { + children + .iter() + .any(|&c| reaches(&OperatorNode::reachable(&memo.groups[c.0 as usize].node), g)) + }; + // Parents before children, so the highest group read is taken first. + for h in (0..memo.groups.len()).rev() { + let h = GroupId(h as u32); + if h == group || candidate.implements.contains(&h) || children.contains(&h) { + continue; + } + if !Rc::ptr_eq(&memo.groups[h.0 as usize].node, &candidate.node) + && memo.groups[h.0 as usize].node.operator != candidate.node.operator + && reaches(&reachable, h) + && !below(h, &children) + { + children.push(h); + } + } + children +} + +fn is_build(payload: &Operator) -> bool { + matches!( + payload, + Operator::ASAP(ASAPOp::SummaryAgg { .. } | ASAPOp::PaneChain { .. }) + ) +} + +struct Coefficients { + coefficients: ChoiceCoefficients, + lazy: LazyTerms, + root_key: u32, + root_producer_keys: Vec, + slot_keys: Vec, + part_keys: Vec>, +} + +/// Every node of `dag` measured for one computation, children first. +fn per_computation<'a>( + dag: &'a PhysicalASAPDAG, + entries: &[QueryWorkloadEntry], + scan_lookback: Option, + horizons: &HashMap, + data: &DataWorkload, + models: &'a PlanningModels<'a>, +) -> Result>, SelectionError> { + let nodes = node_map(dag); + let kept = kept_pane_chains(dag); + let reach = query_reach(dag, &nodes); + let ranges: Vec<_> = entries.iter().map(|e| e.range_evaluation).collect(); + let retention = pane_retention_ms(dag, &reach, &ranges); + 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 = &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) { + (Operator::NonASAP(NonASAPOp::Scan { source, .. }), ExecutionTiming::QueryTime) => { + match horizons.get(&super::scans::raw_horizon_key(source)) { + Some(&horizon) => horizon, + None => windows_read(dag, &nodes, node.id, &kept) + .into_iter() + .max() + .or(scan_lookback) + .unwrap_or(DEFAULT_LOOKBACK_MS), + } + } + _ => 0, + }; + let node_priced = price_node( + node, + &views, + (1, 1.0), + lookback_ms, + retained_panes(node, &retention), + (node.output_state.timing == ExecutionTiming::IngestionTime) + .then(|| maintained_window_ms(node.id, &nodes, &producers_of)) + .flatten(), + data, + models.plan_selection, + 1.0, + ) + .map_err(|error| SelectionError::Cost { + candidate: "memo choice".into(), + node: node.id, + error, + })?; + priced.insert(node.id, node_priced); + } + Ok(priced) +} + +/// Each node's shallow key, by position: its payload with its producers +/// left out, and its output schema. Unlike [`keys_of`] it does not depend +/// on the implementations below a node. +pub(super) fn shallow_keys(dag: &PhysicalASAPDAG, interned: &mut HashMap) -> Vec { + dag.nodes + .iter() + .map(|node| { + let payload = node.payload.map_children(|_| ()); + let text = serde_json::to_string(&("shallow", &payload, &node.output_schema)) + .expect("a physical node serializes"); + let next = interned.len() as u32; + *interned.entry(text).or_insert(next) + }) + .collect() +} + +#[allow(clippy::too_many_arguments)] +fn coefficients_of( + memo: &Memo, + group: GroupId, + _index: usize, + candidate: &ImplementationCandidate, + physical: &PhysicalCandidate, + keys: &[u32], + shallow: &[u32], + private: &[bool], + child_roots: &[Option], + horizons: &HashMap, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, +) -> Result, SelectionError> { + let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { + return Ok(None); + }; + let dag = &physical.dag; + let Some(&root) = dag.roots.first() else { + return Ok(None); + }; + let position: HashMap = dag + .nodes + .iter() + .enumerate() + .map(|(i, n)| (n.id, i)) + .collect(); + if !private[position[&root]] { + return Ok(None); + } + let scan_lookback = scan_lookback_ms(memo, group); + let priced = per_computation(dag, entries, scan_lookback, horizons, data, models)?; + let mut producers_of: HashMap> = HashMap::new(); + for e in &dag.edges { + producers_of.entry(e.consumer).or_default().push(e.producer); + } + let is_private = |id: PhysicalASAPNodeId| private[position[&id]]; + let materialized = |id: PhysicalASAPNodeId| { + dag.nodes[position[&id]].materialization != Materialization::NotMaterialized + }; + // A part's nodes: from `start`, private nodes, not past a materialized + // one (other than `start`); and the non-private nodes it reads. + let walk = |start: PhysicalASAPNodeId| { + let mut part = Vec::new(); + let mut reads = Vec::new(); + let mut seen = HashSet::new(); + let mut stack = vec![start]; + while let Some(id) = stack.pop() { + if !seen.insert(id) { + continue; + } + if !is_private(id) { + reads.push(id); + continue; + } + if id != start && materialized(id) { + continue; + } + part.push(id); + stack.extend(producers_of.get(&id).into_iter().flatten().copied()); + } + (part, reads) + }; + let per_rate = |id: PhysicalASAPNodeId| { + let p = &priced[&id]; + p.cpu_ops_per_second * prices.cost_per_cpu_op + + p.scanned_bytes_per_second * prices.cost_per_scanned_byte + }; + let kept_state = |id: PhysicalASAPNodeId| { + let p = &priced[&id]; + p.average_memory_bytes * prices.cost_per_memory_byte_second + + p.disk_bytes as f64 * prices.cost_per_disk_byte_second + }; + let query_time = |id: PhysicalASAPNodeId| { + dag.nodes[position[&id]].output_state.timing == ExecutionTiming::QueryTime + }; + let work = |part: &[PhysicalASAPNodeId]| -> f64 { + part.iter() + .filter(|&&id| query_time(id)) + .map(|id| priced[id].work_ms) + .sum() + }; + let mut held = Vec::new(); + let mut raw_scans = Vec::new(); + let mut lazy_of = |part: &[PhysicalASAPNodeId], via: Via| { + for &id in part { + let p = &priced[&id]; + if let Some((bytes, _, per_reader)) = p.held { + held.push(HeldTerm { + via, + bytes, + per_reader, + own_node: id == root, + }); + } + if let Some((source, bytes)) = p.raw_scan { + raw_scans.push((source_name(source).to_string(), bytes, via)); + } + } + }; + let (direct, direct_reads) = walk(root); + lazy_of(&direct, Via::Direct); + let direct_coefficients = PartCoefficients { + cost_per_unit_rate: direct.iter().map(|&id| per_rate(id) + kept_state(id)).sum(), + work_ms: work(&direct), + }; + let mut part_keys = vec![direct.iter().map(|id| keys[position[id]]).collect()]; + let mut slots = Vec::new(); + let mut slot_keys = Vec::new(); + let mut slot_reads = Vec::new(); + for node in &dag.nodes { + if !is_private(node.id) || node.materialization == Materialization::NotMaterialized { + continue; + } + let (part, reads) = walk(node.id); + let via = Via::Slot(slots.len()); + lazy_of(&part, via); + let at_ingestion = matches!(node.materialization, Materialization::MatAtIngestion(_)); + let rate_cost: f64 = part.iter().map(|&id| per_rate(id)).sum(); + let state: f64 = part.iter().map(|&id| kept_state(id)).sum(); + slots.push(SlotCoefficients { + materialization: node.materialization, + fixed_cost: match at_ingestion { + true => rate_cost + state, + false => state, + }, + cost_per_unit_rate: match at_ingestion { + true => 0.0, + false => rate_cost, + }, + work_ms: work(&part), + }); + slot_keys.push(shallow[position[&node.id]]); + part_keys.push(part.iter().map(|id| keys[position[id]]).collect()); + slot_reads.push(reads); + } + let child_via: Vec = child_roots + .iter() + .map(|child| { + let reads_child = |reads: &[PhysicalASAPNodeId]| { + reads.iter().any(|id| Some(keys[position[id]]) == *child) + }; + match reads_child(&direct_reads) { + true => Via::Direct, + false => slot_reads + .iter() + .position(|reads| reads_child(reads)) + .map_or(Via::Direct, Via::Slot), + } + }) + .collect(); + // The direct part at ingestion time: the whole sub-DAG placed at + // ingestion time, its direct nodes counted. + let mut at_ingestion = dag.clone(); + for node in &mut at_ingestion.nodes { + node.output_state.timing = ExecutionTiming::IngestionTime; + } + let direct_at_ingestion = per_computation( + &at_ingestion, + entries, + scan_lookback, + horizons, + data, + models, + ) + .ok() + .map(|ingestion_priced| PartCoefficients { + cost_per_unit_rate: direct + .iter() + .map(|id| { + let p = &ingestion_priced[id]; + p.cpu_ops_per_second * prices.cost_per_cpu_op + + p.scanned_bytes_per_second * prices.cost_per_scanned_byte + + p.average_memory_bytes * prices.cost_per_memory_byte_second + + p.disk_bytes as f64 * prices.cost_per_disk_byte_second + }) + .sum(), + work_ms: 0.0, + }); + // The group's own node built once at query time, in memory: the state + // it holds instead of its working memory. + let mut built_once = dag.nodes[position[&root]].clone(); + built_once.materialization = Materialization::MatAtQuery(Medium::Memory); + let built_once_state_bytes = { + let mut once_dag = dag.clone(); + once_dag.nodes[position[&root]] = built_once; + let once_priced = + per_computation(&once_dag, entries, scan_lookback, horizons, data, models)?; + once_priced[&root].held.map_or(0, |(bytes, _, _)| bytes) + }; + let _ = candidate; + Ok(Some(Coefficients { + coefficients: ChoiceCoefficients { + direct: direct_coefficients, + slots, + child_via, + }, + lazy: LazyTerms { + held, + raw_scans, + direct_at_ingestion, + built_once_state_bytes, + thresholds: Vec::new(), + }, + root_key: shallow[position[&root]], + root_producer_keys: producers_of + .get(&root) + .into_iter() + .flatten() + .filter(|p| is_private(**p)) + .map(|p| shallow[position[p]]) + .collect(), + slot_keys, + part_keys, + })) +} + +/// The plan-level charges of `lazy` for one choice in a plan (cost model +/// 1): each held term at its part's rate and its readers' latencies. +/// `part_rate(via)` is the rate of the part's consumers, `readers(via)` the +/// queries reaching it with their rates, `latency_ms(q)` query `q`'s +/// latency, `built_once` whether the group's own node is built once. +pub(crate) fn held_cost( + lazy: &LazyTerms, + built_once: bool, + part_rate: impl Fn(Via) -> f64, + readers: impl Fn(Via) -> Vec<(usize, f64)>, + latency_ms: impl Fn(usize) -> f64, + cost_per_memory_byte_second: f64, +) -> f64 { + let mut cost = 0.0; + let built_once_term = HeldTerm { + via: Via::Direct, + bytes: lazy.built_once_state_bytes, + per_reader: false, + own_node: true, + }; + let terms = lazy + .held + .iter() + .filter(|term| !(built_once && term.own_node)) + .chain(built_once.then_some(&built_once_term)); + for term in terms { + let (bytes, per_reader) = (term.bytes, term.per_reader); + let readers = readers(term.via); + let longest = readers + .iter() + .map(|&(q, _)| latency_ms(q)) + .fold(0.0, f64::max); + let rate: f64 = readers.iter().map(|&(_, r)| r).sum(); + let evaluation_ms = match per_reader && rate > 0.0 { + true => readers.iter().map(|&(q, r)| r * latency_ms(q)).sum::() / rate, + false => longest, + }; + cost += cost_per_memory_byte_second * bytes as f64 * part_rate(term.via) * evaluation_ms + / 1_000.0; + } + cost +} + +/// Each source's longest raw scan among `scans` (Stage 3 keeps it on disk +/// once). +pub(crate) fn raw_data_bytes<'s>( + scans: impl IntoIterator, +) -> BTreeMap<&'s str, u64> { + let mut longest: BTreeMap<&str, u64> = BTreeMap::new(); + for (source, bytes) in scans { + let entry = longest.entry(source.as_str()).or_default(); + *entry = (*entry).max(*bytes); + } + longest +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pass::subexpr::search::{ + at_thresholds, composed_roots, group_choices, memo, resized_builds, shared, + }; + use crate::{computations, price_at_scale}; + 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::ir::schema::schema_support::with_promql_series_identity; + use asap_types::workload::{ + DataArrival, Predictability, 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(intervals_ms: &[u32], targets: &[AccuracyTarget]) -> Vec { + intervals_ms + .iter() + .zip(targets) + .map(|(&ms, target)| { + let mut entry = entry_with( + target.clone(), + QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval( + ms, + ))), + ); + entry.predictability = Predictability::Predictable { known_at: None }; + entry + }) + .collect() + } + + /// For every plan of `queries`' memo (every choice and threshold, every + /// Stage 2 materialization), Σ of its choices' coefficients times the + /// true consumer rates, plus the lazy terms, is Stage 3's price. + fn check(queries: &[&str], intervals_ms: &[u32]) -> Vec> { + let exact = vec![AccuracyTarget::exact(); queries.len()]; + check_with(queries, intervals_ms, &exact) + } + + /// [`check`] with each query's accuracy target; the Stage 2 + /// assignments of the plans checked. + fn check_with( + queries: &[&str], + intervals_ms: &[u32], + targets: &[AccuracyTarget], + ) -> Vec> { + let roots: Vec<(usize, QueryRoot)> = queries + .iter() + .zip(targets) + .map(|(q, target)| { + let root = lower_promql(q, target.clone()); + QueryRoot::Operator(with_promql_series_identity(&root).unwrap()) + }) + .enumerate() + .collect(); + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let entries = entries(intervals_ms, targets); + let data = data(); + let models = PlanningModels::builtin(); + let settings = models.plan_selection; + let CostModel::UsageOverTime(prices) = &settings.cost_model else { + unreachable!() + }; + let memo = memo(&per_computation, &entries); + let mut cache = PriceCache::new(1.0); + let parts: Vec> = (0..memo.groups.len()) + .map(|g| { + choice_parts( + &memo, + GroupId(g as u32), + &entries, + &data, + models, + &mut cache, + ) + .unwrap() + }) + .collect(); + let mut checked = Vec::new(); + for choice in group_choices(&memo).unwrap() { + for plan in at_thresholds(&memo, &choice) { + let Some(roots) = composed_roots(&memo, &per_computation, &plan) else { + continue; + }; + let roots = shared(&roots, &resized_builds(&memo, &choice, &plan)); + let Ok(physicals) = optimize_physical_asap( + "L", + &roots, + &entries, + &data, + models.deployment_capabilities, + ) else { + continue; + }; + for physical in physicals { + if let Some((expected, actual)) = priced_both( + &memo, &choice, &parts, &physical, &entries, &data, &models, prices, + &mut cache, + ) { + assert!( + (actual - expected).abs() <= 1e-9 * expected.abs(), + "{queries:?} {:?}: {actual} vs {expected}", + physical.assignment + ); + checked.push(physical.assignment.clone()); + } + } + } + } + checked + } + + /// Stage 3's price of `physical`, and the coefficients' (`None` when + /// a group's choice is not found in it). + #[allow(clippy::too_many_arguments)] + fn priced_both( + _memo: &Memo, + choice: &[Option], + parts: &[Vec], + physical: &PhysicalCandidate, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: &PlanningModels<'_>, + prices: &crate::ResourceUnitCost, + cache: &mut PriceCache, + ) -> Option<(f64, f64)> { + let settings = models.plan_selection; + let dag = &physical.dag; + let recurrence: Vec<_> = entries.iter().map(|e| e.recurrence.clone()).collect(); + let ranges: Vec<_> = entries.iter().map(|e| e.range_evaluation).collect(); + let nodes = node_map(dag); + let (cost, _) = price_at_scale( + dag, + &nodes, + &kept_pane_chains(dag), + &recurrence, + &ranges, + data, + settings, + 1.0, + ) + .ok()?; + let expected = cost.cost_per_second?; + + let rates = computations( + dag, + &recurrence, + &ranges, + settings.non_repeating_query_horizon_seconds, + None, + ) + .ok()?; + let reach = query_reach(dag, &nodes); + let keys = shallow_keys(dag, cache.keys_mut()); + let position_of: HashMap = dag + .nodes + .iter() + .enumerate() + .map(|(i, n)| (n.id, i)) + .collect(); + let by_key = |key: u32| -> Vec { + dag.nodes + .iter() + .zip(&keys) + .filter(|(_, k)| **k == key) + .map(|(n, _)| n.id) + .collect() + }; + let query_rate = + |q: usize| crate::shared_evaluation_rate(&[&entries[q]], settings).unwrap(); + // Each needed group's choice in this plan: its candidate, with the + // plan's materialization of its private builds. + struct Placed<'p> { + part: &'p ChoiceParts, + root: PhysicalASAPNodeId, + /// Whether `root` is its query-time copy. + at_query_time: bool, + /// Its copies at ingestion time (under ingestion-time builds). + at_ingestion: Vec, + slots: Vec, + } + let mut placed: Vec = Vec::new(); + let mut used: HashSet = HashSet::new(); + // The plan's groups: the chosen ones, and the children their + // choices read that the memo's candidates do not list (their first + // candidate). + let mut chosen: Vec> = choice.to_vec(); + let mut queue: Vec = (0..chosen.len()).filter(|&g| chosen[g].is_some()).collect(); + let mut order = Vec::new(); + while let Some(g) = queue.pop() { + order.push(g); + let c = chosen[g].expect("queued groups are chosen"); + for part in parts[g].iter().filter(|p| p.choice.candidate == c).take(1) { + for child in &part.choice.children { + let h = child.0 as usize; + if chosen[h].is_none() { + chosen[h] = Some(0); + queue.push(h); + } + } + } + } + order.sort_unstable(); + for g in order { + let c = chosen[g].expect("chosen"); + let candidates: Vec<&ChoiceParts> = parts[g] + .iter() + .filter(|p| p.choice.candidate == c) + .collect(); + let part = candidates.into_iter().find(|p| { + p.build_keys + .iter() + .zip(&p.choice.assignment) + .all(|(key, m)| { + by_key(*key) + .iter() + .any(|id| nodes[id].materialization == *m) + }) + }); + let part = part?; + // The plan's copies of its node: equal shallow keys, reading + // equal private producers. + let copies: Vec = by_key(part.root_key) + .into_iter() + .filter(|id| !used.contains(id)) + .filter(|id| { + let producers: Vec = dag + .edges + .iter() + .filter(|e| e.consumer == *id) + .map(|e| keys[position_of[&e.producer]]) + .collect(); + part.root_producer_keys + .iter() + .all(|k| producers.contains(k)) + }) + .collect(); + used.extend(copies.iter().copied()); + let timing = |id: &PhysicalASAPNodeId| nodes[id].output_state.timing; + let query_time = copies + .iter() + .copied() + .find(|id| timing(id) == ExecutionTiming::QueryTime); + let at_ingestion: Vec = copies + .iter() + .copied() + .filter(|id| timing(id) == ExecutionTiming::IngestionTime) + .collect(); + let root = query_time.or(at_ingestion.first().copied())?; + let slots = part + .slot_keys + .iter() + .map(|key| by_key(*key).first().copied()) + .collect::>>()?; + placed.push(Placed { + part, + root, + at_query_time: query_time.is_some(), + at_ingestion, + slots, + }); + } + let part_root = |p: &Placed, via: Via| match via { + Via::Direct => p.root, + Via::Slot(i) => p.slots[i], + }; + // Latency: each query's sum of the work of the parts it reaches. + let mut latency = vec![0.0; entries.len()]; + for p in &placed { + let vias = std::iter::once(Via::Direct).chain((0..p.slots.len()).map(Via::Slot)); + for via in vias { + let work = match via { + Via::Direct if !p.at_query_time => 0.0, + Via::Direct => p.part.choice.coefficients.direct.work_ms, + Via::Slot(i) => p.part.choice.coefficients.slots[i].work_ms, + }; + for (q, reached) in reach.iter().enumerate() { + if reached.contains(&part_root(p, via)) { + latency[q] += work; + } + } + } + } + let mut actual = 0.0; + let mut scans = Vec::new(); + for p in &placed { + let coefficients = &p.part.choice.coefficients; + if p.at_query_time { + actual += coefficients.direct.cost_per_unit_rate * rates[&p.root].1; + } + for copy in &p.at_ingestion { + let at_ingestion = p.part.lazy.direct_at_ingestion?; + actual += at_ingestion.cost_per_unit_rate * rates[copy].1; + } + for (i, slot) in coefficients.slots.iter().enumerate() { + actual += slot.fixed_cost + slot.cost_per_unit_rate * rates[&p.slots[i]].1; + } + let built_once = nodes[&p.root].materialization != Materialization::NotMaterialized + && !is_build(&nodes[&p.root].payload); + // Only a query-time direct part holds memory or scans raw data. + let mut lazy = p.part.lazy.clone(); + if !p.at_query_time { + lazy.held.retain(|term| term.via != Via::Direct); + lazy.raw_scans.retain(|(_, _, via)| *via != Via::Direct); + } + actual += held_cost( + &lazy, + built_once, + |via| rates[&part_root(p, via)].1, + |via| { + reach + .iter() + .enumerate() + .filter(|(_, r)| r.contains(&part_root(p, via))) + .map(|(q, _)| (q, query_rate(q))) + .collect() + }, + |q| latency[q], + prices.cost_per_memory_byte_second, + ); + for (source, bytes, via) in &lazy.raw_scans { + if nodes[&part_root(p, *via)].output_state.timing == ExecutionTiming::QueryTime { + scans.push((source.clone(), *bytes)); + } + } + } + actual += raw_data_bytes(&scans) + .values() + .map(|&bytes| prices.cost_per_disk_byte_second * bytes as f64) + .sum::(); + Some((expected, actual)) + } + + #[test] + fn coefficients_reproduce_stage_3_on_one_query_plans() { + for query in [ + "sum by (job) (rate(m[5m]))", + "max(avg_over_time(m[10m]))", + "sum_over_time(m[5m])", + ] { + assert!(!check(&[query], &[60_000]).is_empty(), "{query}"); + } + } + + /// A group both queries read: not materialized, built once at query + /// time, and its builds at ingestion time, in every plan. + #[test] + fn coefficients_reproduce_stage_3_with_a_shared_group() { + let checked = check( + &[ + "sum by (job) (rate(m[5m]))", + "sum by (job) (rate(m[5m])) * 2", + ], + &[60_000, 30_000], + ); + let at = |m: fn(&Materialization) -> bool| checked.iter().flatten().any(m); + assert!(at(|m| matches!(m, Materialization::MatAtQuery(_)))); + assert!(at(|m| matches!(m, Materialization::MatAtIngestion(_)))); + } + + /// A 5-minute sum read from its own panes or from a 1-hour sum's + /// chain (the chain site's thresholds), built at query or ingestion + /// time. + #[test] + #[ignore = "a pane chain over a raw scan is no raw-horizon reader: the shared \ + scan's look-back (the longest window of its readers) is not a site \ + threshold, so a chain's input rows depend on other groups"] + fn coefficients_reproduce_stage_3_with_a_shared_pane_chain() { + let checked = check( + &["sum_over_time(m[5m])", "sum_over_time(m[1h])"], + &[60_000, 60_000], + ); + assert!(checked.len() > 2); + } + + /// Two quantiles within 1% and 2% sharing one KLL build at a size + /// threshold. + #[test] + fn coefficients_reproduce_stage_3_with_a_shared_build() { + let checked = check_with( + &[ + "quantile_over_time(0.9, m[5m])", + "quantile_over_time(0.9, m[5m])", + ], + &[60_000, 60_000], + &[ + AccuracyTarget::max_error(0.01), + AccuracyTarget::max_error(0.02), + ], + ); + assert!(!checked.is_empty()); + } +} diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 6a82c2fe8..108331bc1 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -23,6 +23,7 @@ #![allow(dead_code)] pub(crate) mod capability; +pub(crate) mod coefficients; pub(crate) mod groups; pub(crate) mod milp; pub(crate) mod pricing; diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index 9c200c032..1b492553f 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -151,6 +151,11 @@ impl PriceCache { ..Self::default() } } + + /// The interned structural node keys. + pub(super) fn keys_mut(&mut self) -> &mut HashMap { + &mut self.keys + } } /// A candidate's sub-DAG as Stage 2 builds it, with its nodes' keys. @@ -162,9 +167,9 @@ struct Part { /// A child group's sub-DAG's node keys, and its root's. #[derive(Debug, Default)] -struct ChildKeys { - nodes: HashSet, - root: Option, +pub(super) struct ChildKeys { + pub(super) nodes: HashSet, + pub(super) root: Option, } /// The usage of `candidate`'s private nodes (from its node down to, @@ -427,7 +432,7 @@ fn built( /// 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( +pub(super) fn child_keys( memo: &Memo, child: GroupId, entries: &[QueryWorkloadEntry], @@ -470,7 +475,7 @@ fn child_keys( /// 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 { +pub(super) fn keys_of(dag: &PhysicalASAPDAG, interned: &mut HashMap) -> Vec { let mut key_of: HashMap = HashMap::new(); dag.nodes .iter() @@ -632,7 +637,7 @@ fn price_part( /// 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 { +pub(super) 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; diff --git a/crates/plan-selection/src/pass/subexpr/scans.rs b/crates/plan-selection/src/pass/subexpr/scans.rs index 2188adaa7..a9d08590b 100644 --- a/crates/plan-selection/src/pass/subexpr/scans.rs +++ b/crates/plan-selection/src/pass/subexpr/scans.rs @@ -415,7 +415,7 @@ fn raw_reader(node: &Rc) -> Option { return None; }; Some(RawReader { - key: hash_of(&("RawHorizon", format!("{source:?}"))), + key: raw_horizon_key(source), horizon_ms: u64::try_from(range.as_millis()).unwrap_or(u64::MAX), }) } @@ -479,6 +479,11 @@ fn child_groups(boundary: &mut Boundary, node: &Rc) -> Vec u64 { + hash_of(&("RawHorizon", format!("{source:?}"))) +} + fn hash_of(value: &T) -> u64 { let mut hasher = DefaultHasher::new(); value.hash(&mut hasher); diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index 39a37a693..7bb8eade3 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -67,12 +67,12 @@ pub(crate) fn stats(memo: &Memo) -> String { /// 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>; +pub(super) 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> { +pub(super) 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(); @@ -137,7 +137,10 @@ fn group_choices(memo: &Memo) -> Option> { /// of its coupling sites: every combination of thresholds every reader may /// 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>> { +pub(super) fn at_thresholds( + memo: &Memo, + choice: &GroupChoice, +) -> Vec>> { let chosen: Vec> = choice .iter() .enumerate() @@ -246,7 +249,7 @@ fn at_thresholds(memo: &Memo, choice: &GroupChoice) -> Vec( +pub(super) fn composed_roots( memo: &Memo, per_computation: &PerComputationCandidates, chosen: &[Option], @@ -340,7 +343,7 @@ fn substituted( /// `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 { +pub(super) fn shared(roots: &[QueryRoot], resized: &[Rc]) -> Vec { let identical = |roots: Vec| { let shareable = sharing_candidates(&roots); match shareable.is_empty() { @@ -357,7 +360,7 @@ fn shared(roots: &[QueryRoot], resized: &[Rc]) -> Vec { /// 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( +pub(super) fn resized_builds( memo: &Memo, choice: &GroupChoice, plan: &[Option], From 338325d796433852827bbc1c7717fa7b642fd378 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:54:48 +0000 Subject: [PATCH 26/67] plan-selection: SearchAlgorithm setting, MILP stub, search report fields, stage_pipeline --search, three-way harness Stage1ToStage3Search.memo_groups becomes search: SearchAlgorithm (Legacy, MemoBruteForce, MemoMilp; default Legacy). MemoMilp calls subexpr::milp::search, a stub returning None until the model is built; under cost models 2 and 3, or while the stub returns None, the legacy search runs and SearchReport.routed_to_legacy says why. SearchReport gains algorithm, routed_to_legacy, solver_gap, solver_iterations and solver_time_ms (omitted in JSON at their defaults, so fixtures are unchanged). stage_pipeline takes --search legacy|memo-brute-force|memo-milp. joint_search_equivalence gains assert_three_way (exhaustive, memo brute force, memo MILP), random-workload tests under cost models 1, 2 and 3 (ignored: minutes in a debug build) and three targeted workloads: peak cores deciding under cost model 2, a memory limit ruling out the best plan under cost model 3, and a raw-data horizon linking two windows over one source. --- crates/devtools/src/bin/stage_pipeline.rs | 44 ++- crates/devtools/src/planning_stages.rs | 10 +- crates/devtools/tests/stage_pipeline.rs | 80 ++++ .../tests/joint_search_equivalence.rs | 367 +++++++++++++++++- crates/plan-selection/src/lib.rs | 65 ++++ .../src/pass/stage1_to_stage3.rs | 144 +++++-- .../plan-selection/src/pass/subexpr/milp.rs | 31 ++ crates/plan-selection/src/pass/subexpr/mod.rs | 7 +- docs/user_guide_docs/stage-viewer.md | 1 + 9 files changed, 690 insertions(+), 59 deletions(-) create mode 100644 crates/plan-selection/src/pass/subexpr/milp.rs diff --git a/crates/devtools/src/bin/stage_pipeline.rs b/crates/devtools/src/bin/stage_pipeline.rs index 02961938f..2a1d6e78f 100644 --- a/crates/devtools/src/bin/stage_pipeline.rs +++ b/crates/devtools/src/bin/stage_pipeline.rs @@ -33,6 +33,13 @@ // `"default"`. `--exhaustive` instead composes and prices every candidate // (the search the default one is tested against, for small workloads and // comparison) and the document lists them; `search` is `"exhaustive"`. +// `--search legacy|memo-brute-force|memo-milp` picks the default search's +// algorithm (`SearchAlgorithm`): `legacy` (the default, `search` is +// `"default"`), the memo over logical sub-expressions searched by brute +// force (the reference), or solved as a MILP; `search` names it, and the +// search report says which ran (`algorithm`, omitted for legacy) and, +// when the legacy search ran instead, why (`routed_to_legacy`). Not with +// `--exhaustive`. // Unlike `plan_installable`'s default settings, Stage 3 ignores accuracy // guarantees unless `--enforce-accuracy-guarantees` is given (below); // `accuracy_guarantees` is `"ignored"` or `"enforced"`. @@ -202,7 +209,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; @@ -213,8 +220,8 @@ use asap_logical_asap_optimization::per_computation_candidates::{ 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, + PlanSelectionSettings, PlanningModels, Rejection, ResourceUnitCost, ResourceUsage, + SearchAlgorithm, Selection, SelectionError, }; use asap_planner::{installable_plan_of, PlannerExecutorCompileCheck}; use asap_types::ir::flat::{flatten, FlatDag}; @@ -233,7 +240,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|memo-milp] [--max-displayed ] --out "; /// The default `--max-combinations`: the most Stage 1 candidates (Pass 2's /// included) an exhaustive run enumerates. @@ -334,6 +341,7 @@ fn run(args: Vec) -> Result<(), Failure> { let mut plan_selection_settings_path: Option = None; let mut enforce_accuracy_guarantees = false; let mut exhaustive = false; + let mut search_algorithm: Option = None; let mut max_combinations_given = false; let mut args = args.into_iter(); while let Some(flag) = args.next() { @@ -393,6 +401,19 @@ 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" => { + search_algorithm = Some(match value()?.as_str() { + "legacy" => SearchAlgorithm::Legacy, + "memo-brute-force" => SearchAlgorithm::MemoBruteForce, + "memo-milp" => SearchAlgorithm::MemoMilp, + other => { + return Err(format!( + "--search {other}: expected legacy, memo-brute-force or memo-milp" + ) + .into()) + } + }) + } other => return Err(format!("unknown argument {other}").into()), } } @@ -403,6 +424,12 @@ fn run(args: Vec) -> Result<(), Failure> { return Err("--table needs --sql".into()); } let out = out.ok_or("--out is required")?; + if exhaustive && search_algorithm.is_some() { + return Err( + "--search picks the default search's algorithm: drop it with --exhaustive".into(), + ); + } + let search_algorithm = search_algorithm.unwrap_or_default(); if max_combinations_given && !exhaustive { return Err("--max-combinations limits the exhaustive search: it needs --exhaustive (the default search has no limit)".into()); } @@ -614,7 +641,7 @@ 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, search_algorithm), }; // A selected plan's very large peak memory, when the deployment // declares no memory capacity to check it against. @@ -636,7 +663,12 @@ 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, search_algorithm) { + (true, _) => "exhaustive", + (false, SearchAlgorithm::Legacy) => "default", + (false, SearchAlgorithm::MemoBruteForce) => "memo-brute-force", + (false, SearchAlgorithm::MemoMilp) => "memo-milp", + }); 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..d411894bf 100644 --- a/crates/devtools/src/planning_stages.rs +++ b/crates/devtools/src/planning_stages.rs @@ -11,7 +11,7 @@ use asap_plan_selection::pass::{ run_stage1_to_stage3_from_pass1, run_stage2_and_stage3, Stage1ToStage3Search, }; pub use asap_plan_selection::pass::{LogicalCandidate, PrunedCandidate}; -use asap_plan_selection::PlanningModels; +use asap_plan_selection::{PlanningModels, SearchAlgorithm}; use asap_plan_selection::{Selection, SelectionError}; use asap_types::ir::schema::schema_support::with_promql_series_identity; use asap_types::ir::QueryRoot; @@ -220,12 +220,14 @@ pub fn run_planning_stages_over( /// fails), and Stage 3 records the first rejection of each cause and the /// search's report (`Selection::search`). `models` should be the ones /// `plan_installable` plans with (the executor's summary families and -/// compile check) to see what it selects. No limit applies. -pub fn run_default_search_over( +/// compile check) to see what it selects. No limit applies. `algorithm` +/// picks the search ([`SearchAlgorithm::Legacy`] is the default search). +pub fn run_search_over( roots: Vec, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, + algorithm: SearchAlgorithm, ) -> Result { validate(&roots)?; let per_computation = @@ -236,7 +238,7 @@ pub fn run_default_search_over( entries, data, models, - Stage1ToStage3Search::default(), + Stage1ToStage3Search::with_algorithm(algorithm), ); if let Some(Ok(roots)) = run.stage1_selected.as_ref().map(|c| &c.roots) { validate(roots)?; diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index 8e4b116c6..a36dfdd45 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1688,3 +1688,83 @@ fn scans_projecting_different_labels_are_shared() { "{rolled_up} vs {without_rollup}" ); } + +/// `--search` picks the default search's algorithm. Example 3b's p99 panel +/// under cost model 1: the memo searched by brute force selects a plan of +/// the legacy search's price; the MILP, not built yet, runs the legacy +/// search and says why; under cost model 2 the MILP runs the legacy search +/// because it prices cost model 1 only. +#[test] +fn search_flag_picks_the_search_algorithm() { + let run = |extra: &[&str]| { + let mut args = vec!["--example", "planner-layering-3b"]; + args.extend_from_slice(extra); + document_of(&args) + }; + let price = |document: &Value| { + let selection = &document["stage3_plan_selection"]; + let selected = selection["selected"].as_str().unwrap(); + selection["costs"][selected]["cost_per_second"] + .as_f64() + .unwrap() + }; + let legacy = run(&["--search", "legacy"]); + assert_eq!(legacy["search"], "default"); + let report = &legacy["stage3_plan_selection"]["search"]; + // The legacy search's report omits the algorithm and the solver's + // fields, which only the MILP fills. + for field in [ + "algorithm", + "routed_to_legacy", + "solver_gap", + "solver_iterations", + "solver_time_ms", + ] { + assert!(report.get(field).is_none(), "{field}: {report}"); + } + + let brute_force = run(&["--search", "memo-brute-force"]); + assert_eq!(brute_force["search"], "memo-brute-force"); + let report = &brute_force["stage3_plan_selection"]["search"]; + assert_eq!(report["algorithm"], "memo-brute-force"); + assert!((price(&brute_force) - price(&legacy)).abs() <= 1e-9 * price(&legacy)); + + let milp = run(&["--search", "memo-milp"]); + assert_eq!(milp["search"], "memo-milp"); + let report = &milp["stage3_plan_selection"]["search"]; + assert!(report.get("algorithm").is_none(), "{report}"); + assert_eq!(report["routed_to_legacy"], "milp_not_built"); + assert_eq!(price(&milp), price(&legacy)); + + let cost_model_2 = run(&[ + "--search", + "memo-milp", + "--cost-model", + "provisioned-capacity", + ]); + assert_eq!( + cost_model_2["stage3_plan_selection"]["search"]["routed_to_legacy"], + "cost_model_not_supported" + ); +} + +/// `--search` takes one of its three algorithms, and not with +/// `--exhaustive`, which is its own search. +#[test] +fn search_flag_is_checked() { + let (status, stderr) = failure(&["--promql", "sum(m)", "--search", "milp"]); + assert_eq!(status, Some(2)); + assert!( + stderr.contains("expected legacy, memo-brute-force or memo-milp"), + "{stderr}" + ); + let (status, stderr) = failure(&[ + "--promql", + "sum(m)", + "--exhaustive", + "--search", + "memo-milp", + ]); + assert_eq!(status, Some(2)); + assert!(stderr.contains("drop it with --exhaustive"), "{stderr}"); +} diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index b9e88ca5d..cd34d4bb4 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -10,7 +10,8 @@ use asap_devtools::planning_stages::{lower_promql_roots, lower_sql_roots}; use asap_plan_selection::{ run_stage1_to_stage3, run_stage1_to_stage3_with, AccuracyModel, BuiltinAccuracyModel, ClusterLimits, ClusterObjective, CostModel, PlanSelectionSettings, PlanningModels, - RejectionCause, ResourceUnitCost, SelectionError, Stage1ToStage3Run, Stage1ToStage3Search, + RejectionCause, ResourceUnitCost, RoutedToLegacy, SearchAlgorithm, SelectionError, + Stage1ToStage3Run, Stage1ToStage3Search, }; use asap_types::ir::flat::flatten; use asap_types::ir::properties::{AccuracyGuarantee, AccuracyMetric, AccuracyTarget}; @@ -503,7 +504,7 @@ fn a_known_bound_that_misses_is_not_discarded_before_composition() { Stage1ToStage3Search { keep_every_candidate: true, exhaustive: false, - memo_groups: false, + search: SearchAlgorithm::Legacy, }, ) .unwrap(); @@ -1203,10 +1204,7 @@ fn memo_group_search_selects_the_exhaustive_price() { &entries, &data, models, - Stage1ToStage3Search { - memo_groups: true, - ..Stage1ToStage3Search::default() - }, + Stage1ToStage3Search::with_algorithm(SearchAlgorithm::MemoBruteForce), ) .unwrap(); let price = |run: &Stage1ToStage3Run| { @@ -1329,10 +1327,7 @@ fn debug_memo_seed() { ("exhaustive", Stage1ToStage3Search::exhaustive()), ( "memo", - Stage1ToStage3Search { - memo_groups: true, - ..Stage1ToStage3Search::default() - }, + Stage1ToStage3Search::with_algorithm(SearchAlgorithm::MemoBruteForce), ), ] { let run = run_stage1_to_stage3_with(indexed(), &entries, &data, models, search).unwrap(); @@ -1355,3 +1350,355 @@ fn debug_memo_seed() { eprintln!("{out}"); } } + +/// What a search's selection is compared by: under cost models 1 and 2 +/// its price per second; under cost model 3 the key its objective ranks +/// by (worst-query latency in ms, peak memory and disk in bytes, in the +/// objective's order; the Pareto front selects its lowest-latency +/// member). `None` when no candidate is valid. +fn objective_of(run: &Stage1ToStage3Run, cost_model: &CostModel) -> Option> { + let selection = run.stage3_plan_selection.as_ref().ok()?; + let cost = &selection.costs[&selection.selected]; + let usage = &cost.usage; + let worst_latency_ms = usage + .latency_per_query + .iter() + .map(|l| l.total_ms()) + .fold(0.0, f64::max); + let (memory, disk) = (usage.peak_memory_bytes as f64, usage.disk_bytes as f64); + Some(match cost_model { + CostModel::ClusterCapacity(limits) => match limits.objective { + ClusterObjective::LeastMemory => vec![memory, disk, worst_latency_ms], + ClusterObjective::LowestLatency | ClusterObjective::ParetoFront => { + vec![worst_latency_ms, memory, disk] + } + }, + _ => vec![cost + .cost_per_second + .expect("cost models 1 and 2 price every valid candidate")], + }) +} + +/// Whether `a` and `b` are equal within 1e-9, relative (absolute for 0). +fn same_objective(a: &Option>, b: &Option>) -> bool { + match (a, b) { + (Some(a), Some(b)) => a + .iter() + .zip(b) + .all(|(a, b)| (a - b).abs() <= 1e-9 * a.abs().max(b.abs()).max(1.0)), + (None, None) => true, + _ => false, + } +} + +/// `roots` planned three ways under `models`: the exhaustive search, the +/// memo searched by brute force ([`SearchAlgorithm::MemoBruteForce`]) and +/// the memo solved as a MILP ([`SearchAlgorithm::MemoMilp`]). The +/// exhaustive search and the brute force select plans of the same price +/// (within 1e-9; [`objective_of`] under cost model 3), or both find none. +/// While the MILP runs the legacy search instead (it prices cost model 1 +/// only, and is not built yet), the report says why and only the first two +/// are compared; once it runs, it selects that price too, with a solver gap +/// of 0. Returns whether the MILP ran. +fn assert_three_way( + label: &str, + roots: &[QueryRoot], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> bool { + three_way(label, roots, entries, data, models).unwrap_or_else(|mismatch| panic!("{mismatch}")) +} + +/// [`assert_three_way`], with a mismatch returned rather than panicking. +fn three_way( + label: &str, + roots: &[QueryRoot], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result { + let run = |search| { + run_stage1_to_stage3_with( + roots.iter().cloned().enumerate().collect(), + entries, + data, + models, + search, + ) + .unwrap_or_else(|e| panic!("{label}: {e}")) + }; + let cost_model = &models.plan_selection.cost_model; + let exhaustive = objective_of(&run(Stage1ToStage3Search::exhaustive()), cost_model); + let brute_force_run = run(Stage1ToStage3Search::with_algorithm( + SearchAlgorithm::MemoBruteForce, + )); + if let Ok(selection) = &brute_force_run.stage3_plan_selection { + assert_eq!( + selection.search.algorithm, + SearchAlgorithm::MemoBruteForce, + "{label}" + ); + } + let brute_force = objective_of(&brute_force_run, cost_model); + if !same_objective(&exhaustive, &brute_force) { + return Err(format!( + "{label}: exhaustive {exhaustive:?}, memo brute force {brute_force:?}" + )); + } + let milp_run = run(Stage1ToStage3Search::with_algorithm( + SearchAlgorithm::MemoMilp, + )); + let Ok(selection) = &milp_run.stage3_plan_selection else { + // The legacy search or the MILP found no valid plan. + return match exhaustive { + None => Ok(false), + Some(_) => Err(format!( + "{label}: exhaustive {exhaustive:?}, memo MILP {:?}", + milp_run.stage3_plan_selection + )), + }; + }; + let report = &selection.search; + if report.algorithm == SearchAlgorithm::Legacy { + let expected = match cost_model { + CostModel::UsageOverTime(_) => RoutedToLegacy::MilpNotBuilt, + _ => RoutedToLegacy::CostModelNotSupported, + }; + assert_eq!(report.routed_to_legacy, Some(expected), "{label}"); + return Ok(false); + } + assert_eq!(report.algorithm, SearchAlgorithm::MemoMilp, "{label}"); + assert_eq!(report.routed_to_legacy, None, "{label}"); + let milp = objective_of(&milp_run, cost_model); + if !same_objective(&exhaustive, &milp) { + return Err(format!( + "{label}: exhaustive {exhaustive:?}, memo MILP {milp:?}" + )); + } + let gap = report.solver_gap.expect("the MILP reports its gap"); + if gap.abs() > 1e-9 { + return Err(format!("{label}: MILP gap {gap}")); + } + Ok(true) +} + +/// [`assert_three_way`] on seeds 1 to 60 of the random workloads small +/// enough to search exhaustively, over one or two gauges, with raw data +/// kept or not, under each of `settings` in turn (every one with `every`). +fn assert_three_way_on_random_workloads(settings: &[PlanSelectionSettings], every: bool) { + 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 turn = seed as usize % settings.len(); + let chosen = match every { + true => settings, + false => &settings[turn..=turn], + }; + for settings in chosen { + let models = PlanningModels::builtin() + .with_plan_selection(settings) + .with_deployment_capabilities(&deployment); + if let Err(mismatch) = three_way( + &format!("seed {seed} under {:?}", settings.cost_model), + &roots, + &entries, + &data, + models, + ) { + mismatches.push(mismatch); + } + } + } + assert!( + searched >= 20, + "only {searched} workloads were small enough" + ); + assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); +} + +/// The three searches agree on the random workloads under cost model 1. +#[test] +#[ignore = "the exhaustive search and the brute force over every seed take minutes in a debug build; run with --release --ignored"] +fn three_way_random_workloads_under_cost_model_1() { + assert_three_way_on_random_workloads(&cost_models(1)[..1], true); +} + +/// The three searches agree on the random workloads under cost model 2. +#[test] +#[ignore = "the exhaustive search and the brute force over every seed under cost model 2 take minutes in a debug build; run with --release --ignored"] +fn three_way_random_workloads_under_cost_model_2() { + assert_three_way_on_random_workloads(&cost_models(1)[1..2], true); +} + +/// The three searches agree on the random workloads under cost model 3, +/// each objective in turn, with 1 GB of cluster memory. +#[test] +#[ignore = "the exhaustive search and the brute force over every seed under cost model 3 take minutes in a debug build; run with --release --ignored"] +fn three_way_random_workloads_under_cost_model_3() { + assert_three_way_on_random_workloads(&cost_models(1)[2..], false); +} + +/// `queries`, each repeated every minute within 5% error, over the gauge +/// `m` with `rows_per_second` series each sampled every second. +fn workload_at_rate(queries: &[&str], rows_per_second: u64) -> PlanningWorkload { + let queries: Vec<(&str, u32)> = queries.iter().map(|&q| (q, 60_000)).collect(); + let mut workload = workload_of(&queries, AccuracyTarget::max_error(0.05)); + let data = workload.data_workload.as_mut().unwrap(); + data.arrival = DataArrival::Ingesting { + rows_per_second: rows_per_second as f64, + }; + for facts in data.sources.values_mut() { + facts.series_count = Some(rows_per_second); + facts.sample_interval = Some(DurationMs(1_000)); + } + workload +} + +/// `settings` with `cost_model`. +fn under(cost_model: CostModel) -> PlanSelectionSettings { + PlanSelectionSettings { + cost_model, + ..PlanSelectionSettings::default() + } +} + +/// The exhaustive search's selected candidate's usage under `settings`. +fn exhaustive_selection( + workload: &PlanningWorkload, + settings: &PlanSelectionSettings, +) -> asap_plan_selection::ResourceUsage { + let roots = lower_promql_roots(workload).unwrap(); + let entries: Vec<_> = workload.query_workload.entries().collect(); + let data = workload.data_workload.clone().unwrap(); + let run = run_stage1_to_stage3_with( + roots.into_iter().enumerate().collect(), + &entries, + &data, + PlanningModels::builtin().with_plan_selection(settings), + Stage1ToStage3Search::exhaustive(), + ) + .unwrap(); + let selection = run.stage3_plan_selection.unwrap(); + selection.costs[&selection.selected].usage.clone() +} + +/// [`assert_three_way`] on `workload` under `settings`. +fn assert_three_way_on(label: &str, workload: &PlanningWorkload, settings: &PlanSelectionSettings) { + let roots = lower_promql_roots(workload).unwrap(); + let entries: Vec<_> = workload.query_workload.entries().collect(); + let data = workload.data_workload.clone().unwrap(); + assert_three_way( + label, + &roots, + &entries, + &data, + PlanningModels::builtin().with_plan_selection(settings), + ); +} + +/// A peak case: under cost model 2, a p99 and a sum over 5 minutes of a +/// million rows per second, raw data on disk priced at 1e-7 per byte-second. +/// Scanning the raw window at query time needs one core per query and no +/// ingestion core; maintaining state at ingestion time needs fewer disk +/// bytes and one more core. At 1,000 per core-second the query-time plan is +/// selected; with cores free, a plan with an ingestion core is: the peak +/// cores decide. The three searches agree. +#[test] +fn three_way_peak_cores_decide_under_cost_model_2() { + let workload = workload_at_rate( + &["quantile_over_time(0.99, m[5m])", "sum_over_time(m[5m])"], + 1_000_000, + ); + let prices = |cost_per_provisioned_cpu_core_second| ResourceUnitCost { + cost_per_disk_byte_second: 1e-7, + cost_per_provisioned_cpu_core_second, + ..ResourceUnitCost::illustrative() + }; + let priced = under(CostModel::ProvisionedCapacity(prices(1_000.0))); + let free_cores = under(CostModel::ProvisionedCapacity(prices(0.0))); + let (with_core_price, without) = ( + exhaustive_selection(&workload, &priced), + exhaustive_selection(&workload, &free_cores), + ); + assert!( + with_core_price.peak_cpu_cores < without.peak_cpu_cores, + "{with_core_price:?} {without:?}" + ); + assert_three_way_on("peak cores", &workload, &priced); + assert_three_way_on("peak cores, free", &workload, &free_cores); +} + +/// A capacity case: under cost model 3 with the lowest-latency objective, +/// a p99 and a sum over 5 minutes of a million rows per second. Without a +/// binding limit the lowest-latency plan holds more than 5 GB; with 5 GB +/// of memory it is ruled out and another plan is selected. The three +/// searches agree. +#[test] +fn three_way_memory_limit_rules_out_the_best_plan_under_cost_model_3() { + let workload = workload_at_rate( + &["quantile_over_time(0.99, m[5m])", "sum_over_time(m[5m])"], + 1_000_000, + ); + let limited = |memory_bytes| { + under(CostModel::ClusterCapacity(ClusterLimits { + memory_bytes, + disk_bytes: 1_000_000_000_000, + cpu_cores: None, + objective: ClusterObjective::LowestLatency, + })) + }; + let (unlimited, five_gb) = (limited(1_000_000_000_000_000), limited(5_000_000_000)); + let best = exhaustive_selection(&workload, &unlimited); + assert!(best.peak_memory_bytes > 5_000_000_000, "{best:?}"); + let within = exhaustive_selection(&workload, &five_gb); + assert!(within.peak_memory_bytes <= 5_000_000_000, "{within:?}"); + assert_three_way_on("memory limit", &workload, &five_gb); +} + +/// A raw-horizon case: raw data kept, on disk priced at 1e-8 per +/// byte-second under cost model 1, and a p99 over 5 minutes and a sum over +/// 10 minutes of one source at 100,000 rows per second. Raw data is kept +/// for the longest window any query scans at query time, so the queries' +/// choices are linked through the source's horizon: the selection scans +/// the 5-minute window and keeps no 10-minute horizon. The three searches +/// agree. +#[test] +fn three_way_raw_horizon_links_two_windows_over_one_source() { + let workload = workload_at_rate( + &["quantile_over_time(0.99, m[5m])", "sum_over_time(m[10m])"], + 100_000, + ); + let settings = under(CostModel::UsageOverTime(ResourceUnitCost { + cost_per_disk_byte_second: 1e-8, + ..ResourceUnitCost::illustrative() + })); + let selected = exhaustive_selection(&workload, &settings); + let ten_minutes_of_raw_bytes = 2 * 960_000_000; + assert!( + selected.disk_bytes > 0 && selected.disk_bytes < ten_minutes_of_raw_bytes, + "{selected:?}" + ); + assert_three_way_on("raw horizon", &workload, &settings); +} diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index aae16cc0c..6362399ea 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1028,6 +1028,71 @@ 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 search algorithm that ran; omitted in JSON for the legacy + /// search. + #[serde(default, skip_serializing_if = "SearchAlgorithm::is_legacy")] + pub algorithm: SearchAlgorithm, + /// Why the legacy search ran instead of the requested memo search; + /// `None` when the requested algorithm ran. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub routed_to_legacy: Option, + /// The MILP solver's relative gap between its best plan and its bound + /// at the end ([`SearchAlgorithm::MemoMilp`]); 0 when the plan is + /// proven optimal. `None` for every other search. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub solver_gap: Option, + /// The MILP solver's simplex iterations; 0 (omitted in JSON) for every + /// other search. + #[serde(default, skip_serializing_if = "is_zero")] + pub solver_iterations: u32, + /// The MILP solver's wall time in milliseconds; 0 (omitted in JSON) for + /// every other search. + #[serde(default, skip_serializing_if = "is_zero")] + pub solver_time_ms: u64, +} + +/// How the Stage 1 to Stage 3 search explores the plans +/// ([`Stage1ToStage3Search::search`]). +#[derive( + Debug, Clone, Copy, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize, +)] +#[serde(rename_all = "kebab-case")] +pub enum SearchAlgorithm { + /// The search over whole-workload candidates: each query's own + /// candidates checked first, independent components searched + /// separately, only the best so far kept. + #[default] + Legacy, + /// Every plan of the memo over logical sub-expressions, each priced by + /// Stage 3: the reference the MILP is checked against. + MemoBruteForce, + /// The memo over logical sub-expressions solved as a mixed-integer + /// linear program. + MemoMilp, +} + +impl SearchAlgorithm { + fn is_legacy(&self) -> bool { + *self == SearchAlgorithm::Legacy + } +} + +fn is_zero(value: &T) -> bool { + *value == T::default() +} + +/// Why the legacy search ran instead of the requested memo search. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RoutedToLegacy { + /// The caller asked for the exhaustive search, which the memo searches + /// do not replace. + ExhaustiveRequested, + /// The MILP's objective is cost model 1's sum of local prices; cost + /// models 2 and 3 are searched by the legacy search. + CostModelNotSupported, + /// The MILP model is not built yet. + MilpNotBuilt, } /// 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..486d80377 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -57,8 +57,8 @@ use super::query_candidates::QueryCandidates; use super::{OptimizationInput, OptimizationPass, OptimizeError, PlanOutput, QueryPlan}; use crate::PlanningModels; use crate::{ - CalibrationFit, PlanSelector, QuerySearchReport, Rejection, RejectionCause, Selection, - SelectionError, + CalibrationFit, CostModel, PlanSelector, QuerySearchReport, Rejection, RejectionCause, + RoutedToLegacy, SearchAlgorithm, Selection, SelectionError, }; use asap_logical_asap_optimization::per_computation_candidates::{ generate_per_computation_candidates, PerComputationCandidateError, PerComputationCandidates, @@ -159,11 +159,12 @@ 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, + /// The search algorithm (joint-planning-search design: the legacy + /// search, or the memo over logical sub-expressions, "Memo-group + /// search"); ignored with `exhaustive`. Default + /// [`SearchAlgorithm::Legacy`] until the MILP selects the exhaustive + /// search's price on every equivalence test. + pub search: SearchAlgorithm, } impl Stage1ToStage3Search { @@ -173,7 +174,15 @@ impl Stage1ToStage3Search { Self { keep_every_candidate: true, exhaustive: true, - memo_groups: false, + search: SearchAlgorithm::Legacy, + } + } + + /// The default search run with `algorithm`. + pub fn with_algorithm(algorithm: SearchAlgorithm) -> Self { + Self { + search: algorithm, + ..Self::default() } } } @@ -244,36 +253,97 @@ 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 routed_to_legacy = match (search.exhaustive, search.search) { + (_, SearchAlgorithm::Legacy) => None, + (true, _) => Some(RoutedToLegacy::ExhaustiveRequested), + (false, SearchAlgorithm::MemoBruteForce) => { + let pruned = + per_computation.without_unsupported_summaries(models.deployment_capabilities); + let found = super::subexpr::search::brute_force(&pruned, entries, data, models); + return memo_run( + pruned, + found, + entries, + data, + SearchAlgorithm::MemoBruteForce, + ); + } + (false, SearchAlgorithm::MemoMilp) => { + if !matches!( + models.plan_selection.cost_model, + CostModel::UsageOverTime(_) + ) { + Some(RoutedToLegacy::CostModelNotSupported) + } else { + let pruned = per_computation + .clone() + .without_unsupported_summaries(models.deployment_capabilities); + match super::subexpr::milp::search(&pruned, entries, data, models) { + Some(found) => { + return memo_run(pruned, found, entries, data, SearchAlgorithm::MemoMilp) + } + None => Some(RoutedToLegacy::MilpNotBuilt), + } + } + } + }; + let mut run = legacy_search(per_computation, entries, data, models, search); + if let Ok(selection) = &mut run.stage3_plan_selection { + selection.search.algorithm = SearchAlgorithm::Legacy; + selection.search.routed_to_legacy = routed_to_legacy; } + run +} + +/// The run of a memo search (`algorithm`) over `per_computation`, pruned +/// for the deployment, that `found` the selection, its physical candidate +/// and its logical roots, or why none is valid. +fn memo_run( + per_computation: PerComputationCandidates, + found: super::subexpr::milp::MemoSearchResult, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + algorithm: SearchAlgorithm, +) -> Stage1ToStage3Run { + let combinations = per_computation.combinations(); + let (selection, selected, logical) = match found { + Ok((mut selection, selected, logical)) => { + selection.search.algorithm = algorithm; + (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)); + Stage1ToStage3Run { + stage1_pass1_per_computation: per_computation, + 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, + }), + } +} + +/// The legacy search ([`SearchAlgorithm::Legacy`]) of +/// [`run_stage1_to_stage3_from_pass1`]: independent components searched +/// separately when that is proven exact, else the workload as one. +fn legacy_search( + per_computation: PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + search: Stage1ToStage3Search, +) -> Stage1ToStage3Run { 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/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs new file mode 100644 index 000000000..0f3732633 --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -0,0 +1,31 @@ +//! The memo over logical sub-expressions solved as a mixed-integer linear +//! program (joint-planning-search design, "Memo-group search"), behind +//! [`crate::SearchAlgorithm::MemoMilp`]. + +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; +use asap_physical_asap_optimization::PhysicalCandidate; +use asap_types::ir::QueryRoot; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; + +use crate::{PlanningModels, Selection, SelectionError}; + +/// A memo search's result: the selection, its physical candidate and its +/// logical roots, or why no plan is valid. +pub(crate) type MemoSearchResult = + Result<(Selection, PhysicalCandidate, Vec), SelectionError>; + +/// The cheapest valid plan of `per_computation`'s memo under cost model 1, +/// as [`super::search::brute_force`] returns it, with the solver's gap, +/// iterations and time in the selection's report +/// ([`crate::SearchReport::solver_gap`]). `None` while the model is not +/// built: the caller runs the legacy search and reports +/// [`crate::RoutedToLegacy::MilpNotBuilt`]. +pub(crate) fn search( + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Option { + let _ = (per_computation, entries, data, models); + None +} diff --git a/crates/plan-selection/src/pass/subexpr/mod.rs b/crates/plan-selection/src/pass/subexpr/mod.rs index 61dcea355..54a8bbaf0 100644 --- a/crates/plan-selection/src/pass/subexpr/mod.rs +++ b/crates/plan-selection/src/pass/subexpr/mod.rs @@ -5,7 +5,7 @@ //! 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 +//! Behind [`crate::pass::Stage1ToStage3Search::search`] until it //! selects the exhaustive search's price on every equivalence test. //! //! - [`groups`] (W1): the hash-consed group DAG and Pass 1's @@ -19,11 +19,14 @@ //! - [`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. +//! bound, the dynamic program, build-late; +//! - [`milp`] (integrator): the memo solved as a mixed-integer linear +//! program. #![allow(dead_code)] pub(crate) mod capability; pub(crate) mod groups; +pub(crate) mod milp; pub(crate) mod pricing; pub(crate) mod scans; pub(crate) mod search; diff --git a/docs/user_guide_docs/stage-viewer.md b/docs/user_guide_docs/stage-viewer.md index c13758ca2..fb33a0712 100644 --- a/docs/user_guide_docs/stage-viewer.md +++ b/docs/user_guide_docs/stage-viewer.md @@ -330,6 +330,7 @@ cargo run -p asap-devtools --bin stage_pipeline -- \ | `--measured-operation-costs ` | A JSON table of measured CPU time per sketch operation (`[{"sketch": …, "operation": "update", "cpu_ns_per_operation": …}]`, as `convert_rqe_atomic_costs` writes it). A sketch node whose parameters match an entry exactly is priced from it; each node's `price_source` in the document says `measured` or `analytical`. | | `--enforce-accuracy-guarantees` | Reject a plan whose accuracy guarantee does not meet its query's target, as `plan_installable` does with the default settings. Without it (the default), every query output is accepted regardless of its guarantee, the plans are compared on cost alone, and each output's line says `regardless of guarantee`, plus `ignoring a failed guarantee check ()` when its guarantee would not have met the target; the document's `accuracy_guarantees` is `ignored` or `enforced`. With `--plan-selection-settings`, this flag decides, whatever the file's `ignore_accuracy_guarantees` says. | | `--exhaustive` | Compose and price every candidate instead of running the default search, and list them all. Use it for small workloads or to compare with the default search. | +| `--search ` | The default search's algorithm: `legacy` (the default), `memo-brute-force` (every plan of the memo over logical sub-expressions, each priced by Stage 3: the reference) or `memo-milp` (the memo solved as a mixed-integer linear program). The document's `search` names it (`default` for `legacy`), and the search report says which algorithm ran (`algorithm`, omitted for `legacy`) and, when the legacy search ran instead, why (`routed_to_legacy`: the MILP prices cost model 1 only, or is not built yet). Not with `--exhaustive`. | | `--max-displayed ` | With `--exhaustive`, how many plans the document shows (default 64), in enumeration order, plus the selected plan. The planner still optimizes and prices every plan, so the selection does not depend on it. | | `--max-combinations ` | With `--exhaustive` only: the most candidates a run plans (default 10000), counted at Stage 1 and again at Stage 2; a larger workload is refused with exit status 3 and no document. The default search has no limit. | | `--out ` | Where to write the document | From 3ffb7411234c415431f4106b6731ed19f0fc8f98 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:36:49 +0000 Subject: [PATCH 27/67] Memo MILP search: lazy loop over the model with exact pricing of each group choice Solve the model for a lower bound, price the solution's group choice exactly as the reference search does, cut that group choice, and stop when the bound reaches the incumbent. --- crates/plan-selection/src/lib.rs | 9 + .../plan-selection/src/pass/subexpr/milp.rs | 271 +++++++++++++++--- .../plan-selection/src/pass/subexpr/search.rs | 149 ++++++---- 3 files changed, 328 insertions(+), 101 deletions(-) diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index e9244c094..320928334 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1207,6 +1207,9 @@ pub enum SelectionError { }, #[error("no valid candidate: {0:?}")] NoValidCandidate(Vec), + /// The memo search's MILP solver failed. + #[error("MILP solver: {0}")] + Solver(String), } /// Stage 3 (plan selection): reject candidates that the executor cannot @@ -2068,6 +2071,12 @@ impl<'a> PlanSelector<'a> { /// Under a price ranking, fold the price of valid candidate `cost` /// into the cheapest price and the runner-up ([`ROUNDING`]). + /// Under a price ranking, the cheapest price of the valid candidates + /// offered so far. + pub(crate) fn cheapest_price(&self) -> Option { + self.cheapest + } + fn track_runner_up(&mut self, cost: &CandidateCost) { let (Ranking::Price, Some(price)) = (self.ranking, cost.cost_per_second) else { return; diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 8eab9a195..4358dc375 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -18,8 +18,13 @@ use good_lp::{ constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, WithMipGap, }; -use super::GroupId; +use super::{GroupId, Memo}; +use crate::{CostModel, PlanSelector, PlanningModels, Selection, SelectionError}; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; +use asap_physical_asap_optimization::PhysicalCandidate; +use asap_types::ir::QueryRoot; use asap_types::physical::Materialization; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; /// How a child group of a choice is reached: by the consumers reaching the /// choice's own node, or by one of its materialized private slots. @@ -109,8 +114,9 @@ pub(crate) struct ModelInput { /// The union rate of the classes in each set, by sorted class set: how /// often a node built once for queries of those classes is computed. pub union_rate: HashMap, f64>, - /// Choice sets excluded by earlier cuts (no-good cuts). - pub excluded: Vec>, + /// No-good cuts from earlier iterations: at most `bound` of these + /// choices together (a plan whose every group took one of them). + pub excluded: Vec<(Vec, f64)>, /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all /// of a choice set is chosen (local cuts). pub extra_bounds: Vec<(Vec, f64)>, @@ -172,10 +178,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> // it reaches the group and the group is not built once (or it is the // group's own built-once consumer). let mut passes: HashMap<(Consumer, usize), Variable> = HashMap::new(); - for (&(k, g), _) in &reach { + for &(k, g) in reach.keys() { passes.insert((k, g), vars.add(variable().min(0).max(1))); } - for (&g, _) in &once { + for &g in once.keys() { let k = Consumer::Group(g); passes .entry((k, g.0 as usize)) @@ -185,7 +191,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let mut y: HashMap<(Consumer, usize), Variable> = HashMap::new(); for (c, choice) in input.choices.iter().enumerate() { let g = choice.group.0 as usize; - for (&(k, h), _) in &passes { + for &(k, h) in passes.keys() { if h == g { y.insert((k, c), vars.add(variable().min(0).max(1))); } @@ -308,7 +314,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> model = model.with(constraint!(need[root.0 as usize] == 1)); } } - for g in 0..n { + for (g, &needed) in need.iter().enumerate() { // One choice implementing it (its own, or a fused one covering // it) when needed, none otherwise. let mut implementing = Expression::from(0.0); @@ -317,7 +323,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> implementing += x[c]; } } - model = model.with(constraint!(implementing == need[g])); + model = model.with(constraint!(implementing == needed)); } for (c, choice) in input.choices.iter().enumerate() { for child in &choice.children { @@ -414,7 +420,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let h = child.0 as usize; match via { Via::Direct => { - for (&(k, g), _) in &passes { + for &(k, g) in passes.keys() { if g != choice.group.0 as usize { continue; } @@ -452,12 +458,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> model = model.with(constraint!(latency <= bound)); } // Cuts. - for set in &input.excluded { + for (set, bound) in &input.excluded { let mut chosen = Expression::from(0.0); for &c in set { chosen += x[c]; } - model = model.with(constraint!(chosen <= set.len() as f64 - 1.0)); + model = model.with(constraint!(chosen <= *bound)); } for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { let mut chosen = Expression::from(0.0); @@ -542,6 +548,221 @@ fn possible_consumers(input: &ModelInput) -> BTreeMap> consumers } +/// A memo search's result: the selection, its physical candidate and its +/// logical roots, or why no plan is valid. +pub(crate) type MemoSearchResult = + Result<(Selection, PhysicalCandidate, Vec), SelectionError>; + +/// The cheapest valid plan of `per_computation`'s memo under cost model 1 +/// (`None` under cost models 2 and 3, which the caller routes to the +/// legacy search). +/// +/// Lazy loop: solve the model (a lower bound: it omits memory held while a +/// query runs, and checks composed guarantees only on its incumbents); +/// price the solution's group choice exactly, as the reference search +/// prices one ([`offer_group_choice`]: every threshold, identical-expression +/// variant and Stage 2 materialization), which updates the incumbent; add +/// a no-good cut excluding that group choice; stop when the model's bound +/// reaches the incumbent. Each group choice is priced at most once, so the +/// loop ends, and the incumbent is the reference search's optimum. +pub(crate) fn search( + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Option { + if !matches!( + models.plan_selection.cost_model, + CostModel::UsageOverTime(_) + ) { + return None; + } + Some(solve_memo(per_computation, entries, data, models)) +} + +fn solve_memo( + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> MemoSearchResult { + let started = std::time::Instant::now(); + let memo = super::search::memo(per_computation, entries); + let memo_stats = super::search::stats(&memo, started.elapsed().as_secs_f64() * 1e3); + let scale = super::pricing::frozen_calibration(&memo, entries, data, models)?; + let mut cache = super::pricing::PriceCache::new(scale); + let mut choices: Vec = Vec::new(); + for g in 0..memo.groups.len() { + let (group_choices, _lazy) = super::coefficients::choices( + &memo, + GroupId(g as u32), + entries, + data, + models, + &mut cache, + )?; + choices.extend(group_choices); + } + let input = model_input(&memo, choices, entries, models)?; + let mut input = input; + let mut selector = PlanSelector::new(entries, data, models, false)?; + let mut logical: HashMap> = HashMap::new(); + let mut number = 0usize; + let mut iterations = 0u32; + let mut bound = None; + loop { + iterations += 1; + let Some(solution) = solve(&input).map_err(SelectionError::Solver)? else { + break; + }; + bound = Some(solution.objective); + if let Some(incumbent) = selector.cheapest_price() { + if solution.objective >= incumbent - 1e-9 * incumbent.abs() { + break; + } + } + // The solution's group choice: the candidate each chosen choice + // takes for its group. + let mut group_choice: super::search::GroupChoice = vec![None; memo.groups.len()]; + for &c in &solution.chosen { + group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); + } + super::search::offer_group_choice( + &memo, + per_computation, + &group_choice, + entries, + data, + models, + &mut selector, + &mut logical, + &mut number, + )?; + // Exclude it: every choice of each chosen (group, candidate). + let groups = group_choice.iter().flatten().count(); + let excluded: Vec = input + .choices + .iter() + .enumerate() + .filter(|(_, choice)| group_choice[choice.group.0 as usize] == Some(choice.candidate)) + .map(|(c, _)| c) + .collect(); + input.excluded.push((excluded, groups as f64 - 1.0)); + } + let (mut selection, selected) = selector.finish()?; + let price = selection.costs[&selection.selected].cost_per_second; + selection.search.memo = Some(memo_stats); + selection.search.solver_iterations = iterations; + selection.search.solver_time_ms = started.elapsed().as_millis() as u64; + selection.search.solver_gap = match (price, bound) { + (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), + (Some(_), Some(_)) => Some(0.0), + _ => None, + }; + let roots = logical + .remove(&selected.from_logical) + .expect("the selection's logical roots are kept"); + Ok((selection, selected, roots)) +} + +/// The model's input for `memo`'s `choices`. +fn model_input( + memo: &Memo, + choices: Vec, + entries: &[QueryWorkloadEntry], + models: PlanningModels<'_>, +) -> Result { + let settings = models.plan_selection; + let rate = |entries: &[&QueryWorkloadEntry]| { + crate::shared_evaluation_rate(entries, settings).map_err(|error| SelectionError::Cost { + candidate: "memo".to_string(), + node: Default::default(), + error, + }) + }; + // Schedule classes: queries with one recurrence and range evaluation. + let mut classes: Vec<&QueryWorkloadEntry> = Vec::new(); + let schedule: Vec = entries + .iter() + .map(|entry| { + match classes.iter().position(|c| { + c.recurrence == entry.recurrence && c.range_evaluation == entry.range_evaluation + }) { + Some(class) => class, + None => { + classes.push(entry); + classes.len() - 1 + } + } + }) + .collect(); + let mut union_rate = HashMap::new(); + let mut queries = Vec::with_capacity(entries.len()); + for (q, entry) in entries.iter().enumerate() { + let roots = memo + .groups + .iter() + .enumerate() + .filter(|(_, group)| group.roots_of.contains(&q)) + .map(|(g, _)| GroupId(g as u32)) + .collect(); + // A query's evaluation must fit its latency bound and its repeat + // interval. + let interval = entry.recurrence.fixed_interval_ms().map(|ms| ms as f64); + let max_latency_ms = match (entry.requirements.max_latency_ms, interval) { + (Some(a), Some(b)) => Some(a.min(b)), + (a, b) => a.or(b), + }; + queries.push(ModelQuery { + rate: rate(&[entry])?, + roots, + max_latency_ms, + }); + } + let group_queries: Vec> = memo.groups.iter().map(|g| g.queries.clone()).collect(); + for queries_of in &group_queries { + let mut present: Vec = queries_of.iter().map(|&q| schedule[q]).collect(); + present.sort_unstable(); + present.dedup(); + for mask in 1u64..(1 << present.len()) { + let set: Vec = (0..present.len()) + .filter(|i| mask >> i & 1 == 1) + .map(|i| present[i]) + .collect(); + if let std::collections::hash_map::Entry::Vacant(slot) = union_rate.entry(set) { + let members: Vec<&QueryWorkloadEntry> = + slot.key().iter().map(|&c| classes[c]).collect(); + slot.insert(rate(&members)?); + } + } + } + let group_slots = memo + .groups + .iter() + .enumerate() + .filter(|(_, group)| group.parents.len() > 1 || group.roots_of.len() > 1) + .map(|(g, _)| { + let allowed = &memo.candidates[g][0].materializations; + GroupSlot { + group: GroupId(g as u32), + may_recompute: allowed.is_empty() + || allowed.contains(&Materialization::NotMaterialized), + } + }) + .collect(); + Ok(ModelInput { + group_count: memo.groups.len(), + group_queries, + choices, + group_slots, + queries, + schedule, + union_rate, + excluded: Vec::new(), + extra_bounds: Vec::new(), + }) +} + #[cfg(test)] mod tests { use super::*; @@ -618,31 +839,3 @@ mod tests { ); } } - -use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; -use asap_physical_asap_optimization::PhysicalCandidate; -use asap_types::ir::QueryRoot; -use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; - -use crate::{PlanningModels, Selection, SelectionError}; - -/// A memo search's result: the selection, its physical candidate and its -/// logical roots, or why no plan is valid. -pub(crate) type MemoSearchResult = - Result<(Selection, PhysicalCandidate, Vec), SelectionError>; - -/// The cheapest valid plan of `per_computation`'s memo under cost model 1, -/// as [`super::search::brute_force`] returns it, with the solver's gap, -/// iterations and time in the selection's report -/// ([`crate::SearchReport::solver_gap`]). `None` while the model is not -/// built: the caller runs the legacy search and reports -/// [`crate::RoutedToLegacy::MilpNotBuilt`]. -pub(crate) fn search( - per_computation: &PerComputationCandidates, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: PlanningModels<'_>, -) -> Option { - let _ = (per_computation, entries, data, models); - None -} diff --git a/crates/plan-selection/src/pass/subexpr/search.rs b/crates/plan-selection/src/pass/subexpr/search.rs index df2d84b29..27ec1177f 100644 --- a/crates/plan-selection/src/pass/subexpr/search.rs +++ b/crates/plan-selection/src/pass/subexpr/search.rs @@ -459,11 +459,83 @@ pub(super) fn resized_builds( .collect() } +/// Offer `selector` every plan of the group choice `choice`: each +/// combination of its sites' thresholds, composed, its guarantees +/// recomposed, with each identical-expression and scan-widening variant +/// the exhaustive search lists, in every Stage 2 materialization. Logical +/// roots of the plans `selector` keeps are kept in `logical`, by id +/// (`L`). `false` when the choice breaks a joint rule (nothing +/// offered). +#[allow(clippy::too_many_arguments)] +pub(super) fn offer_group_choice( + memo: &Memo, + per_computation: &PerComputationCandidates, + choice: &GroupChoice, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + selector: &mut PlanSelector<'_>, + logical: &mut HashMap>, + number: &mut usize, +) -> Result { + 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) { + return Ok(false); + } + // 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) + { + return Ok(false); + } + for plan in at_thresholds(memo, choice) { + 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; + }; + // 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; + }; + 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)); + } + } + Ok(true) +} + /// 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. +/// [`at_thresholds`]), each offered as [`offer_group_choice`] offers it; +/// the cheapest valid one, with its logical roots. pub(crate) fn brute_force( per_computation: &PerComputationCandidates, entries: &[QueryWorkloadEntry], @@ -478,64 +550,17 @@ pub(crate) fn brute_force( 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; - } - // 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; - }; - // 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; - }; - // 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; - }; - 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)); - } - } + offer_group_choice( + &memo, + per_computation, + choice, + entries, + data, + models, + &mut selector, + &mut logical, + &mut number, + )?; } let (mut selection, selected) = selector.finish()?; selection.search.memo = Some(memo_stats); From 2a1bca09e14a4b88f37b2daab6a64ec4fc880e24 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:38:06 +0000 Subject: [PATCH 28/67] wip: memo choices shared or apart per group, scan children apart or shared at a horizon (tests not yet passing) --- .../src/pass/subexpr/coefficients.rs | 578 +++++++++++++----- .../plan-selection/src/pass/subexpr/milp.rs | 14 + 2 files changed, 424 insertions(+), 168 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index fceaf7a17..aae394c25 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -34,7 +34,7 @@ //! chain's slot is charged per unit of its group's union rate, while //! Stage 3 computes it once per request (range steps not counted). -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::rc::Rc; use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate}; @@ -136,9 +136,24 @@ pub(super) fn choice_parts( models: PlanningModels<'_>, cache: &mut PriceCache, ) -> Result, SelectionError> { + let node = &memo.groups[group.0 as usize].node; let queries = &memo.groups[group.0 as usize].queries; - let group_entries: Vec = - queries.iter().map(|&q| entries[q].clone()).collect(); + let is_scan = matches!(node.non_asap(), Some(NonASAPOp::Scan { .. })); + // Shared (one node for every query), and for a group several queries + // read, apart (a copy per query); a scan apart is its reader's. + let mut modes: Vec> = vec![None]; + if queries.len() > 1 && !is_scan { + modes.extend(queries.iter().map(|&q| Some(q))); + } + // A shared scan reads the longest window of its readers: one choice + // per horizon a reader may read. + let own_horizons: Vec> = match node.non_asap() { + Some(NonASAPOp::Scan { source, .. }) => raw_windows(memo) + .get(&super::scans::raw_horizon_key(source)) + .map(|windows| windows.iter().map(|&w| Some(w)).collect()) + .unwrap_or_else(|| vec![None]), + _ => vec![None], + }; let mut out = Vec::new(); let versions: Vec<(usize, ImplementationCandidate, SiteThresholds)> = memo.candidates [group.0 as usize] @@ -150,89 +165,168 @@ pub(super) fn choice_parts( .map(move |(version, thresholds)| (index, version, thresholds)) }) .collect(); - for (index, candidate, thresholds) in &versions { - let (index, candidate) = (*index, candidate); - let roots = vec![QueryRoot::Operator(Rc::clone(&candidate.node)); queries.len().max(1)]; - let Ok(physicals) = optimize_physical_asap( - "S", - &roots, - &group_entries, - data, - models.deployment_capabilities, - ) else { - continue; + for copy_of in modes { + let built_for: Vec = match copy_of { + Some(q) => vec![q], + None => queries.clone(), }; - let child_groups = children_of(memo, group, candidate); - - let mut children: HashSet = HashSet::new(); - let mut child_roots: Vec> = Vec::new(); - for &child in &child_groups { - let keys = child_keys(memo, child, entries, data, &models, cache); - children.extend(keys.nodes.iter().copied()); - child_roots.push(keys.root); - } - let mut seen: Vec> = Vec::new(); - for physical in &physicals { - let keys = keys_of(&physical.dag, cache.keys_mut()); - let shallow = shallow_keys(&physical.dag, cache.keys_mut()); - let private: Vec = keys.iter().map(|k| !children.contains(k)).collect(); - let builds: Vec = (0..physical.dag.nodes.len()) - .filter(|&i| private[i] && is_build(&physical.dag.nodes[i].payload)) - .collect(); - // Only the private builds are the choice's; any other - // materialized node (the group's own node built once, a - // child's build) is another variable of the model. - let foreign = physical.dag.nodes.iter().enumerate().any(|(i, n)| { - n.materialization != Materialization::NotMaterialized && !builds.contains(&i) - }); - if foreign { + let part_entries: Vec = + built_for.iter().map(|&q| entries[q].clone()).collect(); + for (index, candidate, thresholds) in &versions { + let (index, candidate) = (*index, candidate); + let roots = + vec![QueryRoot::Operator(Rc::clone(&candidate.node)); built_for.len().max(1)]; + let Ok(physicals) = optimize_physical_asap( + "S", + &roots, + &part_entries, + data, + models.deployment_capabilities, + ) else { continue; - } - let assignment: Vec = builds + }; + let child_groups = children_of(memo, group, candidate); + let child_key_sets: Vec> = child_groups .iter() - .map(|&i| physical.dag.nodes[i].materialization) + .map(|&child| child_keys(memo, child, entries, data, &models, cache)) .collect(); - if seen.contains(&assignment) { - continue; + // Each scan child: apart (its nodes this choice's) or shared at + // a horizon no shorter than this choice's window on it. + let scan_children: Vec<(usize, u64, Vec>)> = child_groups + .iter() + .enumerate() + .filter_map( + |(i, child)| match memo.groups[child.0 as usize].node.non_asap() { + Some(NonASAPOp::Scan { source, .. }) => { + let key = super::scans::raw_horizon_key(source); + let mut variants = vec![None]; + variants.extend( + raw_windows(memo) + .get(&key) + .into_iter() + .flatten() + .map(|&w| Some(w)), + ); + Some((i, key, variants)) + } + _ => None, + }, + ) + .collect(); + let mut combinations: Vec>> = vec![Vec::new()]; + for (_, _, variants) in &scan_children { + combinations = combinations + .into_iter() + .flat_map(|chosen| { + variants.iter().map(move |v| { + let mut chosen = chosen.clone(); + chosen.push(*v); + chosen + }) + }) + .collect(); + } + for own_horizon in &own_horizons { + for combination in &combinations { + // Children whose nodes are not this choice's, and the + // horizon each shared scan child reads. + let mut horizons: HashMap = HashMap::new(); + let mut absorbed: Vec = Vec::new(); + for ((i, key, _), variant) in scan_children.iter().zip(combination) { + match variant { + Some(h) => { + horizons.insert(*key, *h); + } + None => absorbed.push(*i), + } + } + let mut children: HashSet = HashSet::new(); + let mut child_roots: Vec> = Vec::new(); + let mut kept_children: Vec = Vec::new(); + for (i, keys) in child_key_sets.iter().enumerate() { + if absorbed.contains(&i) { + continue; + } + children.extend(keys.nodes.iter().copied()); + child_roots.push(keys.root); + kept_children.push(child_groups[i]); + } + let mut seen: Vec> = Vec::new(); + for physical in &physicals { + let keys = keys_of(&physical.dag, cache.keys_mut()); + let shallow = shallow_keys(&physical.dag, cache.keys_mut()); + let private: Vec = + keys.iter().map(|k| !children.contains(k)).collect(); + let builds: Vec = (0..physical.dag.nodes.len()) + .filter(|&i| private[i] && is_build(&physical.dag.nodes[i].payload)) + .collect(); + // Only the private builds are the choice's; any + // other materialized node (the group's own node + // built once, a child's build) is another variable + // of the model. + let foreign = physical.dag.nodes.iter().enumerate().any(|(i, n)| { + n.materialization != Materialization::NotMaterialized + && !builds.contains(&i) + }); + if foreign { + continue; + } + let assignment: Vec = builds + .iter() + .map(|&i| physical.dag.nodes[i].materialization) + .collect(); + if seen.contains(&assignment) { + continue; + } + seen.push(assignment.clone()); + let Some(parts) = coefficients_of( + memo, + group, + index, + candidate, + physical, + &keys, + &shallow, + &private, + &child_roots, + &horizons, + own_horizon.or_else(|| scan_lookback_ms(memo, group)), + &part_entries, + data, + &models, + )? + else { + continue; + }; + out.push(ChoiceParts { + choice: Choice { + group, + candidate: index, + implements: candidate.implements.clone(), + children: kept_children.clone(), + assignment, + coefficients: parts.coefficients, + copy_of, + scan_children: scan_children + .iter() + .zip(combination) + .map(|((i, _, _), v)| (child_groups[*i], *v)) + .collect(), + scan_horizon_ms: *own_horizon, + }, + lazy: LazyTerms { + thresholds: thresholds.clone(), + ..parts.lazy + }, + root_key: parts.root_key, + root_producer_keys: parts.root_producer_keys, + slot_keys: parts.slot_keys, + build_keys: builds.iter().map(|&i| shallow[i]).collect(), + part_keys: parts.part_keys, + }); + } + } } - seen.push(assignment.clone()); - let Some(parts) = coefficients_of( - memo, - group, - index, - candidate, - physical, - &keys, - &shallow, - &private, - &child_roots, - &horizons_of(memo, thresholds), - &group_entries, - data, - &models, - )? - else { - continue; - }; - out.push(ChoiceParts { - choice: Choice { - group, - candidate: index, - implements: candidate.implements.clone(), - children: child_groups.clone(), - assignment, - coefficients: parts.coefficients, - }, - lazy: LazyTerms { - thresholds: thresholds.clone(), - ..parts.lazy - }, - root_key: parts.root_key, - root_producer_keys: parts.root_producer_keys, - slot_keys: parts.slot_keys, - build_keys: builds.iter().map(|&i| shallow[i]).collect(), - part_keys: parts.part_keys, - }); } } Ok(out) @@ -240,10 +334,9 @@ pub(super) fn choice_parts( /// `candidate` rewritten at each combination of thresholds of its coupling /// sites that it may read, as the brute-force search rewrites a plan's -/// candidates, with the index of each threshold. At a raw-horizon -/// threshold the candidate is unchanged and its query-time scans of the -/// site's source read the threshold's horizon ([`horizons_of`]). A -/// candidate without sites is itself. +/// candidates, with the index of each threshold. A raw-horizon site has +/// no threshold here: scan sharing is part of each choice +/// ([`Choice::scan_children`]). A candidate without sites is itself. fn at_site_thresholds( memo: &Memo, candidate: &ImplementationCandidate, @@ -304,19 +397,10 @@ fn at_site_thresholds( next.extend(at(index, rewritten)); } } - CouplingKind::RawHorizon => { - let own: Vec = scans::raw_readers(memo, &version) - .into_iter() - .filter(|r| r.key == site.key) - .collect(); - for (index, threshold) in - scans::horizon_thresholds(memo, site).iter().enumerate() - { - if own.iter().all(|r| scans::may_read_horizon(threshold, r)) { - next.extend(at(index, Some(version.clone()))); - } - } - } + // A raw horizon is the look-back of a shared scan: each + // choice prices its scan children shared at a horizon or + // apart ([`Choice::scan_children`]). + CouplingKind::RawHorizon => next.push((version, chosen)), } } versions = next; @@ -324,17 +408,56 @@ fn at_site_thresholds( versions } -/// The raw-data horizon of each source key `thresholds` fixes, in ms. -fn horizons_of(memo: &Memo, thresholds: &[(CouplingSite, usize)]) -> HashMap { - thresholds - .iter() - .filter(|(site, _)| site.kind == CouplingKind::RawHorizon) - .filter_map(|(site, index)| { - super::scans::horizon_thresholds(memo, site) - .get(*index) - .map(|t| (site.key, t.horizon_ms)) - }) - .collect() +/// Every window a raw scan of each source may be read for, by the source's +/// raw-horizon key: each range over it (through time shifts), and each +/// pane chain over it, kept across evaluations (its pane width) or not +/// (its panes), in every candidate and every chain threshold of the memo. +fn raw_windows(memo: &Memo) -> HashMap> { + use asap_types::ir::OperatorNode; + let mut windows: HashMap> = HashMap::new(); + let scan_below = |node: &Rc| -> Option { + let mut node = Rc::clone(node); + loop { + match node.non_asap() { + Some(NonASAPOp::Scan { source, .. }) => { + return Some(super::scans::raw_horizon_key(source)) + } + Some(NonASAPOp::TimeShift { child, .. }) => node = Rc::clone(child), + _ => return None, + } + } + }; + for candidate in memo.candidates.iter().flatten() { + for (version, _) in at_site_thresholds(memo, candidate) { + for node in OperatorNode::reachable(&version.node) { + match (node.non_asap(), node.asap()) { + (Some(NonASAPOp::TimeRange { range, child, .. }), _) => { + if let Some(key) = scan_below(child) { + let ms = u64::try_from(range.as_millis()).unwrap_or(u64::MAX); + windows.entry(key).or_default().insert(ms); + } + } + ( + _, + Some(ASAPOp::PaneChain { + child, + pane_width_ms, + pane_count, + .. + }), + ) => { + if let Some(key) = scan_below(child) { + let entry = windows.entry(key).or_default(); + entry.insert(*pane_width_ms); + entry.insert(pane_width_ms.saturating_mul(u64::from(*pane_count))); + } + } + _ => {} + } + } + } + } + windows } /// The child groups of `candidate`: its own, and every other group whose @@ -489,6 +612,7 @@ fn coefficients_of( private: &[bool], child_roots: &[Option], horizons: &HashMap, + scan_lookback: Option, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: &PlanningModels<'_>, @@ -509,7 +633,6 @@ fn coefficients_of( if !private[position[&root]] { return Ok(None); } - let scan_lookback = scan_lookback_ms(memo, group); let priced = per_computation(dag, entries, scan_lookback, horizons, data, models)?; let mut producers_of: HashMap> = HashMap::new(); for e in &dag.edges { @@ -920,6 +1043,19 @@ mod tests { ) .ok()?; let expected = cost.cost_per_second?; + let debug = std::env::var("W2DBG").is_ok(); + if debug { + for (id, n) in &cost.per_node { + eprintln!( + "stage3 {id} {:?} {:?} {:?} x{} {}", + n.cost_per_second, + nodes[id].output_state.timing, + nodes[id].materialization, + n.multiplier, + n.detail + ); + } + } let rates = computations( dag, @@ -947,6 +1083,49 @@ mod tests { }; let query_rate = |q: usize| crate::shared_evaluation_rate(&[&entries[q]], settings).unwrap(); + // Each node's producers, and the nodes below a node. + let producers = |id: PhysicalASAPNodeId| -> Vec { + dag.edges + .iter() + .filter(|e| e.consumer == id) + .map(|e| e.producer) + .collect() + }; + let consumers = + |id: PhysicalASAPNodeId| dag.edges.iter().filter(|e| e.producer == id).count(); + let below = |root: PhysicalASAPNodeId| -> Vec { + let mut seen = HashSet::new(); + let mut stack = vec![root]; + let mut out = Vec::new(); + while let Some(id) = stack.pop() { + if seen.insert(id) { + out.push(id); + stack.extend(producers(id)); + } + } + out + }; + // The queries whose plan reaches each node at all. + let starts = crate::query_starts(dag); + let reaches_at_all = |id: PhysicalASAPNodeId| -> Vec { + starts + .iter() + .enumerate() + .filter(|(_, start)| start.iter().any(|s| below(*s).contains(&id))) + .map(|(q, _)| q) + .collect() + }; + let kept = kept_pane_chains(dag); + // A scan's horizon in the plan, or `None` when one reader has it + // (apart, its reader's). + let scan_variant = |s: PhysicalASAPNodeId| -> Option { + (consumers(s) > 1).then(|| { + windows_read(dag, &nodes, s, &kept) + .into_iter() + .max() + .unwrap_or(0) + }) + }; // Each needed group's choice in this plan: its candidate, with the // plan's materialization of its private builds. struct Placed<'p> { @@ -960,84 +1139,147 @@ mod tests { } let mut placed: Vec = Vec::new(); let mut used: HashSet = HashSet::new(); - // The plan's groups: the chosen ones, and the children their - // choices read that the memo's candidates do not list (their first - // candidate). + // The plan's groups, parents first: the chosen ones, and the + // children their choices read that the memo's candidates do not + // list (their first candidate). let mut chosen: Vec> = choice.to_vec(); - let mut queue: Vec = (0..chosen.len()).filter(|&g| chosen[g].is_some()).collect(); - let mut order = Vec::new(); - while let Some(g) = queue.pop() { - order.push(g); - let c = chosen[g].expect("queued groups are chosen"); - for part in parts[g].iter().filter(|p| p.choice.candidate == c).take(1) { - for child in &part.choice.children { - let h = child.0 as usize; - if chosen[h].is_none() { - chosen[h] = Some(0); - queue.push(h); - } + for g in (0..chosen.len()).rev() { + let Some(c) = chosen[g] else { continue }; + for part in parts[g].iter().filter(|p| p.choice.candidate == c) { + for child in part + .choice + .children + .iter() + .chain(part.choice.scan_children.iter().map(|(h, _)| h)) + { + chosen[child.0 as usize].get_or_insert(0); } } } - order.sort_unstable(); - for g in order { - let c = chosen[g].expect("chosen"); - let candidates: Vec<&ChoiceParts> = parts[g] - .iter() - .filter(|p| p.choice.candidate == c) - .collect(); - let part = candidates.into_iter().find(|p| { - p.build_keys - .iter() - .zip(&p.choice.assignment) - .all(|(key, m)| { - by_key(*key) - .iter() - .any(|id| nodes[id].materialization == *m) - }) - }); - let part = part?; + for g in (0..chosen.len()).rev() { + let Some(c) = chosen[g] else { continue }; + let Some(sample) = parts[g].iter().find(|p| p.choice.candidate == c) else { + return None; + }; // The plan's copies of its node: equal shallow keys, reading // equal private producers. - let copies: Vec = by_key(part.root_key) + let copies: Vec = by_key(sample.root_key) .into_iter() .filter(|id| !used.contains(id)) .filter(|id| { - let producers: Vec = dag - .edges + let keys_below: Vec = producers(*id) .iter() - .filter(|e| e.consumer == *id) - .map(|e| keys[position_of[&e.producer]]) + .map(|p| keys[position_of[p]]) .collect(); - part.root_producer_keys + sample + .root_producer_keys .iter() - .all(|k| producers.contains(k)) + .all(|k| keys_below.contains(k)) }) .collect(); + if copies.is_empty() { + // A scan every reader has apart. + if sample.choice.group.0 as usize == g + && parts[g] + .iter() + .all(|p| p.choice.scan_horizon_ms.is_some() || p.choice.copy_of.is_none()) + { + continue; + } + return None; + } used.extend(copies.iter().copied()); let timing = |id: &PhysicalASAPNodeId| nodes[id].output_state.timing; - let query_time = copies + let query_time: Vec = copies .iter() .copied() - .find(|id| timing(id) == ExecutionTiming::QueryTime); + .filter(|id| timing(id) == ExecutionTiming::QueryTime) + .collect(); let at_ingestion: Vec = copies .iter() .copied() .filter(|id| timing(id) == ExecutionTiming::IngestionTime) .collect(); - let root = query_time.or(at_ingestion.first().copied())?; - let slots = part - .slot_keys - .iter() - .map(|key| by_key(*key).first().copied()) - .collect::>>()?; - placed.push(Placed { - part, - root, - at_query_time: query_time.is_some(), - at_ingestion, - slots, - }); + let apart = query_time.len() > 1; + let roots: Vec<(PhysicalASAPNodeId, bool)> = match query_time.is_empty() { + true => vec![(at_ingestion[0], false)], + false => query_time.iter().map(|&id| (id, true)).collect(), + }; + for (k, &(root, at_query_time)) in roots.iter().enumerate() { + let copy_of = match apart { + true => Some(*reaches_at_all(root).first()?), + false => None, + }; + let mine = below(root); + let scans_below: Vec> = mine + .iter() + .filter(|id| { + matches!(nodes[id].payload, Operator::NonASAP(NonASAPOp::Scan { .. })) + }) + .filter(|id| { + producers(root).contains(id) + || mine.iter().any(|m| producers(*m).contains(id)) + }) + .map(|&s| scan_variant(s)) + .collect(); + let part = parts[g].iter().find(|p| { + p.choice.candidate == c + && p.choice.copy_of == copy_of + && p.build_keys + .iter() + .zip(&p.choice.assignment) + .all(|(key, m)| { + mine.iter().any(|id| { + keys[position_of[id]] == *key && nodes[id].materialization == *m + }) + }) + && p.choice + .scan_children + .iter() + .all(|(_, v)| scans_below.contains(v)) + && match p.choice.scan_horizon_ms { + None => true, + Some(h) => { + scan_variant(root).unwrap_or(h) == h + || windows_read(dag, &nodes, root, &kept).into_iter().max() + == Some(h) + } + } + })?; + // Apart scans below are this choice's. + for (_, v) in &part.choice.scan_children { + if v.is_none() { + for id in &mine { + if matches!( + nodes[id].payload, + Operator::NonASAP(NonASAPOp::Scan { .. }) + ) && consumers(*id) == 1 + { + used.insert(*id); + } + } + } + } + let slots = part + .slot_keys + .iter() + .map(|key| { + mine.iter() + .copied() + .find(|id| keys[position_of[id]] == *key) + }) + .collect::>>()?; + placed.push(Placed { + part, + root, + at_query_time, + at_ingestion: match k { + 0 => at_ingestion.clone(), + _ => Vec::new(), + }, + slots, + }); + } } let part_root = |p: &Placed, via: Via| match via { Via::Direct => p.root, @@ -1063,6 +1305,9 @@ mod tests { let mut actual = 0.0; let mut scans = Vec::new(); for p in &placed { + if debug { + eprintln!("placed g{} c{} copy {:?} scans {:?} h {:?} root {} qt {} ing {:?} slots {:?} direct {:?} slotsc {:?} ingc {:?} rate {}", p.part.choice.group.0, p.part.choice.candidate, p.part.choice.copy_of, p.part.choice.scan_children, p.part.choice.scan_horizon_ms, p.root, p.at_query_time, p.at_ingestion, p.slots, p.part.choice.coefficients.direct, p.part.choice.coefficients.slots, p.part.lazy.direct_at_ingestion, rates[&p.root].1); + } let coefficients = &p.part.choice.coefficients; if p.at_query_time { actual += coefficients.direct.cost_per_unit_rate * rates[&p.root].1; @@ -1141,9 +1386,6 @@ mod tests { /// chain (the chain site's thresholds), built at query or ingestion /// time. #[test] - #[ignore = "a pane chain over a raw scan is no raw-horizon reader: the shared \ - scan's look-back (the longest window of its readers) is not a site \ - threshold, so a chain's input rows depend on other groups"] fn coefficients_reproduce_stage_3_with_a_shared_pane_chain() { let checked = check( &["sum_over_time(m[5m])", "sum_over_time(m[1h])"], diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 8eab9a195..d6ea3b8e9 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -73,6 +73,17 @@ pub(crate) struct Choice { pub children: Vec, pub assignment: Vec, pub coefficients: ChoiceCoefficients, + /// `None`: the group's one node, shared by every query reading it. + /// `Some(q)`: query `q`'s own copy of it (the group apart); a group's + /// copies are all apart or all shared. + pub copy_of: Option, + /// Each child group that is a scan: `None` when it is apart and its + /// copy is this choice's own (not in `children`); `Some(h)` when it is + /// the scan shared at horizon `h` ms (its choice's `scan_horizon_ms`). + pub scan_children: Vec<(GroupId, Option)>, + /// For a scan group's choice: the horizon it reads, the longest window + /// among its readers. + pub scan_horizon_ms: Option, } /// A shared group's own node, built once per evaluation when chosen. @@ -561,6 +572,9 @@ mod tests { slots: Vec::new(), child_via: children.iter().map(|_| Via::Direct).collect(), }, + copy_of: None, + scan_children: Vec::new(), + scan_horizon_ms: None, } } From f901bdecfd127eae85d4600f1bce189cdb87f1c9 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:49:52 +0000 Subject: [PATCH 29/67] Memo MILP search: report the model size, the first round's bound and the solver's time; cost model 1 always runs the MILP --- crates/devtools/tests/stage_pipeline.rs | 11 ++- .../tests/joint_search_equivalence.rs | 14 +-- crates/plan-selection/src/lib.rs | 22 +++-- .../src/pass/stage1_to_stage3.rs | 8 +- .../plan-selection/src/pass/subexpr/milp.rs | 90 ++++++++++--------- 5 files changed, 83 insertions(+), 62 deletions(-) diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index 4ccc01fe1..1c72bfa8b 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1749,9 +1749,14 @@ fn search_flag_picks_the_search_algorithm() { let milp = run(&["--search", "memo-milp"]); assert_eq!(milp["search"], "memo-milp"); let report = &milp["stage3_plan_selection"]["search"]; - assert!(report.get("algorithm").is_none(), "{report}"); - assert_eq!(report["routed_to_legacy"], "milp_not_built"); - assert_eq!(price(&milp), price(&legacy)); + assert_eq!(report["algorithm"], "memo-milp"); + assert!(report.get("routed_to_legacy").is_none(), "{report}"); + assert!( + report["solver_iterations"].as_u64().unwrap() >= 1, + "{report}" + ); + assert!(report["solver_gap"].as_f64().unwrap() <= 1e-9, "{report}"); + assert!((price(&milp) - price(&legacy)).abs() <= 1e-9 * price(&legacy)); let cost_model_2 = run(&[ "--search", diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index cd34d4bb4..e494d0114 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1461,11 +1461,15 @@ fn three_way( }; let report = &selection.search; if report.algorithm == SearchAlgorithm::Legacy { - let expected = match cost_model { - CostModel::UsageOverTime(_) => RoutedToLegacy::MilpNotBuilt, - _ => RoutedToLegacy::CostModelNotSupported, - }; - assert_eq!(report.routed_to_legacy, Some(expected), "{label}"); + assert!( + !matches!(cost_model, CostModel::UsageOverTime(_)), + "{label}: cost model 1 is searched by the MILP" + ); + assert_eq!( + report.routed_to_legacy, + Some(RoutedToLegacy::CostModelNotSupported), + "{label}" + ); return Ok(false); } assert_eq!(report.algorithm, SearchAlgorithm::MemoMilp, "{label}"); diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 320928334..e372abd65 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1040,14 +1040,26 @@ pub struct SearchReport { /// proven optimal. `None` for every other search. #[serde(default, skip_serializing_if = "Option::is_none")] pub solver_gap: Option, - /// The MILP solver's simplex iterations; 0 (omitted in JSON) for every - /// other search. + /// The MILP solves of the lazy loop (one per round: a solve, the exact + /// price of its plan, a cut); 0 (omitted in JSON) for every other + /// search. #[serde(default, skip_serializing_if = "is_zero")] pub solver_iterations: u32, - /// The MILP solver's wall time in milliseconds; 0 (omitted in JSON) for - /// every other search. + /// The wall time in milliseconds spent in the MILP solver, every round + /// together; 0 (omitted in JSON) for every other search. #[serde(default, skip_serializing_if = "is_zero")] pub solver_time_ms: u64, + /// The MILP's variables and constraints in its last round; 0 (omitted + /// in JSON) for every other search. + #[serde(default, skip_serializing_if = "is_zero")] + pub solver_variables: u64, + #[serde(default, skip_serializing_if = "is_zero")] + pub solver_constraints: u64, + /// The first round's lower bound on the cost per second, before any + /// cut: how tight the model is against the selected price; `None` for + /// every other search. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub solver_first_bound_cost_per_second: Option, } /// The size of a sub-expression memo (joint-planning-search design, @@ -1110,8 +1122,6 @@ pub enum RoutedToLegacy { /// The MILP's objective is cost model 1's sum of local prices; cost /// models 2 and 3 are searched by the legacy search. CostModelNotSupported, - /// The MILP model is not built yet. - MilpNotBuilt, } /// 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 7861e10a6..5a80128f7 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -268,12 +268,8 @@ pub fn run_stage1_to_stage3_from_pass1( let pruned = per_computation .clone() .without_unsupported_summaries(models.deployment_capabilities); - match super::subexpr::milp::search(&pruned, entries, data, models) { - Some(found) => { - return memo_run(pruned, found, entries, data, SearchAlgorithm::MemoMilp) - } - None => Some(RoutedToLegacy::MilpNotBuilt), - } + let found = super::subexpr::milp::search(&pruned, entries, data, models); + return memo_run(pruned, found, entries, data, SearchAlgorithm::MemoMilp); } } }; diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 4358dc375..6e8271f87 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -19,7 +19,7 @@ use good_lp::{ }; use super::{GroupId, Memo}; -use crate::{CostModel, PlanSelector, PlanningModels, Selection, SelectionError}; +use crate::{PlanSelector, PlanningModels, Selection, SelectionError}; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; use asap_physical_asap_optimization::PhysicalCandidate; use asap_types::ir::QueryRoot; @@ -131,6 +131,9 @@ pub(crate) struct ModelSolution { pub built_once: BTreeSet, /// The model's objective: a lower bound on the plan's price. pub objective: f64, + /// The model's size. + pub variables: u64, + pub constraints: u64, } /// A consumer of the model: a query, a materialized private slot of a @@ -307,11 +310,13 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } - let mut model = vars.minimise(objective.clone()).using(highs); + let variable_count = vars.len() as u64; + let model = vars.minimise(objective.clone()).using(highs); + let mut constraints = Vec::new(); // Needed groups: every query's roots; children of chosen choices. for query in &input.queries { for root in &query.roots { - model = model.with(constraint!(need[root.0 as usize] == 1)); + constraints.push(constraint!(need[root.0 as usize] == 1)); } } for (g, &needed) in need.iter().enumerate() { @@ -323,18 +328,18 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> implementing += x[c]; } } - model = model.with(constraint!(implementing == needed)); + constraints.push(constraint!(implementing == needed)); } for (c, choice) in input.choices.iter().enumerate() { for child in &choice.children { - model = model.with(constraint!(need[child.0 as usize] >= x[c])); + constraints.push(constraint!(need[child.0 as usize] >= x[c])); } } for (&g, &v) in &once { - model = model.with(constraint!(v <= need[g.0 as usize])); + constraints.push(constraint!(v <= need[g.0 as usize])); let slot = input.group_slots.iter().find(|s| s.group == g).unwrap(); if !slot.may_recompute { - model = model.with(constraint!(v == need[g.0 as usize])); + constraints.push(constraint!(v == need[g.0 as usize])); } } // Reach at all: a query reaches its roots, and every child of a chosen @@ -342,7 +347,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for (q, query) in input.queries.iter().enumerate() { for root in &query.roots { if let Some(&v) = reaches_at_all.get(&(q, root.0 as usize)) { - model = model.with(constraint!(v == 1)); + constraints.push(constraint!(v == 1)); } } } @@ -354,7 +359,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, child.0 as usize)), ) { - model = model.with(constraint!(below >= above + x[c] - 1)); + constraints.push(constraint!(below >= above + x[c] - 1)); } } } @@ -367,7 +372,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for (_, selector) in sets { any += *selector; } - model = model.with(constraint!(any <= 1)); + constraints.push(constraint!(any <= 1)); for class in classes_of(g) { let mut containing = Expression::from(0.0); for (set, selector) in sets { @@ -377,7 +382,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } for &q in &input.group_queries[g] { if input.schedule[q] == class { - model = model.with(constraint!(containing.clone() >= reaches_at_all[&(q, g)])); + constraints.push(constraint!(containing.clone() >= reaches_at_all[&(q, g)])); } } let mut members = Expression::from(0.0); @@ -386,34 +391,34 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> members += reaches_at_all[&(q, g)]; } } - model = model.with(constraint!(containing <= members)); + constraints.push(constraint!(containing <= members)); } } // Consumer reach. for (q, query) in input.queries.iter().enumerate() { for root in &query.roots { if let Some(&v) = reach.get(&(Consumer::Query(q), root.0 as usize)) { - model = model.with(constraint!(v >= need[root.0 as usize])); + constraints.push(constraint!(v >= need[root.0 as usize])); } } } for (&(k, g), &p) in &passes { match k { Consumer::Group(own) if own.0 as usize == g => { - model = model.with(constraint!(p >= once[&own])); + constraints.push(constraint!(p >= once[&own])); } _ => { let r = reach[&(k, g)]; match once.get(&GroupId(g as u32)) { - Some(&o) => model = model.with(constraint!(p >= r - o)), - None => model = model.with(constraint!(p >= r)), + Some(&o) => constraints.push(constraint!(p >= r - o)), + None => constraints.push(constraint!(p >= r)), } } } } for (&(k, c), &v) in &y { let g = input.choices[c].group.0 as usize; - model = model.with(constraint!(v >= passes[&(k, g)] + x[c] - 1)); + constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { @@ -425,7 +430,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> continue; } if let Some(&r) = reach.get(&(k, h)) { - model = model.with(constraint!(r >= y[&(k, c)])); + constraints.push(constraint!(r >= y[&(k, c)])); } } } @@ -435,14 +440,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> slot: *slot, }; if let Some(&r) = reach.get(&(k, h)) { - model = model.with(constraint!(r >= x[c])); + constraints.push(constraint!(r >= x[c])); } } } } } for &(z, a, b) in &products { - model = model.with(constraint!(z >= a + b - 1)); + constraints.push(constraint!(z >= a + b - 1)); } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { @@ -455,7 +460,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> latency += input.choices[c].coefficients.direct.work_ms * v; } } - model = model.with(constraint!(latency <= bound)); + constraints.push(constraint!(latency <= bound)); } // Cuts. for (set, bound) in &input.excluded { @@ -463,18 +468,20 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &c in set { chosen += x[c]; } - model = model.with(constraint!(chosen <= *bound)); + constraints.push(constraint!(chosen <= *bound)); } for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { let mut chosen = Expression::from(0.0); for &c in set { chosen += x[c]; } - model = model.with(constraint!( + constraints.push(constraint!( e >= *bound * (chosen - (set.len() as f64 - 1.0)) )); } + let constraint_count = constraints.len() as u64; let model = model + .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; let solution = match model.solve() { @@ -497,6 +504,8 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> chosen, built_once, objective: solution.eval(&objective), + variables: variable_count, + constraints: constraint_count, })) } @@ -554,8 +563,7 @@ pub(crate) type MemoSearchResult = Result<(Selection, PhysicalCandidate, Vec), SelectionError>; /// The cheapest valid plan of `per_computation`'s memo under cost model 1 -/// (`None` under cost models 2 and 3, which the caller routes to the -/// legacy search). +/// (the caller routes cost models 2 and 3 to the legacy search). /// /// Lazy loop: solve the model (a lower bound: it omits memory held while a /// query runs, and checks composed guarantees only on its incumbents); @@ -570,21 +578,6 @@ pub(crate) fn search( entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, -) -> Option { - if !matches!( - models.plan_selection.cost_model, - CostModel::UsageOverTime(_) - ) { - return None; - } - Some(solve_memo(per_computation, entries, data, models)) -} - -fn solve_memo( - per_computation: &PerComputationCandidates, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: PlanningModels<'_>, ) -> MemoSearchResult { let started = std::time::Instant::now(); let memo = super::search::memo(per_computation, entries); @@ -610,12 +603,20 @@ fn solve_memo( let mut number = 0usize; let mut iterations = 0u32; let mut bound = None; + let mut first_bound = None; + let mut solver_ms = 0.0; + let mut size = (0, 0); loop { iterations += 1; - let Some(solution) = solve(&input).map_err(SelectionError::Solver)? else { + let solve_started = std::time::Instant::now(); + let solved = solve(&input).map_err(SelectionError::Solver)?; + solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; + let Some(solution) = solved else { break; }; bound = Some(solution.objective); + first_bound.get_or_insert(solution.objective); + size = (solution.variables, solution.constraints); if let Some(incumbent) = selector.cheapest_price() { if solution.objective >= incumbent - 1e-9 * incumbent.abs() { break; @@ -653,7 +654,12 @@ fn solve_memo( let price = selection.costs[&selection.selected].cost_per_second; selection.search.memo = Some(memo_stats); selection.search.solver_iterations = iterations; - selection.search.solver_time_ms = started.elapsed().as_millis() as u64; + selection.search.solver_time_ms = solver_ms as u64; + ( + selection.search.solver_variables, + selection.search.solver_constraints, + ) = size; + selection.search.solver_first_bound_cost_per_second = first_bound; selection.search.solver_gap = match (price, bound) { (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), (Some(_), Some(_)) => Some(0.0), 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 30/67] 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 c8e3d31c2ce72e99f69dff28106ac8c63c113876 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:58:59 +0000 Subject: [PATCH 31/67] plan-selection: memo choices shared or apart, shared-node and scan-horizon coefficients Each group's choices come shared (one node for its queries) and, for a group several queries read, apart (one copy per query, Choice.copy_of). A scan child is either apart (counted in its reader's choice) or shared at a horizon, the longest window among its readers (Choice.scan_children; a scan group's shared choice per horizon, Choice.scan_horizon_ms). Every shareable node of a choice's direct part gets SharedNodeCoefficients: what it and the direct nodes below it cost per unit of rate when built once at query time, its work and the state it holds. The tests take every plan of the brute force (identical-expression variants included, widened scans excepted: TODO, no choices yet): the model's terms are a lower bound of price_at_scale, and with held memory and raw data equal to it within 1e-9. --- .../src/pass/subexpr/coefficients.rs | 336 ++++++++++++++---- .../plan-selection/src/pass/subexpr/milp.rs | 18 + 2 files changed, 275 insertions(+), 79 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index aae394c25..07172f095 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -44,7 +44,9 @@ use asap_types::ir::{ASAPOp, NonASAPOp, Operator, QueryRoot}; use asap_types::physical::{Materialization, Medium}; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; -use super::milp::{Choice, ChoiceCoefficients, PartCoefficients, SlotCoefficients, Via}; +use super::milp::{ + Choice, ChoiceCoefficients, PartCoefficients, SharedNodeCoefficients, SlotCoefficients, Via, +}; use super::pricing::{child_keys, keys_of, scan_lookback_ms, PriceCache}; use super::{CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo}; use crate::{ @@ -88,6 +90,8 @@ pub(crate) struct HeldTerm { pub per_reader: bool, /// Whether it is the group's own node (replaced when it is built once). pub own_node: bool, + /// Its node's shallow key ([`shallow_keys`]). + pub node_key: u32, } /// Every choice of `group` with its lazy terms, in candidate order and, @@ -120,6 +124,9 @@ pub(super) struct ChoiceParts { pub root_key: u32, /// The shallow keys of the candidate node's private producers. pub root_producer_keys: Vec, + /// Per shared node of the coefficients, the shallow keys of the nodes + /// it computes when built once, itself first. + pub shared_node_keys: Vec>, /// Per slot of the coefficients, the shallow key of its build. pub slot_keys: Vec, /// The shallow keys of the private builds, in assignment order. @@ -166,9 +173,12 @@ pub(super) fn choice_parts( }) .collect(); for copy_of in modes { - let built_for: Vec = match copy_of { - Some(q) => vec![q], - None => queries.clone(), + // A copy for one query is built for it twice, so that Stage 2 + // offers its builds built once for several readers (`MatAtQuery`): + // another group's or copy's identical build may share it. + let built_for: Vec = match (copy_of, queries.as_slice()) { + (Some(q), _) | (None, &[q]) => vec![q, q], + (None, _) => queries.clone(), }; let part_entries: Vec = built_for.iter().map(|&q| entries[q].clone()).collect(); @@ -280,10 +290,6 @@ pub(super) fn choice_parts( } seen.push(assignment.clone()); let Some(parts) = coefficients_of( - memo, - group, - index, - candidate, physical, &keys, &shallow, @@ -320,6 +326,7 @@ pub(super) fn choice_parts( }, root_key: parts.root_key, root_producer_keys: parts.root_producer_keys, + shared_node_keys: parts.shared_node_keys, slot_keys: parts.slot_keys, build_keys: builds.iter().map(|&i| shallow[i]).collect(), part_keys: parts.part_keys, @@ -510,6 +517,7 @@ struct Coefficients { lazy: LazyTerms, root_key: u32, root_producer_keys: Vec, + shared_node_keys: Vec>, slot_keys: Vec, part_keys: Vec>, } @@ -602,10 +610,6 @@ pub(super) fn shallow_keys(dag: &PhysicalASAPDAG, interned: &mut HashMap = direct.iter().copied().collect(); + for &id in &direct { + let node = &dag.nodes[position[&id]]; + let shareable = + matches!(node.payload, Operator::ASAP(_)) || !node.output_schema.unique_keys.is_empty(); + if id == root || !shareable || is_build(&node.payload) || !query_time(id) { + continue; + } + let below: Vec = walk(id) + .0 + .into_iter() + .filter(|n| direct_set.contains(n)) + .collect(); + let mut once = dag.clone(); + once.nodes[position[&id]].materialization = Materialization::MatAtQuery(Medium::Memory); + let state_bytes = per_computation(&once, entries, scan_lookback, horizons, data, models)? + [&id] + .held + .map_or(0, |(bytes, _, _)| bytes); + shared_nodes.push(SharedNodeCoefficients { + cost_per_unit_rate: below.iter().map(|&n| per_rate(n)).sum(), + work_ms: work(&below), + state_bytes, + }); + shared_node_keys.push( + below + .iter() + .map(|n| shallow[position[n]]) + .collect::>(), + ); + } let mut part_keys = vec![direct.iter().map(|id| keys[position[id]]).collect()]; let mut slots = Vec::new(); let mut slot_keys = Vec::new(); @@ -792,12 +833,12 @@ fn coefficients_of( per_computation(&once_dag, entries, scan_lookback, horizons, data, models)?; once_priced[&root].held.map_or(0, |(bytes, _, _)| bytes) }; - let _ = candidate; Ok(Some(Coefficients { coefficients: ChoiceCoefficients { direct: direct_coefficients, slots, child_via, + shared_nodes, }, lazy: LazyTerms { held, @@ -807,6 +848,7 @@ fn coefficients_of( thresholds: Vec::new(), }, root_key: shallow[position[&root]], + shared_node_keys, root_producer_keys: producers_of .get(&root) .into_iter() @@ -838,6 +880,7 @@ pub(crate) fn held_cost( bytes: lazy.built_once_state_bytes, per_reader: false, own_node: true, + node_key: u32::MAX, }; let terms = lazy .held @@ -974,7 +1017,22 @@ mod tests { }) .collect(); let mut checked = Vec::new(); + let mut unplaced: Vec> = Vec::new(); for choice in group_choices(&memo).unwrap() { + // The plans the brute force prices: requirements met, scans + // widened only through the identical-expression variants. + let chosen: Vec<(GroupId, &ImplementationCandidate)> = choice + .iter() + .enumerate() + .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) + .collect(); + if !crate::pass::subexpr::windows::requirements_met(&chosen) + || chosen + .iter() + .any(|(_, c)| c.origin == crate::pass::subexpr::CandidateOrigin::WidenedScans) + { + continue; + } for plan in at_thresholds(&memo, &choice) { let Some(roots) = composed_roots(&memo, &per_computation, &plan) else { continue; @@ -983,7 +1041,18 @@ mod tests { &roots, &resized_builds(&memo, &choice, &plan), ); - for roots in identical_expression_variants(roots) { + // Widened variants (every metric's scans widened to one + // label set) come after the others; they have no choices + // yet (TODO: choices of widened candidates), and are only + // priced exactly. + let unwidened = 1 + asap_logical_asap_optimization::asap_aware_cse::identical_expression_pass2_candidates( + &roots, + &asap_logical_asap_optimization::asap_aware_cse::sharing_candidates(&roots), + ) + .len(); + for (variant, roots) in identical_expression_variants(roots).into_iter().enumerate() + { + let widened = variant >= unwidened; let Ok(physicals) = optimize_physical_asap( "L", &roots, @@ -994,21 +1063,37 @@ mod tests { continue; }; for physical in physicals { - if let Some((expected, actual)) = priced_both( + match priced_both( &memo, &choice, &parts, &physical, &entries, &data, &models, prices, &mut cache, ) { - assert!( - (actual - expected).abs() <= 1e-9 * expected.abs(), - "{queries:?} {:?}: {actual} vs {expected}", - physical.assignment - ); - checked.push(physical.assignment.clone()); + Some((expected, model, full)) => { + // The model's terms are a lower bound, and + // with the lazy terms the price itself. + assert!( + model <= expected * (1.0 + 1e-9), + "{queries:?} {:?}: model {model} over {expected}", + physical.assignment + ); + assert!( + (full - expected).abs() <= 1e-9 * expected.abs(), + "{queries:?} {:?}: {full} vs {expected}", + physical.assignment + ); + checked.push(physical.assignment.clone()); + } + None if widened => {} + None => unplaced.push(physical.assignment.clone()), } } } } } + assert!( + unplaced.is_empty(), + "{queries:?}: {} plans whose choices are not found", + unplaced.len() + ); checked } @@ -1016,7 +1101,7 @@ mod tests { /// a group's choice is not found in it). #[allow(clippy::too_many_arguments)] fn priced_both( - _memo: &Memo, + memo: &Memo, choice: &[Option], parts: &[Vec], physical: &PhysicalCandidate, @@ -1025,7 +1110,7 @@ mod tests { models: &PlanningModels<'_>, prices: &crate::ResourceUnitCost, cache: &mut PriceCache, - ) -> Option<(f64, f64)> { + ) -> Option<(f64, f64, f64)> { let settings = models.plan_selection; let dag = &physical.dag; let recurrence: Vec<_> = entries.iter().map(|e| e.recurrence.clone()).collect(); @@ -1043,19 +1128,6 @@ mod tests { ) .ok()?; let expected = cost.cost_per_second?; - let debug = std::env::var("W2DBG").is_ok(); - if debug { - for (id, n) in &cost.per_node { - eprintln!( - "stage3 {id} {:?} {:?} {:?} x{} {}", - n.cost_per_second, - nodes[id].output_state.timing, - nodes[id].materialization, - n.multiplier, - n.detail - ); - } - } let rates = computations( dag, @@ -1158,37 +1230,38 @@ mod tests { } for g in (0..chosen.len()).rev() { let Some(c) = chosen[g] else { continue }; - let Some(sample) = parts[g].iter().find(|p| p.choice.candidate == c) else { - return None; - }; + let sample = parts[g].iter().find(|p| p.choice.candidate == c)?; // The plan's copies of its node: equal shallow keys, reading // equal private producers. - let copies: Vec = by_key(sample.root_key) - .into_iter() - .filter(|id| !used.contains(id)) - .filter(|id| { - let keys_below: Vec = producers(*id) - .iter() - .map(|p| keys[position_of[p]]) - .collect(); - sample - .root_producer_keys - .iter() - .all(|k| keys_below.contains(k)) - }) - .collect(); + let group_queries = &memo.groups[g].queries; + // The plan's copies of its node, by any of its candidate's + // versions (thresholds): equal shallow keys, reading equal + // private producers. + let mut copies: Vec = Vec::new(); + for part in parts[g].iter().filter(|p| p.choice.candidate == c) { + for id in by_key(part.root_key) { + let keys_below: Vec = + producers(id).iter().map(|p| keys[position_of[p]]).collect(); + if !used.contains(&id) + && !copies.contains(&id) + && (part.choice.scan_horizon_ms.is_some() + || reaches_at_all(id).iter().any(|q| group_queries.contains(q))) + && part + .root_producer_keys + .iter() + .all(|k| keys_below.contains(k)) + { + copies.push(id); + } + } + } if copies.is_empty() { - // A scan every reader has apart. - if sample.choice.group.0 as usize == g - && parts[g] - .iter() - .all(|p| p.choice.scan_horizon_ms.is_some() || p.choice.copy_of.is_none()) - { + // A scan every reader has apart (its readers'). + if sample.choice.scan_horizon_ms.is_some() { continue; } return None; } - used.extend(copies.iter().copied()); let timing = |id: &PhysicalASAPNodeId| nodes[id].output_state.timing; let query_time: Vec = copies .iter() @@ -1207,7 +1280,11 @@ mod tests { }; for (k, &(root, at_query_time)) in roots.iter().enumerate() { let copy_of = match apart { - true => Some(*reaches_at_all(root).first()?), + true => Some( + *reaches_at_all(root) + .iter() + .find(|q| group_queries.contains(q))?, + ), false => None, }; let mine = below(root); @@ -1240,12 +1317,11 @@ mod tests { && match p.choice.scan_horizon_ms { None => true, Some(h) => { - scan_variant(root).unwrap_or(h) == h - || windows_read(dag, &nodes, root, &kept).into_iter().max() - == Some(h) + windows_read(dag, &nodes, root, &kept).into_iter().max() == Some(h) } } - })?; + }); + let part = part?; // Apart scans below are this choice's. for (_, v) in &part.choice.scan_children { if v.is_none() { @@ -1268,7 +1344,8 @@ mod tests { .copied() .find(|id| keys[position_of[id]] == *key) }) - .collect::>>()?; + .collect::>>(); + let slots = slots?; placed.push(Placed { part, root, @@ -1287,9 +1364,15 @@ mod tests { }; // Latency: each query's sum of the work of the parts it reaches. let mut latency = vec![0.0; entries.len()]; + let mut counted_slots: HashSet = HashSet::new(); for p in &placed { let vias = std::iter::once(Via::Direct).chain((0..p.slots.len()).map(Via::Slot)); for via in vias { + // A part's node several copies or groups share is one + // node, its work counted once. + if !counted_slots.insert(part_root(p, via)) { + continue; + } let work = match via { Via::Direct if !p.at_query_time => 0.0, Via::Direct => p.part.choice.coefficients.direct.work_ms, @@ -1304,30 +1387,107 @@ mod tests { } let mut actual = 0.0; let mut scans = Vec::new(); + // Shared nodes built once at query time: charged once, at their + // rate, in place of each copy's per-consumer charge. + let mut built_once_nodes: HashSet = HashSet::new(); + let mut held_extra = 0.0; + let mut charged_slots: HashSet = HashSet::new(); + let mut charged_roots: HashSet = HashSet::new(); for p in &placed { - if debug { - eprintln!("placed g{} c{} copy {:?} scans {:?} h {:?} root {} qt {} ing {:?} slots {:?} direct {:?} slotsc {:?} ingc {:?} rate {}", p.part.choice.group.0, p.part.choice.candidate, p.part.choice.copy_of, p.part.choice.scan_children, p.part.choice.scan_horizon_ms, p.root, p.at_query_time, p.at_ingestion, p.slots, p.part.choice.coefficients.direct, p.part.choice.coefficients.slots, p.part.lazy.direct_at_ingestion, rates[&p.root].1); - } let coefficients = &p.part.choice.coefficients; - if p.at_query_time { + let mut lazy = p.part.lazy.clone(); + // Two groups whose chosen nodes are one node (identical after a + // shared build is resized): its consumers' rate is both groups', + // charged once. + let first_root = charged_roots.insert(p.root); + if !first_root { + lazy.held.retain(|term| term.via != Via::Direct); + lazy.raw_scans.retain(|(_, _, via)| *via != Via::Direct); + } + if p.at_query_time && first_root { actual += coefficients.direct.cost_per_unit_rate * rates[&p.root].1; + let mine = below(p.root); + for (shared, sub) in coefficients + .shared_nodes + .iter() + .zip(&p.part.shared_node_keys) + { + let Some(node) = mine.iter().copied().find(|id| { + keys[position_of[id]] == sub[0] + && nodes[id].materialization != Materialization::NotMaterialized + }) else { + continue; + }; + actual -= shared.cost_per_unit_rate * rates[&p.root].1; + // Its nodes' held memory, at the shared node's rate. + let below_held: Vec = lazy + .held + .iter() + .filter(|t| sub[1..].contains(&t.node_key)) + .copied() + .collect(); + lazy.held.retain(|term| !sub.contains(&term.node_key)); + if built_once_nodes.insert(node) { + actual += shared.cost_per_unit_rate * rates[&node].1; + let readers: Vec<(usize, f64)> = reach + .iter() + .enumerate() + .filter(|(_, r)| r.contains(&node)) + .map(|(q, _)| (q, query_rate(q))) + .collect(); + let rate: f64 = readers.iter().map(|(_, r)| r).sum(); + for term in below_held { + let evaluation_ms = match term.per_reader && rate > 0.0 { + true => { + readers.iter().map(|&(q, r)| r * latency[q]).sum::() / rate + } + false => { + readers.iter().map(|&(q, _)| latency[q]).fold(0.0, f64::max) + } + }; + held_extra += prices.cost_per_memory_byte_second + * term.bytes as f64 + * rates[&node].1 + * evaluation_ms + / 1_000.0; + } + let longest = reach + .iter() + .enumerate() + .filter(|(_, r)| r.contains(&node)) + .map(|(q, _)| latency[q]) + .fold(0.0, f64::max); + held_extra += prices.cost_per_memory_byte_second + * shared.state_bytes as f64 + * rates[&node].1 + * longest + / 1_000.0; + } + } } for copy in &p.at_ingestion { let at_ingestion = p.part.lazy.direct_at_ingestion?; actual += at_ingestion.cost_per_unit_rate * rates[copy].1; } + // A slot two copies share is one node, charged once, its held + // memory too. for (i, slot) in coefficients.slots.iter().enumerate() { - actual += slot.fixed_cost + slot.cost_per_unit_rate * rates[&p.slots[i]].1; + if charged_slots.insert(p.slots[i]) { + actual += slot.fixed_cost + slot.cost_per_unit_rate * rates[&p.slots[i]].1; + } else { + lazy.held.retain(|term| term.via != Via::Slot(i)); + lazy.raw_scans.retain(|(_, _, via)| *via != Via::Slot(i)); + } } let built_once = nodes[&p.root].materialization != Materialization::NotMaterialized - && !is_build(&nodes[&p.root].payload); + && !is_build(&nodes[&p.root].payload) + && first_root; // Only a query-time direct part holds memory or scans raw data. - let mut lazy = p.part.lazy.clone(); if !p.at_query_time { lazy.held.retain(|term| term.via != Via::Direct); lazy.raw_scans.retain(|(_, _, via)| *via != Via::Direct); } - actual += held_cost( + held_extra += held_cost( &lazy, built_once, |via| rates[&part_root(p, via)].1, @@ -1348,11 +1508,14 @@ mod tests { } } } - actual += raw_data_bytes(&scans) - .values() - .map(|&bytes| prices.cost_per_disk_byte_second * bytes as f64) - .sum::(); - Some((expected, actual)) + let model = actual; + let full = model + + held_extra + + raw_data_bytes(&scans) + .values() + .map(|&bytes| prices.cost_per_disk_byte_second * bytes as f64) + .sum::(); + Some((expected, model, full)) } #[test] @@ -1394,6 +1557,21 @@ mod tests { assert!(checked.len() > 2); } + /// Three queries over one metric: a shared rate under two + /// aggregations, and a sum over a longer window of the same scan. + #[test] + fn coefficients_reproduce_stage_3_on_a_mixed_workload() { + let checked = check( + &[ + "sum by (job) (rate(m[5m]))", + "max(rate(m[5m]))", + "sum_over_time(m[10m])", + ], + &[60_000, 30_000, 60_000], + ); + assert!(!checked.is_empty()); + } + /// Two quantiles within 1% and 2% sharing one KLL build at a size /// threshold. #[test] diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index d6ea3b8e9..e442fad54 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -60,6 +60,23 @@ pub(crate) struct ChoiceCoefficients { pub slots: Vec, /// Per child group of the candidate, in order. pub child_via: Vec, + /// The shareable nodes of the direct part (an ASAP node, or one whose + /// schema has a unique key) of a group two or more queries read, other + /// than its own node: each may be one node for every copy of the group + /// (the identical-expression rule), and then built once at query time. + pub shared_nodes: Vec, +} + +/// A shareable node of a choice's direct part ([`ChoiceCoefficients::shared_nodes`]). +/// Its price is in `direct.cost_per_unit_rate`, charged per consumer; built +/// once at query time it costs `cost_per_unit_rate` per unit of the union +/// rate of the queries reaching it instead, and holds `state_bytes` while +/// its longest reader runs. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct SharedNodeCoefficients { + pub cost_per_unit_rate: f64, + pub work_ms: f64, + pub state_bytes: u64, } /// One choice of a group: a candidate and an option per private slot. @@ -571,6 +588,7 @@ mod tests { }, slots: Vec::new(), child_via: children.iter().map(|_| Via::Direct).collect(), + shared_nodes: Vec::new(), }, copy_of: None, scan_children: Vec::new(), From 633c82ea1bf7e4bc751c9b80c0f16594a225709a Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:01:55 +0000 Subject: [PATCH 32/67] Memo MILP model: shared choices only (a lower bound on a group's copies), shareable nodes charged per consumer or once, the cheaper --- .../plan-selection/src/pass/subexpr/milp.rs | 103 ++++++++++++++++-- 1 file changed, 92 insertions(+), 11 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 2173fa8e5..e482623e1 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -254,11 +254,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> }; let mut groups_with_rates: BTreeSet = once.keys().map(|g| g.0 as usize).collect(); for choice in &input.choices { - if choice - .coefficients - .slots - .iter() - .any(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) + if !choice.coefficients.shared_nodes.is_empty() + || choice + .coefficients + .slots + .iter() + .any(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) { groups_with_rates.insert(choice.group.0 as usize); } @@ -273,6 +274,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> // The objective. let mut objective = Expression::from(0.0); + let mut products_pending: Vec<(Variable, Variable, Variable)> = Vec::new(); // A consumer's constant rate, or `None` for one built once at query // time, whose rate is its group's union rate. let rate = |k: Consumer| -> Option { @@ -287,13 +289,44 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); + // The shareable nodes of a choice's direct part cost, per unit of + // rate, `shared_cost[c]`, charged either per consumer (the rest of the + // direct part's charge) or once at the union rate of the queries + // reaching the group when built once: the model charges the smaller + // (`per_consumer_shared[c]` against `built_once_shared[c]`). + let shared_cost: Vec = input + .choices + .iter() + .map(|choice| { + choice + .coefficients + .shared_nodes + .iter() + .map(|node| node.cost_per_unit_rate) + .sum() + }) + .collect(); + let mut per_consumer_shared: Vec> = vec![Vec::new(); input.choices.len()]; + let mut built_once_shared: Vec> = vec![Vec::new(); input.choices.len()]; + let mut product = |vars: &mut ProblemVariables, a: Variable, b: Variable| { + let z = vars.add(variable().min(0).max(1)); + products_pending.push((z, a, b)); + z + }; for (&(k, c), &v) in &y { - let cost = input.choices[c].coefficients.direct.cost_per_unit_rate; - if cost == 0.0 { + let total = input.choices[c].coefficients.direct.cost_per_unit_rate; + let shared = shared_cost[c].min(total); + let cost = total - shared; + if total == 0.0 { continue; } match rate(k) { - Some(r) => objective += r * cost * v, + Some(r) => { + objective += r * cost * v; + if shared != 0.0 { + per_consumer_shared[c].push((v, r * shared)); + } + } None => { let g = match k { Consumer::Group(g) => g.0 as usize, @@ -301,11 +334,41 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Query(_) => unreachable!("a query's rate is constant"), }; for (set, selector) in &rate_of[&g] { - rated.push((*selector, v, input.union_rate[set] * cost)); + let union = input.union_rate[set]; + if cost != 0.0 { + rated.push((*selector, v, union * cost)); + } + if shared != 0.0 { + let z = product(&mut vars, *selector, v); + per_consumer_shared[c].push((z, union * shared)); + } } } } } + for (c, choice) in input.choices.iter().enumerate() { + if shared_cost[c] == 0.0 { + continue; + } + for (set, selector) in &rate_of[&(choice.group.0 as usize)] { + let z = product(&mut vars, *selector, x[c]); + built_once_shared[c].push((z, input.union_rate[set] * shared_cost[c])); + } + } + // min(A, B) as a + b with a ≥ A − M·u, b ≥ B − M·(1 − u), u binary. + let mut cheaper_of: Vec<(Variable, Variable, Variable, usize, f64)> = Vec::new(); + for c in 0..input.choices.len() { + if per_consumer_shared[c].is_empty() && built_once_shared[c].is_empty() { + continue; + } + let big_m: f64 = per_consumer_shared[c].iter().map(|(_, w)| w).sum::() + + built_once_shared[c].iter().map(|(_, w)| w).sum::(); + let a = vars.add(variable().min(0)); + let b = vars.add(variable().min(0)); + let u = vars.add(variable().binary()); + objective += a + b; + cheaper_of.push((a, b, u, c, big_m)); + } for (c, choice) in input.choices.iter().enumerate() { for slot in &choice.coefficients.slots { objective += slot.fixed_cost * x[c]; @@ -474,9 +537,21 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } - for &(z, a, b) in &products { + for &(z, a, b) in products.iter().chain(&products_pending) { constraints.push(constraint!(z >= a + b - 1)); } + for &(a, b, u, c, big_m) in &cheaper_of { + let mut per_consumer = Expression::from(0.0); + for &(v, w) in &per_consumer_shared[c] { + per_consumer += w * v; + } + let mut built_once = Expression::from(0.0); + for &(v, w) in &built_once_shared[c] { + built_once += w * v; + } + constraints.push(constraint!(a >= per_consumer - big_m * u)); + constraints.push(constraint!(b >= built_once - big_m * (1 - u))); + } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { let Some(bound) = query.max_latency_ms else { @@ -622,7 +697,13 @@ pub(crate) fn search( models, &mut cache, )?; - choices.extend(group_choices); + // A group's copies (apart) cost at least its shared choice with its + // node not built once: each copy is charged per consumer as the + // shared node is, its own scans read at least the shortest horizon, + // and its builds are sized for fewer readers but built per copy. So + // the model keeps the shared choices (a lower bound), and the exact + // pricing of a group choice prices both. + choices.extend(group_choices.into_iter().filter(|c| c.copy_of.is_none())); } let input = model_input(&memo, choices, entries, models)?; let mut input = input; From 62216320bcc8d08c9612d503d67ea22ae6ea97ea Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:22:12 +0000 Subject: [PATCH 33/67] Memo MILP model: requirements and pane/top-k exclusivity as constraints, readers agree on a shared scan's horizon, raw data on disk and kept-or-not, held-memory lower bound; widened scans and copies left to exact pricing --- .../plan-selection/src/pass/subexpr/milp.rs | 298 +++++++++++++++++- 1 file changed, 292 insertions(+), 6 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index e482623e1..9bf7cf260 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -18,8 +18,9 @@ use good_lp::{ constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, WithMipGap, }; +use super::coefficients::LazyTerms; use super::{GroupId, Memo}; -use crate::{PlanSelector, PlanningModels, Selection, SelectionError}; +use crate::{CostModel, PlanSelector, PlanningModels, ResourceUnitCost, Selection, SelectionError}; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; use asap_physical_asap_optimization::PhysicalCandidate; use asap_types::ir::QueryRoot; @@ -148,6 +149,59 @@ pub(crate) struct ModelInput { /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all /// of a choice set is chosen (local cuts). pub extra_bounds: Vec<(Vec, f64)>, + /// Requirements of a choice on another group's choice: when the first + /// is chosen, one of the second (an exact implementation, or one + /// carrying a coupling site) is. + pub requires: Vec<(usize, Vec)>, + /// Choices that may not be chosen together with any of the second set + /// (panes and an exact top-k reuse in one plan). + pub exclusive: (Vec, Vec), + /// Each query-time raw scan of a choice's part: each source's longest + /// one is kept on disk. + pub raw_scans: Vec, + /// Whether the deployment keeps raw data: without it, no part with a + /// raw scan may run at query time. + pub raw_data_retained: bool, +} + +/// A raw scan of one part of a choice, when that part runs at query time. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct RawScan { + pub choice: usize, + /// Index of its source. + pub source: usize, + /// The price per second of keeping its look-back on disk. + pub cost: f64, + pub via: Via, +} + +/// Adds to `choice`'s coefficients a lower bound on the memory its parts +/// hold while an evaluation runs ([`LazyTerms::held`]): Stage 3 charges +/// `bytes · rate · E / 1000` at the memory price, `E` the latency of the +/// queries reaching the part, at least the part's own work. The group's +/// own node's working memory is left out (replaced when it is built +/// once). +fn add_lazy_lower_bounds(choice: &mut Choice, lazy: &LazyTerms, prices: &ResourceUnitCost) { + let coefficients = &mut choice.coefficients; + for term in &lazy.held { + if term.own_node { + continue; + } + let per_unit_rate = |work_ms: f64| { + prices.cost_per_memory_byte_second * term.bytes as f64 * work_ms / 1000.0 + }; + match term.via { + Via::Direct => { + coefficients.direct.cost_per_unit_rate += per_unit_rate(coefficients.direct.work_ms) + } + Via::Slot(i) => { + let slot = &mut coefficients.slots[i]; + if !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { + slot.cost_per_unit_rate += per_unit_rate(slot.work_ms); + } + } + } + } } /// The model's optimum. @@ -401,6 +455,19 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } + let vars_exclusive = (vars.add(variable().binary()), vars.add(variable().binary())); + let source_count = input + .raw_scans + .iter() + .map(|s| s.source + 1) + .max() + .unwrap_or(0); + let kept_on_disk: Vec = (0..source_count) + .map(|_| vars.add(variable().min(0))) + .collect(); + for &v in &kept_on_disk { + objective += v; + } let variable_count = vars.len() as u64; let model = vars.minimise(objective.clone()).using(highs); let mut constraints = Vec::new(); @@ -582,6 +649,60 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> e >= *bound * (chosen - (set.len() as f64 - 1.0)) )); } + for (c, allowed) in &input.requires { + let mut any = Expression::from(0.0); + for &a in allowed { + any += x[a]; + } + constraints.push(constraint!(x[*c] <= any)); + } + let (first, second) = &input.exclusive; + if !first.is_empty() && !second.is_empty() { + let (p, t) = (vars_exclusive.0, vars_exclusive.1); + for &c in first { + constraints.push(constraint!(x[c] <= p)); + } + for &c in second { + constraints.push(constraint!(x[c] <= t)); + } + constraints.push(constraint!(p + t <= 1)); + } + // Raw data: each source's longest query-time scan is kept on disk; a + // deployment without raw data runs no raw scan at query time. A part + // runs at query time for every consumer but an ingestion-time slot. + let at_query_time = |k: Consumer| match k { + Consumer::Slot { choice, slot } => !matches!( + input.choices[choice].coefficients.slots[slot].materialization, + Materialization::MatAtIngestion(_) + ), + _ => true, + }; + for scan in &input.raw_scans { + let c = scan.choice; + let mut runs: Vec = Vec::new(); + match scan.via { + Via::Direct => { + for (&(k, d), &v) in &y { + if d == c && at_query_time(k) { + runs.push(v); + } + } + } + Via::Slot(i) => { + let m = input.choices[c].coefficients.slots[i].materialization; + if !matches!(m, Materialization::MatAtIngestion(_)) { + runs.push(x[c]); + } + } + } + for v in runs { + if !input.raw_data_retained { + constraints.push(constraint!(v == 0)); + } else if scan.cost > 0.0 { + constraints.push(constraint!(kept_on_disk[scan.source] >= scan.cost * v)); + } + } + } let constraint_count = constraints.len() as u64; let model = model .with_all(constraints) @@ -592,12 +713,29 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Err(good_lp::ResolutionError::Infeasible) => return Ok(None), Err(error) => return Err(error.to_string()), }; - let chosen = x + let chosen: Vec = x .iter() .enumerate() .filter(|(_, v)| solution.value(**v) > 0.5) .map(|(c, _)| c) .collect(); + if std::env::var_os("JP_MILP_TRACE").is_some() { + for &c in &chosen { + let choice: &Choice = &input.choices[c]; + let reach: Vec = y + .iter() + .filter(|((_, d), v)| *d == c && solution.value(**v) > 1e-9) + .map(|((k, _), v)| format!("{k:?}={:.2}", solution.value(*v))) + .collect(); + eprintln!( + " g{} cand{} direct {:.4e}/rate work {:.2}ms shared {:?} slots {:?} scans {:?} horizon {:?} assignment {:?} y {:?}", + choice.group.0, choice.candidate, choice.coefficients.direct.cost_per_unit_rate, + choice.coefficients.direct.work_ms, choice.coefficients.shared_nodes.len(), + choice.coefficients.slots.iter().map(|s| (s.fixed_cost, s.cost_per_unit_rate)).collect::>(), + choice.scan_children, choice.scan_horizon_ms, choice.assignment, reach + ); + } + } let built_once = once .iter() .filter(|(_, v)| solution.value(**v) > 0.5) @@ -688,8 +826,12 @@ pub(crate) fn search( let scale = super::pricing::frozen_calibration(&memo, entries, data, models)?; let mut cache = super::pricing::PriceCache::new(scale); let mut choices: Vec = Vec::new(); + let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); + let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { + unreachable!("the caller routes cost models 2 and 3 to the legacy search"); + }; for g in 0..memo.groups.len() { - let (group_choices, _lazy) = super::coefficients::choices( + let (group_choices, lazy) = super::coefficients::choices( &memo, GroupId(g as u32), entries, @@ -703,10 +845,51 @@ pub(crate) fn search( // and its builds are sized for fewer readers but built per copy. So // the model keeps the shared choices (a lower bound), and the exact // pricing of a group choice prices both. - choices.extend(group_choices.into_iter().filter(|c| c.copy_of.is_none())); + // Widened scans are reached through the identical-expression + // variants of the exact pricing, never chosen directly. + for (mut choice, lazy) in group_choices.into_iter().zip(lazy) { + if choice.copy_of.is_some() + || memo.candidates[g][choice.candidate].origin + == super::CandidateOrigin::WidenedScans + { + continue; + } + add_lazy_lower_bounds(&mut choice, &lazy, prices); + raw_scans.extend( + lazy.raw_scans + .iter() + .map(|(source, bytes, via)| (choices.len(), source.clone(), *bytes, *via)), + ); + choices.push(choice); + } + } + let trace = std::env::var_os("JP_MILP_TRACE").is_some(); + if trace { + eprintln!( + "milp: memo {} groups, {} choices in {:.1}s", + memo.groups.len(), + choices.len(), + started.elapsed().as_secs_f64() + ); + } + let mut input = model_input(&memo, choices, entries, models)?; + let mut sources: Vec = Vec::new(); + for (c, source, bytes, via) in raw_scans { + let index = match sources.iter().position(|s| *s == source) { + Some(i) => i, + None => { + sources.push(source); + sources.len() - 1 + } + }; + input.raw_scans.push(RawScan { + choice: c, + source: index, + cost: prices.cost_per_disk_byte_second * bytes as f64, + via, + }); } - let input = model_input(&memo, choices, entries, models)?; - let mut input = input; + input.raw_data_retained = models.deployment_capabilities.raw_data_retained; let mut selector = PlanSelector::new(entries, data, models, false)?; let mut logical: HashMap> = HashMap::new(); let mut number = 0usize; @@ -725,6 +908,9 @@ pub(crate) fn search( }; bound = Some(solution.objective); first_bound.get_or_insert(solution.objective); + if trace { + eprintln!("milp: round {iterations} bound {:.6e} incumbent {:?} solve {:.0}ms vars {} cons {} at {:.1}s", solution.objective, selector.cheapest_price(), solve_started.elapsed().as_secs_f64()*1e3, solution.variables, solution.constraints, started.elapsed().as_secs_f64()); + } size = (solution.variables, solution.constraints); if let Some(incumbent) = selector.cheapest_price() { if solution.objective >= incumbent - 1e-9 * incumbent.abs() { @@ -737,6 +923,53 @@ pub(crate) fn search( for &c in &solution.chosen { group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); } + let chosen_candidates: Vec<(GroupId, &super::ImplementationCandidate)> = group_choice + .iter() + .enumerate() + .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) + .collect(); + if !super::windows::requirements_met(&chosen_candidates) { + // A rollup's requirement: cut the smallest failing pairs of + // (group, candidate), else the whole group choice below. + let mut cut = false; + for (i, a) in chosen_candidates.iter().enumerate() { + if a.1.rollup_of.is_none() { + continue; + } + for (j, b) in chosen_candidates.iter().enumerate() { + if i != j && !super::windows::requirements_met(&[*a, *b]) { + let members: Vec = (0..input.choices.len()) + .filter(|&c| { + let choice = &input.choices[c]; + (choice.group == a.0 + && choice.candidate == group_choice[a.0 .0 as usize].unwrap()) + || (choice.group == b.0 + && choice.candidate + == group_choice[b.0 .0 as usize].unwrap()) + }) + .collect(); + input.excluded.push((members, 1.0)); + cut = true; + } + } + if !super::windows::requirements_met(&[*a]) { + let members: Vec = (0..input.choices.len()) + .filter(|&c| { + input.choices[c].group == a.0 + && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] + }) + .collect(); + input.excluded.push((members, 0.0)); + cut = true; + } + } + if trace { + eprintln!("milp: round {iterations} requirements not met (pair cut {cut})"); + } + if cut { + continue; + } + } super::search::offer_group_choice( &memo, per_computation, @@ -851,6 +1084,51 @@ fn model_input( } } } + // Requirements the reference search checks on a group choice + // ([`super::windows::requirements_met`]) other than rollups, which the + // loop cuts lazily. + let candidate_of = |c: &Choice| &memo.candidates[c.group.0 as usize][c.candidate]; + let mut requires = Vec::new(); + for (c, choice) in choices.iter().enumerate() { + let candidate = candidate_of(choice); + for &g in &candidate.requires_exact_children { + let allowed = (0..choices.len()) + .filter(|&a| { + choices[a].group == g && super::windows::is_exact(candidate_of(&choices[a])) + }) + .collect(); + requires.push((c, allowed)); + } + for (g, site) in &candidate.requires_sites { + let allowed = (0..choices.len()) + .filter(|&a| { + choices[a].group == *g + && candidate_of(&choices[a]).coupling_sites.contains(site) + }) + .collect(); + requires.push((c, allowed)); + } + } + // A reader of a scan shared at horizon `h` reads the scan group's + // choice at that horizon, so every reader of one scan agrees on it. + for (c, choice) in choices.iter().enumerate() { + for &(scan, horizon) in &choice.scan_children { + let Some(h) = horizon else { continue }; + let allowed = (0..choices.len()) + .filter(|&a| choices[a].group == scan && choices[a].scan_horizon_ms == Some(h)) + .collect(); + requires.push((c, allowed)); + } + } + let of_origin = |origin: super::CandidateOrigin| -> Vec { + (0..choices.len()) + .filter(|&c| candidate_of(&choices[c]).origin == origin) + .collect() + }; + let exclusive = ( + of_origin(super::CandidateOrigin::Panes), + of_origin(super::CandidateOrigin::TopKReuse), + ); let group_slots = memo .groups .iter() @@ -875,6 +1153,10 @@ fn model_input( union_rate, excluded: Vec::new(), extra_bounds: Vec::new(), + requires, + exclusive, + raw_scans: Vec::new(), + raw_data_retained: true, }) } @@ -937,6 +1219,10 @@ mod tests { union_rate: [(vec![0], 1.0)].into(), excluded: Vec::new(), extra_bounds: Vec::new(), + requires: Vec::new(), + exclusive: (Vec::new(), Vec::new()), + raw_scans: Vec::new(), + raw_data_retained: true, }; let solution = solve(&input).unwrap().unwrap(); assert_eq!(solution.built_once, [GroupId(0)].into()); 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 34/67] 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 470feed6455e05be3b46ea9f159f4e6e1418e3f1 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:31:39 +0000 Subject: [PATCH 35/67] Memo MILP model: aggregated reach constraints and linear floors on rate-dependent charges for a tight relaxation (solves ~50x faster); trace hooks --- .../plan-selection/src/pass/subexpr/milp.rs | 182 ++++++++++++++---- 1 file changed, 144 insertions(+), 38 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 9bf7cf260..61e3748d2 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -231,23 +231,27 @@ enum Consumer { pub(crate) fn solve(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let x: Vec = input - .choices - .iter() - .map(|_| vars.add(variable().binary())) - .collect(); - let need: Vec = (0..n).map(|_| vars.add(variable().binary())).collect(); + let relaxed = std::env::var_os("JP_MILP_LP").is_some(); + let binary = || { + if relaxed { + variable().min(0).max(1) + } else { + variable().binary() + } + }; + let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); + let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); let once: BTreeMap = input .group_slots .iter() - .map(|slot| (slot.group, vars.add(variable().binary()))) + .map(|slot| (slot.group, vars.add(binary()))) .collect(); // Which queries' plans reach each group at all (any materialization): // the union rate of a node built once reads them. let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); for (g, queries) in input.group_queries.iter().enumerate() { for &q in queries { - reaches_at_all.insert((q, g), vars.add(variable().binary())); + reaches_at_all.insert((q, g), vars.add(binary())); } } // Consumers that may reach each group's own node, and through it its @@ -321,7 +325,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &g in &groups_with_rates { let sets = class_sets(g) .into_iter() - .map(|set| (set, vars.add(variable().binary()))) + .map(|set| (set, vars.add(binary()))) .collect(); rate_of.insert(g, sets); } @@ -341,6 +345,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Group(_) => None, } }; + let slowest_rate = |g: usize| -> f64 { + rate_of[&g] + .iter() + .map(|(set, _)| input.union_rate[set]) + .fold(f64::INFINITY, f64::min) + }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); // The shareable nodes of a choice's direct part cost, per unit of @@ -387,10 +397,15 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, Consumer::Query(_) => unreachable!("a query's rate is constant"), }; + // At least the slowest class's rate is charged linearly + // (some query reaches the group), the rest through the + // selector products: a tighter relaxation. + let slowest = slowest_rate(g); + objective += slowest * cost * v; for (set, selector) in &rate_of[&g] { let union = input.union_rate[set]; - if cost != 0.0 { - rated.push((*selector, v, union * cost)); + if cost != 0.0 && union > slowest { + rated.push((*selector, v, (union - slowest) * cost)); } if shared != 0.0 { let z = product(&mut vars, *selector, v); @@ -419,8 +434,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> + built_once_shared[c].iter().map(|(_, w)| w).sum::(); let a = vars.add(variable().min(0)); let b = vars.add(variable().min(0)); - let u = vars.add(variable().binary()); + let u = vars.add(binary()); objective += a + b; + objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; + if std::env::var_os("JP_MILP_NOMIN").is_some() { + continue; + } cheaper_of.push((a, b, u, c, big_m)); } for (c, choice) in input.choices.iter().enumerate() { @@ -429,12 +448,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if slot.cost_per_unit_rate != 0.0 && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - for (set, selector) in &rate_of[&(choice.group.0 as usize)] { - rated.push(( - *selector, - x[c], - input.union_rate[set] * slot.cost_per_unit_rate, - )); + let g = choice.group.0 as usize; + let slowest = slowest_rate(g); + objective += slowest * slot.cost_per_unit_rate * x[c]; + for (set, selector) in &rate_of[&g] { + let union = input.union_rate[set]; + if union > slowest { + rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); + } } } } @@ -455,7 +476,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } - let vars_exclusive = (vars.add(variable().binary()), vars.add(variable().binary())); + let vars_exclusive = (vars.add(binary()), vars.add(binary())); let source_count = input .raw_scans .iter() @@ -509,15 +530,29 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + // Aggregated over a group's choices (one is chosen when it is + // needed), which keeps the LP relaxation tight. + let mut choices_of: Vec> = vec![Vec::new(); n]; for (c, choice) in input.choices.iter().enumerate() { - let g = choice.group.0 as usize; - for child in &choice.children { + choices_of[choice.group.0 as usize].push(c); + } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); + } + for h in children { + let mut with_child = Expression::from(0.0); + for &c in of_g { + if input.choices[c].children.iter().any(|k| k.0 as usize == h) { + with_child += x[c]; + } + } for &q in &input.group_queries[g] { - if let (Some(&above), Some(&below)) = ( - reaches_at_all.get(&(q, g)), - reaches_at_all.get(&(q, child.0 as usize)), - ) { - constraints.push(constraint!(below >= above + x[c] - 1)); + if let (Some(&above), Some(&below)) = + (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) + { + constraints.push(constraint!(below >= above + with_child.clone() - 1)); } } } @@ -578,19 +613,26 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let g = input.choices[c].group.0 as usize; constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); } + for (&(k, g), &p) in &passes { + let mut into = Expression::from(0.0); + for &c in &choices_of[g] { + into += y[&(k, c)]; + } + // Unless a fused choice of another group implements it. + for (c, choice) in input.choices.iter().enumerate() { + if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { + into += x[c]; + } + } + constraints.push(constraint!(into >= p)); + } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { let h = child.0 as usize; match via { Via::Direct => { - for &(k, g) in passes.keys() { - if g != choice.group.0 as usize { - continue; - } - if let Some(&r) = reach.get(&(k, h)) { - constraints.push(constraint!(r >= y[&(k, c)])); - } - } + // Aggregated below, over the group's choices. + let _ = h; } Via::Slot(slot) => { let k = Consumer::Slot { @@ -604,6 +646,43 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + for (child, via) in input.choices[c] + .children + .iter() + .zip(&input.choices[c].coefficients.child_via) + { + if *via == Via::Direct { + children.insert(child.0 as usize); + } + } + } + for h in children { + for &(k, pg) in passes.keys() { + if pg != g { + continue; + } + let Some(&r) = reach.get(&(k, h)) else { + continue; + }; + let mut into = Expression::from(0.0); + for &c in of_g { + let choice = &input.choices[c]; + if choice + .children + .iter() + .zip(&choice.coefficients.child_via) + .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) + { + into += y[&(k, c)]; + } + } + constraints.push(constraint!(r >= into)); + } + } + } for &(z, a, b) in products.iter().chain(&products_pending) { constraints.push(constraint!(z >= a + b - 1)); } @@ -616,8 +695,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &(v, w) in &built_once_shared[c] { built_once += w * v; } - constraints.push(constraint!(a >= per_consumer - big_m * u)); - constraints.push(constraint!(b >= built_once - big_m * (1 - u))); + // Both are at least the slowest class's rate (some query reaches + // the chosen choice): charged linearly, which keeps the relaxation + // tight; only the excess goes through the binary choice. + let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); + constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); + constraints.push(constraint!( + b >= built_once - floor * x[c] - big_m * (1 - u) + )); } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { @@ -704,10 +789,13 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } let constraint_count = constraints.len() as u64; - let model = model + let mut model = model .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; + if std::env::var_os("JP_MILP_VERBOSE").is_some() { + model.set_verbose(true); + } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), @@ -720,7 +808,25 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> .map(|(c, _)| c) .collect(); if std::env::var_os("JP_MILP_TRACE").is_some() { - for &c in &chosen { + if relaxed { + for (g, v) in once.iter() { + eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); + } + for (g, v) in need.iter().enumerate() { + eprintln!(" need g{} = {:.3}", g, solution.value(*v)); + } + } + let shown: Vec = if relaxed { + (0..x.len()) + .filter(|&c| solution.value(x[c]) > 1e-6) + .collect() + } else { + chosen.clone() + }; + for &c in &shown { + if relaxed { + eprintln!(" x = {:.3}", solution.value(x[c])); + } let choice: &Choice = &input.choices[c]; let reach: Vec = y .iter() From fe1b31065d88f9037e7842d851e490f7cecd6a6b Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:33:44 +0000 Subject: [PATCH 36/67] plan-selection: a reader's choice at a shared scan's horizon only when the horizon covers its own windows A shared scan reads the longest window among its readers, so a reader's choice at horizon h exists only for h no shorter than every window it reads the scan through (time shifts included). Choices at shorter horizons let the model price a 5-minute quantile's scan at 10 seconds. --- .../src/pass/subexpr/coefficients.rs | 21 ++ .../plan-selection/src/pass/subexpr/milp.rs | 182 ++++-------------- 2 files changed, 59 insertions(+), 144 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 07172f095..bea3da6e7 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -263,6 +263,27 @@ pub(super) fn choice_parts( } let mut seen: Vec> = Vec::new(); for physical in &physicals { + // A shared scan's horizon is the longest window of + // its readers: no shorter than this choice's own + // window on it (through time shifts). + if !horizons.is_empty() { + let dag = &physical.dag; + let nodes = node_map(dag); + let kept = kept_pane_chains(dag); + let too_short = dag.nodes.iter().any(|n| match &n.payload { + Operator::NonASAP(NonASAPOp::Scan { source, .. }) => horizons + .get(&super::scans::raw_horizon_key(source)) + .is_some_and(|&h| { + windows_read(dag, &nodes, n.id, &kept) + .into_iter() + .any(|w| w > h) + }), + _ => false, + }); + if too_short { + continue; + } + } let keys = keys_of(&physical.dag, cache.keys_mut()); let shallow = shallow_keys(&physical.dag, cache.keys_mut()); let private: Vec = diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 61e3748d2..9bf7cf260 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -231,27 +231,23 @@ enum Consumer { pub(crate) fn solve(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let relaxed = std::env::var_os("JP_MILP_LP").is_some(); - let binary = || { - if relaxed { - variable().min(0).max(1) - } else { - variable().binary() - } - }; - let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); - let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); + let x: Vec = input + .choices + .iter() + .map(|_| vars.add(variable().binary())) + .collect(); + let need: Vec = (0..n).map(|_| vars.add(variable().binary())).collect(); let once: BTreeMap = input .group_slots .iter() - .map(|slot| (slot.group, vars.add(binary()))) + .map(|slot| (slot.group, vars.add(variable().binary()))) .collect(); // Which queries' plans reach each group at all (any materialization): // the union rate of a node built once reads them. let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); for (g, queries) in input.group_queries.iter().enumerate() { for &q in queries { - reaches_at_all.insert((q, g), vars.add(binary())); + reaches_at_all.insert((q, g), vars.add(variable().binary())); } } // Consumers that may reach each group's own node, and through it its @@ -325,7 +321,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &g in &groups_with_rates { let sets = class_sets(g) .into_iter() - .map(|set| (set, vars.add(binary()))) + .map(|set| (set, vars.add(variable().binary()))) .collect(); rate_of.insert(g, sets); } @@ -345,12 +341,6 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Group(_) => None, } }; - let slowest_rate = |g: usize| -> f64 { - rate_of[&g] - .iter() - .map(|(set, _)| input.union_rate[set]) - .fold(f64::INFINITY, f64::min) - }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); // The shareable nodes of a choice's direct part cost, per unit of @@ -397,15 +387,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, Consumer::Query(_) => unreachable!("a query's rate is constant"), }; - // At least the slowest class's rate is charged linearly - // (some query reaches the group), the rest through the - // selector products: a tighter relaxation. - let slowest = slowest_rate(g); - objective += slowest * cost * v; for (set, selector) in &rate_of[&g] { let union = input.union_rate[set]; - if cost != 0.0 && union > slowest { - rated.push((*selector, v, (union - slowest) * cost)); + if cost != 0.0 { + rated.push((*selector, v, union * cost)); } if shared != 0.0 { let z = product(&mut vars, *selector, v); @@ -434,12 +419,8 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> + built_once_shared[c].iter().map(|(_, w)| w).sum::(); let a = vars.add(variable().min(0)); let b = vars.add(variable().min(0)); - let u = vars.add(binary()); + let u = vars.add(variable().binary()); objective += a + b; - objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; - if std::env::var_os("JP_MILP_NOMIN").is_some() { - continue; - } cheaper_of.push((a, b, u, c, big_m)); } for (c, choice) in input.choices.iter().enumerate() { @@ -448,14 +429,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if slot.cost_per_unit_rate != 0.0 && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - let g = choice.group.0 as usize; - let slowest = slowest_rate(g); - objective += slowest * slot.cost_per_unit_rate * x[c]; - for (set, selector) in &rate_of[&g] { - let union = input.union_rate[set]; - if union > slowest { - rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); - } + for (set, selector) in &rate_of[&(choice.group.0 as usize)] { + rated.push(( + *selector, + x[c], + input.union_rate[set] * slot.cost_per_unit_rate, + )); } } } @@ -476,7 +455,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } - let vars_exclusive = (vars.add(binary()), vars.add(binary())); + let vars_exclusive = (vars.add(variable().binary()), vars.add(variable().binary())); let source_count = input .raw_scans .iter() @@ -530,29 +509,15 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } - // Aggregated over a group's choices (one is chosen when it is - // needed), which keeps the LP relaxation tight. - let mut choices_of: Vec> = vec![Vec::new(); n]; for (c, choice) in input.choices.iter().enumerate() { - choices_of[choice.group.0 as usize].push(c); - } - for (g, of_g) in choices_of.iter().enumerate() { - let mut children: BTreeSet = BTreeSet::new(); - for &c in of_g { - children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); - } - for h in children { - let mut with_child = Expression::from(0.0); - for &c in of_g { - if input.choices[c].children.iter().any(|k| k.0 as usize == h) { - with_child += x[c]; - } - } + let g = choice.group.0 as usize; + for child in &choice.children { for &q in &input.group_queries[g] { - if let (Some(&above), Some(&below)) = - (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) - { - constraints.push(constraint!(below >= above + with_child.clone() - 1)); + if let (Some(&above), Some(&below)) = ( + reaches_at_all.get(&(q, g)), + reaches_at_all.get(&(q, child.0 as usize)), + ) { + constraints.push(constraint!(below >= above + x[c] - 1)); } } } @@ -613,26 +578,19 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let g = input.choices[c].group.0 as usize; constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); } - for (&(k, g), &p) in &passes { - let mut into = Expression::from(0.0); - for &c in &choices_of[g] { - into += y[&(k, c)]; - } - // Unless a fused choice of another group implements it. - for (c, choice) in input.choices.iter().enumerate() { - if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { - into += x[c]; - } - } - constraints.push(constraint!(into >= p)); - } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { let h = child.0 as usize; match via { Via::Direct => { - // Aggregated below, over the group's choices. - let _ = h; + for &(k, g) in passes.keys() { + if g != choice.group.0 as usize { + continue; + } + if let Some(&r) = reach.get(&(k, h)) { + constraints.push(constraint!(r >= y[&(k, c)])); + } + } } Via::Slot(slot) => { let k = Consumer::Slot { @@ -646,43 +604,6 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } - for (g, of_g) in choices_of.iter().enumerate() { - let mut children: BTreeSet = BTreeSet::new(); - for &c in of_g { - for (child, via) in input.choices[c] - .children - .iter() - .zip(&input.choices[c].coefficients.child_via) - { - if *via == Via::Direct { - children.insert(child.0 as usize); - } - } - } - for h in children { - for &(k, pg) in passes.keys() { - if pg != g { - continue; - } - let Some(&r) = reach.get(&(k, h)) else { - continue; - }; - let mut into = Expression::from(0.0); - for &c in of_g { - let choice = &input.choices[c]; - if choice - .children - .iter() - .zip(&choice.coefficients.child_via) - .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) - { - into += y[&(k, c)]; - } - } - constraints.push(constraint!(r >= into)); - } - } - } for &(z, a, b) in products.iter().chain(&products_pending) { constraints.push(constraint!(z >= a + b - 1)); } @@ -695,14 +616,8 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &(v, w) in &built_once_shared[c] { built_once += w * v; } - // Both are at least the slowest class's rate (some query reaches - // the chosen choice): charged linearly, which keeps the relaxation - // tight; only the excess goes through the binary choice. - let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); - constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); - constraints.push(constraint!( - b >= built_once - floor * x[c] - big_m * (1 - u) - )); + constraints.push(constraint!(a >= per_consumer - big_m * u)); + constraints.push(constraint!(b >= built_once - big_m * (1 - u))); } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { @@ -789,13 +704,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } let constraint_count = constraints.len() as u64; - let mut model = model + let model = model .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; - if std::env::var_os("JP_MILP_VERBOSE").is_some() { - model.set_verbose(true); - } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), @@ -808,25 +720,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> .map(|(c, _)| c) .collect(); if std::env::var_os("JP_MILP_TRACE").is_some() { - if relaxed { - for (g, v) in once.iter() { - eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); - } - for (g, v) in need.iter().enumerate() { - eprintln!(" need g{} = {:.3}", g, solution.value(*v)); - } - } - let shown: Vec = if relaxed { - (0..x.len()) - .filter(|&c| solution.value(x[c]) > 1e-6) - .collect() - } else { - chosen.clone() - }; - for &c in &shown { - if relaxed { - eprintln!(" x = {:.3}", solution.value(x[c])); - } + for &c in &chosen { let choice: &Choice = &input.choices[c]; let reach: Vec = y .iter() From 8ea3ee19dcfb38310eea425fbca9f8e85009ee74 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:41:12 +0000 Subject: [PATCH 37/67] Restore the MILP model's tight relaxation, which the horizon fix's rebase dropped --- .../plan-selection/src/pass/subexpr/milp.rs | 182 ++++++++++++++---- 1 file changed, 144 insertions(+), 38 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 9bf7cf260..61e3748d2 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -231,23 +231,27 @@ enum Consumer { pub(crate) fn solve(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let x: Vec = input - .choices - .iter() - .map(|_| vars.add(variable().binary())) - .collect(); - let need: Vec = (0..n).map(|_| vars.add(variable().binary())).collect(); + let relaxed = std::env::var_os("JP_MILP_LP").is_some(); + let binary = || { + if relaxed { + variable().min(0).max(1) + } else { + variable().binary() + } + }; + let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); + let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); let once: BTreeMap = input .group_slots .iter() - .map(|slot| (slot.group, vars.add(variable().binary()))) + .map(|slot| (slot.group, vars.add(binary()))) .collect(); // Which queries' plans reach each group at all (any materialization): // the union rate of a node built once reads them. let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); for (g, queries) in input.group_queries.iter().enumerate() { for &q in queries { - reaches_at_all.insert((q, g), vars.add(variable().binary())); + reaches_at_all.insert((q, g), vars.add(binary())); } } // Consumers that may reach each group's own node, and through it its @@ -321,7 +325,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &g in &groups_with_rates { let sets = class_sets(g) .into_iter() - .map(|set| (set, vars.add(variable().binary()))) + .map(|set| (set, vars.add(binary()))) .collect(); rate_of.insert(g, sets); } @@ -341,6 +345,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Group(_) => None, } }; + let slowest_rate = |g: usize| -> f64 { + rate_of[&g] + .iter() + .map(|(set, _)| input.union_rate[set]) + .fold(f64::INFINITY, f64::min) + }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); // The shareable nodes of a choice's direct part cost, per unit of @@ -387,10 +397,15 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, Consumer::Query(_) => unreachable!("a query's rate is constant"), }; + // At least the slowest class's rate is charged linearly + // (some query reaches the group), the rest through the + // selector products: a tighter relaxation. + let slowest = slowest_rate(g); + objective += slowest * cost * v; for (set, selector) in &rate_of[&g] { let union = input.union_rate[set]; - if cost != 0.0 { - rated.push((*selector, v, union * cost)); + if cost != 0.0 && union > slowest { + rated.push((*selector, v, (union - slowest) * cost)); } if shared != 0.0 { let z = product(&mut vars, *selector, v); @@ -419,8 +434,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> + built_once_shared[c].iter().map(|(_, w)| w).sum::(); let a = vars.add(variable().min(0)); let b = vars.add(variable().min(0)); - let u = vars.add(variable().binary()); + let u = vars.add(binary()); objective += a + b; + objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; + if std::env::var_os("JP_MILP_NOMIN").is_some() { + continue; + } cheaper_of.push((a, b, u, c, big_m)); } for (c, choice) in input.choices.iter().enumerate() { @@ -429,12 +448,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if slot.cost_per_unit_rate != 0.0 && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - for (set, selector) in &rate_of[&(choice.group.0 as usize)] { - rated.push(( - *selector, - x[c], - input.union_rate[set] * slot.cost_per_unit_rate, - )); + let g = choice.group.0 as usize; + let slowest = slowest_rate(g); + objective += slowest * slot.cost_per_unit_rate * x[c]; + for (set, selector) in &rate_of[&g] { + let union = input.union_rate[set]; + if union > slowest { + rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); + } } } } @@ -455,7 +476,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } - let vars_exclusive = (vars.add(variable().binary()), vars.add(variable().binary())); + let vars_exclusive = (vars.add(binary()), vars.add(binary())); let source_count = input .raw_scans .iter() @@ -509,15 +530,29 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + // Aggregated over a group's choices (one is chosen when it is + // needed), which keeps the LP relaxation tight. + let mut choices_of: Vec> = vec![Vec::new(); n]; for (c, choice) in input.choices.iter().enumerate() { - let g = choice.group.0 as usize; - for child in &choice.children { + choices_of[choice.group.0 as usize].push(c); + } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); + } + for h in children { + let mut with_child = Expression::from(0.0); + for &c in of_g { + if input.choices[c].children.iter().any(|k| k.0 as usize == h) { + with_child += x[c]; + } + } for &q in &input.group_queries[g] { - if let (Some(&above), Some(&below)) = ( - reaches_at_all.get(&(q, g)), - reaches_at_all.get(&(q, child.0 as usize)), - ) { - constraints.push(constraint!(below >= above + x[c] - 1)); + if let (Some(&above), Some(&below)) = + (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) + { + constraints.push(constraint!(below >= above + with_child.clone() - 1)); } } } @@ -578,19 +613,26 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let g = input.choices[c].group.0 as usize; constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); } + for (&(k, g), &p) in &passes { + let mut into = Expression::from(0.0); + for &c in &choices_of[g] { + into += y[&(k, c)]; + } + // Unless a fused choice of another group implements it. + for (c, choice) in input.choices.iter().enumerate() { + if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { + into += x[c]; + } + } + constraints.push(constraint!(into >= p)); + } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { let h = child.0 as usize; match via { Via::Direct => { - for &(k, g) in passes.keys() { - if g != choice.group.0 as usize { - continue; - } - if let Some(&r) = reach.get(&(k, h)) { - constraints.push(constraint!(r >= y[&(k, c)])); - } - } + // Aggregated below, over the group's choices. + let _ = h; } Via::Slot(slot) => { let k = Consumer::Slot { @@ -604,6 +646,43 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + for (child, via) in input.choices[c] + .children + .iter() + .zip(&input.choices[c].coefficients.child_via) + { + if *via == Via::Direct { + children.insert(child.0 as usize); + } + } + } + for h in children { + for &(k, pg) in passes.keys() { + if pg != g { + continue; + } + let Some(&r) = reach.get(&(k, h)) else { + continue; + }; + let mut into = Expression::from(0.0); + for &c in of_g { + let choice = &input.choices[c]; + if choice + .children + .iter() + .zip(&choice.coefficients.child_via) + .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) + { + into += y[&(k, c)]; + } + } + constraints.push(constraint!(r >= into)); + } + } + } for &(z, a, b) in products.iter().chain(&products_pending) { constraints.push(constraint!(z >= a + b - 1)); } @@ -616,8 +695,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &(v, w) in &built_once_shared[c] { built_once += w * v; } - constraints.push(constraint!(a >= per_consumer - big_m * u)); - constraints.push(constraint!(b >= built_once - big_m * (1 - u))); + // Both are at least the slowest class's rate (some query reaches + // the chosen choice): charged linearly, which keeps the relaxation + // tight; only the excess goes through the binary choice. + let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); + constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); + constraints.push(constraint!( + b >= built_once - floor * x[c] - big_m * (1 - u) + )); } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { @@ -704,10 +789,13 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } let constraint_count = constraints.len() as u64; - let model = model + let mut model = model .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; + if std::env::var_os("JP_MILP_VERBOSE").is_some() { + model.set_verbose(true); + } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), @@ -720,7 +808,25 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> .map(|(c, _)| c) .collect(); if std::env::var_os("JP_MILP_TRACE").is_some() { - for &c in &chosen { + if relaxed { + for (g, v) in once.iter() { + eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); + } + for (g, v) in need.iter().enumerate() { + eprintln!(" need g{} = {:.3}", g, solution.value(*v)); + } + } + let shown: Vec = if relaxed { + (0..x.len()) + .filter(|&c| solution.value(x[c]) > 1e-6) + .collect() + } else { + chosen.clone() + }; + for &c in &shown { + if relaxed { + eprintln!(" x = {:.3}", solution.value(x[c])); + } let choice: &Choice = &input.choices[c]; let reach: Vec = y .iter() From a64198fae680e781e923dd15fe8da5a766ff931a Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:44:24 +0000 Subject: [PATCH 38/67] Memo MILP search: a no-good cut excludes only its own group choice (a plan needing more groups is another plan); a debugging check of the model's lower bound per group choice --- .../plan-selection/src/pass/subexpr/milp.rs | 91 ++++++++++++++++++- 1 file changed, 86 insertions(+), 5 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 61e3748d2..f41161575 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -145,7 +145,9 @@ pub(crate) struct ModelInput { pub union_rate: HashMap, f64>, /// No-good cuts from earlier iterations: at most `bound` of these /// choices together (a plan whose every group took one of them). - pub excluded: Vec<(Vec, f64)>, + /// With the groups that plan does not need: a plan needing any of them + /// is another plan, not cut. + pub excluded: Vec<(Vec, Vec, f64)>, /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all /// of a choice set is chosen (local cuts). pub extra_bounds: Vec<(Vec, f64)>, @@ -718,11 +720,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> constraints.push(constraint!(latency <= bound)); } // Cuts. - for (set, bound) in &input.excluded { + for (set, outside, bound) in &input.excluded { let mut chosen = Expression::from(0.0); for &c in set { chosen += x[c]; } + for &g in outside { + chosen -= need[g]; + } constraints.push(constraint!(chosen <= *bound)); } for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { @@ -1054,7 +1059,7 @@ pub(crate) fn search( == group_choice[b.0 .0 as usize].unwrap()) }) .collect(); - input.excluded.push((members, 1.0)); + input.excluded.push((members, Vec::new(), 1.0)); cut = true; } } @@ -1065,7 +1070,7 @@ pub(crate) fn search( && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] }) .collect(); - input.excluded.push((members, 0.0)); + input.excluded.push((members, Vec::new(), 0.0)); cut = true; } } @@ -1096,7 +1101,65 @@ pub(crate) fn search( .filter(|(_, choice)| group_choice[choice.group.0 as usize] == Some(choice.candidate)) .map(|(c, _)| c) .collect(); - input.excluded.push((excluded, groups as f64 - 1.0)); + let mut covered = vec![false; memo.groups.len()]; + for (g, c) in group_choice.iter().enumerate() { + if let Some(c) = c { + for i in &memo.candidates[g][*c].implements { + covered[i.0 as usize] = true; + } + } + } + let outside = (0..memo.groups.len()).filter(|&g| !covered[g]).collect(); + input + .excluded + .push((excluded, outside, groups as f64 - 1.0)); + } + if std::env::var_os("JP_MILP_CHECK").is_some() { + // Debugging aid: every group choice's exact price against the + // model's value restricted to it. + let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); + for group_choice in super::search::group_choices(&memo).unwrap_or_default() { + let mut one = PlanSelector::new(entries, data, models, false)?; + let mut scratch = HashMap::new(); + let mut k = 0usize; + super::search::offer_group_choice( + &memo, + per_computation, + &group_choice, + entries, + data, + models, + &mut one, + &mut scratch, + &mut k, + )?; + let Some(exact) = one.cheapest_price() else { + continue; + }; + let mut restricted = model_input_clone(&input); + restricted.excluded.clear(); + let outside: Vec = (0..restricted.choices.len()) + .filter(|&c| { + group_choice[restricted.choices[c].group.0 as usize] + != Some(restricted.choices[c].candidate) + }) + .collect(); + restricted.excluded.push((outside, Vec::new(), 0.0)); + let value = solve(&restricted) + .map_err(SelectionError::Solver)? + .map(|s| s.objective); + let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); + if bad || exact < incumbent * (1.0 - 1e-9) { + eprintln!( + "check: {group_choice:?} exact {exact:.7e} model {value:?}{}", + if bad { + " NOT A LOWER BOUND" + } else { + " BELOW INCUMBENT" + } + ); + } + } } let (mut selection, selected) = selector.finish()?; let price = selection.costs[&selection.selected].cost_per_second; @@ -1119,6 +1182,24 @@ pub(crate) fn search( Ok((selection, selected, roots)) } +fn model_input_clone(input: &ModelInput) -> ModelInput { + ModelInput { + group_count: input.group_count, + group_queries: input.group_queries.clone(), + choices: input.choices.clone(), + group_slots: input.group_slots.clone(), + queries: input.queries.clone(), + schedule: input.schedule.clone(), + union_rate: input.union_rate.clone(), + excluded: input.excluded.clone(), + extra_bounds: input.extra_bounds.clone(), + requires: input.requires.clone(), + exclusive: input.exclusive.clone(), + raw_scans: input.raw_scans.clone(), + raw_data_retained: input.raw_data_retained, + } +} + /// The model's input for `memo`'s `choices`. fn model_input( memo: &Memo, From 0ca983521d1d64bc5609c9275e2f4685ac6abfdc Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:45:16 +0000 Subject: [PATCH 39/67] plan-selection: memo MILP encodes the rollup requirements as constraints Each rollup needs a finer group candidate whose build is its coarse build's source; a rollup, its finer candidate and a third group's candidate whose build beats that source are never chosen together; for a sketch rollup, every unrolled coarse build whose source is the rolled state needs a candidate whose build beats it. The lazy pair cut stays as a safety net. --- .../src/pass/subexpr/coefficients.rs | 21 - .../plan-selection/src/pass/subexpr/milp.rs | 497 +++++++++--------- .../src/pass/subexpr/windows.rs | 155 ++++++ 3 files changed, 418 insertions(+), 255 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index bea3da6e7..07172f095 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -263,27 +263,6 @@ pub(super) fn choice_parts( } let mut seen: Vec> = Vec::new(); for physical in &physicals { - // A shared scan's horizon is the longest window of - // its readers: no shorter than this choice's own - // window on it (through time shifts). - if !horizons.is_empty() { - let dag = &physical.dag; - let nodes = node_map(dag); - let kept = kept_pane_chains(dag); - let too_short = dag.nodes.iter().any(|n| match &n.payload { - Operator::NonASAP(NonASAPOp::Scan { source, .. }) => horizons - .get(&super::scans::raw_horizon_key(source)) - .is_some_and(|&h| { - windows_read(dag, &nodes, n.id, &kept) - .into_iter() - .any(|w| w > h) - }), - _ => false, - }); - if too_short { - continue; - } - } let keys = keys_of(&physical.dag, cache.keys_mut()); let shallow = shallow_keys(&physical.dag, cache.keys_mut()); let private: Vec = diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index f41161575..0f55b614c 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -145,9 +145,7 @@ pub(crate) struct ModelInput { pub union_rate: HashMap, f64>, /// No-good cuts from earlier iterations: at most `bound` of these /// choices together (a plan whose every group took one of them). - /// With the groups that plan does not need: a plan needing any of them - /// is another plan, not cut. - pub excluded: Vec<(Vec, Vec, f64)>, + pub excluded: Vec<(Vec, f64)>, /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all /// of a choice set is chosen (local cuts). pub extra_bounds: Vec<(Vec, f64)>, @@ -158,6 +156,10 @@ pub(crate) struct ModelInput { /// Choices that may not be chosen together with any of the second set /// (panes and an exact top-k reuse in one plan). pub exclusive: (Vec, Vec), + /// The rollup requirements ([`super::windows::rollup_rules`]): when a + /// choice of each set of the first list is chosen (each set a + /// candidate's choices), one of the second is. + pub rollup_rules: Vec<(Vec>, Vec)>, /// Each query-time raw scan of a choice's part: each source's longest /// one is kept on disk. pub raw_scans: Vec, @@ -233,27 +235,23 @@ enum Consumer { pub(crate) fn solve(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let relaxed = std::env::var_os("JP_MILP_LP").is_some(); - let binary = || { - if relaxed { - variable().min(0).max(1) - } else { - variable().binary() - } - }; - let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); - let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); + let x: Vec = input + .choices + .iter() + .map(|_| vars.add(variable().binary())) + .collect(); + let need: Vec = (0..n).map(|_| vars.add(variable().binary())).collect(); let once: BTreeMap = input .group_slots .iter() - .map(|slot| (slot.group, vars.add(binary()))) + .map(|slot| (slot.group, vars.add(variable().binary()))) .collect(); // Which queries' plans reach each group at all (any materialization): // the union rate of a node built once reads them. let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); for (g, queries) in input.group_queries.iter().enumerate() { for &q in queries { - reaches_at_all.insert((q, g), vars.add(binary())); + reaches_at_all.insert((q, g), vars.add(variable().binary())); } } // Consumers that may reach each group's own node, and through it its @@ -327,7 +325,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &g in &groups_with_rates { let sets = class_sets(g) .into_iter() - .map(|set| (set, vars.add(binary()))) + .map(|set| (set, vars.add(variable().binary()))) .collect(); rate_of.insert(g, sets); } @@ -347,12 +345,6 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Group(_) => None, } }; - let slowest_rate = |g: usize| -> f64 { - rate_of[&g] - .iter() - .map(|(set, _)| input.union_rate[set]) - .fold(f64::INFINITY, f64::min) - }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); // The shareable nodes of a choice's direct part cost, per unit of @@ -399,15 +391,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, Consumer::Query(_) => unreachable!("a query's rate is constant"), }; - // At least the slowest class's rate is charged linearly - // (some query reaches the group), the rest through the - // selector products: a tighter relaxation. - let slowest = slowest_rate(g); - objective += slowest * cost * v; for (set, selector) in &rate_of[&g] { let union = input.union_rate[set]; - if cost != 0.0 && union > slowest { - rated.push((*selector, v, (union - slowest) * cost)); + if cost != 0.0 { + rated.push((*selector, v, union * cost)); } if shared != 0.0 { let z = product(&mut vars, *selector, v); @@ -436,12 +423,8 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> + built_once_shared[c].iter().map(|(_, w)| w).sum::(); let a = vars.add(variable().min(0)); let b = vars.add(variable().min(0)); - let u = vars.add(binary()); + let u = vars.add(variable().binary()); objective += a + b; - objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; - if std::env::var_os("JP_MILP_NOMIN").is_some() { - continue; - } cheaper_of.push((a, b, u, c, big_m)); } for (c, choice) in input.choices.iter().enumerate() { @@ -450,14 +433,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if slot.cost_per_unit_rate != 0.0 && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - let g = choice.group.0 as usize; - let slowest = slowest_rate(g); - objective += slowest * slot.cost_per_unit_rate * x[c]; - for (set, selector) in &rate_of[&g] { - let union = input.union_rate[set]; - if union > slowest { - rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); - } + for (set, selector) in &rate_of[&(choice.group.0 as usize)] { + rated.push(( + *selector, + x[c], + input.union_rate[set] * slot.cost_per_unit_rate, + )); } } } @@ -478,7 +459,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } - let vars_exclusive = (vars.add(binary()), vars.add(binary())); + let vars_exclusive = (vars.add(variable().binary()), vars.add(variable().binary())); let source_count = input .raw_scans .iter() @@ -532,29 +513,15 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } - // Aggregated over a group's choices (one is chosen when it is - // needed), which keeps the LP relaxation tight. - let mut choices_of: Vec> = vec![Vec::new(); n]; for (c, choice) in input.choices.iter().enumerate() { - choices_of[choice.group.0 as usize].push(c); - } - for (g, of_g) in choices_of.iter().enumerate() { - let mut children: BTreeSet = BTreeSet::new(); - for &c in of_g { - children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); - } - for h in children { - let mut with_child = Expression::from(0.0); - for &c in of_g { - if input.choices[c].children.iter().any(|k| k.0 as usize == h) { - with_child += x[c]; - } - } + let g = choice.group.0 as usize; + for child in &choice.children { for &q in &input.group_queries[g] { - if let (Some(&above), Some(&below)) = - (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) - { - constraints.push(constraint!(below >= above + with_child.clone() - 1)); + if let (Some(&above), Some(&below)) = ( + reaches_at_all.get(&(q, g)), + reaches_at_all.get(&(q, child.0 as usize)), + ) { + constraints.push(constraint!(below >= above + x[c] - 1)); } } } @@ -615,26 +582,19 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let g = input.choices[c].group.0 as usize; constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); } - for (&(k, g), &p) in &passes { - let mut into = Expression::from(0.0); - for &c in &choices_of[g] { - into += y[&(k, c)]; - } - // Unless a fused choice of another group implements it. - for (c, choice) in input.choices.iter().enumerate() { - if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { - into += x[c]; - } - } - constraints.push(constraint!(into >= p)); - } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { let h = child.0 as usize; match via { Via::Direct => { - // Aggregated below, over the group's choices. - let _ = h; + for &(k, g) in passes.keys() { + if g != choice.group.0 as usize { + continue; + } + if let Some(&r) = reach.get(&(k, h)) { + constraints.push(constraint!(r >= y[&(k, c)])); + } + } } Via::Slot(slot) => { let k = Consumer::Slot { @@ -648,43 +608,6 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } - for (g, of_g) in choices_of.iter().enumerate() { - let mut children: BTreeSet = BTreeSet::new(); - for &c in of_g { - for (child, via) in input.choices[c] - .children - .iter() - .zip(&input.choices[c].coefficients.child_via) - { - if *via == Via::Direct { - children.insert(child.0 as usize); - } - } - } - for h in children { - for &(k, pg) in passes.keys() { - if pg != g { - continue; - } - let Some(&r) = reach.get(&(k, h)) else { - continue; - }; - let mut into = Expression::from(0.0); - for &c in of_g { - let choice = &input.choices[c]; - if choice - .children - .iter() - .zip(&choice.coefficients.child_via) - .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) - { - into += y[&(k, c)]; - } - } - constraints.push(constraint!(r >= into)); - } - } - } for &(z, a, b) in products.iter().chain(&products_pending) { constraints.push(constraint!(z >= a + b - 1)); } @@ -697,14 +620,8 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &(v, w) in &built_once_shared[c] { built_once += w * v; } - // Both are at least the slowest class's rate (some query reaches - // the chosen choice): charged linearly, which keeps the relaxation - // tight; only the excess goes through the binary choice. - let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); - constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); - constraints.push(constraint!( - b >= built_once - floor * x[c] - big_m * (1 - u) - )); + constraints.push(constraint!(a >= per_consumer - big_m * u)); + constraints.push(constraint!(b >= built_once - big_m * (1 - u))); } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { @@ -720,14 +637,11 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> constraints.push(constraint!(latency <= bound)); } // Cuts. - for (set, outside, bound) in &input.excluded { + for (set, bound) in &input.excluded { let mut chosen = Expression::from(0.0); for &c in set { chosen += x[c]; } - for &g in outside { - chosen -= need[g]; - } constraints.push(constraint!(chosen <= *bound)); } for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { @@ -746,6 +660,18 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } constraints.push(constraint!(x[*c] <= any)); } + for (all_of, one_of) in &input.rollup_rules { + let mut chosen = Expression::from(0.0); + for set in all_of { + for &c in set { + chosen += x[c]; + } + } + for &c in one_of { + chosen -= x[c]; + } + constraints.push(constraint!(chosen <= all_of.len() as f64 - 1.0)); + } let (first, second) = &input.exclusive; if !first.is_empty() && !second.is_empty() { let (p, t) = (vars_exclusive.0, vars_exclusive.1); @@ -794,13 +720,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } let constraint_count = constraints.len() as u64; - let mut model = model + let model = model .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; - if std::env::var_os("JP_MILP_VERBOSE").is_some() { - model.set_verbose(true); - } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), @@ -813,25 +736,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> .map(|(c, _)| c) .collect(); if std::env::var_os("JP_MILP_TRACE").is_some() { - if relaxed { - for (g, v) in once.iter() { - eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); - } - for (g, v) in need.iter().enumerate() { - eprintln!(" need g{} = {:.3}", g, solution.value(*v)); - } - } - let shown: Vec = if relaxed { - (0..x.len()) - .filter(|&c| solution.value(x[c]) > 1e-6) - .collect() - } else { - chosen.clone() - }; - for &c in &shown { - if relaxed { - eprintln!(" x = {:.3}", solution.value(x[c])); - } + for &c in &chosen { let choice: &Choice = &input.choices[c]; let reach: Vec = y .iter() @@ -1040,8 +945,9 @@ pub(crate) fn search( .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) .collect(); if !super::windows::requirements_met(&chosen_candidates) { - // A rollup's requirement: cut the smallest failing pairs of - // (group, candidate), else the whole group choice below. + // The model encodes every requirement (`rollup_rules`), so this + // is a safety net: cut the smallest failing pairs of (group, + // candidate), else the whole group choice below. let mut cut = false; for (i, a) in chosen_candidates.iter().enumerate() { if a.1.rollup_of.is_none() { @@ -1059,7 +965,7 @@ pub(crate) fn search( == group_choice[b.0 .0 as usize].unwrap()) }) .collect(); - input.excluded.push((members, Vec::new(), 1.0)); + input.excluded.push((members, 1.0)); cut = true; } } @@ -1070,7 +976,7 @@ pub(crate) fn search( && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] }) .collect(); - input.excluded.push((members, Vec::new(), 0.0)); + input.excluded.push((members, 0.0)); cut = true; } } @@ -1101,65 +1007,7 @@ pub(crate) fn search( .filter(|(_, choice)| group_choice[choice.group.0 as usize] == Some(choice.candidate)) .map(|(c, _)| c) .collect(); - let mut covered = vec![false; memo.groups.len()]; - for (g, c) in group_choice.iter().enumerate() { - if let Some(c) = c { - for i in &memo.candidates[g][*c].implements { - covered[i.0 as usize] = true; - } - } - } - let outside = (0..memo.groups.len()).filter(|&g| !covered[g]).collect(); - input - .excluded - .push((excluded, outside, groups as f64 - 1.0)); - } - if std::env::var_os("JP_MILP_CHECK").is_some() { - // Debugging aid: every group choice's exact price against the - // model's value restricted to it. - let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); - for group_choice in super::search::group_choices(&memo).unwrap_or_default() { - let mut one = PlanSelector::new(entries, data, models, false)?; - let mut scratch = HashMap::new(); - let mut k = 0usize; - super::search::offer_group_choice( - &memo, - per_computation, - &group_choice, - entries, - data, - models, - &mut one, - &mut scratch, - &mut k, - )?; - let Some(exact) = one.cheapest_price() else { - continue; - }; - let mut restricted = model_input_clone(&input); - restricted.excluded.clear(); - let outside: Vec = (0..restricted.choices.len()) - .filter(|&c| { - group_choice[restricted.choices[c].group.0 as usize] - != Some(restricted.choices[c].candidate) - }) - .collect(); - restricted.excluded.push((outside, Vec::new(), 0.0)); - let value = solve(&restricted) - .map_err(SelectionError::Solver)? - .map(|s| s.objective); - let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); - if bad || exact < incumbent * (1.0 - 1e-9) { - eprintln!( - "check: {group_choice:?} exact {exact:.7e} model {value:?}{}", - if bad { - " NOT A LOWER BOUND" - } else { - " BELOW INCUMBENT" - } - ); - } - } + input.excluded.push((excluded, groups as f64 - 1.0)); } let (mut selection, selected) = selector.finish()?; let price = selection.costs[&selection.selected].cost_per_second; @@ -1182,24 +1030,6 @@ pub(crate) fn search( Ok((selection, selected, roots)) } -fn model_input_clone(input: &ModelInput) -> ModelInput { - ModelInput { - group_count: input.group_count, - group_queries: input.group_queries.clone(), - choices: input.choices.clone(), - group_slots: input.group_slots.clone(), - queries: input.queries.clone(), - schedule: input.schedule.clone(), - union_rate: input.union_rate.clone(), - excluded: input.excluded.clone(), - extra_bounds: input.extra_bounds.clone(), - requires: input.requires.clone(), - exclusive: input.exclusive.clone(), - raw_scans: input.raw_scans.clone(), - raw_data_retained: input.raw_data_retained, - } -} - /// The model's input for `memo`'s `choices`. fn model_input( memo: &Memo, @@ -1272,8 +1102,7 @@ fn model_input( } } // Requirements the reference search checks on a group choice - // ([`super::windows::requirements_met`]) other than rollups, which the - // loop cuts lazily. + // ([`super::windows::requirements_met`]); rollups below. let candidate_of = |c: &Choice| &memo.candidates[c.group.0 as usize][c.candidate]; let mut requires = Vec::new(); for (c, choice) in choices.iter().enumerate() { @@ -1307,6 +1136,28 @@ fn model_input( requires.push((c, allowed)); } } + // Rollups: every rule over the candidates with a choice, each + // candidate standing for its choices (one per group is chosen). + let mut choices_of: BTreeMap<(usize, usize), Vec> = BTreeMap::new(); + for (c, choice) in choices.iter().enumerate() { + choices_of + .entry((choice.group.0 as usize, choice.candidate)) + .or_default() + .push(c); + } + let present: BTreeSet<(usize, usize)> = choices_of.keys().copied().collect(); + let rollup_rules = super::windows::rollup_rules(memo, &present) + .into_iter() + .map(|rule| { + ( + rule.all_of.iter().map(|k| choices_of[k].clone()).collect(), + rule.one_of + .iter() + .flat_map(|k| choices_of[k].clone()) + .collect(), + ) + }) + .collect(); let of_origin = |origin: super::CandidateOrigin| -> Vec { (0..choices.len()) .filter(|&c| candidate_of(&choices[c]).origin == origin) @@ -1342,6 +1193,7 @@ fn model_input( extra_bounds: Vec::new(), requires, exclusive, + rollup_rules, raw_scans: Vec::new(), raw_data_retained: true, }) @@ -1350,6 +1202,7 @@ fn model_input( #[cfg(test)] mod tests { use super::*; + use asap_types::ir::properties::AccuracyTarget; fn choice(group: u32, children: &[u32], cost: f64) -> Choice { Choice { @@ -1408,6 +1261,7 @@ mod tests { extra_bounds: Vec::new(), requires: Vec::new(), exclusive: (Vec::new(), Vec::new()), + rollup_rules: Vec::new(), raw_scans: Vec::new(), raw_data_retained: true, }; @@ -1430,4 +1284,179 @@ mod tests { solution.objective ); } + + /// A group choice: per group, its candidate. + type Plan = Vec>; + + /// The plans of `queries`' memo (each query every minute, at `targets`) + /// the model may return, as group choices: the model of every choice + /// with no latency bound, solved with a no-good cut on each group + /// choice it returns until it is infeasible; and the group choices of + /// the same candidates the reference search accepts + /// ([`super::super::windows::requirements_met`]). + fn model_and_reference_plans( + queries: &[&str], + targets: &[AccuracyTarget], + ) -> (BTreeSet, BTreeSet) { + 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::schema::schema_support::with_promql_series_identity; + use asap_types::workload::{ + DataArrival, Predictability, QueryRecurrence, RepeatedDemand, RepetitionInterval, + SourceFacts, + }; + let roots: Vec<(usize, QueryRoot)> = queries + .iter() + .zip(targets) + .map(|(q, target)| { + let root = lower_promql(q, target.clone()); + QueryRoot::Operator(with_promql_series_identity(&root).unwrap()) + }) + .enumerate() + .collect(); + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let entries: Vec = targets + .iter() + .map(|target| { + let mut entry = entry_with( + target.clone(), + QueryRecurrence::Repeated(RepeatedDemand::FixedInterval(RepetitionInterval( + 60_000, + ))), + ); + entry.predictability = Predictability::Predictable { known_at: None }; + entry + }) + .collect(); + let data = 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(), + }; + let models = PlanningModels::builtin(); + let memo = super::super::search::memo(&per_computation, &entries); + let scale = + super::super::pricing::frozen_calibration(&memo, &entries, &data, models).unwrap(); + let mut cache = super::super::pricing::PriceCache::new(scale); + let mut choices = Vec::new(); + for g in 0..memo.groups.len() { + let (group_choices, _) = super::super::coefficients::choices( + &memo, + GroupId(g as u32), + &entries, + &data, + models, + &mut cache, + ) + .unwrap(); + choices.extend(group_choices.into_iter().filter(|choice| { + choice.copy_of.is_none() + && memo.candidates[g][choice.candidate].origin + != super::super::CandidateOrigin::WidenedScans + })); + } + let mut input = model_input(&memo, choices, &entries, models).unwrap(); + assert!( + input + .choices + .iter() + .any(|c| memo.candidates[c.group.0 as usize][c.candidate] + .rollup_of + .is_some()), + "a rollup is a choice" + ); + for query in &mut input.queries { + query.max_latency_ms = None; + } + let present: BTreeSet<(usize, usize)> = input + .choices + .iter() + .map(|c| (c.group.0 as usize, c.candidate)) + .collect(); + let mut model: BTreeSet>> = BTreeSet::new(); + while let Some(solution) = solve(&input).unwrap() { + let mut plan = vec![None; memo.groups.len()]; + for &c in &solution.chosen { + plan[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); + } + let excluded: Vec = (0..input.choices.len()) + .filter(|&c| { + plan[input.choices[c].group.0 as usize] == Some(input.choices[c].candidate) + }) + .collect(); + let groups = plan.iter().flatten().count(); + input.excluded.push((excluded, groups as f64 - 1.0)); + assert!(model.insert(plan), "a group choice returned twice"); + } + // Only groups with a choice to make: one with a single candidate + // is in a plan or not by how its readers reach it (a scan read + // apart is each reader's own copy), not by a choice. + let decided: Vec = (0..memo.groups.len()) + .map(|g| present.iter().filter(|(h, _)| *h == g).count() > 1) + .collect(); + let project = |plan: Vec>| -> Vec> { + plan.into_iter() + .zip(&decided) + .map(|(c, &decided)| c.filter(|_| decided)) + .collect() + }; + let model = model.into_iter().map(project).collect(); + let reference = super::super::search::group_choices(&memo) + .unwrap() + .into_iter() + .filter(|plan| { + plan.iter() + .enumerate() + .all(|(g, c)| c.is_none_or(|c| present.contains(&(g, c)))) + }) + .filter(|plan| { + let chosen: Vec<(GroupId, &super::super::ImplementationCandidate)> = plan + .iter() + .enumerate() + .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) + .collect(); + super::super::windows::requirements_met(&chosen) + }) + .map(project) + .collect(); + (model, reference) + } + + /// Exact sums by `job`, `(job, instance)` and `(job, instance, zone)`: + /// a coarser sum may roll up from either finer one, but the joint rule + /// rolls it up from the coarsest finer build chosen, which must be exact. + /// The group choices the model may return are exactly those the + /// reference search accepts. + #[test] + fn the_model_admits_exactly_the_plans_whose_rollups_meet_the_joint_rule() { + let (model, reference) = model_and_reference_plans( + &[ + "sum by (job) (m)", + "sum by (job, instance) (m)", + "sum by (job, instance, zone) (m)", + ], + &[ + AccuracyTarget::exact(), + AccuracyTarget::exact(), + AccuracyTarget::exact(), + ], + ); + assert!(!reference.is_empty()); + let only_model: Vec<_> = model.difference(&reference).collect(); + let only_reference: Vec<_> = reference.difference(&model).collect(); + assert!( + only_model.is_empty() && only_reference.is_empty(), + "model only {only_model:?}, reference only {only_reference:?} ({} plans)", + reference.len() + ); + } } diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index 6196df94b..183012553 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -943,6 +943,161 @@ pub(crate) fn requirements_met(chosen: &[(GroupId, &ImplementationCandidate)]) - .all(|(coarse, _)| rolled_coarse.contains(coarse)) } +/// One rollup requirement of [`requirements_met`] as a rule over +/// candidates, by `(group, candidate index)`: when every candidate of +/// `all_of` is chosen, one of `one_of` is (none: they may not be chosen +/// together). +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RollupRule { + pub all_of: Vec<(usize, usize)>, + pub one_of: Vec<(usize, usize)>, +} + +/// The rollup requirements of [`requirements_met`] over the candidates in +/// `present` (one per group in a plan), as rules: a plan of one present +/// candidate per needed group meets every rollup requirement iff it meets +/// every rule. The joint rule's source of a coarse build is the best build +/// over the plan's nodes (each rollup undone, in group order): fewest keys +/// for an exact accumulator, most for a sketch, the first in walk order +/// among equals, so a total order over the builds present, and a build is +/// the source iff no single other candidate's build beats it. For a rollup +/// `R` of coarse build `C` reading group `f`: +/// +/// - `R` needs one of `f`'s candidates whose build is `C`'s source with `R` +/// alone, exposing the state at `R`'s site for a sketch; +/// - `R`, such a candidate `F` and a candidate `Y` of a third group whose +/// build beats `F`'s for `C` are never chosen together; +/// - for a sketch, every unrolled coarse build `C'` (of `R`'s rewritten +/// candidate, of `F`, or of a third group's `Y`) whose source would be +/// `F`'s build needs a candidate `Z` of a fourth group whose build beats +/// it for `C'`. +pub(crate) fn rollup_rules(memo: &Memo, present: &BTreeSet<(usize, usize)>) -> Vec { + let candidate = |(g, c): (usize, usize)| &memo.candidates[g][c]; + // A candidate's node in the joint rule's plan: a rollup undone. + let plan_node = |key: (usize, usize)| match &candidate(key).rollup_of { + Some(read) => Rc::clone(&read.coarse_candidate), + None => Rc::clone(&candidate(key).node), + }; + // The joint rule's source of each coarse build over `nodes`, in group + // order. + let sources = + |nodes: &[(usize, Rc)]| -> Vec<(*const OperatorNode, Rc)> { + let mut nodes = nodes.to_vec(); + nodes.sort_by_key(|(g, _)| *g); + let roots: Vec = nodes + .into_iter() + .map(|(_, node)| QueryRoot::Operator(node)) + .collect(); + rollup_candidates(&roots) + .into_iter() + .map(|r| (Rc::as_ptr(&r.coarse), r.finer)) + .collect() + }; + let source_of = |coarse: *const OperatorNode, nodes: &[(usize, Rc)]| { + sources(nodes) + .into_iter() + .find(|(c, _)| *c == coarse) + .map(|(_, finer)| finer) + }; + let mut rules = Vec::new(); + for &rollup in present { + let Some(read) = &candidate(rollup).rollup_of else { + continue; + }; + let (coarse_group, finer_group) = (rollup.0, read.finer_group.0 as usize); + let coarse = Rc::as_ptr(&read.coarse_build); + let site = read.finer_state.as_ref().map(|state| CouplingSite { + key: state_key(state), + kind: CouplingKind::SharedBuildSize, + }); + let r_node = (coarse_group, plan_node(rollup)); + let mut allowed = Vec::new(); + for &finer in present.iter().filter(|(g, _)| *g == finer_group) { + let f_node = (finer_group, plan_node(finer)); + let Some(build) = source_of(coarse, &[r_node.clone(), f_node.clone()]) else { + continue; + }; + let in_finer = reached(&candidate(finer).node).contains(&Rc::as_ptr(&build)); + let state_ok = site.is_none_or(|site| { + candidate(finer).coupling_sites.contains(&site) && state_key(&build) == site.key + }); + if !in_finer || !state_ok { + continue; + } + allowed.push(finer); + let others = || { + present + .iter() + .copied() + .filter(move |(g, _)| *g != coarse_group && *g != finer_group) + }; + // A third group's build that beats the finer build for `C`. + for other in others() { + let nodes = [r_node.clone(), f_node.clone(), (other.0, plan_node(other))]; + if source_of(coarse, &nodes).is_none_or(|b| !Rc::ptr_eq(&b, &build)) { + rules.push(RollupRule { + all_of: vec![rollup, finer, other], + one_of: Vec::new(), + }); + } + } + if site.is_none() { + continue; + } + // Every unrolled coarse build whose source the state is rolls up + // too, unless a fourth group's build beats it. + let mut holders: Vec> = vec![None, Some(finer)]; + holders.extend( + others() + .filter(|&o| candidate(o).rollup_of.is_none()) + .map(Some), + ); + for holder in holders { + let mut nodes = vec![r_node.clone(), f_node.clone()]; + if let Some(h) = holder.filter(|&h| h != finer) { + nodes.push((h.0, plan_node(h))); + } + let holder_node = match holder { + None => Rc::clone(&r_node.1), + Some(h) => plan_node(h), + }; + let in_holder = reached(&holder_node); + for (other_coarse, finer_build) in sources(&nodes) { + if other_coarse == coarse + || !Rc::ptr_eq(&finer_build, &build) + || !in_holder.contains(&other_coarse) + { + continue; + } + let mut all_of = vec![rollup, finer]; + if let Some(h) = holder.filter(|&h| h != finer) { + all_of.push(h); + } + let groups: Vec = all_of.iter().map(|(g, _)| *g).collect(); + let one_of = present + .iter() + .copied() + .filter(|(g, _)| !groups.contains(g)) + .filter(|&z| { + let mut with = nodes.clone(); + with.push((z.0, plan_node(z))); + source_of(other_coarse, &with).is_none_or(|b| !Rc::ptr_eq(&b, &build)) + }) + .collect(); + rules.push(RollupRule { all_of, one_of }); + } + } + } + rules.push(RollupRule { + all_of: vec![rollup], + one_of: allowed, + }); + } + rules.sort_by(|a, b| (&a.all_of, &a.one_of).cmp(&(&b.all_of, &b.one_of))); + rules.dedup(); + rules +} + /// 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 { From fee059477afb1ca0a35ccbafd00bc143a89fea6e Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:02:03 +0000 Subject: [PATCH 40/67] plan-selection: memo shared-node coefficients are disjoint Each direct node of a choice belongs to its nearest shareable node (itself or above it), so the shared nodes' costs add up to at most the direct part's and the model, which charges their sum at least once, stays a lower bound. Before, a shareable node's part included the shareable nodes below it (an estimate over its merge), and the sum overcounted. --- .../src/pass/subexpr/coefficients.rs | 84 ++++++++++++------- 1 file changed, 54 insertions(+), 30 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 596298045..9ae6caebe 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -720,17 +720,38 @@ fn coefficients_of( let mut shared_nodes = Vec::new(); let mut shared_node_keys = Vec::new(); let direct_set: HashSet = direct.iter().copied().collect(); - for &id in &direct { - let node = &dag.nodes[position[&id]]; - let shareable = - matches!(node.payload, Operator::ASAP(_)) || !node.output_schema.unique_keys.is_empty(); - if id == root || !shareable || is_build(&node.payload) || !query_time(id) { - continue; - } - let below: Vec = walk(id) - .0 - .into_iter() - .filter(|n| direct_set.contains(n)) + // Each shareable node with the direct nodes below it. + let shareable: Vec<(PhysicalASAPNodeId, Vec)> = direct + .iter() + .copied() + .filter(|&id| { + let node = &dag.nodes[position[&id]]; + let shareable = matches!(node.payload, Operator::ASAP(_)) + || !node.output_schema.unique_keys.is_empty(); + id != root && shareable && !is_build(&node.payload) && query_time(id) + }) + .map(|id| { + let below = walk(id) + .0 + .into_iter() + .filter(|n| direct_set.contains(n)) + .collect(); + (id, below) + }) + .collect(); + // Each direct node belongs to its nearest shareable node (itself or + // above it), so that no node is in two of them: the model charges + // their sum once. + for (id, below) in &shareable { + let id = *id; + let below: Vec = below + .iter() + .copied() + .filter(|n| { + !shareable.iter().any(|(other, other_below)| { + *other != id && below.contains(other) && other_below.contains(n) + }) + }) .collect(); let mut once = dag.clone(); once.nodes[position[&id]].materialization = Materialization::MatAtQuery(Medium::Memory); @@ -1408,23 +1429,31 @@ mod tests { if p.at_query_time && first_root { actual += coefficients.direct.cost_per_unit_rate * rates[&p.root].1; let mine = below(p.root); + let mut matched: HashSet = HashSet::new(); for (shared, sub) in coefficients .shared_nodes .iter() .zip(&p.part.shared_node_keys) { + // Charged once per plan node at its rate: the node is + // one node for the copies or groups sharing it, built + // once when it (or a node above it) is materialized. let Some(node) = mine.iter().copied().find(|id| { keys[position_of[id]] == sub[0] - && nodes[id].materialization != Materialization::NotMaterialized + && !matched.contains(id) + && nodes[id].output_state.timing == ExecutionTiming::QueryTime }) else { continue; }; + matched.insert(node); actual -= shared.cost_per_unit_rate * rates[&p.root].1; - // Its nodes' held memory, at the shared node's rate. - let below_held: Vec = lazy + let materialized = + nodes[&node].materialization != Materialization::NotMaterialized; + let part_held: Vec = lazy .held .iter() - .filter(|t| sub[1..].contains(&t.node_key)) + .filter(|t| sub.contains(&t.node_key)) + .filter(|t| !(materialized && t.node_key == sub[0])) .copied() .collect(); lazy.held.retain(|term| !sub.contains(&term.node_key)); @@ -1437,14 +1466,13 @@ mod tests { .map(|(q, _)| (q, query_rate(q))) .collect(); let rate: f64 = readers.iter().map(|(_, r)| r).sum(); - for term in below_held { + let longest = readers.iter().map(|&(q, _)| latency[q]).fold(0.0, f64::max); + for term in part_held { let evaluation_ms = match term.per_reader && rate > 0.0 { true => { readers.iter().map(|&(q, r)| r * latency[q]).sum::() / rate } - false => { - readers.iter().map(|&(q, _)| latency[q]).fold(0.0, f64::max) - } + false => longest, }; held_extra += prices.cost_per_memory_byte_second * term.bytes as f64 @@ -1452,17 +1480,13 @@ mod tests { * evaluation_ms / 1_000.0; } - let longest = reach - .iter() - .enumerate() - .filter(|(_, r)| r.contains(&node)) - .map(|(q, _)| latency[q]) - .fold(0.0, f64::max); - held_extra += prices.cost_per_memory_byte_second - * shared.state_bytes as f64 - * rates[&node].1 - * longest - / 1_000.0; + if materialized { + held_extra += prices.cost_per_memory_byte_second + * shared.state_bytes as f64 + * rates[&node].1 + * longest + / 1_000.0; + } } } } From fcabf39fa86c51b51b1a693be69732b044ce10d0 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:02:57 +0000 Subject: [PATCH 41/67] Restore the MILP model's tight relaxation, cut form and lower-bound check, and the scan-horizon filter, which the rollup-constraints commit dropped; keep its rollup rules --- .../src/pass/subexpr/coefficients.rs | 21 + .../plan-selection/src/pass/subexpr/milp.rs | 287 +++- .../src/pass/subexpr/milp.rs.orig | 1433 +++++++++++++++++ 3 files changed, 1697 insertions(+), 44 deletions(-) create mode 100644 crates/plan-selection/src/pass/subexpr/milp.rs.orig diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 9ae6caebe..3135a5908 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -263,6 +263,27 @@ pub(super) fn choice_parts( } let mut seen: Vec> = Vec::new(); for physical in &physicals { + // A shared scan's horizon is the longest window of + // its readers: no shorter than this choice's own + // window on it (through time shifts). + if !horizons.is_empty() { + let dag = &physical.dag; + let nodes = node_map(dag); + let kept = kept_pane_chains(dag); + let too_short = dag.nodes.iter().any(|n| match &n.payload { + Operator::NonASAP(NonASAPOp::Scan { source, .. }) => horizons + .get(&super::scans::raw_horizon_key(source)) + .is_some_and(|&h| { + windows_read(dag, &nodes, n.id, &kept) + .into_iter() + .any(|w| w > h) + }), + _ => false, + }); + if too_short { + continue; + } + } let keys = keys_of(&physical.dag, cache.keys_mut()); let shallow = shallow_keys(&physical.dag, cache.keys_mut()); let private: Vec = diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 0f55b614c..1488e0dd8 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -145,7 +145,9 @@ pub(crate) struct ModelInput { pub union_rate: HashMap, f64>, /// No-good cuts from earlier iterations: at most `bound` of these /// choices together (a plan whose every group took one of them). - pub excluded: Vec<(Vec, f64)>, + /// With the groups that plan does not need: a plan needing any of them + /// is another plan, not cut. + pub excluded: Vec<(Vec, Vec, f64)>, /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all /// of a choice set is chosen (local cuts). pub extra_bounds: Vec<(Vec, f64)>, @@ -235,23 +237,27 @@ enum Consumer { pub(crate) fn solve(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let x: Vec = input - .choices - .iter() - .map(|_| vars.add(variable().binary())) - .collect(); - let need: Vec = (0..n).map(|_| vars.add(variable().binary())).collect(); + let relaxed = std::env::var_os("JP_MILP_LP").is_some(); + let binary = || { + if relaxed { + variable().min(0).max(1) + } else { + variable().binary() + } + }; + let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); + let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); let once: BTreeMap = input .group_slots .iter() - .map(|slot| (slot.group, vars.add(variable().binary()))) + .map(|slot| (slot.group, vars.add(binary()))) .collect(); // Which queries' plans reach each group at all (any materialization): // the union rate of a node built once reads them. let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); for (g, queries) in input.group_queries.iter().enumerate() { for &q in queries { - reaches_at_all.insert((q, g), vars.add(variable().binary())); + reaches_at_all.insert((q, g), vars.add(binary())); } } // Consumers that may reach each group's own node, and through it its @@ -325,7 +331,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &g in &groups_with_rates { let sets = class_sets(g) .into_iter() - .map(|set| (set, vars.add(variable().binary()))) + .map(|set| (set, vars.add(binary()))) .collect(); rate_of.insert(g, sets); } @@ -345,6 +351,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Group(_) => None, } }; + let slowest_rate = |g: usize| -> f64 { + rate_of[&g] + .iter() + .map(|(set, _)| input.union_rate[set]) + .fold(f64::INFINITY, f64::min) + }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); // The shareable nodes of a choice's direct part cost, per unit of @@ -391,10 +403,15 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, Consumer::Query(_) => unreachable!("a query's rate is constant"), }; + // At least the slowest class's rate is charged linearly + // (some query reaches the group), the rest through the + // selector products: a tighter relaxation. + let slowest = slowest_rate(g); + objective += slowest * cost * v; for (set, selector) in &rate_of[&g] { let union = input.union_rate[set]; - if cost != 0.0 { - rated.push((*selector, v, union * cost)); + if cost != 0.0 && union > slowest { + rated.push((*selector, v, (union - slowest) * cost)); } if shared != 0.0 { let z = product(&mut vars, *selector, v); @@ -423,8 +440,12 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> + built_once_shared[c].iter().map(|(_, w)| w).sum::(); let a = vars.add(variable().min(0)); let b = vars.add(variable().min(0)); - let u = vars.add(variable().binary()); + let u = vars.add(binary()); objective += a + b; + objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; + if std::env::var_os("JP_MILP_NOMIN").is_some() { + continue; + } cheaper_of.push((a, b, u, c, big_m)); } for (c, choice) in input.choices.iter().enumerate() { @@ -433,12 +454,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if slot.cost_per_unit_rate != 0.0 && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - for (set, selector) in &rate_of[&(choice.group.0 as usize)] { - rated.push(( - *selector, - x[c], - input.union_rate[set] * slot.cost_per_unit_rate, - )); + let g = choice.group.0 as usize; + let slowest = slowest_rate(g); + objective += slowest * slot.cost_per_unit_rate * x[c]; + for (set, selector) in &rate_of[&g] { + let union = input.union_rate[set]; + if union > slowest { + rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); + } } } } @@ -459,7 +482,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } - let vars_exclusive = (vars.add(variable().binary()), vars.add(variable().binary())); + let vars_exclusive = (vars.add(binary()), vars.add(binary())); let source_count = input .raw_scans .iter() @@ -513,15 +536,29 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + // Aggregated over a group's choices (one is chosen when it is + // needed), which keeps the LP relaxation tight. + let mut choices_of: Vec> = vec![Vec::new(); n]; for (c, choice) in input.choices.iter().enumerate() { - let g = choice.group.0 as usize; - for child in &choice.children { + choices_of[choice.group.0 as usize].push(c); + } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); + } + for h in children { + let mut with_child = Expression::from(0.0); + for &c in of_g { + if input.choices[c].children.iter().any(|k| k.0 as usize == h) { + with_child += x[c]; + } + } for &q in &input.group_queries[g] { - if let (Some(&above), Some(&below)) = ( - reaches_at_all.get(&(q, g)), - reaches_at_all.get(&(q, child.0 as usize)), - ) { - constraints.push(constraint!(below >= above + x[c] - 1)); + if let (Some(&above), Some(&below)) = + (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) + { + constraints.push(constraint!(below >= above + with_child.clone() - 1)); } } } @@ -582,19 +619,26 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let g = input.choices[c].group.0 as usize; constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); } + for (&(k, g), &p) in &passes { + let mut into = Expression::from(0.0); + for &c in &choices_of[g] { + into += y[&(k, c)]; + } + // Unless a fused choice of another group implements it. + for (c, choice) in input.choices.iter().enumerate() { + if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { + into += x[c]; + } + } + constraints.push(constraint!(into >= p)); + } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { let h = child.0 as usize; match via { Via::Direct => { - for &(k, g) in passes.keys() { - if g != choice.group.0 as usize { - continue; - } - if let Some(&r) = reach.get(&(k, h)) { - constraints.push(constraint!(r >= y[&(k, c)])); - } - } + // Aggregated below, over the group's choices. + let _ = h; } Via::Slot(slot) => { let k = Consumer::Slot { @@ -608,6 +652,43 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + for (child, via) in input.choices[c] + .children + .iter() + .zip(&input.choices[c].coefficients.child_via) + { + if *via == Via::Direct { + children.insert(child.0 as usize); + } + } + } + for h in children { + for &(k, pg) in passes.keys() { + if pg != g { + continue; + } + let Some(&r) = reach.get(&(k, h)) else { + continue; + }; + let mut into = Expression::from(0.0); + for &c in of_g { + let choice = &input.choices[c]; + if choice + .children + .iter() + .zip(&choice.coefficients.child_via) + .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) + { + into += y[&(k, c)]; + } + } + constraints.push(constraint!(r >= into)); + } + } + } for &(z, a, b) in products.iter().chain(&products_pending) { constraints.push(constraint!(z >= a + b - 1)); } @@ -620,8 +701,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &(v, w) in &built_once_shared[c] { built_once += w * v; } - constraints.push(constraint!(a >= per_consumer - big_m * u)); - constraints.push(constraint!(b >= built_once - big_m * (1 - u))); + // Both are at least the slowest class's rate (some query reaches + // the chosen choice): charged linearly, which keeps the relaxation + // tight; only the excess goes through the binary choice. + let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); + constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); + constraints.push(constraint!( + b >= built_once - floor * x[c] - big_m * (1 - u) + )); } // Latency bounds: the work each query reaches at query time. for (q, query) in input.queries.iter().enumerate() { @@ -637,11 +724,14 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> constraints.push(constraint!(latency <= bound)); } // Cuts. - for (set, bound) in &input.excluded { + for (set, outside, bound) in &input.excluded { let mut chosen = Expression::from(0.0); for &c in set { chosen += x[c]; } + for &g in outside { + chosen -= need[g]; + } constraints.push(constraint!(chosen <= *bound)); } for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { @@ -720,10 +810,13 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } let constraint_count = constraints.len() as u64; - let model = model + let mut model = model .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; + if std::env::var_os("JP_MILP_VERBOSE").is_some() { + model.set_verbose(true); + } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), @@ -736,7 +829,25 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> .map(|(c, _)| c) .collect(); if std::env::var_os("JP_MILP_TRACE").is_some() { - for &c in &chosen { + if relaxed { + for (g, v) in once.iter() { + eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); + } + for (g, v) in need.iter().enumerate() { + eprintln!(" need g{} = {:.3}", g, solution.value(*v)); + } + } + let shown: Vec = if relaxed { + (0..x.len()) + .filter(|&c| solution.value(x[c]) > 1e-6) + .collect() + } else { + chosen.clone() + }; + for &c in &shown { + if relaxed { + eprintln!(" x = {:.3}", solution.value(x[c])); + } let choice: &Choice = &input.choices[c]; let reach: Vec = y .iter() @@ -965,7 +1076,7 @@ pub(crate) fn search( == group_choice[b.0 .0 as usize].unwrap()) }) .collect(); - input.excluded.push((members, 1.0)); + input.excluded.push((members, Vec::new(), 1.0)); cut = true; } } @@ -976,7 +1087,7 @@ pub(crate) fn search( && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] }) .collect(); - input.excluded.push((members, 0.0)); + input.excluded.push((members, Vec::new(), 0.0)); cut = true; } } @@ -1007,7 +1118,65 @@ pub(crate) fn search( .filter(|(_, choice)| group_choice[choice.group.0 as usize] == Some(choice.candidate)) .map(|(c, _)| c) .collect(); - input.excluded.push((excluded, groups as f64 - 1.0)); + let mut covered = vec![false; memo.groups.len()]; + for (g, c) in group_choice.iter().enumerate() { + if let Some(c) = c { + for i in &memo.candidates[g][*c].implements { + covered[i.0 as usize] = true; + } + } + } + let outside = (0..memo.groups.len()).filter(|&g| !covered[g]).collect(); + input + .excluded + .push((excluded, outside, groups as f64 - 1.0)); + } + if std::env::var_os("JP_MILP_CHECK").is_some() { + // Debugging aid: every group choice's exact price against the + // model's value restricted to it. + let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); + for group_choice in super::search::group_choices(&memo).unwrap_or_default() { + let mut one = PlanSelector::new(entries, data, models, false)?; + let mut scratch = HashMap::new(); + let mut k = 0usize; + super::search::offer_group_choice( + &memo, + per_computation, + &group_choice, + entries, + data, + models, + &mut one, + &mut scratch, + &mut k, + )?; + let Some(exact) = one.cheapest_price() else { + continue; + }; + let mut restricted = model_input_clone(&input); + restricted.excluded.clear(); + let outside: Vec = (0..restricted.choices.len()) + .filter(|&c| { + group_choice[restricted.choices[c].group.0 as usize] + != Some(restricted.choices[c].candidate) + }) + .collect(); + restricted.excluded.push((outside, Vec::new(), 0.0)); + let value = solve(&restricted) + .map_err(SelectionError::Solver)? + .map(|s| s.objective); + let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); + if bad || exact < incumbent * (1.0 - 1e-9) { + eprintln!( + "check: {group_choice:?} exact {exact:.7e} model {value:?}{}", + if bad { + " NOT A LOWER BOUND" + } else { + " BELOW INCUMBENT" + } + ); + } + } } let (mut selection, selected) = selector.finish()?; let price = selection.costs[&selection.selected].cost_per_second; @@ -1030,6 +1199,25 @@ pub(crate) fn search( Ok((selection, selected, roots)) } +fn model_input_clone(input: &ModelInput) -> ModelInput { + ModelInput { + group_count: input.group_count, + group_queries: input.group_queries.clone(), + choices: input.choices.clone(), + group_slots: input.group_slots.clone(), + queries: input.queries.clone(), + schedule: input.schedule.clone(), + union_rate: input.union_rate.clone(), + excluded: input.excluded.clone(), + extra_bounds: input.extra_bounds.clone(), + requires: input.requires.clone(), + exclusive: input.exclusive.clone(), + raw_scans: input.raw_scans.clone(), + raw_data_retained: input.raw_data_retained, + rollup_rules: input.rollup_rules.clone(), + } +} + /// The model's input for `memo`'s `choices`. fn model_input( memo: &Memo, @@ -1394,7 +1582,18 @@ mod tests { }) .collect(); let groups = plan.iter().flatten().count(); - input.excluded.push((excluded, groups as f64 - 1.0)); + let mut covered = vec![false; memo.groups.len()]; + for (g, c) in plan.iter().enumerate() { + if let Some(c) = c { + for i in &memo.candidates[g][*c].implements { + covered[i.0 as usize] = true; + } + } + } + let outside = (0..memo.groups.len()).filter(|&g| !covered[g]).collect(); + input + .excluded + .push((excluded, outside, groups as f64 - 1.0)); assert!(model.insert(plan), "a group choice returned twice"); } // Only groups with a choice to make: one with a single candidate diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs.orig b/crates/plan-selection/src/pass/subexpr/milp.rs.orig new file mode 100644 index 000000000..f41161575 --- /dev/null +++ b/crates/plan-selection/src/pass/subexpr/milp.rs.orig @@ -0,0 +1,1433 @@ +//! The memo's solver: one MILP over the memo's choices (design "The +//! solver: MILP over the memo"). +//! +//! A choice is one candidate of one group with an option for each of its +//! private materialization slots ([`Choice`]); its coefficients come from +//! Stage 3's per-node pricing ([`ChoiceCoefficients`]). The model picks one +//! choice per needed group, a materialization for each shared group's own +//! node, and propagates which consumers (queries, materialized slots) reach +//! which group, so each private node is charged once per consumer reaching +//! it, at that consumer's rate. Terms the model omits (memory held while a +//! query runs, composed accuracy) are added back by the caller's lazy cuts: +//! the model's optimum is a lower bound on the plan's Stage 3 price. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use good_lp::solvers::highs::highs; +use good_lp::{ + constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, WithMipGap, +}; + +use super::coefficients::LazyTerms; +use super::{GroupId, Memo}; +use crate::{CostModel, PlanSelector, PlanningModels, ResourceUnitCost, Selection, SelectionError}; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; +use asap_physical_asap_optimization::PhysicalCandidate; +use asap_types::ir::QueryRoot; +use asap_types::physical::Materialization; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; + +/// How a child group of a choice is reached: by the consumers reaching the +/// choice's own node, or by one of its materialized private slots. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Via { + Direct, + Slot(usize), +} + +/// What one part of a choice costs, from Stage 3's per-node pricing. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub(crate) struct PartCoefficients { + /// Price per second per unit of the rate of a consumer reaching it. + pub cost_per_unit_rate: f64, + /// Work in ms per computation, which a query reaching it at query time + /// adds to its latency. + pub work_ms: f64, +} + +/// One materialized private slot of a choice. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct SlotCoefficients { + pub materialization: Materialization, + /// Price per second that does not depend on a rate (ingestion-time + /// work and kept state). + pub fixed_cost: f64, + /// Price per second per unit of the slot's own rate (query-time slots: + /// the union rate of the queries reaching the group). + pub cost_per_unit_rate: f64, + /// Work in ms per computation at query time. + pub work_ms: f64, +} + +/// One choice's coefficients. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ChoiceCoefficients { + pub direct: PartCoefficients, + pub slots: Vec, + /// Per child group of the candidate, in order. + pub child_via: Vec, + /// The shareable nodes of the direct part (an ASAP node, or one whose + /// schema has a unique key) of a group two or more queries read, other + /// than its own node: each may be one node for every copy of the group + /// (the identical-expression rule), and then built once at query time. + pub shared_nodes: Vec, +} + +/// A shareable node of a choice's direct part ([`ChoiceCoefficients::shared_nodes`]). +/// Its price is in `direct.cost_per_unit_rate`, charged per consumer; built +/// once at query time it costs `cost_per_unit_rate` per unit of the union +/// rate of the queries reaching it instead, and holds `state_bytes` while +/// its longest reader runs. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct SharedNodeCoefficients { + pub cost_per_unit_rate: f64, + pub work_ms: f64, + pub state_bytes: u64, +} + +/// One choice of a group: a candidate and an option per private slot. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Choice { + pub group: GroupId, + pub candidate: usize, + /// The groups it implements (the first is `group`). + pub implements: Vec, + /// The child groups it reads. + pub children: Vec, + pub assignment: Vec, + pub coefficients: ChoiceCoefficients, + /// `None`: the group's one node, shared by every query reading it. + /// `Some(q)`: query `q`'s own copy of it (the group apart); a group's + /// copies are all apart or all shared. + pub copy_of: Option, + /// Each child group that is a scan: `None` when it is apart and its + /// copy is this choice's own (not in `children`); `Some(h)` when it is + /// the scan shared at horizon `h` ms (its choice's `scan_horizon_ms`). + pub scan_children: Vec<(GroupId, Option)>, + /// For a scan group's choice: the horizon it reads, the longest window + /// among its readers. + pub scan_horizon_ms: Option, +} + +/// A shared group's own node, built once per evaluation when chosen. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct GroupSlot { + pub group: GroupId, + /// Whether it may stay `NotMaterialized` (recomputed per consumer). + pub may_recompute: bool, +} + +/// One query of the model. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ModelQuery { + /// Its evaluations per second. + pub rate: f64, + /// The groups its root reads directly (one for an operator root). + pub roots: Vec, + /// Its latency bound in ms, if any (the tighter of its own bound and + /// its repeat interval). + pub max_latency_ms: Option, +} + +/// Everything the model reads. +pub(crate) struct ModelInput { + pub group_count: usize, + /// Per group, the queries whose plan may reach it. + pub group_queries: Vec>, + pub choices: Vec, + pub group_slots: Vec, + pub queries: Vec, + /// Per query, its schedule class: queries of one class run at the same + /// instants (one recurrence and range evaluation). + pub schedule: Vec, + /// The union rate of the classes in each set, by sorted class set: how + /// often a node built once for queries of those classes is computed. + pub union_rate: HashMap, f64>, + /// No-good cuts from earlier iterations: at most `bound` of these + /// choices together (a plan whose every group took one of them). + /// With the groups that plan does not need: a plan needing any of them + /// is another plan, not cut. + pub excluded: Vec<(Vec, Vec, f64)>, + /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all + /// of a choice set is chosen (local cuts). + pub extra_bounds: Vec<(Vec, f64)>, + /// Requirements of a choice on another group's choice: when the first + /// is chosen, one of the second (an exact implementation, or one + /// carrying a coupling site) is. + pub requires: Vec<(usize, Vec)>, + /// Choices that may not be chosen together with any of the second set + /// (panes and an exact top-k reuse in one plan). + pub exclusive: (Vec, Vec), + /// Each query-time raw scan of a choice's part: each source's longest + /// one is kept on disk. + pub raw_scans: Vec, + /// Whether the deployment keeps raw data: without it, no part with a + /// raw scan may run at query time. + pub raw_data_retained: bool, +} + +/// A raw scan of one part of a choice, when that part runs at query time. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct RawScan { + pub choice: usize, + /// Index of its source. + pub source: usize, + /// The price per second of keeping its look-back on disk. + pub cost: f64, + pub via: Via, +} + +/// Adds to `choice`'s coefficients a lower bound on the memory its parts +/// hold while an evaluation runs ([`LazyTerms::held`]): Stage 3 charges +/// `bytes · rate · E / 1000` at the memory price, `E` the latency of the +/// queries reaching the part, at least the part's own work. The group's +/// own node's working memory is left out (replaced when it is built +/// once). +fn add_lazy_lower_bounds(choice: &mut Choice, lazy: &LazyTerms, prices: &ResourceUnitCost) { + let coefficients = &mut choice.coefficients; + for term in &lazy.held { + if term.own_node { + continue; + } + let per_unit_rate = |work_ms: f64| { + prices.cost_per_memory_byte_second * term.bytes as f64 * work_ms / 1000.0 + }; + match term.via { + Via::Direct => { + coefficients.direct.cost_per_unit_rate += per_unit_rate(coefficients.direct.work_ms) + } + Via::Slot(i) => { + let slot = &mut coefficients.slots[i]; + if !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { + slot.cost_per_unit_rate += per_unit_rate(slot.work_ms); + } + } + } + } +} + +/// The model's optimum. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ModelSolution { + /// The chosen choices, by index in [`ModelInput::choices`]. + pub chosen: Vec, + /// The shared groups built once per evaluation. + pub built_once: BTreeSet, + /// The model's objective: a lower bound on the plan's price. + pub objective: f64, + /// The model's size. + pub variables: u64, + pub constraints: u64, +} + +/// A consumer of the model: a query, a materialized private slot of a +/// choice, or a shared group built once. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +enum Consumer { + Query(usize), + Slot { choice: usize, slot: usize }, + Group(GroupId), +} + +/// Build and solve the model. `None` when it is infeasible. +pub(crate) fn solve(input: &ModelInput) -> Result, String> { + let mut vars = ProblemVariables::new(); + let n = input.group_count; + let relaxed = std::env::var_os("JP_MILP_LP").is_some(); + let binary = || { + if relaxed { + variable().min(0).max(1) + } else { + variable().binary() + } + }; + let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); + let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); + let once: BTreeMap = input + .group_slots + .iter() + .map(|slot| (slot.group, vars.add(binary()))) + .collect(); + // Which queries' plans reach each group at all (any materialization): + // the union rate of a node built once reads them. + let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); + for (g, queries) in input.group_queries.iter().enumerate() { + for &q in queries { + reaches_at_all.insert((q, g), vars.add(binary())); + } + } + // Consumers that may reach each group's own node, and through it its + // direct part. + let consumers = possible_consumers(input); + let mut reach: HashMap<(Consumer, usize), Variable> = HashMap::new(); + for (&g, ks) in &consumers { + for &k in ks { + reach.insert((k, g), vars.add(variable().min(0).max(1))); + } + } + // A consumer passes through a group into its chosen direct part when + // it reaches the group and the group is not built once (or it is the + // group's own built-once consumer). + let mut passes: HashMap<(Consumer, usize), Variable> = HashMap::new(); + for &(k, g) in reach.keys() { + passes.insert((k, g), vars.add(variable().min(0).max(1))); + } + for &g in once.keys() { + let k = Consumer::Group(g); + passes + .entry((k, g.0 as usize)) + .or_insert_with(|| vars.add(variable().min(0).max(1))); + } + // y[k, c]: consumer k passes into choice c's direct part. + let mut y: HashMap<(Consumer, usize), Variable> = HashMap::new(); + for (c, choice) in input.choices.iter().enumerate() { + let g = choice.group.0 as usize; + for &(k, h) in passes.keys() { + if h == g { + y.insert((k, c), vars.add(variable().min(0).max(1))); + } + } + } + // Union-rate selectors for built-once consumers (shared groups and + // query-time private slots): one per set of schedule classes among the + // group's queries, exactly the classes of the queries reaching it. + let mut rate_of: HashMap, Variable)>> = HashMap::new(); + let classes_of = |g: usize| -> Vec { + let mut classes: Vec = input.group_queries[g] + .iter() + .map(|&q| input.schedule[q]) + .collect(); + classes.sort_unstable(); + classes.dedup(); + classes + }; + let class_sets = |g: usize| -> Vec> { + let classes = classes_of(g); + (1u64..(1 << classes.len())) + .map(|mask| { + (0..classes.len()) + .filter(|i| mask >> i & 1 == 1) + .map(|i| classes[i]) + .collect() + }) + .collect() + }; + let mut groups_with_rates: BTreeSet = once.keys().map(|g| g.0 as usize).collect(); + for choice in &input.choices { + if !choice.coefficients.shared_nodes.is_empty() + || choice + .coefficients + .slots + .iter() + .any(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) + { + groups_with_rates.insert(choice.group.0 as usize); + } + } + for &g in &groups_with_rates { + let sets = class_sets(g) + .into_iter() + .map(|set| (set, vars.add(binary()))) + .collect(); + rate_of.insert(g, sets); + } + + // The objective. + let mut objective = Expression::from(0.0); + let mut products_pending: Vec<(Variable, Variable, Variable)> = Vec::new(); + // A consumer's constant rate, or `None` for one built once at query + // time, whose rate is its group's union rate. + let rate = |k: Consumer| -> Option { + match k { + Consumer::Query(q) => Some(input.queries[q].rate), + Consumer::Slot { choice, slot } => { + let m = input.choices[choice].coefficients.slots[slot].materialization; + matches!(m, Materialization::MatAtIngestion(_)).then_some(1.0) + } + Consumer::Group(_) => None, + } + }; + let slowest_rate = |g: usize| -> f64 { + rate_of[&g] + .iter() + .map(|(set, _)| input.union_rate[set]) + .fold(f64::INFINITY, f64::min) + }; + // Products of a variable rate (a set selector) and a reach variable. + let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); + // The shareable nodes of a choice's direct part cost, per unit of + // rate, `shared_cost[c]`, charged either per consumer (the rest of the + // direct part's charge) or once at the union rate of the queries + // reaching the group when built once: the model charges the smaller + // (`per_consumer_shared[c]` against `built_once_shared[c]`). + let shared_cost: Vec = input + .choices + .iter() + .map(|choice| { + choice + .coefficients + .shared_nodes + .iter() + .map(|node| node.cost_per_unit_rate) + .sum() + }) + .collect(); + let mut per_consumer_shared: Vec> = vec![Vec::new(); input.choices.len()]; + let mut built_once_shared: Vec> = vec![Vec::new(); input.choices.len()]; + let mut product = |vars: &mut ProblemVariables, a: Variable, b: Variable| { + let z = vars.add(variable().min(0).max(1)); + products_pending.push((z, a, b)); + z + }; + for (&(k, c), &v) in &y { + let total = input.choices[c].coefficients.direct.cost_per_unit_rate; + let shared = shared_cost[c].min(total); + let cost = total - shared; + if total == 0.0 { + continue; + } + match rate(k) { + Some(r) => { + objective += r * cost * v; + if shared != 0.0 { + per_consumer_shared[c].push((v, r * shared)); + } + } + None => { + let g = match k { + Consumer::Group(g) => g.0 as usize, + Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, + Consumer::Query(_) => unreachable!("a query's rate is constant"), + }; + // At least the slowest class's rate is charged linearly + // (some query reaches the group), the rest through the + // selector products: a tighter relaxation. + let slowest = slowest_rate(g); + objective += slowest * cost * v; + for (set, selector) in &rate_of[&g] { + let union = input.union_rate[set]; + if cost != 0.0 && union > slowest { + rated.push((*selector, v, (union - slowest) * cost)); + } + if shared != 0.0 { + let z = product(&mut vars, *selector, v); + per_consumer_shared[c].push((z, union * shared)); + } + } + } + } + } + for (c, choice) in input.choices.iter().enumerate() { + if shared_cost[c] == 0.0 { + continue; + } + for (set, selector) in &rate_of[&(choice.group.0 as usize)] { + let z = product(&mut vars, *selector, x[c]); + built_once_shared[c].push((z, input.union_rate[set] * shared_cost[c])); + } + } + // min(A, B) as a + b with a ≥ A − M·u, b ≥ B − M·(1 − u), u binary. + let mut cheaper_of: Vec<(Variable, Variable, Variable, usize, f64)> = Vec::new(); + for c in 0..input.choices.len() { + if per_consumer_shared[c].is_empty() && built_once_shared[c].is_empty() { + continue; + } + let big_m: f64 = per_consumer_shared[c].iter().map(|(_, w)| w).sum::() + + built_once_shared[c].iter().map(|(_, w)| w).sum::(); + let a = vars.add(variable().min(0)); + let b = vars.add(variable().min(0)); + let u = vars.add(binary()); + objective += a + b; + objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; + if std::env::var_os("JP_MILP_NOMIN").is_some() { + continue; + } + cheaper_of.push((a, b, u, c, big_m)); + } + for (c, choice) in input.choices.iter().enumerate() { + for slot in &choice.coefficients.slots { + objective += slot.fixed_cost * x[c]; + if slot.cost_per_unit_rate != 0.0 + && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) + { + let g = choice.group.0 as usize; + let slowest = slowest_rate(g); + objective += slowest * slot.cost_per_unit_rate * x[c]; + for (set, selector) in &rate_of[&g] { + let union = input.union_rate[set]; + if union > slowest { + rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); + } + } + } + } + } + let mut products = Vec::new(); + for &(a, b, coefficient) in &rated { + let z = vars.add(variable().min(0).max(1)); + objective += coefficient * z; + products.push((z, a, b)); + } + // Lower bounds on the omitted terms from earlier cuts. + let extra: Vec = input + .extra_bounds + .iter() + .map(|_| vars.add(variable().min(0))) + .collect(); + for &e in &extra { + objective += e; + } + + let vars_exclusive = (vars.add(binary()), vars.add(binary())); + let source_count = input + .raw_scans + .iter() + .map(|s| s.source + 1) + .max() + .unwrap_or(0); + let kept_on_disk: Vec = (0..source_count) + .map(|_| vars.add(variable().min(0))) + .collect(); + for &v in &kept_on_disk { + objective += v; + } + let variable_count = vars.len() as u64; + let model = vars.minimise(objective.clone()).using(highs); + let mut constraints = Vec::new(); + // Needed groups: every query's roots; children of chosen choices. + for query in &input.queries { + for root in &query.roots { + constraints.push(constraint!(need[root.0 as usize] == 1)); + } + } + for (g, &needed) in need.iter().enumerate() { + // One choice implementing it (its own, or a fused one covering + // it) when needed, none otherwise. + let mut implementing = Expression::from(0.0); + for (c, choice) in input.choices.iter().enumerate() { + if choice.implements.iter().any(|i| i.0 as usize == g) { + implementing += x[c]; + } + } + constraints.push(constraint!(implementing == needed)); + } + for (c, choice) in input.choices.iter().enumerate() { + for child in &choice.children { + constraints.push(constraint!(need[child.0 as usize] >= x[c])); + } + } + for (&g, &v) in &once { + constraints.push(constraint!(v <= need[g.0 as usize])); + let slot = input.group_slots.iter().find(|s| s.group == g).unwrap(); + if !slot.may_recompute { + constraints.push(constraint!(v == need[g.0 as usize])); + } + } + // Reach at all: a query reaches its roots, and every child of a chosen + // choice of a group it reaches. + for (q, query) in input.queries.iter().enumerate() { + for root in &query.roots { + if let Some(&v) = reaches_at_all.get(&(q, root.0 as usize)) { + constraints.push(constraint!(v == 1)); + } + } + } + // Aggregated over a group's choices (one is chosen when it is + // needed), which keeps the LP relaxation tight. + let mut choices_of: Vec> = vec![Vec::new(); n]; + for (c, choice) in input.choices.iter().enumerate() { + choices_of[choice.group.0 as usize].push(c); + } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); + } + for h in children { + let mut with_child = Expression::from(0.0); + for &c in of_g { + if input.choices[c].children.iter().any(|k| k.0 as usize == h) { + with_child += x[c]; + } + } + for &q in &input.group_queries[g] { + if let (Some(&above), Some(&below)) = + (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) + { + constraints.push(constraint!(below >= above + with_child.clone() - 1)); + } + } + } + } + // Rate selectors: exactly the classes of the queries reaching the + // group (at most one selector; a class is in it iff one of its + // queries reaches the group). + for (&g, sets) in &rate_of { + let mut any = Expression::from(0.0); + for (_, selector) in sets { + any += *selector; + } + constraints.push(constraint!(any <= 1)); + for class in classes_of(g) { + let mut containing = Expression::from(0.0); + for (set, selector) in sets { + if set.contains(&class) { + containing += *selector; + } + } + for &q in &input.group_queries[g] { + if input.schedule[q] == class { + constraints.push(constraint!(containing.clone() >= reaches_at_all[&(q, g)])); + } + } + let mut members = Expression::from(0.0); + for &q in &input.group_queries[g] { + if input.schedule[q] == class { + members += reaches_at_all[&(q, g)]; + } + } + constraints.push(constraint!(containing <= members)); + } + } + // Consumer reach. + for (q, query) in input.queries.iter().enumerate() { + for root in &query.roots { + if let Some(&v) = reach.get(&(Consumer::Query(q), root.0 as usize)) { + constraints.push(constraint!(v >= need[root.0 as usize])); + } + } + } + for (&(k, g), &p) in &passes { + match k { + Consumer::Group(own) if own.0 as usize == g => { + constraints.push(constraint!(p >= once[&own])); + } + _ => { + let r = reach[&(k, g)]; + match once.get(&GroupId(g as u32)) { + Some(&o) => constraints.push(constraint!(p >= r - o)), + None => constraints.push(constraint!(p >= r)), + } + } + } + } + for (&(k, c), &v) in &y { + let g = input.choices[c].group.0 as usize; + constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); + } + for (&(k, g), &p) in &passes { + let mut into = Expression::from(0.0); + for &c in &choices_of[g] { + into += y[&(k, c)]; + } + // Unless a fused choice of another group implements it. + for (c, choice) in input.choices.iter().enumerate() { + if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { + into += x[c]; + } + } + constraints.push(constraint!(into >= p)); + } + for (c, choice) in input.choices.iter().enumerate() { + for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { + let h = child.0 as usize; + match via { + Via::Direct => { + // Aggregated below, over the group's choices. + let _ = h; + } + Via::Slot(slot) => { + let k = Consumer::Slot { + choice: c, + slot: *slot, + }; + if let Some(&r) = reach.get(&(k, h)) { + constraints.push(constraint!(r >= x[c])); + } + } + } + } + } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + for (child, via) in input.choices[c] + .children + .iter() + .zip(&input.choices[c].coefficients.child_via) + { + if *via == Via::Direct { + children.insert(child.0 as usize); + } + } + } + for h in children { + for &(k, pg) in passes.keys() { + if pg != g { + continue; + } + let Some(&r) = reach.get(&(k, h)) else { + continue; + }; + let mut into = Expression::from(0.0); + for &c in of_g { + let choice = &input.choices[c]; + if choice + .children + .iter() + .zip(&choice.coefficients.child_via) + .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) + { + into += y[&(k, c)]; + } + } + constraints.push(constraint!(r >= into)); + } + } + } + for &(z, a, b) in products.iter().chain(&products_pending) { + constraints.push(constraint!(z >= a + b - 1)); + } + for &(a, b, u, c, big_m) in &cheaper_of { + let mut per_consumer = Expression::from(0.0); + for &(v, w) in &per_consumer_shared[c] { + per_consumer += w * v; + } + let mut built_once = Expression::from(0.0); + for &(v, w) in &built_once_shared[c] { + built_once += w * v; + } + // Both are at least the slowest class's rate (some query reaches + // the chosen choice): charged linearly, which keeps the relaxation + // tight; only the excess goes through the binary choice. + let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); + constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); + constraints.push(constraint!( + b >= built_once - floor * x[c] - big_m * (1 - u) + )); + } + // Latency bounds: the work each query reaches at query time. + for (q, query) in input.queries.iter().enumerate() { + let Some(bound) = query.max_latency_ms else { + continue; + }; + let mut latency = Expression::from(0.0); + for (&(k, c), &v) in &y { + if k == Consumer::Query(q) { + latency += input.choices[c].coefficients.direct.work_ms * v; + } + } + constraints.push(constraint!(latency <= bound)); + } + // Cuts. + for (set, outside, bound) in &input.excluded { + let mut chosen = Expression::from(0.0); + for &c in set { + chosen += x[c]; + } + for &g in outside { + chosen -= need[g]; + } + constraints.push(constraint!(chosen <= *bound)); + } + for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { + let mut chosen = Expression::from(0.0); + for &c in set { + chosen += x[c]; + } + constraints.push(constraint!( + e >= *bound * (chosen - (set.len() as f64 - 1.0)) + )); + } + for (c, allowed) in &input.requires { + let mut any = Expression::from(0.0); + for &a in allowed { + any += x[a]; + } + constraints.push(constraint!(x[*c] <= any)); + } + let (first, second) = &input.exclusive; + if !first.is_empty() && !second.is_empty() { + let (p, t) = (vars_exclusive.0, vars_exclusive.1); + for &c in first { + constraints.push(constraint!(x[c] <= p)); + } + for &c in second { + constraints.push(constraint!(x[c] <= t)); + } + constraints.push(constraint!(p + t <= 1)); + } + // Raw data: each source's longest query-time scan is kept on disk; a + // deployment without raw data runs no raw scan at query time. A part + // runs at query time for every consumer but an ingestion-time slot. + let at_query_time = |k: Consumer| match k { + Consumer::Slot { choice, slot } => !matches!( + input.choices[choice].coefficients.slots[slot].materialization, + Materialization::MatAtIngestion(_) + ), + _ => true, + }; + for scan in &input.raw_scans { + let c = scan.choice; + let mut runs: Vec = Vec::new(); + match scan.via { + Via::Direct => { + for (&(k, d), &v) in &y { + if d == c && at_query_time(k) { + runs.push(v); + } + } + } + Via::Slot(i) => { + let m = input.choices[c].coefficients.slots[i].materialization; + if !matches!(m, Materialization::MatAtIngestion(_)) { + runs.push(x[c]); + } + } + } + for v in runs { + if !input.raw_data_retained { + constraints.push(constraint!(v == 0)); + } else if scan.cost > 0.0 { + constraints.push(constraint!(kept_on_disk[scan.source] >= scan.cost * v)); + } + } + } + let constraint_count = constraints.len() as u64; + let mut model = model + .with_all(constraints) + .with_mip_gap(0.0) + .map_err(|e| format!("MIP gap: {e:?}"))?; + if std::env::var_os("JP_MILP_VERBOSE").is_some() { + model.set_verbose(true); + } + let solution = match model.solve() { + Ok(solution) => solution, + Err(good_lp::ResolutionError::Infeasible) => return Ok(None), + Err(error) => return Err(error.to_string()), + }; + let chosen: Vec = x + .iter() + .enumerate() + .filter(|(_, v)| solution.value(**v) > 0.5) + .map(|(c, _)| c) + .collect(); + if std::env::var_os("JP_MILP_TRACE").is_some() { + if relaxed { + for (g, v) in once.iter() { + eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); + } + for (g, v) in need.iter().enumerate() { + eprintln!(" need g{} = {:.3}", g, solution.value(*v)); + } + } + let shown: Vec = if relaxed { + (0..x.len()) + .filter(|&c| solution.value(x[c]) > 1e-6) + .collect() + } else { + chosen.clone() + }; + for &c in &shown { + if relaxed { + eprintln!(" x = {:.3}", solution.value(x[c])); + } + let choice: &Choice = &input.choices[c]; + let reach: Vec = y + .iter() + .filter(|((_, d), v)| *d == c && solution.value(**v) > 1e-9) + .map(|((k, _), v)| format!("{k:?}={:.2}", solution.value(*v))) + .collect(); + eprintln!( + " g{} cand{} direct {:.4e}/rate work {:.2}ms shared {:?} slots {:?} scans {:?} horizon {:?} assignment {:?} y {:?}", + choice.group.0, choice.candidate, choice.coefficients.direct.cost_per_unit_rate, + choice.coefficients.direct.work_ms, choice.coefficients.shared_nodes.len(), + choice.coefficients.slots.iter().map(|s| (s.fixed_cost, s.cost_per_unit_rate)).collect::>(), + choice.scan_children, choice.scan_horizon_ms, choice.assignment, reach + ); + } + } + let built_once = once + .iter() + .filter(|(_, v)| solution.value(**v) > 0.5) + .map(|(g, _)| *g) + .collect(); + Ok(Some(ModelSolution { + chosen, + built_once, + objective: solution.eval(&objective), + variables: variable_count, + constraints: constraint_count, + })) +} + +/// For each group, the consumers that may reach its node: its queries +/// (that reach it without passing a materialized node), the materialized +/// private slots of choices of groups above it, and the shared groups +/// above it (or itself) that may be built once. +fn possible_consumers(input: &ModelInput) -> BTreeMap> { + let mut consumers: BTreeMap> = BTreeMap::new(); + for (g, queries) in input.group_queries.iter().enumerate() { + let entry = consumers.entry(g).or_default(); + for &q in queries { + entry.insert(Consumer::Query(q)); + } + } + // Propagate downwards until nothing changes. + let shared: BTreeSet = input + .group_slots + .iter() + .map(|s| s.group.0 as usize) + .collect(); + loop { + let mut changed = false; + for (c, choice) in input.choices.iter().enumerate() { + let g = choice.group.0 as usize; + let mut above: BTreeSet = consumers.get(&g).cloned().unwrap_or_default(); + if shared.contains(&g) { + above.insert(Consumer::Group(GroupId(g as u32))); + } + for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { + let entry = consumers.entry(child.0 as usize).or_default(); + let before = entry.len(); + match via { + Via::Direct => entry.extend(above.iter().copied()), + Via::Slot(slot) => { + entry.insert(Consumer::Slot { + choice: c, + slot: *slot, + }); + } + } + changed |= entry.len() != before; + } + } + if !changed { + break; + } + } + consumers +} + +/// A memo search's result: the selection, its physical candidate and its +/// logical roots, or why no plan is valid. +pub(crate) type MemoSearchResult = + Result<(Selection, PhysicalCandidate, Vec), SelectionError>; + +/// The cheapest valid plan of `per_computation`'s memo under cost model 1 +/// (the caller routes cost models 2 and 3 to the legacy search). +/// +/// Lazy loop: solve the model (a lower bound: it omits memory held while a +/// query runs, and checks composed guarantees only on its incumbents); +/// price the solution's group choice exactly, as the reference search +/// prices one ([`offer_group_choice`]: every threshold, identical-expression +/// variant and Stage 2 materialization), which updates the incumbent; add +/// a no-good cut excluding that group choice; stop when the model's bound +/// reaches the incumbent. Each group choice is priced at most once, so the +/// loop ends, and the incumbent is the reference search's optimum. +pub(crate) fn search( + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> MemoSearchResult { + let started = std::time::Instant::now(); + let memo = super::search::memo(per_computation, entries); + let memo_stats = super::search::stats(&memo, started.elapsed().as_secs_f64() * 1e3); + let scale = super::pricing::frozen_calibration(&memo, entries, data, models)?; + let mut cache = super::pricing::PriceCache::new(scale); + let mut choices: Vec = Vec::new(); + let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); + let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { + unreachable!("the caller routes cost models 2 and 3 to the legacy search"); + }; + for g in 0..memo.groups.len() { + let (group_choices, lazy) = super::coefficients::choices( + &memo, + GroupId(g as u32), + entries, + data, + models, + &mut cache, + )?; + // A group's copies (apart) cost at least its shared choice with its + // node not built once: each copy is charged per consumer as the + // shared node is, its own scans read at least the shortest horizon, + // and its builds are sized for fewer readers but built per copy. So + // the model keeps the shared choices (a lower bound), and the exact + // pricing of a group choice prices both. + // Widened scans are reached through the identical-expression + // variants of the exact pricing, never chosen directly. + for (mut choice, lazy) in group_choices.into_iter().zip(lazy) { + if choice.copy_of.is_some() + || memo.candidates[g][choice.candidate].origin + == super::CandidateOrigin::WidenedScans + { + continue; + } + add_lazy_lower_bounds(&mut choice, &lazy, prices); + raw_scans.extend( + lazy.raw_scans + .iter() + .map(|(source, bytes, via)| (choices.len(), source.clone(), *bytes, *via)), + ); + choices.push(choice); + } + } + let trace = std::env::var_os("JP_MILP_TRACE").is_some(); + if trace { + eprintln!( + "milp: memo {} groups, {} choices in {:.1}s", + memo.groups.len(), + choices.len(), + started.elapsed().as_secs_f64() + ); + } + let mut input = model_input(&memo, choices, entries, models)?; + let mut sources: Vec = Vec::new(); + for (c, source, bytes, via) in raw_scans { + let index = match sources.iter().position(|s| *s == source) { + Some(i) => i, + None => { + sources.push(source); + sources.len() - 1 + } + }; + input.raw_scans.push(RawScan { + choice: c, + source: index, + cost: prices.cost_per_disk_byte_second * bytes as f64, + via, + }); + } + input.raw_data_retained = models.deployment_capabilities.raw_data_retained; + let mut selector = PlanSelector::new(entries, data, models, false)?; + let mut logical: HashMap> = HashMap::new(); + let mut number = 0usize; + let mut iterations = 0u32; + let mut bound = None; + let mut first_bound = None; + let mut solver_ms = 0.0; + let mut size = (0, 0); + loop { + iterations += 1; + let solve_started = std::time::Instant::now(); + let solved = solve(&input).map_err(SelectionError::Solver)?; + solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; + let Some(solution) = solved else { + break; + }; + bound = Some(solution.objective); + first_bound.get_or_insert(solution.objective); + if trace { + eprintln!("milp: round {iterations} bound {:.6e} incumbent {:?} solve {:.0}ms vars {} cons {} at {:.1}s", solution.objective, selector.cheapest_price(), solve_started.elapsed().as_secs_f64()*1e3, solution.variables, solution.constraints, started.elapsed().as_secs_f64()); + } + size = (solution.variables, solution.constraints); + if let Some(incumbent) = selector.cheapest_price() { + if solution.objective >= incumbent - 1e-9 * incumbent.abs() { + break; + } + } + // The solution's group choice: the candidate each chosen choice + // takes for its group. + let mut group_choice: super::search::GroupChoice = vec![None; memo.groups.len()]; + for &c in &solution.chosen { + group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); + } + let chosen_candidates: Vec<(GroupId, &super::ImplementationCandidate)> = group_choice + .iter() + .enumerate() + .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) + .collect(); + if !super::windows::requirements_met(&chosen_candidates) { + // A rollup's requirement: cut the smallest failing pairs of + // (group, candidate), else the whole group choice below. + let mut cut = false; + for (i, a) in chosen_candidates.iter().enumerate() { + if a.1.rollup_of.is_none() { + continue; + } + for (j, b) in chosen_candidates.iter().enumerate() { + if i != j && !super::windows::requirements_met(&[*a, *b]) { + let members: Vec = (0..input.choices.len()) + .filter(|&c| { + let choice = &input.choices[c]; + (choice.group == a.0 + && choice.candidate == group_choice[a.0 .0 as usize].unwrap()) + || (choice.group == b.0 + && choice.candidate + == group_choice[b.0 .0 as usize].unwrap()) + }) + .collect(); + input.excluded.push((members, Vec::new(), 1.0)); + cut = true; + } + } + if !super::windows::requirements_met(&[*a]) { + let members: Vec = (0..input.choices.len()) + .filter(|&c| { + input.choices[c].group == a.0 + && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] + }) + .collect(); + input.excluded.push((members, Vec::new(), 0.0)); + cut = true; + } + } + if trace { + eprintln!("milp: round {iterations} requirements not met (pair cut {cut})"); + } + if cut { + continue; + } + } + super::search::offer_group_choice( + &memo, + per_computation, + &group_choice, + entries, + data, + models, + &mut selector, + &mut logical, + &mut number, + )?; + // Exclude it: every choice of each chosen (group, candidate). + let groups = group_choice.iter().flatten().count(); + let excluded: Vec = input + .choices + .iter() + .enumerate() + .filter(|(_, choice)| group_choice[choice.group.0 as usize] == Some(choice.candidate)) + .map(|(c, _)| c) + .collect(); + let mut covered = vec![false; memo.groups.len()]; + for (g, c) in group_choice.iter().enumerate() { + if let Some(c) = c { + for i in &memo.candidates[g][*c].implements { + covered[i.0 as usize] = true; + } + } + } + let outside = (0..memo.groups.len()).filter(|&g| !covered[g]).collect(); + input + .excluded + .push((excluded, outside, groups as f64 - 1.0)); + } + if std::env::var_os("JP_MILP_CHECK").is_some() { + // Debugging aid: every group choice's exact price against the + // model's value restricted to it. + let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); + for group_choice in super::search::group_choices(&memo).unwrap_or_default() { + let mut one = PlanSelector::new(entries, data, models, false)?; + let mut scratch = HashMap::new(); + let mut k = 0usize; + super::search::offer_group_choice( + &memo, + per_computation, + &group_choice, + entries, + data, + models, + &mut one, + &mut scratch, + &mut k, + )?; + let Some(exact) = one.cheapest_price() else { + continue; + }; + let mut restricted = model_input_clone(&input); + restricted.excluded.clear(); + let outside: Vec = (0..restricted.choices.len()) + .filter(|&c| { + group_choice[restricted.choices[c].group.0 as usize] + != Some(restricted.choices[c].candidate) + }) + .collect(); + restricted.excluded.push((outside, Vec::new(), 0.0)); + let value = solve(&restricted) + .map_err(SelectionError::Solver)? + .map(|s| s.objective); + let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); + if bad || exact < incumbent * (1.0 - 1e-9) { + eprintln!( + "check: {group_choice:?} exact {exact:.7e} model {value:?}{}", + if bad { + " NOT A LOWER BOUND" + } else { + " BELOW INCUMBENT" + } + ); + } + } + } + let (mut selection, selected) = selector.finish()?; + let price = selection.costs[&selection.selected].cost_per_second; + selection.search.memo = Some(memo_stats); + selection.search.solver_iterations = iterations; + selection.search.solver_time_ms = solver_ms as u64; + ( + selection.search.solver_variables, + selection.search.solver_constraints, + ) = size; + selection.search.solver_first_bound_cost_per_second = first_bound; + selection.search.solver_gap = match (price, bound) { + (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), + (Some(_), Some(_)) => Some(0.0), + _ => None, + }; + let roots = logical + .remove(&selected.from_logical) + .expect("the selection's logical roots are kept"); + Ok((selection, selected, roots)) +} + +fn model_input_clone(input: &ModelInput) -> ModelInput { + ModelInput { + group_count: input.group_count, + group_queries: input.group_queries.clone(), + choices: input.choices.clone(), + group_slots: input.group_slots.clone(), + queries: input.queries.clone(), + schedule: input.schedule.clone(), + union_rate: input.union_rate.clone(), + excluded: input.excluded.clone(), + extra_bounds: input.extra_bounds.clone(), + requires: input.requires.clone(), + exclusive: input.exclusive.clone(), + raw_scans: input.raw_scans.clone(), + raw_data_retained: input.raw_data_retained, + } +} + +/// The model's input for `memo`'s `choices`. +fn model_input( + memo: &Memo, + choices: Vec, + entries: &[QueryWorkloadEntry], + models: PlanningModels<'_>, +) -> Result { + let settings = models.plan_selection; + let rate = |entries: &[&QueryWorkloadEntry]| { + crate::shared_evaluation_rate(entries, settings).map_err(|error| SelectionError::Cost { + candidate: "memo".to_string(), + node: Default::default(), + error, + }) + }; + // Schedule classes: queries with one recurrence and range evaluation. + let mut classes: Vec<&QueryWorkloadEntry> = Vec::new(); + let schedule: Vec = entries + .iter() + .map(|entry| { + match classes.iter().position(|c| { + c.recurrence == entry.recurrence && c.range_evaluation == entry.range_evaluation + }) { + Some(class) => class, + None => { + classes.push(entry); + classes.len() - 1 + } + } + }) + .collect(); + let mut union_rate = HashMap::new(); + let mut queries = Vec::with_capacity(entries.len()); + for (q, entry) in entries.iter().enumerate() { + let roots = memo + .groups + .iter() + .enumerate() + .filter(|(_, group)| group.roots_of.contains(&q)) + .map(|(g, _)| GroupId(g as u32)) + .collect(); + // A query's evaluation must fit its latency bound and its repeat + // interval. + let interval = entry.recurrence.fixed_interval_ms().map(|ms| ms as f64); + let max_latency_ms = match (entry.requirements.max_latency_ms, interval) { + (Some(a), Some(b)) => Some(a.min(b)), + (a, b) => a.or(b), + }; + queries.push(ModelQuery { + rate: rate(&[entry])?, + roots, + max_latency_ms, + }); + } + let group_queries: Vec> = memo.groups.iter().map(|g| g.queries.clone()).collect(); + for queries_of in &group_queries { + let mut present: Vec = queries_of.iter().map(|&q| schedule[q]).collect(); + present.sort_unstable(); + present.dedup(); + for mask in 1u64..(1 << present.len()) { + let set: Vec = (0..present.len()) + .filter(|i| mask >> i & 1 == 1) + .map(|i| present[i]) + .collect(); + if let std::collections::hash_map::Entry::Vacant(slot) = union_rate.entry(set) { + let members: Vec<&QueryWorkloadEntry> = + slot.key().iter().map(|&c| classes[c]).collect(); + slot.insert(rate(&members)?); + } + } + } + // Requirements the reference search checks on a group choice + // ([`super::windows::requirements_met`]) other than rollups, which the + // loop cuts lazily. + let candidate_of = |c: &Choice| &memo.candidates[c.group.0 as usize][c.candidate]; + let mut requires = Vec::new(); + for (c, choice) in choices.iter().enumerate() { + let candidate = candidate_of(choice); + for &g in &candidate.requires_exact_children { + let allowed = (0..choices.len()) + .filter(|&a| { + choices[a].group == g && super::windows::is_exact(candidate_of(&choices[a])) + }) + .collect(); + requires.push((c, allowed)); + } + for (g, site) in &candidate.requires_sites { + let allowed = (0..choices.len()) + .filter(|&a| { + choices[a].group == *g + && candidate_of(&choices[a]).coupling_sites.contains(site) + }) + .collect(); + requires.push((c, allowed)); + } + } + // A reader of a scan shared at horizon `h` reads the scan group's + // choice at that horizon, so every reader of one scan agrees on it. + for (c, choice) in choices.iter().enumerate() { + for &(scan, horizon) in &choice.scan_children { + let Some(h) = horizon else { continue }; + let allowed = (0..choices.len()) + .filter(|&a| choices[a].group == scan && choices[a].scan_horizon_ms == Some(h)) + .collect(); + requires.push((c, allowed)); + } + } + let of_origin = |origin: super::CandidateOrigin| -> Vec { + (0..choices.len()) + .filter(|&c| candidate_of(&choices[c]).origin == origin) + .collect() + }; + let exclusive = ( + of_origin(super::CandidateOrigin::Panes), + of_origin(super::CandidateOrigin::TopKReuse), + ); + let group_slots = memo + .groups + .iter() + .enumerate() + .filter(|(_, group)| group.parents.len() > 1 || group.roots_of.len() > 1) + .map(|(g, _)| { + let allowed = &memo.candidates[g][0].materializations; + GroupSlot { + group: GroupId(g as u32), + may_recompute: allowed.is_empty() + || allowed.contains(&Materialization::NotMaterialized), + } + }) + .collect(); + Ok(ModelInput { + group_count: memo.groups.len(), + group_queries, + choices, + group_slots, + queries, + schedule, + union_rate, + excluded: Vec::new(), + extra_bounds: Vec::new(), + requires, + exclusive, + raw_scans: Vec::new(), + raw_data_retained: true, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn choice(group: u32, children: &[u32], cost: f64) -> Choice { + Choice { + group: GroupId(group), + candidate: 0, + implements: vec![GroupId(group)], + children: children.iter().map(|&c| GroupId(c)).collect(), + assignment: Vec::new(), + coefficients: ChoiceCoefficients { + direct: PartCoefficients { + cost_per_unit_rate: cost, + work_ms: 1.0, + }, + slots: Vec::new(), + child_via: children.iter().map(|_| Via::Direct).collect(), + shared_nodes: Vec::new(), + }, + copy_of: None, + scan_children: Vec::new(), + scan_horizon_ms: None, + } + } + + /// A scan two queries read costs once per query recomputing it, or once + /// per evaluation when built once: the model builds it once when that + /// is cheaper, and its objective is the plan's price. + #[test] + fn a_shared_group_is_built_once_when_that_is_cheaper() { + let input = ModelInput { + group_count: 3, + group_queries: vec![vec![0, 1], vec![0], vec![1]], + choices: vec![ + choice(0, &[], 10.0), + choice(1, &[0], 1.0), + choice(2, &[0], 1.0), + ], + group_slots: vec![GroupSlot { + group: GroupId(0), + may_recompute: true, + }], + queries: vec![ + ModelQuery { + rate: 1.0, + roots: vec![GroupId(1)], + max_latency_ms: None, + }, + ModelQuery { + rate: 1.0, + roots: vec![GroupId(2)], + max_latency_ms: None, + }, + ], + schedule: vec![0, 0], + union_rate: [(vec![0], 1.0)].into(), + excluded: Vec::new(), + extra_bounds: Vec::new(), + requires: Vec::new(), + exclusive: (Vec::new(), Vec::new()), + raw_scans: Vec::new(), + raw_data_retained: true, + }; + let solution = solve(&input).unwrap().unwrap(); + assert_eq!(solution.built_once, [GroupId(0)].into()); + assert!( + (solution.objective - 12.0).abs() < 1e-9, + "{}", + solution.objective + ); + // Recomputing is cheaper when the scan is cheap. + let mut cheap = input; + cheap.choices[0].coefficients.direct.cost_per_unit_rate = 0.1; + cheap.union_rate.insert(vec![0], 1.0); + // Built once still costs its union rate × 0.1 = 0.1 < 0.2. + let solution = solve(&cheap).unwrap().unwrap(); + assert!( + (solution.objective - 2.1).abs() < 1e-9, + "{}", + solution.objective + ); + } +} From 7694939226e00c8c7cd0285e13381fe50e5ecad3 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:03:22 +0000 Subject: [PATCH 42/67] plan-selection: memo choices at a sketch rollup's state-size thresholds A candidate carrying a rollup-state site is priced at each threshold its readers may read, its finer build resized as the brute force resizes it. Test: an exact rollup's plans reproduce Stage 3. --- .../src/pass/subexpr/coefficients.rs | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 3135a5908..6aadab79b 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -410,6 +410,23 @@ fn at_site_thresholds( next.extend(at(index, rewritten)); } } + // A sketch rollup's finer build, sized for its readers, as + // the brute force sizes it. + CouplingKind::SharedBuildSize if windows::is_rollup_state(memo, site) => { + for (index, threshold) in windows::state_size_thresholds(memo, site) + .iter() + .enumerate() + { + let rewritten = windows::may_read_state(threshold, &version, site) + .then(|| windows::resize_state_at(&version, site, threshold).ok()) + .flatten() + .map(|node| ImplementationCandidate { + node, + ..version.clone() + }); + next.extend(at(index, rewritten)); + } + } CouplingKind::SharedBuildSize => { for (index, threshold) in capability::size_thresholds(memo, site).iter().enumerate() @@ -1618,6 +1635,18 @@ mod tests { assert!(!checked.is_empty()); } + /// A per-job sum and the total rolled up from it. (TODO: a sketch + /// rollup, which reads its finer group's build rather than its node; + /// the test does not place those plans yet.) + #[test] + fn coefficients_reproduce_stage_3_with_a_rollup() { + assert!(!check( + &["sum by (job) (rate(m[5m]))", "sum(rate(m[5m]))"], + &[60_000, 60_000], + ) + .is_empty()); + } + /// Two quantiles within 1% and 2% sharing one KLL build at a size /// threshold. #[test] From 6769081fbff9f68bb7989c91e6132133be9f4d4d Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:28:24 +0000 Subject: [PATCH 43/67] Memo MILP model: held memory over each query's query-time latency (linearized), latency from query-time reach through built-once nodes and query-time slots --- .../plan-selection/src/pass/subexpr/milp.rs | 223 +++++++++++++++--- 1 file changed, 191 insertions(+), 32 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 1488e0dd8..5f89ac016 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -168,6 +168,21 @@ pub(crate) struct ModelInput { /// Whether the deployment keeps raw data: without it, no part with a /// raw scan may run at query time. pub raw_data_retained: bool, + /// Per choice, the memory its parts hold while a query's evaluation + /// runs ([`HeldAtQueryTime`]). + pub held: Vec, +} + +/// The memory a choice's parts hold while an evaluation runs, priced per +/// ms of latency: Stage 3 charges `bytes · rate · E / 1000` at the memory +/// price, `E` the latency of the queries reaching the part. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub(crate) struct HeldAtQueryTime { + /// Per unit of a reading query's rate: each reader holds its own copy + /// over its own latency. + pub per_reader_per_ms: f64, + /// Per unit of rate: one copy held over the longest reader's latency. + pub longest_per_ms: f64, } /// A raw scan of one part of a choice, when that part runs at query time. @@ -181,33 +196,36 @@ pub(crate) struct RawScan { pub via: Via, } -/// Adds to `choice`'s coefficients a lower bound on the memory its parts -/// hold while an evaluation runs ([`LazyTerms::held`]): Stage 3 charges -/// `bytes · rate · E / 1000` at the memory price, `E` the latency of the -/// queries reaching the part, at least the part's own work. The group's -/// own node's working memory is left out (replaced when it is built -/// once). -fn add_lazy_lower_bounds(choice: &mut Choice, lazy: &LazyTerms, prices: &ResourceUnitCost) { - let coefficients = &mut choice.coefficients; +/// The memory `choice`'s parts hold while an evaluation runs +/// ([`LazyTerms::held`]), but its ingestion-time slots'. Its own node, +/// when built once, holds the smaller of its working memory and its built +/// state. +fn held_at_query_time( + choice: &Choice, + lazy: &LazyTerms, + prices: &ResourceUnitCost, +) -> HeldAtQueryTime { + let mut held = HeldAtQueryTime::default(); for term in &lazy.held { - if term.own_node { - continue; + if let Via::Slot(i) = term.via { + if matches!( + choice.coefficients.slots[i].materialization, + Materialization::MatAtIngestion(_) + ) { + continue; + } } - let per_unit_rate = |work_ms: f64| { - prices.cost_per_memory_byte_second * term.bytes as f64 * work_ms / 1000.0 + let bytes = match term.own_node { + true => term.bytes.min(lazy.built_once_state_bytes), + false => term.bytes, }; - match term.via { - Via::Direct => { - coefficients.direct.cost_per_unit_rate += per_unit_rate(coefficients.direct.work_ms) - } - Via::Slot(i) => { - let slot = &mut coefficients.slots[i]; - if !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - slot.cost_per_unit_rate += per_unit_rate(slot.work_ms); - } - } + let per_ms = prices.cost_per_memory_byte_second * bytes as f64 / 1000.0; + match term.per_reader { + true => held.per_reader_per_ms += per_ms, + false => held.longest_per_ms += per_ms, } } + held } /// The model's optimum. @@ -482,6 +500,54 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> objective += e; } + // Query-time reach and latency (for held memory and latency bounds): + // `at_query_time[q, g]`, query `q` reaches group `g` without passing an + // ingestion-time slot; `work_reached[q, c]`, it reaches choice `c` so. + let mut at_query_time: HashMap<(usize, usize), Variable> = HashMap::new(); + for (g, queries) in input.group_queries.iter().enumerate() { + for &q in queries { + at_query_time.insert((q, g), vars.add(variable().min(0).max(1))); + } + } + let query_time_work = |c: usize| -> f64 { + let coefficients = &input.choices[c].coefficients; + coefficients.direct.work_ms + + coefficients + .slots + .iter() + .filter(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) + .map(|s| s.work_ms) + .sum::() + }; + let held_of = |c: usize| input.held.get(c).copied().unwrap_or_default(); + let mut work_reached: HashMap<(usize, usize), Variable> = HashMap::new(); + for (c, choice) in input.choices.iter().enumerate() { + let held = held_of(c); + if query_time_work(c) <= 0.0 && held.per_reader_per_ms <= 0.0 && held.longest_per_ms <= 0.0 + { + continue; + } + for &q in &input.group_queries[choice.group.0 as usize] { + work_reached.insert((q, c), vars.add(variable().min(0).max(1))); + } + } + // Held memory per (choice, query) for a reader's own copy, and per + // choice for one copy over the longest reader. + let mut held_per_reader: HashMap<(usize, usize), Variable> = HashMap::new(); + let mut held_longest: HashMap = HashMap::new(); + for &(q, c) in work_reached.keys() { + let held = held_of(c); + if held.per_reader_per_ms > 0.0 { + let h = vars.add(variable().min(0)); + objective += h; + held_per_reader.insert((q, c), h); + } + if held.longest_per_ms > 0.0 && !held_longest.contains_key(&c) { + let h = vars.add(variable().min(0)); + objective += h; + held_longest.insert(c, h); + } + } let vars_exclusive = (vars.add(binary()), vars.add(binary())); let source_count = input .raw_scans @@ -710,18 +776,103 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> b >= built_once - floor * x[c] - big_m * (1 - u) )); } - // Latency bounds: the work each query reaches at query time. + // Query-time reach: a query reaches its roots, and through a chosen + // choice its children read directly or by a query-time slot. for (q, query) in input.queries.iter().enumerate() { - let Some(bound) = query.max_latency_ms else { - continue; - }; + for root in &query.roots { + if let Some(&v) = at_query_time.get(&(q, root.0 as usize)) { + constraints.push(constraint!(v == 1)); + } + } + } + for (g, of_g) in choices_of.iter().enumerate() { + let mut children: BTreeSet = BTreeSet::new(); + for &c in of_g { + children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); + } + for h in children { + let mut with_child = Expression::from(0.0); + for &c in of_g { + let choice = &input.choices[c]; + let at_query = choice + .children + .iter() + .zip(&choice.coefficients.child_via) + .any(|(child, via)| { + child.0 as usize == h + && match via { + Via::Direct => true, + Via::Slot(i) => !matches!( + choice.coefficients.slots[*i].materialization, + Materialization::MatAtIngestion(_) + ), + } + }); + if at_query { + with_child += x[c]; + } + } + for &q in &input.group_queries[g] { + if let (Some(&above), Some(&below)) = + (at_query_time.get(&(q, g)), at_query_time.get(&(q, h))) + { + constraints.push(constraint!(below >= above + with_child.clone() - 1)); + } + } + } + } + for (&(q, c), &w) in &work_reached { + let g = input.choices[c].group.0 as usize; + constraints.push(constraint!(w >= at_query_time[&(q, g)] + x[c] - 1)); + } + // Each query's latency: the query-time work of the choices it reaches + // (a lower bound: compaction is left out). + let latency_of = |q: usize| { let mut latency = Expression::from(0.0); - for (&(k, c), &v) in &y { - if k == Consumer::Query(q) { - latency += input.choices[c].coefficients.direct.work_ms * v; + for (&(r, c), &w) in &work_reached { + if r == q { + latency += query_time_work(c) * w; } } - constraints.push(constraint!(latency <= bound)); + latency + }; + for (q, query) in input.queries.iter().enumerate() { + if let Some(bound) = query.max_latency_ms { + constraints.push(constraint!(latency_of(q) <= bound)); + } + } + // Held memory, linearized: h ≥ coefficient · (latency − L · (1 − w)), + // L the longest the query's latency can be; exact when w is 0 or 1. + let longest_latency: Vec = (0..input.queries.len()) + .map(|q| { + choices_of + .iter() + .enumerate() + .filter(|(g, _)| input.group_queries[*g].contains(&q)) + .map(|(_, of_g)| of_g.iter().map(|&c| query_time_work(c)).fold(0.0, f64::max)) + .sum() + }) + .collect(); + for (&(q, c), &w) in &work_reached { + let held = held_of(c); + let longest = longest_latency[q]; + if let Some(&h) = held_per_reader.get(&(q, c)) { + let coefficient = held.per_reader_per_ms * input.queries[q].rate; + constraints.push(constraint!( + h >= coefficient * (latency_of(q) - longest * (1 - w)) + )); + } + if let Some(&h) = held_longest.get(&c) { + let g = input.choices[c].group.0 as usize; + let rate = input.group_queries[g] + .iter() + .map(|&r| input.queries[r].rate) + .fold(f64::INFINITY, f64::min); + let coefficient = held.longest_per_ms * rate; + constraints.push(constraint!( + h >= coefficient * (latency_of(q) - longest * (1 - w)) + )); + } } // Cuts. for (set, outside, bound) in &input.excluded { @@ -954,6 +1105,7 @@ pub(crate) fn search( let mut cache = super::pricing::PriceCache::new(scale); let mut choices: Vec = Vec::new(); let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); + let mut held: Vec = Vec::new(); let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { unreachable!("the caller routes cost models 2 and 3 to the legacy search"); }; @@ -974,14 +1126,14 @@ pub(crate) fn search( // pricing of a group choice prices both. // Widened scans are reached through the identical-expression // variants of the exact pricing, never chosen directly. - for (mut choice, lazy) in group_choices.into_iter().zip(lazy) { + for (choice, lazy) in group_choices.into_iter().zip(lazy) { if choice.copy_of.is_some() || memo.candidates[g][choice.candidate].origin == super::CandidateOrigin::WidenedScans { continue; } - add_lazy_lower_bounds(&mut choice, &lazy, prices); + held.push(held_at_query_time(&choice, &lazy, prices)); raw_scans.extend( lazy.raw_scans .iter() @@ -1017,6 +1169,7 @@ pub(crate) fn search( }); } input.raw_data_retained = models.deployment_capabilities.raw_data_retained; + input.held = held; let mut selector = PlanSelector::new(entries, data, models, false)?; let mut logical: HashMap> = HashMap::new(); let mut number = 0usize; @@ -1050,6 +1203,9 @@ pub(crate) fn search( for &c in &solution.chosen { group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); } + if trace { + eprintln!("milp: round {iterations} group choice {group_choice:?}"); + } let chosen_candidates: Vec<(GroupId, &super::ImplementationCandidate)> = group_choice .iter() .enumerate() @@ -1214,6 +1370,7 @@ fn model_input_clone(input: &ModelInput) -> ModelInput { exclusive: input.exclusive.clone(), raw_scans: input.raw_scans.clone(), raw_data_retained: input.raw_data_retained, + held: input.held.clone(), rollup_rules: input.rollup_rules.clone(), } } @@ -1384,6 +1541,7 @@ fn model_input( rollup_rules, raw_scans: Vec::new(), raw_data_retained: true, + held: Vec::new(), }) } @@ -1452,6 +1610,7 @@ mod tests { rollup_rules: Vec::new(), raw_scans: Vec::new(), raw_data_retained: true, + held: Vec::new(), }; let solution = solve(&input).unwrap().unwrap(); assert_eq!(solution.built_once, [GroupId(0)].into()); From 56759e0c51a4da62732be5b44f13e2ed4db63ffc Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:34:53 +0000 Subject: [PATCH 44/67] Memo MILP search: cut off plans the model prices at or above the incumbent (experimental, off by default); solver option hooks for measurements --- .../plan-selection/src/pass/subexpr/milp.rs | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 5f89ac016..06d34ee56 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -171,6 +171,10 @@ pub(crate) struct ModelInput { /// Per choice, the memory its parts hold while a query's evaluation /// runs ([`HeldAtQueryTime`]). pub held: Vec, + /// Only plans the model prices below this: the incumbent's price. The + /// model is then infeasible exactly when no group choice left can beat + /// the incumbent. + pub cutoff: Option, } /// The memory a choice's parts hold while an evaluation runs, priced per @@ -960,6 +964,9 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } + if let Some(cutoff) = input.cutoff.filter(|_| std::env::var_os("JP_MILP_CUTOFF").is_some()) { + constraints.push(constraint!(objective.clone() <= cutoff)); + } let constraint_count = constraints.len() as u64; let mut model = model .with_all(constraints) @@ -968,6 +975,20 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if std::env::var_os("JP_MILP_VERBOSE").is_some() { model.set_verbose(true); } + if let Ok(options) = std::env::var("JP_HIGHS_OPTIONS") { + for option in options.split(',').filter(|o| !o.is_empty()) { + let (key, value) = option.split_once('=').unwrap_or((option, "")); + if let Ok(v) = value.parse::() { + model = model.set_option(key, v); + } else if let Ok(v) = value.parse::() { + model = model.set_option(key, v); + } else if let Ok(v) = value.parse::() { + model = model.set_option(key, v); + } else { + model = model.set_option(key, value); + } + } + } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), @@ -1181,9 +1202,16 @@ pub(crate) fn search( loop { iterations += 1; let solve_started = std::time::Instant::now(); + input.cutoff = selector + .cheapest_price() + .map(|incumbent| incumbent - 1e-9 * incumbent.abs()); let solved = solve(&input).map_err(SelectionError::Solver)?; solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; let Some(solution) = solved else { + // No plan left below the incumbent: it is optimal. + if let Some(incumbent) = selector.cheapest_price() { + bound = Some(incumbent); + } break; }; bound = Some(solution.objective); @@ -1371,6 +1399,7 @@ fn model_input_clone(input: &ModelInput) -> ModelInput { raw_scans: input.raw_scans.clone(), raw_data_retained: input.raw_data_retained, held: input.held.clone(), + cutoff: input.cutoff, rollup_rules: input.rollup_rules.clone(), } } @@ -1542,6 +1571,7 @@ fn model_input( raw_scans: Vec::new(), raw_data_retained: true, held: Vec::new(), + cutoff: None, }) } @@ -1611,6 +1641,7 @@ mod tests { raw_scans: Vec::new(), raw_data_retained: true, held: Vec::new(), + cutoff: None, }; let solution = solve(&input).unwrap().unwrap(); assert_eq!(solution.built_once, [GroupId(0)].into()); 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 45/67] 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 From a4bea41d62732de7bdac7ac749fda1a3053772f0 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:00:39 +0000 Subject: [PATCH 46/67] Memo MILP search: a time limit that reports its gap and is never called optimal; prices scaled to the solver's tolerances; groups no query reads charge no rate; tracing of each round's exact pricing --- crates/devtools/src/bin/stage_pipeline.rs | 27 ++- crates/devtools/src/planning_stages.rs | 17 +- .../tests/joint_search_equivalence.rs | 1 + crates/plan-selection/src/lib.rs | 9 + .../src/pass/stage1_to_stage3.rs | 15 +- .../plan-selection/src/pass/subexpr/milp.rs | 178 +++++++++++++++++- 6 files changed, 226 insertions(+), 21 deletions(-) diff --git a/crates/devtools/src/bin/stage_pipeline.rs b/crates/devtools/src/bin/stage_pipeline.rs index 2a1d6e78f..d6a8bd6ac 100644 --- a/crates/devtools/src/bin/stage_pipeline.rs +++ b/crates/devtools/src/bin/stage_pipeline.rs @@ -221,7 +221,7 @@ use asap_logical_asap_optimization::DeploymentCapabilities; use asap_plan_selection::{ large_peak_memory_warning, ClusterLimits, ClusterObjective, CostModel, MeasuredOperationCosts, PlanSelectionSettings, PlanningModels, Rejection, ResourceUnitCost, ResourceUsage, - SearchAlgorithm, Selection, SelectionError, + SearchAlgorithm, Selection, SelectionError, Stage1ToStage3Search, }; use asap_planner::{installable_plan_of, PlannerExecutorCompileCheck}; use asap_types::ir::flat::{flatten, FlatDag}; @@ -240,7 +240,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 ] | --search legacy|memo-brute-force|memo-milp] [--max-displayed ] --out "; +[--plan-selection-settings ] [--enforce-accuracy-guarantees] [--deployment-capabilities ] [--exhaustive [--max-combinations ] | --search legacy|memo-brute-force|memo-milp [--milp-time-limit-s ]] [--max-displayed ] --out "; /// The default `--max-combinations`: the most Stage 1 candidates (Pass 2's /// included) an exhaustive run enumerates. @@ -342,6 +342,7 @@ fn run(args: Vec) -> Result<(), Failure> { let mut enforce_accuracy_guarantees = false; let mut exhaustive = false; let mut search_algorithm: Option = None; + let mut milp_time_limit_ms: Option = None; let mut max_combinations_given = false; let mut args = args.into_iter(); while let Some(flag) = args.next() { @@ -401,6 +402,12 @@ 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, + "--milp-time-limit-s" => { + let seconds: f64 = value()? + .parse() + .map_err(|e| format!("--milp-time-limit-s: {e}"))?; + milp_time_limit_ms = Some((seconds * 1e3) as u64); + } "--search" => { search_algorithm = Some(match value()?.as_str() { "legacy" => SearchAlgorithm::Legacy, @@ -430,6 +437,11 @@ fn run(args: Vec) -> Result<(), Failure> { ); } let search_algorithm = search_algorithm.unwrap_or_default(); + if milp_time_limit_ms.is_some() && search_algorithm != SearchAlgorithm::MemoMilp { + return Err( + "--milp-time-limit-s limits the MILP search: it needs --search memo-milp".into(), + ); + } if max_combinations_given && !exhaustive { return Err("--max-combinations limits the exhaustive search: it needs --exhaustive (the default search has no limit)".into()); } @@ -641,7 +653,16 @@ fn run(args: Vec) -> Result<(), Failure> { Some(max_combinations), models, ), - false => run_search_over(roots.clone(), &entries, &data, models, search_algorithm), + false => run_search_over( + roots.clone(), + &entries, + &data, + models, + Stage1ToStage3Search { + milp_time_limit_ms, + ..Stage1ToStage3Search::with_algorithm(search_algorithm) + }, + ), }; // A selected plan's very large peak memory, when the deployment // declares no memory capacity to check it against. diff --git a/crates/devtools/src/planning_stages.rs b/crates/devtools/src/planning_stages.rs index d411894bf..c4569af26 100644 --- a/crates/devtools/src/planning_stages.rs +++ b/crates/devtools/src/planning_stages.rs @@ -11,7 +11,7 @@ use asap_plan_selection::pass::{ run_stage1_to_stage3_from_pass1, run_stage2_and_stage3, Stage1ToStage3Search, }; pub use asap_plan_selection::pass::{LogicalCandidate, PrunedCandidate}; -use asap_plan_selection::{PlanningModels, SearchAlgorithm}; +use asap_plan_selection::PlanningModels; use asap_plan_selection::{Selection, SelectionError}; use asap_types::ir::schema::schema_support::with_promql_series_identity; use asap_types::ir::QueryRoot; @@ -220,26 +220,21 @@ pub fn run_planning_stages_over( /// fails), and Stage 3 records the first rejection of each cause and the /// search's report (`Selection::search`). `models` should be the ones /// `plan_installable` plans with (the executor's summary families and -/// compile check) to see what it selects. No limit applies. `algorithm` -/// picks the search ([`SearchAlgorithm::Legacy`] is the default search). +/// compile check) to see what it selects. `search` picks the algorithm +/// ([`asap_plan_selection::SearchAlgorithm::Legacy`] is the default search) and the MILP +/// search's time limit; nothing else is limited. pub fn run_search_over( roots: Vec, entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, - algorithm: SearchAlgorithm, + 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, - Stage1ToStage3Search::with_algorithm(algorithm), - ); + 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 7e4977104..d0fbd8778 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -505,6 +505,7 @@ fn a_known_bound_that_misses_is_not_discarded_before_composition() { keep_every_candidate: true, exhaustive: false, search: SearchAlgorithm::Legacy, + milp_time_limit_ms: None, }, ) .unwrap(); diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 98f9caa63..ed9cd931d 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1060,6 +1060,11 @@ pub struct SearchReport { /// every other search. #[serde(default, skip_serializing_if = "Option::is_none")] pub solver_first_bound_cost_per_second: Option, + /// Whether the MILP search stopped at its time limit + /// ([`crate::Stage1ToStage3Search::milp_time_limit_ms`]): its plan is + /// then proven optimal only if [`Self::solver_gap`] is 0. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub solver_time_limit_reached: bool, } /// The size of a sub-expression memo (joint-planning-search design, @@ -2111,6 +2116,10 @@ impl<'a> PlanSelector<'a> { /// into the cheapest price and the runner-up ([`ROUNDING`]). /// Under a price ranking, the cheapest price of the valid candidates /// offered so far. + pub(crate) fn rejections(&self) -> &[Rejection] { + &self.rejected + } + pub(crate) fn cheapest_price(&self) -> Option { self.cheapest } diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index 5a80128f7..368810d11 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -165,6 +165,12 @@ pub struct Stage1ToStage3Search { /// [`SearchAlgorithm::Legacy`] until the MILP selects the exhaustive /// search's price on every equivalence test. pub search: SearchAlgorithm, + /// The longest the MILP search may run, in ms; `None` for no limit. + /// Stopped at the limit, it selects its best plan so far and reports + /// the gap to its bound ([`crate::SearchReport::solver_gap`]) and + /// [`crate::SearchReport::solver_time_limit_reached`]: not proven + /// optimal unless the gap is 0. + pub milp_time_limit_ms: Option, } impl Stage1ToStage3Search { @@ -175,6 +181,7 @@ impl Stage1ToStage3Search { keep_every_candidate: true, exhaustive: true, search: SearchAlgorithm::Legacy, + milp_time_limit_ms: None, } } @@ -268,7 +275,13 @@ pub fn run_stage1_to_stage3_from_pass1( let pruned = per_computation .clone() .without_unsupported_summaries(models.deployment_capabilities); - let found = super::subexpr::milp::search(&pruned, entries, data, models); + let found = super::subexpr::milp::search( + &pruned, + entries, + data, + models, + search.milp_time_limit_ms, + ); return memo_run(pruned, found, entries, data, SearchAlgorithm::MemoMilp); } } diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 06d34ee56..15c274853 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -11,7 +11,7 @@ //! query runs, composed accuracy) are added back by the caller's lazy cuts: //! the model's optimum is a lower bound on the plan's Stage 3 price. -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use good_lp::solvers::highs::highs; use good_lp::{ @@ -175,6 +175,8 @@ pub(crate) struct ModelInput { /// model is then infeasible exactly when no group choice left can beat /// the incumbent. pub cutoff: Option, + /// The solver's time limit in seconds, if any. + pub time_limit_s: Option, } /// The memory a choice's parts hold while an evaluation runs, priced per @@ -241,6 +243,9 @@ pub(crate) struct ModelSolution { pub built_once: BTreeSet, /// The model's objective: a lower bound on the plan's price. pub objective: f64, + /// Whether the solver proved it optimal (not stopped at its time + /// limit): only then is `objective` a lower bound. + pub optimal: bool, /// The model's size. pub variables: u64, pub constraints: u64, @@ -257,6 +262,71 @@ enum Consumer { /// Build and solve the model. `None` when it is infeasible. pub(crate) fn solve(input: &ModelInput) -> Result, String> { + // Prices per second are tiny (1e-6 and less) next to the solver's + // feasibility tolerance (1e-7): solve in units where the largest + // per-choice charge is about 1, and convert the objective back. + let scale = price_scale(input); + let scaled = scaled(input, scale); + let solution = solve_scaled(&scaled)?; + Ok(solution.map(|mut solution| { + solution.objective /= scale; + solution + })) +} + +/// The factor that brings `input`'s largest charge to about 1. +fn price_scale(input: &ModelInput) -> f64 { + let fastest = input.queries.iter().map(|q| q.rate).fold(1.0, f64::max); + let largest = input + .choices + .iter() + .flat_map(|choice| { + let k = &choice.coefficients; + std::iter::once(k.direct.cost_per_unit_rate * fastest).chain( + k.slots + .iter() + .map(move |s| s.fixed_cost + s.cost_per_unit_rate * fastest), + ) + }) + .chain(input.raw_scans.iter().map(|s| s.cost)) + .filter(|v| v.is_finite()) + .fold(0.0, f64::max); + if largest > 0.0 { + 1.0 / largest + } else { + 1.0 + } +} + +/// `input` with every price multiplied by `scale`. +fn scaled(input: &ModelInput, scale: f64) -> ModelInput { + let mut scaled = model_input_clone(input); + for choice in &mut scaled.choices { + let k = &mut choice.coefficients; + k.direct.cost_per_unit_rate *= scale; + for slot in &mut k.slots { + slot.fixed_cost *= scale; + slot.cost_per_unit_rate *= scale; + } + for node in &mut k.shared_nodes { + node.cost_per_unit_rate *= scale; + } + } + for scan in &mut scaled.raw_scans { + scan.cost *= scale; + } + for held in &mut scaled.held { + held.per_reader_per_ms *= scale; + held.longest_per_ms *= scale; + } + for (_, bound) in &mut scaled.extra_bounds { + *bound *= scale; + } + scaled.cutoff = input.cutoff.map(|c| c * scale); + scaled +} + +fn solve_scaled(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; let relaxed = std::env::var_os("JP_MILP_LP").is_some(); @@ -373,11 +443,15 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Consumer::Group(_) => None, } }; + // 0 for a group no query reads (no rate is charged linearly then). let slowest_rate = |g: usize| -> f64 { - rate_of[&g] - .iter() + rate_of + .get(&g) + .into_iter() + .flatten() .map(|(set, _)| input.union_rate[set]) - .fold(f64::INFINITY, f64::min) + .reduce(f64::min) + .unwrap_or(0.0) }; // Products of a variable rate (a set selector) and a reach variable. let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); @@ -565,6 +639,22 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> for &v in &kept_on_disk { objective += v; } + if std::env::var_os("JP_MILP_TRACE").is_some() { + use good_lp::IntoAffineExpression; + let negative: Vec<(Variable, f64)> = (&objective) + .linear_coefficients() + .filter(|(_, c)| *c < 0.0 || c.is_nan()) + .collect(); + if !negative.is_empty() { + let x_set: HashSet = x.iter().copied().collect(); + eprintln!( + "milp: {} negative objective coefficients, {} on choices: {:?}", + negative.len(), + negative.iter().filter(|(v, _)| x_set.contains(v)).count(), + &negative[..negative.len().min(5)] + ); + } + } let variable_count = vars.len() as u64; let model = vars.minimise(objective.clone()).using(highs); let mut constraints = Vec::new(); @@ -871,7 +961,8 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> let rate = input.group_queries[g] .iter() .map(|&r| input.queries[r].rate) - .fold(f64::INFINITY, f64::min); + .reduce(f64::min) + .unwrap_or(0.0); let coefficient = held.longest_per_ms * rate; constraints.push(constraint!( h >= coefficient * (latency_of(q) - longest * (1 - w)) @@ -964,7 +1055,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> } } } - if let Some(cutoff) = input.cutoff.filter(|_| std::env::var_os("JP_MILP_CUTOFF").is_some()) { + if let Some(cutoff) = input + .cutoff + .filter(|_| std::env::var_os("JP_MILP_CUTOFF").is_some()) + { constraints.push(constraint!(objective.clone() <= cutoff)); } let constraint_count = constraints.len() as u64; @@ -975,6 +1069,9 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> if std::env::var_os("JP_MILP_VERBOSE").is_some() { model.set_verbose(true); } + if let Some(limit) = input.time_limit_s { + model = model.set_time_limit(limit.max(0.001)); + } if let Ok(options) = std::env::var("JP_HIGHS_OPTIONS") { for option in options.split(',').filter(|o| !o.is_empty()) { let (key, value) = option.split_once('=').unwrap_or((option, "")); @@ -994,6 +1091,21 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> Err(good_lp::ResolutionError::Infeasible) => return Ok(None), Err(error) => return Err(error.to_string()), }; + if std::env::var_os("JP_MILP_TRACE").is_some() { + use good_lp::IntoAffineExpression; + let low: Vec<(f64, f64)> = (&objective) + .linear_coefficients() + .map(|(v, c)| (c, solution.value(v))) + .filter(|(_, value)| *value < -1e-9) + .collect(); + if !low.is_empty() { + eprintln!( + "milp: {} objective variables below 0: {:?}", + low.len(), + &low[..low.len().min(8)] + ); + } + } let chosen: Vec = x .iter() .enumerate() @@ -1044,6 +1156,7 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> chosen, built_once, objective: solution.eval(&objective), + optimal: matches!(solution.status(), good_lp::solvers::SolutionStatus::Optimal), variables: variable_count, constraints: constraint_count, })) @@ -1118,6 +1231,7 @@ pub(crate) fn search( entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, + time_limit_ms: Option, ) -> MemoSearchResult { let started = std::time::Instant::now(); let memo = super::search::memo(per_computation, entries); @@ -1164,6 +1278,23 @@ pub(crate) fn search( } } let trace = std::env::var_os("JP_MILP_TRACE").is_some(); + if trace { + for (c, choice) in choices.iter().enumerate() { + let k = &choice.coefficients; + let negative = k.direct.cost_per_unit_rate < 0.0 + || k.direct.work_ms < 0.0 + || k.slots + .iter() + .any(|s| s.fixed_cost < 0.0 || s.cost_per_unit_rate < 0.0 || s.work_ms < 0.0) + || k.shared_nodes.iter().any(|n| n.cost_per_unit_rate < 0.0); + if negative { + eprintln!( + "milp: NEGATIVE choice {c}: g{} cand{} {:?} assignment {:?}", + choice.group.0, choice.candidate, k, choice.assignment + ); + } + } + } if trace { eprintln!( "milp: memo {} groups, {} choices in {:.1}s", @@ -1199,7 +1330,15 @@ pub(crate) fn search( let mut first_bound = None; let mut solver_ms = 0.0; let mut size = (0, 0); + let mut time_limit_reached = false; loop { + let remaining_s = + time_limit_ms.map(|limit| limit as f64 / 1e3 - started.elapsed().as_secs_f64()); + if remaining_s.is_some_and(|s| s <= 0.0) { + time_limit_reached = true; + break; + } + input.time_limit_s = remaining_s; iterations += 1; let solve_started = std::time::Instant::now(); input.cutoff = selector @@ -1207,6 +1346,11 @@ pub(crate) fn search( .map(|incumbent| incumbent - 1e-9 * incumbent.abs()); let solved = solve(&input).map_err(SelectionError::Solver)?; solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; + if solved.as_ref().is_some_and(|s| !s.optimal) { + // The solver stopped at the time limit: its plan is no bound. + time_limit_reached = true; + break; + } let Some(solution) = solved else { // No plan left below the incumbent: it is optimal. if let Some(incumbent) = selector.cheapest_price() { @@ -1282,6 +1426,8 @@ pub(crate) fn search( continue; } } + let offer_started = std::time::Instant::now(); + let offered_before = number; super::search::offer_group_choice( &memo, per_computation, @@ -1293,6 +1439,17 @@ pub(crate) fn search( &mut logical, &mut number, )?; + if trace { + eprintln!( + "milp: round {iterations} priced {} plans in {:.1}s; last rejection {:?}", + number - offered_before, + offer_started.elapsed().as_secs_f64(), + selector + .rejections() + .last() + .map(|r| (r.cause, r.reason.chars().take(200).collect::())) + ); + } // Exclude it: every choice of each chosen (group, candidate). let groups = group_choice.iter().flatten().count(); let excluded: Vec = input @@ -1362,6 +1519,11 @@ pub(crate) fn search( } } } + if time_limit_reached && selector.cheapest_price().is_none() { + return Err(SelectionError::Solver(format!( + "the MILP search reached its time limit after {iterations} rounds before finding a valid plan" + ))); + } let (mut selection, selected) = selector.finish()?; let price = selection.costs[&selection.selected].cost_per_second; selection.search.memo = Some(memo_stats); @@ -1372,6 +1534,7 @@ pub(crate) fn search( selection.search.solver_constraints, ) = size; selection.search.solver_first_bound_cost_per_second = first_bound; + selection.search.solver_time_limit_reached = time_limit_reached; selection.search.solver_gap = match (price, bound) { (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), (Some(_), Some(_)) => Some(0.0), @@ -1400,6 +1563,7 @@ fn model_input_clone(input: &ModelInput) -> ModelInput { raw_data_retained: input.raw_data_retained, held: input.held.clone(), cutoff: input.cutoff, + time_limit_s: input.time_limit_s, rollup_rules: input.rollup_rules.clone(), } } @@ -1572,6 +1736,7 @@ fn model_input( raw_data_retained: true, held: Vec::new(), cutoff: None, + time_limit_s: None, }) } @@ -1642,6 +1807,7 @@ mod tests { raw_data_retained: true, held: Vec::new(), cutoff: None, + time_limit_s: None, }; let solution = solve(&input).unwrap().unwrap(); assert_eq!(solution.built_once, [GroupId(0)].into()); From 5578cac2899aec37abd904fe4c0107cf1796222f Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:01:16 +0000 Subject: [PATCH 47/67] Test the MILP search's time limit report and its argument check --- crates/devtools/tests/stage_pipeline.rs | 29 +++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index 1c72bfa8b..9ec2ba966 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1790,3 +1790,32 @@ fn search_flag_is_checked() { assert_eq!(status, Some(2)); assert!(stderr.contains("drop it with --exhaustive"), "{stderr}"); } + +/// The MILP search stopped by its time limit before any valid plan says +/// so (never a plan called optimal), and the limit needs the MILP search. +#[test] +fn milp_time_limit_is_reported_and_needs_the_milp_search() { + let document = document_of(&[ + "--example", + "planner-layering-3b", + "--search", + "memo-milp", + "--milp-time-limit-s", + "0", + ]); + let selection = &document["stage3_plan_selection"]; + assert!(selection["selected"].is_null(), "{selection}"); + assert!( + selection["error"] + .as_str() + .is_some_and(|e| e.contains("reached its time limit")), + "{selection}" + ); + let (_, stderr) = failure(&[ + "--example", + "planner-layering-3b", + "--milp-time-limit-s", + "1", + ]); + assert!(stderr.contains("needs --search memo-milp"), "{stderr}"); +} From f152a4775de17285efe7140c290638d51113190c Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:26:11 +0000 Subject: [PATCH 48/67] Restore the stage_pipeline facts flags and the default-search helper the merge had dropped; --search memo-milp and --milp-time-limit-s on top --- crates/devtools/src/bin/stage_pipeline.rs | 82 +- crates/devtools/src/planning_stages.rs | 21 +- .../src/pass/subexpr/milp.rs.orig | 1433 ----------------- 3 files changed, 86 insertions(+), 1450 deletions(-) delete mode 100644 crates/plan-selection/src/pass/subexpr/milp.rs.orig diff --git a/crates/devtools/src/bin/stage_pipeline.rs b/crates/devtools/src/bin/stage_pipeline.rs index 9685209b3..613e233b7 100644 --- a/crates/devtools/src/bin/stage_pipeline.rs +++ b/crates/devtools/src/bin/stage_pipeline.rs @@ -39,7 +39,9 @@ // force (the reference), or solved as a MILP; `search` names it, and the // search report says which ran (`algorithm`, omitted for legacy) and, // when the legacy search ran instead, why (`routed_to_legacy`). Not with -// `--exhaustive`. +// `--exhaustive`. `--milp-time-limit-s ` limits the MILP search: at +// the limit it selects its best plan so far and reports +// `solver_time_limit_reached` with the gap to its bound. // Unlike `plan_installable`'s default settings, Stage 3 ignores accuracy // guarantees unless `--enforce-accuracy-guarantees` is given (below); // `accuracy_guarantees` is `"ignored"` or `"enforced"`. @@ -195,7 +197,13 @@ // `PlanSelectionSettings`. The document's top-level `accuracy_guarantees` is // `"ignored"` or `"enforced"`. // `--sample-interval-ms` is the sample interval declared for every metric the -// PromQL queries read (default 15000); the data declares no other facts. +// PromQL queries read (default 15000). `--series-count `, +// `--rows-per-second ` (the data then ingests at that rate) and +// `--rows-held ` declare those facts for every source the queries read +// (metric or `--table`); nothing else is declared. Plan selection prices +// only from declared facts: a scan whose source has no arrival rate +// (`--rows-per-second`, or `--series-count` with the sample interval) and +// no `--rows-held` fails, naming the source and the missing fact. // `--table ''` (repeatable) declares a table the SQL queries read: `{"name": ..., // "columns": [{"name": ..., "type": "timestamp|utf8|string|float64|double| // int64|bigint", "nullable": true}], "time_index": 0}` (`nullable` defaults @@ -237,7 +245,7 @@ use serde_json::{json, Value}; const USAGE: &str = "usage: stage_pipeline (--example planner-layering-{1,2,3a,3b} | ((--promql | --sql) \ -[--name ]... [--table ...] [--max-error --failure-probability --accuracy-metric ] [--sample-interval-ms ])) \ +[--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|memo-milp [--milp-time-limit-s ]] [--max-displayed ] --out \n\ @@ -330,6 +338,9 @@ fn run(args: Vec) -> Result<(), Failure> { let mut queries: Vec<(QueryLanguage, String, Option)> = Vec::new(); let mut tables = Vec::new(); let (mut max_error, mut failure_probability, mut sample_interval_ms) = (None, None, 15_000u64); + // Facts declared for every source the --promql/--sql queries read. + let (mut series_count, mut rows_per_second, mut rows_held) = + (None::, None::, None::); let mut accuracy_metric: Option = None; let mut max_displayed = 64usize; let mut max_combinations = MAX_COMBINATIONS; @@ -342,6 +353,8 @@ 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`: the default search's algorithm, legacy unless given; + // `--milp-time-limit-s` limits the MILP search. let mut search_algorithm: Option = None; let mut milp_time_limit_ms: Option = None; let mut max_combinations_given = false; @@ -378,6 +391,9 @@ fn run(args: Vec) -> Result<(), Failure> { "--sample-interval-ms" => { sample_interval_ms = value()?.parse().map_err(|e| format!("{e}"))? } + "--series-count" => series_count = Some(value()?.parse().map_err(|e| format!("{e}"))?), + "--rows-per-second" => rows_per_second = Some(number(value()?)?), + "--rows-held" => rows_held = Some(value()?.parse().map_err(|e| format!("{e}"))?), "--max-displayed" => max_displayed = value()?.parse().map_err(|e| format!("{e}"))?, "--max-combinations" => { max_combinations = value()?.parse().map_err(|e| format!("{e}"))?; @@ -403,12 +419,6 @@ 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, - "--milp-time-limit-s" => { - let seconds: f64 = value()? - .parse() - .map_err(|e| format!("--milp-time-limit-s: {e}"))?; - milp_time_limit_ms = Some((seconds * 1e3) as u64); - } "--search" => { search_algorithm = Some(match value()?.as_str() { "legacy" => SearchAlgorithm::Legacy, @@ -422,6 +432,12 @@ fn run(args: Vec) -> Result<(), Failure> { } }) } + "--milp-time-limit-s" => { + let seconds: f64 = value()? + .parse() + .map_err(|e| format!("--milp-time-limit-s: {e}"))?; + milp_time_limit_ms = Some((seconds * 1e3) as u64); + } other => return Err(format!("unknown argument {other}").into()), } } @@ -432,6 +448,18 @@ fn run(args: Vec) -> Result<(), Failure> { return Err("--table needs --sql".into()); } let out = out.ok_or("--out is required")?; + if (series_count.is_some() || rows_per_second.is_some() || rows_held.is_some()) + && (workload_path.is_some() || example.is_some()) + { + return Err( + "--series-count, --rows-per-second and --rows-held declare facts for --promql/--sql \ + queries; a --workload or --example declares its own" + .into(), + ); + } + 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 && search_algorithm.is_some() { return Err( "--search picks the default search's algorithm: drop it with --exhaustive".into(), @@ -443,9 +471,6 @@ fn run(args: Vec) -> Result<(), Failure> { "--milp-time-limit-s limits the MILP search: it needs --search memo-milp".into(), ); } - if max_combinations_given && !exhaustive { - return Err("--max-combinations limits the exhaustive search: it needs --exhaustive (the default search has no limit)".into()); - } if plan_selection_settings_path.is_some() && (cost_model.is_some() || cluster_memory.is_some() @@ -626,7 +651,18 @@ fn run(args: Vec) -> Result<(), Failure> { query["id"] = json!(name); described.push(query); } - let data = DataWorkload::sampled_every(metrics, DurationMs(sample_interval_ms)); + let mut data = DataWorkload::sampled_every(metrics, DurationMs(sample_interval_ms)); + for table in catalog_table_names(&tables)? { + data.sources.entry(table).or_default(); + } + for facts in data.sources.values_mut() { + facts.series_count = series_count; + facts.rows_held = rows_held; + } + if let Some(rows_per_second) = rows_per_second { + data.arrival = DataArrival::Ingesting { rows_per_second }; + } + data.validate().map_err(|e| e.to_string())?; (described, roots, entries, data) } _ => { @@ -1164,6 +1200,20 @@ fn with_default_metrics(mut queries: Vec, roots: &[QueryRoot]) -> Vec Result, String> { + tables + .iter() + .map(|raw| { + let table: Value = serde_json::from_str(raw).map_err(|e| format!("--table: {e}"))?; + table["name"] + .as_str() + .map(str::to_owned) + .ok_or_else(|| "--table: name must be a string".to_string()) + }) + .collect() +} + /// The tables `--table` declares. fn sql_catalog(tables: &[String]) -> Result { let mut catalog = SqlCatalog::new(); @@ -1306,12 +1356,16 @@ mod tests { /// candidate and its cost, chosen among all candidates. #[test] fn display_cap_keeps_the_selected_candidate() { - let workload = query_batch( + let mut workload = query_batch( QueryLanguage::PromQL, &["sum(m)".into()], AccuracyTarget::exact(), 15_000, ); + // `m` holds 1,000 series: pricing uses only declared facts. + workload.data_workload = workload + .data_workload + .map(|data| data.with_series_count(1_000)); let run = run_planning_stages(&workload, lower_promql_roots(&workload).unwrap(), None).unwrap(); let entries: Vec<_> = workload.query_workload.entries().collect(); diff --git a/crates/devtools/src/planning_stages.rs b/crates/devtools/src/planning_stages.rs index c4569af26..c519a322a 100644 --- a/crates/devtools/src/planning_stages.rs +++ b/crates/devtools/src/planning_stages.rs @@ -220,9 +220,24 @@ pub fn run_planning_stages_over( /// fails), and Stage 3 records the first rejection of each cause and the /// search's report (`Selection::search`). `models` should be the ones /// `plan_installable` plans with (the executor's summary families and -/// compile check) to see what it selects. `search` picks the algorithm -/// ([`asap_plan_selection::SearchAlgorithm::Legacy`] is the default search) and the MILP -/// search's time limit; nothing else is limited. +/// compile check) to see what it selects. No limit applies. +pub fn run_default_search_over( + roots: Vec, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result { + 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], diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs.orig b/crates/plan-selection/src/pass/subexpr/milp.rs.orig deleted file mode 100644 index f41161575..000000000 --- a/crates/plan-selection/src/pass/subexpr/milp.rs.orig +++ /dev/null @@ -1,1433 +0,0 @@ -//! The memo's solver: one MILP over the memo's choices (design "The -//! solver: MILP over the memo"). -//! -//! A choice is one candidate of one group with an option for each of its -//! private materialization slots ([`Choice`]); its coefficients come from -//! Stage 3's per-node pricing ([`ChoiceCoefficients`]). The model picks one -//! choice per needed group, a materialization for each shared group's own -//! node, and propagates which consumers (queries, materialized slots) reach -//! which group, so each private node is charged once per consumer reaching -//! it, at that consumer's rate. Terms the model omits (memory held while a -//! query runs, composed accuracy) are added back by the caller's lazy cuts: -//! the model's optimum is a lower bound on the plan's Stage 3 price. - -use std::collections::{BTreeMap, BTreeSet, HashMap}; - -use good_lp::solvers::highs::highs; -use good_lp::{ - constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, WithMipGap, -}; - -use super::coefficients::LazyTerms; -use super::{GroupId, Memo}; -use crate::{CostModel, PlanSelector, PlanningModels, ResourceUnitCost, Selection, SelectionError}; -use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; -use asap_physical_asap_optimization::PhysicalCandidate; -use asap_types::ir::QueryRoot; -use asap_types::physical::Materialization; -use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; - -/// How a child group of a choice is reached: by the consumers reaching the -/// choice's own node, or by one of its materialized private slots. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum Via { - Direct, - Slot(usize), -} - -/// What one part of a choice costs, from Stage 3's per-node pricing. -#[derive(Debug, Clone, Copy, Default, PartialEq)] -pub(crate) struct PartCoefficients { - /// Price per second per unit of the rate of a consumer reaching it. - pub cost_per_unit_rate: f64, - /// Work in ms per computation, which a query reaching it at query time - /// adds to its latency. - pub work_ms: f64, -} - -/// One materialized private slot of a choice. -#[derive(Debug, Clone, Copy, PartialEq)] -pub(crate) struct SlotCoefficients { - pub materialization: Materialization, - /// Price per second that does not depend on a rate (ingestion-time - /// work and kept state). - pub fixed_cost: f64, - /// Price per second per unit of the slot's own rate (query-time slots: - /// the union rate of the queries reaching the group). - pub cost_per_unit_rate: f64, - /// Work in ms per computation at query time. - pub work_ms: f64, -} - -/// One choice's coefficients. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct ChoiceCoefficients { - pub direct: PartCoefficients, - pub slots: Vec, - /// Per child group of the candidate, in order. - pub child_via: Vec, - /// The shareable nodes of the direct part (an ASAP node, or one whose - /// schema has a unique key) of a group two or more queries read, other - /// than its own node: each may be one node for every copy of the group - /// (the identical-expression rule), and then built once at query time. - pub shared_nodes: Vec, -} - -/// A shareable node of a choice's direct part ([`ChoiceCoefficients::shared_nodes`]). -/// Its price is in `direct.cost_per_unit_rate`, charged per consumer; built -/// once at query time it costs `cost_per_unit_rate` per unit of the union -/// rate of the queries reaching it instead, and holds `state_bytes` while -/// its longest reader runs. -#[derive(Debug, Clone, Copy, PartialEq)] -pub(crate) struct SharedNodeCoefficients { - pub cost_per_unit_rate: f64, - pub work_ms: f64, - pub state_bytes: u64, -} - -/// One choice of a group: a candidate and an option per private slot. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct Choice { - pub group: GroupId, - pub candidate: usize, - /// The groups it implements (the first is `group`). - pub implements: Vec, - /// The child groups it reads. - pub children: Vec, - pub assignment: Vec, - pub coefficients: ChoiceCoefficients, - /// `None`: the group's one node, shared by every query reading it. - /// `Some(q)`: query `q`'s own copy of it (the group apart); a group's - /// copies are all apart or all shared. - pub copy_of: Option, - /// Each child group that is a scan: `None` when it is apart and its - /// copy is this choice's own (not in `children`); `Some(h)` when it is - /// the scan shared at horizon `h` ms (its choice's `scan_horizon_ms`). - pub scan_children: Vec<(GroupId, Option)>, - /// For a scan group's choice: the horizon it reads, the longest window - /// among its readers. - pub scan_horizon_ms: Option, -} - -/// A shared group's own node, built once per evaluation when chosen. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct GroupSlot { - pub group: GroupId, - /// Whether it may stay `NotMaterialized` (recomputed per consumer). - pub may_recompute: bool, -} - -/// One query of the model. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct ModelQuery { - /// Its evaluations per second. - pub rate: f64, - /// The groups its root reads directly (one for an operator root). - pub roots: Vec, - /// Its latency bound in ms, if any (the tighter of its own bound and - /// its repeat interval). - pub max_latency_ms: Option, -} - -/// Everything the model reads. -pub(crate) struct ModelInput { - pub group_count: usize, - /// Per group, the queries whose plan may reach it. - pub group_queries: Vec>, - pub choices: Vec, - pub group_slots: Vec, - pub queries: Vec, - /// Per query, its schedule class: queries of one class run at the same - /// instants (one recurrence and range evaluation). - pub schedule: Vec, - /// The union rate of the classes in each set, by sorted class set: how - /// often a node built once for queries of those classes is computed. - pub union_rate: HashMap, f64>, - /// No-good cuts from earlier iterations: at most `bound` of these - /// choices together (a plan whose every group took one of them). - /// With the groups that plan does not need: a plan needing any of them - /// is another plan, not cut. - pub excluded: Vec<(Vec, Vec, f64)>, - /// Lower bounds on the omitted cost, added as `extra ≥ bound` when all - /// of a choice set is chosen (local cuts). - pub extra_bounds: Vec<(Vec, f64)>, - /// Requirements of a choice on another group's choice: when the first - /// is chosen, one of the second (an exact implementation, or one - /// carrying a coupling site) is. - pub requires: Vec<(usize, Vec)>, - /// Choices that may not be chosen together with any of the second set - /// (panes and an exact top-k reuse in one plan). - pub exclusive: (Vec, Vec), - /// Each query-time raw scan of a choice's part: each source's longest - /// one is kept on disk. - pub raw_scans: Vec, - /// Whether the deployment keeps raw data: without it, no part with a - /// raw scan may run at query time. - pub raw_data_retained: bool, -} - -/// A raw scan of one part of a choice, when that part runs at query time. -#[derive(Debug, Clone, Copy, PartialEq)] -pub(crate) struct RawScan { - pub choice: usize, - /// Index of its source. - pub source: usize, - /// The price per second of keeping its look-back on disk. - pub cost: f64, - pub via: Via, -} - -/// Adds to `choice`'s coefficients a lower bound on the memory its parts -/// hold while an evaluation runs ([`LazyTerms::held`]): Stage 3 charges -/// `bytes · rate · E / 1000` at the memory price, `E` the latency of the -/// queries reaching the part, at least the part's own work. The group's -/// own node's working memory is left out (replaced when it is built -/// once). -fn add_lazy_lower_bounds(choice: &mut Choice, lazy: &LazyTerms, prices: &ResourceUnitCost) { - let coefficients = &mut choice.coefficients; - for term in &lazy.held { - if term.own_node { - continue; - } - let per_unit_rate = |work_ms: f64| { - prices.cost_per_memory_byte_second * term.bytes as f64 * work_ms / 1000.0 - }; - match term.via { - Via::Direct => { - coefficients.direct.cost_per_unit_rate += per_unit_rate(coefficients.direct.work_ms) - } - Via::Slot(i) => { - let slot = &mut coefficients.slots[i]; - if !matches!(slot.materialization, Materialization::MatAtIngestion(_)) { - slot.cost_per_unit_rate += per_unit_rate(slot.work_ms); - } - } - } - } -} - -/// The model's optimum. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct ModelSolution { - /// The chosen choices, by index in [`ModelInput::choices`]. - pub chosen: Vec, - /// The shared groups built once per evaluation. - pub built_once: BTreeSet, - /// The model's objective: a lower bound on the plan's price. - pub objective: f64, - /// The model's size. - pub variables: u64, - pub constraints: u64, -} - -/// A consumer of the model: a query, a materialized private slot of a -/// choice, or a shared group built once. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] -enum Consumer { - Query(usize), - Slot { choice: usize, slot: usize }, - Group(GroupId), -} - -/// Build and solve the model. `None` when it is infeasible. -pub(crate) fn solve(input: &ModelInput) -> Result, String> { - let mut vars = ProblemVariables::new(); - let n = input.group_count; - let relaxed = std::env::var_os("JP_MILP_LP").is_some(); - let binary = || { - if relaxed { - variable().min(0).max(1) - } else { - variable().binary() - } - }; - let x: Vec = input.choices.iter().map(|_| vars.add(binary())).collect(); - let need: Vec = (0..n).map(|_| vars.add(binary())).collect(); - let once: BTreeMap = input - .group_slots - .iter() - .map(|slot| (slot.group, vars.add(binary()))) - .collect(); - // Which queries' plans reach each group at all (any materialization): - // the union rate of a node built once reads them. - let mut reaches_at_all: HashMap<(usize, usize), Variable> = HashMap::new(); - for (g, queries) in input.group_queries.iter().enumerate() { - for &q in queries { - reaches_at_all.insert((q, g), vars.add(binary())); - } - } - // Consumers that may reach each group's own node, and through it its - // direct part. - let consumers = possible_consumers(input); - let mut reach: HashMap<(Consumer, usize), Variable> = HashMap::new(); - for (&g, ks) in &consumers { - for &k in ks { - reach.insert((k, g), vars.add(variable().min(0).max(1))); - } - } - // A consumer passes through a group into its chosen direct part when - // it reaches the group and the group is not built once (or it is the - // group's own built-once consumer). - let mut passes: HashMap<(Consumer, usize), Variable> = HashMap::new(); - for &(k, g) in reach.keys() { - passes.insert((k, g), vars.add(variable().min(0).max(1))); - } - for &g in once.keys() { - let k = Consumer::Group(g); - passes - .entry((k, g.0 as usize)) - .or_insert_with(|| vars.add(variable().min(0).max(1))); - } - // y[k, c]: consumer k passes into choice c's direct part. - let mut y: HashMap<(Consumer, usize), Variable> = HashMap::new(); - for (c, choice) in input.choices.iter().enumerate() { - let g = choice.group.0 as usize; - for &(k, h) in passes.keys() { - if h == g { - y.insert((k, c), vars.add(variable().min(0).max(1))); - } - } - } - // Union-rate selectors for built-once consumers (shared groups and - // query-time private slots): one per set of schedule classes among the - // group's queries, exactly the classes of the queries reaching it. - let mut rate_of: HashMap, Variable)>> = HashMap::new(); - let classes_of = |g: usize| -> Vec { - let mut classes: Vec = input.group_queries[g] - .iter() - .map(|&q| input.schedule[q]) - .collect(); - classes.sort_unstable(); - classes.dedup(); - classes - }; - let class_sets = |g: usize| -> Vec> { - let classes = classes_of(g); - (1u64..(1 << classes.len())) - .map(|mask| { - (0..classes.len()) - .filter(|i| mask >> i & 1 == 1) - .map(|i| classes[i]) - .collect() - }) - .collect() - }; - let mut groups_with_rates: BTreeSet = once.keys().map(|g| g.0 as usize).collect(); - for choice in &input.choices { - if !choice.coefficients.shared_nodes.is_empty() - || choice - .coefficients - .slots - .iter() - .any(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) - { - groups_with_rates.insert(choice.group.0 as usize); - } - } - for &g in &groups_with_rates { - let sets = class_sets(g) - .into_iter() - .map(|set| (set, vars.add(binary()))) - .collect(); - rate_of.insert(g, sets); - } - - // The objective. - let mut objective = Expression::from(0.0); - let mut products_pending: Vec<(Variable, Variable, Variable)> = Vec::new(); - // A consumer's constant rate, or `None` for one built once at query - // time, whose rate is its group's union rate. - let rate = |k: Consumer| -> Option { - match k { - Consumer::Query(q) => Some(input.queries[q].rate), - Consumer::Slot { choice, slot } => { - let m = input.choices[choice].coefficients.slots[slot].materialization; - matches!(m, Materialization::MatAtIngestion(_)).then_some(1.0) - } - Consumer::Group(_) => None, - } - }; - let slowest_rate = |g: usize| -> f64 { - rate_of[&g] - .iter() - .map(|(set, _)| input.union_rate[set]) - .fold(f64::INFINITY, f64::min) - }; - // Products of a variable rate (a set selector) and a reach variable. - let mut rated: Vec<(Variable, Variable, f64)> = Vec::new(); - // The shareable nodes of a choice's direct part cost, per unit of - // rate, `shared_cost[c]`, charged either per consumer (the rest of the - // direct part's charge) or once at the union rate of the queries - // reaching the group when built once: the model charges the smaller - // (`per_consumer_shared[c]` against `built_once_shared[c]`). - let shared_cost: Vec = input - .choices - .iter() - .map(|choice| { - choice - .coefficients - .shared_nodes - .iter() - .map(|node| node.cost_per_unit_rate) - .sum() - }) - .collect(); - let mut per_consumer_shared: Vec> = vec![Vec::new(); input.choices.len()]; - let mut built_once_shared: Vec> = vec![Vec::new(); input.choices.len()]; - let mut product = |vars: &mut ProblemVariables, a: Variable, b: Variable| { - let z = vars.add(variable().min(0).max(1)); - products_pending.push((z, a, b)); - z - }; - for (&(k, c), &v) in &y { - let total = input.choices[c].coefficients.direct.cost_per_unit_rate; - let shared = shared_cost[c].min(total); - let cost = total - shared; - if total == 0.0 { - continue; - } - match rate(k) { - Some(r) => { - objective += r * cost * v; - if shared != 0.0 { - per_consumer_shared[c].push((v, r * shared)); - } - } - None => { - let g = match k { - Consumer::Group(g) => g.0 as usize, - Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, - Consumer::Query(_) => unreachable!("a query's rate is constant"), - }; - // At least the slowest class's rate is charged linearly - // (some query reaches the group), the rest through the - // selector products: a tighter relaxation. - let slowest = slowest_rate(g); - objective += slowest * cost * v; - for (set, selector) in &rate_of[&g] { - let union = input.union_rate[set]; - if cost != 0.0 && union > slowest { - rated.push((*selector, v, (union - slowest) * cost)); - } - if shared != 0.0 { - let z = product(&mut vars, *selector, v); - per_consumer_shared[c].push((z, union * shared)); - } - } - } - } - } - for (c, choice) in input.choices.iter().enumerate() { - if shared_cost[c] == 0.0 { - continue; - } - for (set, selector) in &rate_of[&(choice.group.0 as usize)] { - let z = product(&mut vars, *selector, x[c]); - built_once_shared[c].push((z, input.union_rate[set] * shared_cost[c])); - } - } - // min(A, B) as a + b with a ≥ A − M·u, b ≥ B − M·(1 − u), u binary. - let mut cheaper_of: Vec<(Variable, Variable, Variable, usize, f64)> = Vec::new(); - for c in 0..input.choices.len() { - if per_consumer_shared[c].is_empty() && built_once_shared[c].is_empty() { - continue; - } - let big_m: f64 = per_consumer_shared[c].iter().map(|(_, w)| w).sum::() - + built_once_shared[c].iter().map(|(_, w)| w).sum::(); - let a = vars.add(variable().min(0)); - let b = vars.add(variable().min(0)); - let u = vars.add(binary()); - objective += a + b; - objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; - if std::env::var_os("JP_MILP_NOMIN").is_some() { - continue; - } - cheaper_of.push((a, b, u, c, big_m)); - } - for (c, choice) in input.choices.iter().enumerate() { - for slot in &choice.coefficients.slots { - objective += slot.fixed_cost * x[c]; - if slot.cost_per_unit_rate != 0.0 - && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) - { - let g = choice.group.0 as usize; - let slowest = slowest_rate(g); - objective += slowest * slot.cost_per_unit_rate * x[c]; - for (set, selector) in &rate_of[&g] { - let union = input.union_rate[set]; - if union > slowest { - rated.push((*selector, x[c], (union - slowest) * slot.cost_per_unit_rate)); - } - } - } - } - } - let mut products = Vec::new(); - for &(a, b, coefficient) in &rated { - let z = vars.add(variable().min(0).max(1)); - objective += coefficient * z; - products.push((z, a, b)); - } - // Lower bounds on the omitted terms from earlier cuts. - let extra: Vec = input - .extra_bounds - .iter() - .map(|_| vars.add(variable().min(0))) - .collect(); - for &e in &extra { - objective += e; - } - - let vars_exclusive = (vars.add(binary()), vars.add(binary())); - let source_count = input - .raw_scans - .iter() - .map(|s| s.source + 1) - .max() - .unwrap_or(0); - let kept_on_disk: Vec = (0..source_count) - .map(|_| vars.add(variable().min(0))) - .collect(); - for &v in &kept_on_disk { - objective += v; - } - let variable_count = vars.len() as u64; - let model = vars.minimise(objective.clone()).using(highs); - let mut constraints = Vec::new(); - // Needed groups: every query's roots; children of chosen choices. - for query in &input.queries { - for root in &query.roots { - constraints.push(constraint!(need[root.0 as usize] == 1)); - } - } - for (g, &needed) in need.iter().enumerate() { - // One choice implementing it (its own, or a fused one covering - // it) when needed, none otherwise. - let mut implementing = Expression::from(0.0); - for (c, choice) in input.choices.iter().enumerate() { - if choice.implements.iter().any(|i| i.0 as usize == g) { - implementing += x[c]; - } - } - constraints.push(constraint!(implementing == needed)); - } - for (c, choice) in input.choices.iter().enumerate() { - for child in &choice.children { - constraints.push(constraint!(need[child.0 as usize] >= x[c])); - } - } - for (&g, &v) in &once { - constraints.push(constraint!(v <= need[g.0 as usize])); - let slot = input.group_slots.iter().find(|s| s.group == g).unwrap(); - if !slot.may_recompute { - constraints.push(constraint!(v == need[g.0 as usize])); - } - } - // Reach at all: a query reaches its roots, and every child of a chosen - // choice of a group it reaches. - for (q, query) in input.queries.iter().enumerate() { - for root in &query.roots { - if let Some(&v) = reaches_at_all.get(&(q, root.0 as usize)) { - constraints.push(constraint!(v == 1)); - } - } - } - // Aggregated over a group's choices (one is chosen when it is - // needed), which keeps the LP relaxation tight. - let mut choices_of: Vec> = vec![Vec::new(); n]; - for (c, choice) in input.choices.iter().enumerate() { - choices_of[choice.group.0 as usize].push(c); - } - for (g, of_g) in choices_of.iter().enumerate() { - let mut children: BTreeSet = BTreeSet::new(); - for &c in of_g { - children.extend(input.choices[c].children.iter().map(|h| h.0 as usize)); - } - for h in children { - let mut with_child = Expression::from(0.0); - for &c in of_g { - if input.choices[c].children.iter().any(|k| k.0 as usize == h) { - with_child += x[c]; - } - } - for &q in &input.group_queries[g] { - if let (Some(&above), Some(&below)) = - (reaches_at_all.get(&(q, g)), reaches_at_all.get(&(q, h))) - { - constraints.push(constraint!(below >= above + with_child.clone() - 1)); - } - } - } - } - // Rate selectors: exactly the classes of the queries reaching the - // group (at most one selector; a class is in it iff one of its - // queries reaches the group). - for (&g, sets) in &rate_of { - let mut any = Expression::from(0.0); - for (_, selector) in sets { - any += *selector; - } - constraints.push(constraint!(any <= 1)); - for class in classes_of(g) { - let mut containing = Expression::from(0.0); - for (set, selector) in sets { - if set.contains(&class) { - containing += *selector; - } - } - for &q in &input.group_queries[g] { - if input.schedule[q] == class { - constraints.push(constraint!(containing.clone() >= reaches_at_all[&(q, g)])); - } - } - let mut members = Expression::from(0.0); - for &q in &input.group_queries[g] { - if input.schedule[q] == class { - members += reaches_at_all[&(q, g)]; - } - } - constraints.push(constraint!(containing <= members)); - } - } - // Consumer reach. - for (q, query) in input.queries.iter().enumerate() { - for root in &query.roots { - if let Some(&v) = reach.get(&(Consumer::Query(q), root.0 as usize)) { - constraints.push(constraint!(v >= need[root.0 as usize])); - } - } - } - for (&(k, g), &p) in &passes { - match k { - Consumer::Group(own) if own.0 as usize == g => { - constraints.push(constraint!(p >= once[&own])); - } - _ => { - let r = reach[&(k, g)]; - match once.get(&GroupId(g as u32)) { - Some(&o) => constraints.push(constraint!(p >= r - o)), - None => constraints.push(constraint!(p >= r)), - } - } - } - } - for (&(k, c), &v) in &y { - let g = input.choices[c].group.0 as usize; - constraints.push(constraint!(v >= passes[&(k, g)] + x[c] - 1)); - } - for (&(k, g), &p) in &passes { - let mut into = Expression::from(0.0); - for &c in &choices_of[g] { - into += y[&(k, c)]; - } - // Unless a fused choice of another group implements it. - for (c, choice) in input.choices.iter().enumerate() { - if choice.group.0 as usize != g && choice.implements.iter().any(|i| i.0 as usize == g) { - into += x[c]; - } - } - constraints.push(constraint!(into >= p)); - } - for (c, choice) in input.choices.iter().enumerate() { - for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { - let h = child.0 as usize; - match via { - Via::Direct => { - // Aggregated below, over the group's choices. - let _ = h; - } - Via::Slot(slot) => { - let k = Consumer::Slot { - choice: c, - slot: *slot, - }; - if let Some(&r) = reach.get(&(k, h)) { - constraints.push(constraint!(r >= x[c])); - } - } - } - } - } - for (g, of_g) in choices_of.iter().enumerate() { - let mut children: BTreeSet = BTreeSet::new(); - for &c in of_g { - for (child, via) in input.choices[c] - .children - .iter() - .zip(&input.choices[c].coefficients.child_via) - { - if *via == Via::Direct { - children.insert(child.0 as usize); - } - } - } - for h in children { - for &(k, pg) in passes.keys() { - if pg != g { - continue; - } - let Some(&r) = reach.get(&(k, h)) else { - continue; - }; - let mut into = Expression::from(0.0); - for &c in of_g { - let choice = &input.choices[c]; - if choice - .children - .iter() - .zip(&choice.coefficients.child_via) - .any(|(child, via)| child.0 as usize == h && *via == Via::Direct) - { - into += y[&(k, c)]; - } - } - constraints.push(constraint!(r >= into)); - } - } - } - for &(z, a, b) in products.iter().chain(&products_pending) { - constraints.push(constraint!(z >= a + b - 1)); - } - for &(a, b, u, c, big_m) in &cheaper_of { - let mut per_consumer = Expression::from(0.0); - for &(v, w) in &per_consumer_shared[c] { - per_consumer += w * v; - } - let mut built_once = Expression::from(0.0); - for &(v, w) in &built_once_shared[c] { - built_once += w * v; - } - // Both are at least the slowest class's rate (some query reaches - // the chosen choice): charged linearly, which keeps the relaxation - // tight; only the excess goes through the binary choice. - let floor = shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize); - constraints.push(constraint!(a >= per_consumer - floor * x[c] - big_m * u)); - constraints.push(constraint!( - b >= built_once - floor * x[c] - big_m * (1 - u) - )); - } - // Latency bounds: the work each query reaches at query time. - for (q, query) in input.queries.iter().enumerate() { - let Some(bound) = query.max_latency_ms else { - continue; - }; - let mut latency = Expression::from(0.0); - for (&(k, c), &v) in &y { - if k == Consumer::Query(q) { - latency += input.choices[c].coefficients.direct.work_ms * v; - } - } - constraints.push(constraint!(latency <= bound)); - } - // Cuts. - for (set, outside, bound) in &input.excluded { - let mut chosen = Expression::from(0.0); - for &c in set { - chosen += x[c]; - } - for &g in outside { - chosen -= need[g]; - } - constraints.push(constraint!(chosen <= *bound)); - } - for ((set, bound), &e) in input.extra_bounds.iter().zip(&extra) { - let mut chosen = Expression::from(0.0); - for &c in set { - chosen += x[c]; - } - constraints.push(constraint!( - e >= *bound * (chosen - (set.len() as f64 - 1.0)) - )); - } - for (c, allowed) in &input.requires { - let mut any = Expression::from(0.0); - for &a in allowed { - any += x[a]; - } - constraints.push(constraint!(x[*c] <= any)); - } - let (first, second) = &input.exclusive; - if !first.is_empty() && !second.is_empty() { - let (p, t) = (vars_exclusive.0, vars_exclusive.1); - for &c in first { - constraints.push(constraint!(x[c] <= p)); - } - for &c in second { - constraints.push(constraint!(x[c] <= t)); - } - constraints.push(constraint!(p + t <= 1)); - } - // Raw data: each source's longest query-time scan is kept on disk; a - // deployment without raw data runs no raw scan at query time. A part - // runs at query time for every consumer but an ingestion-time slot. - let at_query_time = |k: Consumer| match k { - Consumer::Slot { choice, slot } => !matches!( - input.choices[choice].coefficients.slots[slot].materialization, - Materialization::MatAtIngestion(_) - ), - _ => true, - }; - for scan in &input.raw_scans { - let c = scan.choice; - let mut runs: Vec = Vec::new(); - match scan.via { - Via::Direct => { - for (&(k, d), &v) in &y { - if d == c && at_query_time(k) { - runs.push(v); - } - } - } - Via::Slot(i) => { - let m = input.choices[c].coefficients.slots[i].materialization; - if !matches!(m, Materialization::MatAtIngestion(_)) { - runs.push(x[c]); - } - } - } - for v in runs { - if !input.raw_data_retained { - constraints.push(constraint!(v == 0)); - } else if scan.cost > 0.0 { - constraints.push(constraint!(kept_on_disk[scan.source] >= scan.cost * v)); - } - } - } - let constraint_count = constraints.len() as u64; - let mut model = model - .with_all(constraints) - .with_mip_gap(0.0) - .map_err(|e| format!("MIP gap: {e:?}"))?; - if std::env::var_os("JP_MILP_VERBOSE").is_some() { - model.set_verbose(true); - } - let solution = match model.solve() { - Ok(solution) => solution, - Err(good_lp::ResolutionError::Infeasible) => return Ok(None), - Err(error) => return Err(error.to_string()), - }; - let chosen: Vec = x - .iter() - .enumerate() - .filter(|(_, v)| solution.value(**v) > 0.5) - .map(|(c, _)| c) - .collect(); - if std::env::var_os("JP_MILP_TRACE").is_some() { - if relaxed { - for (g, v) in once.iter() { - eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); - } - for (g, v) in need.iter().enumerate() { - eprintln!(" need g{} = {:.3}", g, solution.value(*v)); - } - } - let shown: Vec = if relaxed { - (0..x.len()) - .filter(|&c| solution.value(x[c]) > 1e-6) - .collect() - } else { - chosen.clone() - }; - for &c in &shown { - if relaxed { - eprintln!(" x = {:.3}", solution.value(x[c])); - } - let choice: &Choice = &input.choices[c]; - let reach: Vec = y - .iter() - .filter(|((_, d), v)| *d == c && solution.value(**v) > 1e-9) - .map(|((k, _), v)| format!("{k:?}={:.2}", solution.value(*v))) - .collect(); - eprintln!( - " g{} cand{} direct {:.4e}/rate work {:.2}ms shared {:?} slots {:?} scans {:?} horizon {:?} assignment {:?} y {:?}", - choice.group.0, choice.candidate, choice.coefficients.direct.cost_per_unit_rate, - choice.coefficients.direct.work_ms, choice.coefficients.shared_nodes.len(), - choice.coefficients.slots.iter().map(|s| (s.fixed_cost, s.cost_per_unit_rate)).collect::>(), - choice.scan_children, choice.scan_horizon_ms, choice.assignment, reach - ); - } - } - let built_once = once - .iter() - .filter(|(_, v)| solution.value(**v) > 0.5) - .map(|(g, _)| *g) - .collect(); - Ok(Some(ModelSolution { - chosen, - built_once, - objective: solution.eval(&objective), - variables: variable_count, - constraints: constraint_count, - })) -} - -/// For each group, the consumers that may reach its node: its queries -/// (that reach it without passing a materialized node), the materialized -/// private slots of choices of groups above it, and the shared groups -/// above it (or itself) that may be built once. -fn possible_consumers(input: &ModelInput) -> BTreeMap> { - let mut consumers: BTreeMap> = BTreeMap::new(); - for (g, queries) in input.group_queries.iter().enumerate() { - let entry = consumers.entry(g).or_default(); - for &q in queries { - entry.insert(Consumer::Query(q)); - } - } - // Propagate downwards until nothing changes. - let shared: BTreeSet = input - .group_slots - .iter() - .map(|s| s.group.0 as usize) - .collect(); - loop { - let mut changed = false; - for (c, choice) in input.choices.iter().enumerate() { - let g = choice.group.0 as usize; - let mut above: BTreeSet = consumers.get(&g).cloned().unwrap_or_default(); - if shared.contains(&g) { - above.insert(Consumer::Group(GroupId(g as u32))); - } - for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { - let entry = consumers.entry(child.0 as usize).or_default(); - let before = entry.len(); - match via { - Via::Direct => entry.extend(above.iter().copied()), - Via::Slot(slot) => { - entry.insert(Consumer::Slot { - choice: c, - slot: *slot, - }); - } - } - changed |= entry.len() != before; - } - } - if !changed { - break; - } - } - consumers -} - -/// A memo search's result: the selection, its physical candidate and its -/// logical roots, or why no plan is valid. -pub(crate) type MemoSearchResult = - Result<(Selection, PhysicalCandidate, Vec), SelectionError>; - -/// The cheapest valid plan of `per_computation`'s memo under cost model 1 -/// (the caller routes cost models 2 and 3 to the legacy search). -/// -/// Lazy loop: solve the model (a lower bound: it omits memory held while a -/// query runs, and checks composed guarantees only on its incumbents); -/// price the solution's group choice exactly, as the reference search -/// prices one ([`offer_group_choice`]: every threshold, identical-expression -/// variant and Stage 2 materialization), which updates the incumbent; add -/// a no-good cut excluding that group choice; stop when the model's bound -/// reaches the incumbent. Each group choice is priced at most once, so the -/// loop ends, and the incumbent is the reference search's optimum. -pub(crate) fn search( - per_computation: &PerComputationCandidates, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: PlanningModels<'_>, -) -> MemoSearchResult { - let started = std::time::Instant::now(); - let memo = super::search::memo(per_computation, entries); - let memo_stats = super::search::stats(&memo, started.elapsed().as_secs_f64() * 1e3); - let scale = super::pricing::frozen_calibration(&memo, entries, data, models)?; - let mut cache = super::pricing::PriceCache::new(scale); - let mut choices: Vec = Vec::new(); - let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); - let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { - unreachable!("the caller routes cost models 2 and 3 to the legacy search"); - }; - for g in 0..memo.groups.len() { - let (group_choices, lazy) = super::coefficients::choices( - &memo, - GroupId(g as u32), - entries, - data, - models, - &mut cache, - )?; - // A group's copies (apart) cost at least its shared choice with its - // node not built once: each copy is charged per consumer as the - // shared node is, its own scans read at least the shortest horizon, - // and its builds are sized for fewer readers but built per copy. So - // the model keeps the shared choices (a lower bound), and the exact - // pricing of a group choice prices both. - // Widened scans are reached through the identical-expression - // variants of the exact pricing, never chosen directly. - for (mut choice, lazy) in group_choices.into_iter().zip(lazy) { - if choice.copy_of.is_some() - || memo.candidates[g][choice.candidate].origin - == super::CandidateOrigin::WidenedScans - { - continue; - } - add_lazy_lower_bounds(&mut choice, &lazy, prices); - raw_scans.extend( - lazy.raw_scans - .iter() - .map(|(source, bytes, via)| (choices.len(), source.clone(), *bytes, *via)), - ); - choices.push(choice); - } - } - let trace = std::env::var_os("JP_MILP_TRACE").is_some(); - if trace { - eprintln!( - "milp: memo {} groups, {} choices in {:.1}s", - memo.groups.len(), - choices.len(), - started.elapsed().as_secs_f64() - ); - } - let mut input = model_input(&memo, choices, entries, models)?; - let mut sources: Vec = Vec::new(); - for (c, source, bytes, via) in raw_scans { - let index = match sources.iter().position(|s| *s == source) { - Some(i) => i, - None => { - sources.push(source); - sources.len() - 1 - } - }; - input.raw_scans.push(RawScan { - choice: c, - source: index, - cost: prices.cost_per_disk_byte_second * bytes as f64, - via, - }); - } - input.raw_data_retained = models.deployment_capabilities.raw_data_retained; - let mut selector = PlanSelector::new(entries, data, models, false)?; - let mut logical: HashMap> = HashMap::new(); - let mut number = 0usize; - let mut iterations = 0u32; - let mut bound = None; - let mut first_bound = None; - let mut solver_ms = 0.0; - let mut size = (0, 0); - loop { - iterations += 1; - let solve_started = std::time::Instant::now(); - let solved = solve(&input).map_err(SelectionError::Solver)?; - solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; - let Some(solution) = solved else { - break; - }; - bound = Some(solution.objective); - first_bound.get_or_insert(solution.objective); - if trace { - eprintln!("milp: round {iterations} bound {:.6e} incumbent {:?} solve {:.0}ms vars {} cons {} at {:.1}s", solution.objective, selector.cheapest_price(), solve_started.elapsed().as_secs_f64()*1e3, solution.variables, solution.constraints, started.elapsed().as_secs_f64()); - } - size = (solution.variables, solution.constraints); - if let Some(incumbent) = selector.cheapest_price() { - if solution.objective >= incumbent - 1e-9 * incumbent.abs() { - break; - } - } - // The solution's group choice: the candidate each chosen choice - // takes for its group. - let mut group_choice: super::search::GroupChoice = vec![None; memo.groups.len()]; - for &c in &solution.chosen { - group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); - } - let chosen_candidates: Vec<(GroupId, &super::ImplementationCandidate)> = group_choice - .iter() - .enumerate() - .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) - .collect(); - if !super::windows::requirements_met(&chosen_candidates) { - // A rollup's requirement: cut the smallest failing pairs of - // (group, candidate), else the whole group choice below. - let mut cut = false; - for (i, a) in chosen_candidates.iter().enumerate() { - if a.1.rollup_of.is_none() { - continue; - } - for (j, b) in chosen_candidates.iter().enumerate() { - if i != j && !super::windows::requirements_met(&[*a, *b]) { - let members: Vec = (0..input.choices.len()) - .filter(|&c| { - let choice = &input.choices[c]; - (choice.group == a.0 - && choice.candidate == group_choice[a.0 .0 as usize].unwrap()) - || (choice.group == b.0 - && choice.candidate - == group_choice[b.0 .0 as usize].unwrap()) - }) - .collect(); - input.excluded.push((members, Vec::new(), 1.0)); - cut = true; - } - } - if !super::windows::requirements_met(&[*a]) { - let members: Vec = (0..input.choices.len()) - .filter(|&c| { - input.choices[c].group == a.0 - && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] - }) - .collect(); - input.excluded.push((members, Vec::new(), 0.0)); - cut = true; - } - } - if trace { - eprintln!("milp: round {iterations} requirements not met (pair cut {cut})"); - } - if cut { - continue; - } - } - super::search::offer_group_choice( - &memo, - per_computation, - &group_choice, - entries, - data, - models, - &mut selector, - &mut logical, - &mut number, - )?; - // Exclude it: every choice of each chosen (group, candidate). - let groups = group_choice.iter().flatten().count(); - let excluded: Vec = input - .choices - .iter() - .enumerate() - .filter(|(_, choice)| group_choice[choice.group.0 as usize] == Some(choice.candidate)) - .map(|(c, _)| c) - .collect(); - let mut covered = vec![false; memo.groups.len()]; - for (g, c) in group_choice.iter().enumerate() { - if let Some(c) = c { - for i in &memo.candidates[g][*c].implements { - covered[i.0 as usize] = true; - } - } - } - let outside = (0..memo.groups.len()).filter(|&g| !covered[g]).collect(); - input - .excluded - .push((excluded, outside, groups as f64 - 1.0)); - } - if std::env::var_os("JP_MILP_CHECK").is_some() { - // Debugging aid: every group choice's exact price against the - // model's value restricted to it. - let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); - for group_choice in super::search::group_choices(&memo).unwrap_or_default() { - let mut one = PlanSelector::new(entries, data, models, false)?; - let mut scratch = HashMap::new(); - let mut k = 0usize; - super::search::offer_group_choice( - &memo, - per_computation, - &group_choice, - entries, - data, - models, - &mut one, - &mut scratch, - &mut k, - )?; - let Some(exact) = one.cheapest_price() else { - continue; - }; - let mut restricted = model_input_clone(&input); - restricted.excluded.clear(); - let outside: Vec = (0..restricted.choices.len()) - .filter(|&c| { - group_choice[restricted.choices[c].group.0 as usize] - != Some(restricted.choices[c].candidate) - }) - .collect(); - restricted.excluded.push((outside, Vec::new(), 0.0)); - let value = solve(&restricted) - .map_err(SelectionError::Solver)? - .map(|s| s.objective); - let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); - if bad || exact < incumbent * (1.0 - 1e-9) { - eprintln!( - "check: {group_choice:?} exact {exact:.7e} model {value:?}{}", - if bad { - " NOT A LOWER BOUND" - } else { - " BELOW INCUMBENT" - } - ); - } - } - } - let (mut selection, selected) = selector.finish()?; - let price = selection.costs[&selection.selected].cost_per_second; - selection.search.memo = Some(memo_stats); - selection.search.solver_iterations = iterations; - selection.search.solver_time_ms = solver_ms as u64; - ( - selection.search.solver_variables, - selection.search.solver_constraints, - ) = size; - selection.search.solver_first_bound_cost_per_second = first_bound; - selection.search.solver_gap = match (price, bound) { - (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), - (Some(_), Some(_)) => Some(0.0), - _ => None, - }; - let roots = logical - .remove(&selected.from_logical) - .expect("the selection's logical roots are kept"); - Ok((selection, selected, roots)) -} - -fn model_input_clone(input: &ModelInput) -> ModelInput { - ModelInput { - group_count: input.group_count, - group_queries: input.group_queries.clone(), - choices: input.choices.clone(), - group_slots: input.group_slots.clone(), - queries: input.queries.clone(), - schedule: input.schedule.clone(), - union_rate: input.union_rate.clone(), - excluded: input.excluded.clone(), - extra_bounds: input.extra_bounds.clone(), - requires: input.requires.clone(), - exclusive: input.exclusive.clone(), - raw_scans: input.raw_scans.clone(), - raw_data_retained: input.raw_data_retained, - } -} - -/// The model's input for `memo`'s `choices`. -fn model_input( - memo: &Memo, - choices: Vec, - entries: &[QueryWorkloadEntry], - models: PlanningModels<'_>, -) -> Result { - let settings = models.plan_selection; - let rate = |entries: &[&QueryWorkloadEntry]| { - crate::shared_evaluation_rate(entries, settings).map_err(|error| SelectionError::Cost { - candidate: "memo".to_string(), - node: Default::default(), - error, - }) - }; - // Schedule classes: queries with one recurrence and range evaluation. - let mut classes: Vec<&QueryWorkloadEntry> = Vec::new(); - let schedule: Vec = entries - .iter() - .map(|entry| { - match classes.iter().position(|c| { - c.recurrence == entry.recurrence && c.range_evaluation == entry.range_evaluation - }) { - Some(class) => class, - None => { - classes.push(entry); - classes.len() - 1 - } - } - }) - .collect(); - let mut union_rate = HashMap::new(); - let mut queries = Vec::with_capacity(entries.len()); - for (q, entry) in entries.iter().enumerate() { - let roots = memo - .groups - .iter() - .enumerate() - .filter(|(_, group)| group.roots_of.contains(&q)) - .map(|(g, _)| GroupId(g as u32)) - .collect(); - // A query's evaluation must fit its latency bound and its repeat - // interval. - let interval = entry.recurrence.fixed_interval_ms().map(|ms| ms as f64); - let max_latency_ms = match (entry.requirements.max_latency_ms, interval) { - (Some(a), Some(b)) => Some(a.min(b)), - (a, b) => a.or(b), - }; - queries.push(ModelQuery { - rate: rate(&[entry])?, - roots, - max_latency_ms, - }); - } - let group_queries: Vec> = memo.groups.iter().map(|g| g.queries.clone()).collect(); - for queries_of in &group_queries { - let mut present: Vec = queries_of.iter().map(|&q| schedule[q]).collect(); - present.sort_unstable(); - present.dedup(); - for mask in 1u64..(1 << present.len()) { - let set: Vec = (0..present.len()) - .filter(|i| mask >> i & 1 == 1) - .map(|i| present[i]) - .collect(); - if let std::collections::hash_map::Entry::Vacant(slot) = union_rate.entry(set) { - let members: Vec<&QueryWorkloadEntry> = - slot.key().iter().map(|&c| classes[c]).collect(); - slot.insert(rate(&members)?); - } - } - } - // Requirements the reference search checks on a group choice - // ([`super::windows::requirements_met`]) other than rollups, which the - // loop cuts lazily. - let candidate_of = |c: &Choice| &memo.candidates[c.group.0 as usize][c.candidate]; - let mut requires = Vec::new(); - for (c, choice) in choices.iter().enumerate() { - let candidate = candidate_of(choice); - for &g in &candidate.requires_exact_children { - let allowed = (0..choices.len()) - .filter(|&a| { - choices[a].group == g && super::windows::is_exact(candidate_of(&choices[a])) - }) - .collect(); - requires.push((c, allowed)); - } - for (g, site) in &candidate.requires_sites { - let allowed = (0..choices.len()) - .filter(|&a| { - choices[a].group == *g - && candidate_of(&choices[a]).coupling_sites.contains(site) - }) - .collect(); - requires.push((c, allowed)); - } - } - // A reader of a scan shared at horizon `h` reads the scan group's - // choice at that horizon, so every reader of one scan agrees on it. - for (c, choice) in choices.iter().enumerate() { - for &(scan, horizon) in &choice.scan_children { - let Some(h) = horizon else { continue }; - let allowed = (0..choices.len()) - .filter(|&a| choices[a].group == scan && choices[a].scan_horizon_ms == Some(h)) - .collect(); - requires.push((c, allowed)); - } - } - let of_origin = |origin: super::CandidateOrigin| -> Vec { - (0..choices.len()) - .filter(|&c| candidate_of(&choices[c]).origin == origin) - .collect() - }; - let exclusive = ( - of_origin(super::CandidateOrigin::Panes), - of_origin(super::CandidateOrigin::TopKReuse), - ); - let group_slots = memo - .groups - .iter() - .enumerate() - .filter(|(_, group)| group.parents.len() > 1 || group.roots_of.len() > 1) - .map(|(g, _)| { - let allowed = &memo.candidates[g][0].materializations; - GroupSlot { - group: GroupId(g as u32), - may_recompute: allowed.is_empty() - || allowed.contains(&Materialization::NotMaterialized), - } - }) - .collect(); - Ok(ModelInput { - group_count: memo.groups.len(), - group_queries, - choices, - group_slots, - queries, - schedule, - union_rate, - excluded: Vec::new(), - extra_bounds: Vec::new(), - requires, - exclusive, - raw_scans: Vec::new(), - raw_data_retained: true, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn choice(group: u32, children: &[u32], cost: f64) -> Choice { - Choice { - group: GroupId(group), - candidate: 0, - implements: vec![GroupId(group)], - children: children.iter().map(|&c| GroupId(c)).collect(), - assignment: Vec::new(), - coefficients: ChoiceCoefficients { - direct: PartCoefficients { - cost_per_unit_rate: cost, - work_ms: 1.0, - }, - slots: Vec::new(), - child_via: children.iter().map(|_| Via::Direct).collect(), - shared_nodes: Vec::new(), - }, - copy_of: None, - scan_children: Vec::new(), - scan_horizon_ms: None, - } - } - - /// A scan two queries read costs once per query recomputing it, or once - /// per evaluation when built once: the model builds it once when that - /// is cheaper, and its objective is the plan's price. - #[test] - fn a_shared_group_is_built_once_when_that_is_cheaper() { - let input = ModelInput { - group_count: 3, - group_queries: vec![vec![0, 1], vec![0], vec![1]], - choices: vec![ - choice(0, &[], 10.0), - choice(1, &[0], 1.0), - choice(2, &[0], 1.0), - ], - group_slots: vec![GroupSlot { - group: GroupId(0), - may_recompute: true, - }], - queries: vec![ - ModelQuery { - rate: 1.0, - roots: vec![GroupId(1)], - max_latency_ms: None, - }, - ModelQuery { - rate: 1.0, - roots: vec![GroupId(2)], - max_latency_ms: None, - }, - ], - schedule: vec![0, 0], - union_rate: [(vec![0], 1.0)].into(), - excluded: Vec::new(), - extra_bounds: Vec::new(), - requires: Vec::new(), - exclusive: (Vec::new(), Vec::new()), - raw_scans: Vec::new(), - raw_data_retained: true, - }; - let solution = solve(&input).unwrap().unwrap(); - assert_eq!(solution.built_once, [GroupId(0)].into()); - assert!( - (solution.objective - 12.0).abs() < 1e-9, - "{}", - solution.objective - ); - // Recomputing is cheaper when the scan is cheap. - let mut cheap = input; - cheap.choices[0].coefficients.direct.cost_per_unit_rate = 0.1; - cheap.union_rate.insert(vec![0], 1.0); - // Built once still costs its union rate × 0.1 = 0.1 < 0.2. - let solution = solve(&cheap).unwrap().unwrap(); - assert!( - (solution.objective - 2.1).abs() < 1e-9, - "{}", - solution.objective - ); - } -} From f02f24dd96f0c6d00d7e03259a5b5b487dfaaf9e Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:02:53 +0000 Subject: [PATCH 49/67] Memo MILP: no absolute MIP gap; a solve stopped within the solver's gap tolerance counts as optimal with its bound lowered by that tolerance (was reported as a time-limit stop) --- .../plan-selection/src/pass/subexpr/milp.rs | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 15c274853..b9aa56cfe 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -274,6 +274,10 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> })) } +/// How far below its objective a solve stopped at the solver's gap +/// tolerance (a relative gap of about 1e-6 at most) may have its optimum. +const GAP_TOLERANCE: f64 = 1e-4; + /// The factor that brings `input`'s largest charge to about 1. fn price_scale(input: &ModelInput) -> f64 { let fastest = input.queries.iter().map(|q| q.rate).fold(1.0, f64::max); @@ -1069,6 +1073,8 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { if std::env::var_os("JP_MILP_VERBOSE").is_some() { model.set_verbose(true); } + // No absolute gap either: prices are scaled to about 1. + model = model.set_option("mip_abs_gap", 0.0); if let Some(limit) = input.time_limit_s { model = model.set_time_limit(limit.max(0.001)); } @@ -1155,8 +1161,19 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { Ok(Some(ModelSolution { chosen, built_once, - objective: solution.eval(&objective), - optimal: matches!(solution.status(), good_lp::solvers::SolutionStatus::Optimal), + // Stopped within the solver's own tolerance of a zero gap, the + // objective is lowered by more than that tolerance to stay a bound. + objective: match solution.status() { + good_lp::solvers::SolutionStatus::GapLimit => { + let value = solution.eval(&objective); + value - value.abs() * GAP_TOLERANCE + } + _ => solution.eval(&objective), + }, + optimal: matches!( + solution.status(), + good_lp::solvers::SolutionStatus::Optimal | good_lp::solvers::SolutionStatus::GapLimit + ), variables: variable_count, constraints: constraint_count, })) From 34bb909d5d095df3076e229926467e7b280f3dc9 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:25:19 +0000 Subject: [PATCH 50/67] Memo MILP: raw windows computed once per memo; a time-limited solve's plan is still priced and the best plan returned with its gap (relaxation bound when no round finished); infeasibility names the queries with no plan --- .../src/pass/subexpr/coefficients.rs | 9 ++- .../plan-selection/src/pass/subexpr/milp.rs | 64 ++++++++++++++++--- .../src/pass/subexpr/pricing.rs | 3 + 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 6aadab79b..ad3481e18 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -145,6 +145,11 @@ pub(super) fn choice_parts( ) -> Result, SelectionError> { let node = &memo.groups[group.0 as usize].node; let queries = &memo.groups[group.0 as usize].queries; + let windows_of_raw = Rc::clone( + cache + .raw_windows + .get_or_insert_with(|| Rc::new(raw_windows(memo))), + ); let is_scan = matches!(node.non_asap(), Some(NonASAPOp::Scan { .. })); // Shared (one node for every query), and for a group several queries // read, apart (a copy per query); a scan apart is its reader's. @@ -155,7 +160,7 @@ pub(super) fn choice_parts( // A shared scan reads the longest window of its readers: one choice // per horizon a reader may read. let own_horizons: Vec> = match node.non_asap() { - Some(NonASAPOp::Scan { source, .. }) => raw_windows(memo) + Some(NonASAPOp::Scan { source, .. }) => windows_of_raw .get(&super::scans::raw_horizon_key(source)) .map(|windows| windows.iter().map(|&w| Some(w)).collect()) .unwrap_or_else(|| vec![None]), @@ -211,7 +216,7 @@ pub(super) fn choice_parts( let key = super::scans::raw_horizon_key(source); let mut variants = vec![None]; variants.extend( - raw_windows(memo) + windows_of_raw .get(&key) .into_iter() .flatten() diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index b9aa56cfe..450b6f666 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -177,6 +177,9 @@ pub(crate) struct ModelInput { pub cutoff: Option, /// The solver's time limit in seconds, if any. pub time_limit_s: Option, + /// Solve the linear relaxation (every binary in [0, 1]): a lower bound + /// on the model's optimum. + pub relaxed: bool, } /// The memory a choice's parts hold while an evaluation runs, priced per @@ -333,7 +336,7 @@ fn scaled(input: &ModelInput, scale: f64) -> ModelInput { fn solve_scaled(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let relaxed = std::env::var_os("JP_MILP_LP").is_some(); + let relaxed = input.relaxed || std::env::var_os("JP_MILP_LP").is_some(); let binary = || { if relaxed { variable().min(0).max(1) @@ -1349,6 +1352,9 @@ pub(crate) fn search( let mut size = (0, 0); let mut time_limit_reached = false; loop { + if time_limit_reached { + break; + } let remaining_s = time_limit_ms.map(|limit| limit as f64 / 1e3 - started.elapsed().as_secs_f64()); if remaining_s.is_some_and(|s| s <= 0.0) { @@ -1363,11 +1369,11 @@ pub(crate) fn search( .map(|incumbent| incumbent - 1e-9 * incumbent.abs()); let solved = solve(&input).map_err(SelectionError::Solver)?; solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; - if solved.as_ref().is_some_and(|s| !s.optimal) { - // The solver stopped at the time limit: its plan is no bound. - time_limit_reached = true; - break; - } + // The solver stopped at the time limit: its plan is no bound, but + // it is still priced, as a last round, and may become the + // incumbent. + let last_round = solved.as_ref().is_some_and(|s| !s.optimal); + time_limit_reached |= last_round; let Some(solution) = solved else { // No plan left below the incumbent: it is optimal. if let Some(incumbent) = selector.cheapest_price() { @@ -1375,14 +1381,16 @@ pub(crate) fn search( } break; }; - bound = Some(solution.objective); - first_bound.get_or_insert(solution.objective); + if !last_round { + bound = Some(solution.objective); + first_bound.get_or_insert(solution.objective); + } if trace { eprintln!("milp: round {iterations} bound {:.6e} incumbent {:?} solve {:.0}ms vars {} cons {} at {:.1}s", solution.objective, selector.cheapest_price(), solve_started.elapsed().as_secs_f64()*1e3, solution.variables, solution.constraints, started.elapsed().as_secs_f64()); } size = (solution.variables, solution.constraints); if let Some(incumbent) = selector.cheapest_price() { - if solution.objective >= incumbent - 1e-9 * incumbent.abs() { + if !last_round && solution.objective >= incumbent - 1e-9 * incumbent.abs() { break; } } @@ -1536,6 +1544,41 @@ pub(crate) fn search( } } } + if time_limit_reached && bound.is_none() { + // No round solved to optimality: the relaxation's optimum (all cuts + // kept) is still a lower bound for the gap. + input.relaxed = true; + input.time_limit_s = None; + if let Some(relaxation) = solve(&input).map_err(SelectionError::Solver)? { + bound = Some(relaxation.objective); + } + } + if !time_limit_reached && iterations == 1 && selector.cheapest_price().is_none() { + // Infeasible from the first round: name the queries with no plan + // of their own. + let mut without_plan = Vec::new(); + for (q, entry) in entries.iter().enumerate().take(input.queries.len()) { + let mut alone = model_input_clone(&input); + for (r, query) in alone.queries.iter_mut().enumerate() { + if r != q { + query.roots.clear(); + } + } + if solve(&alone).map_err(SelectionError::Solver)?.is_none() { + without_plan.push(format!("query {q} ({})", entry.query.0)); + } + } + if !without_plan.is_empty() { + return Err(SelectionError::Solver(format!( + "no valid plan for {}{}", + without_plan.join(", "), + match models.deployment_capabilities.raw_data_retained { + true => "", + false => " in a deployment that keeps no raw data", + } + ))); + } + } if time_limit_reached && selector.cheapest_price().is_none() { return Err(SelectionError::Solver(format!( "the MILP search reached its time limit after {iterations} rounds before finding a valid plan" @@ -1581,6 +1624,7 @@ fn model_input_clone(input: &ModelInput) -> ModelInput { held: input.held.clone(), cutoff: input.cutoff, time_limit_s: input.time_limit_s, + relaxed: input.relaxed, rollup_rules: input.rollup_rules.clone(), } } @@ -1754,6 +1798,7 @@ fn model_input( held: Vec::new(), cutoff: None, time_limit_s: None, + relaxed: false, }) } @@ -1825,6 +1870,7 @@ mod tests { held: Vec::new(), cutoff: None, time_limit_s: None, + relaxed: false, }; let solution = solve(&input).unwrap().unwrap(); assert_eq!(solution.built_once, [GroupId(0)].into()); diff --git a/crates/plan-selection/src/pass/subexpr/pricing.rs b/crates/plan-selection/src/pass/subexpr/pricing.rs index f6063021c..e682dd714 100644 --- a/crates/plan-selection/src/pass/subexpr/pricing.rs +++ b/crates/plan-selection/src/pass/subexpr/pricing.rs @@ -137,6 +137,9 @@ pub(crate) struct PriceCache { /// Local prices by (candidate's structural key, consumers, /// materialization, children's outputs). prices: HashMap>, + /// The raw windows of each scan horizon key of the memo + /// (`coefficients::raw_windows`), computed once per memo. + pub(super) raw_windows: Option>>>, /// Local prices found in the cache. pub hits: u64, /// Local prices computed. From 700a4d042f3ed7cefdf1c49b5ec2348ad4b14400 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:02:30 +0000 Subject: [PATCH 51/67] Memo MILP report: incumbent_cost, model_bound, gap, rounds, model_variables/constraints, first_model_bound, time_limit_reached --- crates/devtools/src/bin/stage_pipeline.rs | 2 +- crates/devtools/tests/stage_pipeline.rs | 11 ++---- .../tests/joint_search_equivalence.rs | 2 +- crates/plan-selection/src/lib.rs | 39 ++++++++++++------- .../src/pass/stage1_to_stage3.rs | 4 +- .../plan-selection/src/pass/subexpr/milp.rs | 22 ++++++----- 6 files changed, 45 insertions(+), 35 deletions(-) diff --git a/crates/devtools/src/bin/stage_pipeline.rs b/crates/devtools/src/bin/stage_pipeline.rs index 613e233b7..232896d57 100644 --- a/crates/devtools/src/bin/stage_pipeline.rs +++ b/crates/devtools/src/bin/stage_pipeline.rs @@ -41,7 +41,7 @@ // when the legacy search ran instead, why (`routed_to_legacy`). Not with // `--exhaustive`. `--milp-time-limit-s ` limits the MILP search: at // the limit it selects its best plan so far and reports -// `solver_time_limit_reached` with the gap to its bound. +// `time_limit_reached` with the gap to its bound. // Unlike `plan_installable`'s default settings, Stage 3 ignores accuracy // guarantees unless `--enforce-accuracy-guarantees` is given (below); // `accuracy_guarantees` is `"ignored"` or `"enforced"`. diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index 66e169b8b..fec162903 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1733,8 +1733,8 @@ fn search_flag_picks_the_search_algorithm() { for field in [ "algorithm", "routed_to_legacy", - "solver_gap", - "solver_iterations", + "gap", + "rounds", "solver_time_ms", ] { assert!(report.get(field).is_none(), "{field}: {report}"); @@ -1751,11 +1751,8 @@ fn search_flag_picks_the_search_algorithm() { let report = &milp["stage3_plan_selection"]["search"]; assert_eq!(report["algorithm"], "memo-milp"); assert!(report.get("routed_to_legacy").is_none(), "{report}"); - assert!( - report["solver_iterations"].as_u64().unwrap() >= 1, - "{report}" - ); - assert!(report["solver_gap"].as_f64().unwrap() <= 1e-9, "{report}"); + assert!(report["rounds"].as_u64().unwrap() >= 1, "{report}"); + assert!(report["gap"].as_f64().unwrap() <= 1e-9, "{report}"); assert!((price(&milp) - price(&legacy)).abs() <= 1e-9 * price(&legacy)); let cost_model_2 = run(&[ diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index d0fbd8778..a18739396 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1481,7 +1481,7 @@ fn three_way( "{label}: exhaustive {exhaustive:?}, memo MILP {milp:?}" )); } - let gap = report.solver_gap.expect("the MILP reports its gap"); + let gap = report.gap.expect("the MILP reports its gap"); if gap.abs() > 1e-9 { return Err(format!("{label}: MILP gap {gap}")); } diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index a101ed6d1..6506d0829 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1035,16 +1035,26 @@ pub struct SearchReport { /// `None` when the requested algorithm ran. #[serde(default, skip_serializing_if = "Option::is_none")] pub routed_to_legacy: Option, - /// The MILP solver's relative gap between its best plan and its bound - /// at the end ([`SearchAlgorithm::MemoMilp`]); 0 when the plan is - /// proven optimal. `None` for every other search. + /// The MILP search's incumbent cost: the Stage 3 cost per second of + /// the best plan its rounds decoded and evaluated + /// ([`SearchAlgorithm::MemoMilp`]); `None` for every other search. #[serde(default, skip_serializing_if = "Option::is_none")] - pub solver_gap: Option, - /// The MILP solves of the lazy loop (one per round: a solve, the exact - /// price of its plan, a cut); 0 (omitted in JSON) for every other + pub incumbent_cost: Option, + /// The MILP search's model bound at the end: a lower bound on the cost + /// per second of every plan it did not evaluate (the last solver + /// objective, or the incumbent cost once no plan below it is left); + /// `None` for every other search. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_bound: Option, + /// (incumbent cost − model bound) / incumbent cost, at least 0: 0 when + /// the incumbent is proven optimal; `None` for every other search. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub gap: Option, + /// The MILP search's rounds (each: a solve, the decoded plan's + /// incumbent evaluation, a cut); 0 (omitted in JSON) for every other /// search. #[serde(default, skip_serializing_if = "is_zero")] - pub solver_iterations: u32, + pub rounds: u32, /// The wall time in milliseconds spent in the MILP solver, every round /// together; 0 (omitted in JSON) for every other search. #[serde(default, skip_serializing_if = "is_zero")] @@ -1052,19 +1062,18 @@ pub struct SearchReport { /// The MILP's variables and constraints in its last round; 0 (omitted /// in JSON) for every other search. #[serde(default, skip_serializing_if = "is_zero")] - pub solver_variables: u64, + pub model_variables: u64, #[serde(default, skip_serializing_if = "is_zero")] - pub solver_constraints: u64, - /// The first round's lower bound on the cost per second, before any - /// cut: how tight the model is against the selected price; `None` for - /// every other search. + pub model_constraints: u64, + /// The first round's model bound, before any cut: how tight the model + /// is against the incumbent cost; `None` for every other search. #[serde(default, skip_serializing_if = "Option::is_none")] - pub solver_first_bound_cost_per_second: Option, + pub first_model_bound: Option, /// Whether the MILP search stopped at its time limit /// ([`crate::Stage1ToStage3Search::milp_time_limit_ms`]): its plan is - /// then proven optimal only if [`Self::solver_gap`] is 0. + /// then proven optimal only if [`Self::gap`] is 0. #[serde(default, skip_serializing_if = "std::ops::Not::not")] - pub solver_time_limit_reached: bool, + pub time_limit_reached: bool, } /// The size of a sub-expression memo (planner-layering design, "The diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index 7f2b39154..021f68050 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -167,8 +167,8 @@ pub struct Stage1ToStage3Search { pub search: SearchAlgorithm, /// The longest the MILP search may run, in ms; `None` for no limit. /// Stopped at the limit, it selects its best plan so far and reports - /// the gap to its bound ([`crate::SearchReport::solver_gap`]) and - /// [`crate::SearchReport::solver_time_limit_reached`]: not proven + /// the gap to its bound ([`crate::SearchReport::gap`]) and + /// [`crate::SearchReport::time_limit_reached`]: not proven /// optimal unless the gap is 0. pub milp_time_limit_ms: Option, } diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 450b6f666..b86091908 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -1587,15 +1587,19 @@ pub(crate) fn search( let (mut selection, selected) = selector.finish()?; let price = selection.costs[&selection.selected].cost_per_second; selection.search.memo = Some(memo_stats); - selection.search.solver_iterations = iterations; - selection.search.solver_time_ms = solver_ms as u64; - ( - selection.search.solver_variables, - selection.search.solver_constraints, - ) = size; - selection.search.solver_first_bound_cost_per_second = first_bound; - selection.search.solver_time_limit_reached = time_limit_reached; - selection.search.solver_gap = match (price, bound) { + let search = &mut selection.search; + search.rounds = iterations; + search.solver_time_ms = solver_ms as u64; + (search.model_variables, search.model_constraints) = size; + search.first_model_bound = first_bound; + search.time_limit_reached = time_limit_reached; + search.incumbent_cost = price; + // Once no plan below the incumbent is left, the bound is the incumbent. + search.model_bound = bound.map(|bound| match price { + Some(price) => bound.min(price), + None => bound, + }); + search.gap = match (price, search.model_bound) { (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), (Some(_), Some(_)) => Some(0.0), _ => None, From b73c99fea33cfda3f1c9bea5f8bee3fc9e7529cf Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:12:57 +0000 Subject: [PATCH 52/67] Selection report: incumbent_cost for every search; the MILP's lower bound, gap, rounds and effort in a nested MilpReport (search.milp) --- crates/devtools/tests/stage_pipeline.rs | 23 +++--- .../tests/joint_search_equivalence.rs | 6 +- crates/plan-selection/src/lib.rs | 76 +++++++++---------- .../src/pass/stage1_to_stage3.rs | 4 +- .../plan-selection/src/pass/subexpr/milp.rs | 29 +++---- 5 files changed, 72 insertions(+), 66 deletions(-) diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index fec162903..429c13d15 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1728,17 +1728,12 @@ fn search_flag_picks_the_search_algorithm() { let legacy = run(&["--search", "legacy"]); assert_eq!(legacy["search"], "default"); let report = &legacy["stage3_plan_selection"]["search"]; - // The legacy search's report omits the algorithm and the solver's - // fields, which only the MILP fills. - for field in [ - "algorithm", - "routed_to_legacy", - "gap", - "rounds", - "solver_time_ms", - ] { + // The legacy search's report omits the algorithm and the MILP's + // report; every search reports its selected plan's cost. + for field in ["algorithm", "routed_to_legacy", "milp"] { assert!(report.get(field).is_none(), "{field}: {report}"); } + assert_eq!(report["incumbent_cost"].as_f64(), Some(price(&legacy))); let brute_force = run(&["--search", "memo-brute-force"]); assert_eq!(brute_force["search"], "memo-brute-force"); @@ -1751,8 +1746,14 @@ fn search_flag_picks_the_search_algorithm() { let report = &milp["stage3_plan_selection"]["search"]; assert_eq!(report["algorithm"], "memo-milp"); assert!(report.get("routed_to_legacy").is_none(), "{report}"); - assert!(report["rounds"].as_u64().unwrap() >= 1, "{report}"); - assert!(report["gap"].as_f64().unwrap() <= 1e-9, "{report}"); + let milp_report = &report["milp"]; + assert!(milp_report["rounds"].as_u64().unwrap() >= 1, "{report}"); + assert!(milp_report["gap"].as_f64().unwrap() <= 1e-9, "{report}"); + assert!( + milp_report["lower_bound"].as_f64().unwrap() <= price(&milp) * (1.0 + 1e-9), + "{report}" + ); + assert_eq!(milp_report["time_limit_reached"], false, "{report}"); assert!((price(&milp) - price(&legacy)).abs() <= 1e-9 * price(&legacy)); let cost_model_2 = run(&[ diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index a18739396..8f0b07795 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1481,7 +1481,11 @@ fn three_way( "{label}: exhaustive {exhaustive:?}, memo MILP {milp:?}" )); } - let gap = report.gap.expect("the MILP reports its gap"); + let gap = report + .milp + .as_ref() + .and_then(|milp| milp.gap) + .expect("the MILP reports its gap"); if gap.abs() > 1e-9 { return Err(format!("{label}: MILP gap {gap}")); } diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 6506d0829..4a0bef45c 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1035,44 +1035,42 @@ pub struct SearchReport { /// `None` when the requested algorithm ran. #[serde(default, skip_serializing_if = "Option::is_none")] pub routed_to_legacy: Option, - /// The MILP search's incumbent cost: the Stage 3 cost per second of - /// the best plan its rounds decoded and evaluated - /// ([`SearchAlgorithm::MemoMilp`]); `None` for every other search. + /// The selected plan's cost per second (Stage 3's exact cost); `None` + /// when no plan was selected. #[serde(default, skip_serializing_if = "Option::is_none")] pub incumbent_cost: Option, - /// The MILP search's model bound at the end: a lower bound on the cost - /// per second of every plan it did not evaluate (the last solver - /// objective, or the incumbent cost once no plan below it is left); - /// `None` for every other search. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model_bound: Option, - /// (incumbent cost − model bound) / incumbent cost, at least 0: 0 when - /// the incumbent is proven optimal; `None` for every other search. + /// What the MILP search ([`SearchAlgorithm::MemoMilp`]) did; `None` + /// for every other search. #[serde(default, skip_serializing_if = "Option::is_none")] + pub milp: Option, +} + +/// The MILP search's report: its bound, gap and effort. +#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct MilpReport { + /// The MILP lower bound at the end: a lower bound on the cost per + /// second of every plan the search did not evaluate (the last solver + /// objective, or the incumbent cost once no plan below it is left); + /// `None` when no round finished and no relaxation was solved. + pub lower_bound: Option, + /// The first round's MILP lower bound, before any cut: how tight the + /// model is against the incumbent cost. + pub first_lower_bound: Option, + /// (incumbent cost − lower bound) / incumbent cost, at least 0: 0 when + /// the incumbent is proven optimal. pub gap: Option, - /// The MILP search's rounds (each: a solve, the decoded plan's - /// incumbent evaluation, a cut); 0 (omitted in JSON) for every other - /// search. - #[serde(default, skip_serializing_if = "is_zero")] + /// Its rounds (each: a solve, the decoded plan's incumbent evaluation, + /// a cut). pub rounds: u32, - /// The wall time in milliseconds spent in the MILP solver, every round - /// together; 0 (omitted in JSON) for every other search. - #[serde(default, skip_serializing_if = "is_zero")] + /// The MILP's variables and constraints in its last round. + pub variable_count: u64, + pub constraint_count: u64, + /// The wall time in milliseconds spent in the solver, every round + /// together. pub solver_time_ms: u64, - /// The MILP's variables and constraints in its last round; 0 (omitted - /// in JSON) for every other search. - #[serde(default, skip_serializing_if = "is_zero")] - pub model_variables: u64, - #[serde(default, skip_serializing_if = "is_zero")] - pub model_constraints: u64, - /// The first round's model bound, before any cut: how tight the model - /// is against the incumbent cost; `None` for every other search. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub first_model_bound: Option, - /// Whether the MILP search stopped at its time limit - /// ([`crate::Stage1ToStage3Search::milp_time_limit_ms`]): its plan is - /// then proven optimal only if [`Self::gap`] is 0. - #[serde(default, skip_serializing_if = "std::ops::Not::not")] + /// Whether it stopped at its time limit + /// ([`crate::Stage1ToStage3Search::milp_time_limit_ms`]): the plan is + /// then proven optimal only if `gap` is 0. pub time_limit_reached: bool, } @@ -1122,9 +1120,6 @@ impl SearchAlgorithm { } } -fn is_zero(value: &T) -> bool { - *value == T::default() -} /// Why the legacy search ran instead of the requested memo search. #[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] @@ -2112,6 +2107,7 @@ impl<'a> PlanSelector<'a> { &query_reach(&sized.dag, &node_map(&sized.dag)), &self.range_evaluations, ); + let incumbent_cost = self.costs.get(&selected).and_then(|c| c.cost_per_second); let selection = Selection { selected, plan, @@ -2123,24 +2119,26 @@ impl<'a> PlanSelector<'a> { search: SearchReport { candidates_checked: self.candidates_checked, runner_up_cost_per_second: self.runner_up, + incumbent_cost, ..SearchReport::default() }, }; Ok((selection, built)) } - /// Under a price ranking, fold the price of valid candidate `cost` - /// into the cheapest price and the runner-up ([`ROUNDING`]). - /// Under a price ranking, the cheapest price of the valid candidates - /// offered so far. + /// The rejections recorded so far. pub(crate) fn rejections(&self) -> &[Rejection] { &self.rejected } + /// Under a price ranking, the cheapest price of the valid candidates + /// offered so far. pub(crate) fn cheapest_price(&self) -> Option { self.cheapest } + /// Under a price ranking, fold the price of valid candidate `cost` + /// into the cheapest price and the runner-up ([`ROUNDING`]). fn track_runner_up(&mut self, cost: &CandidateCost) { let (Ranking::Price, Some(price)) = (self.ranking, cost.cost_per_second) else { return; diff --git a/crates/plan-selection/src/pass/stage1_to_stage3.rs b/crates/plan-selection/src/pass/stage1_to_stage3.rs index 021f68050..252a2df35 100644 --- a/crates/plan-selection/src/pass/stage1_to_stage3.rs +++ b/crates/plan-selection/src/pass/stage1_to_stage3.rs @@ -167,8 +167,8 @@ pub struct Stage1ToStage3Search { pub search: SearchAlgorithm, /// The longest the MILP search may run, in ms; `None` for no limit. /// Stopped at the limit, it selects its best plan so far and reports - /// the gap to its bound ([`crate::SearchReport::gap`]) and - /// [`crate::SearchReport::time_limit_reached`]: not proven + /// the gap to its bound ([`crate::MilpReport::gap`]) and + /// [`crate::MilpReport::time_limit_reached`]: not proven /// optimal unless the gap is 0. pub milp_time_limit_ms: Option, } diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index b86091908..7ec53b54f 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -1587,23 +1587,26 @@ pub(crate) fn search( let (mut selection, selected) = selector.finish()?; let price = selection.costs[&selection.selected].cost_per_second; selection.search.memo = Some(memo_stats); - let search = &mut selection.search; - search.rounds = iterations; - search.solver_time_ms = solver_ms as u64; - (search.model_variables, search.model_constraints) = size; - search.first_model_bound = first_bound; - search.time_limit_reached = time_limit_reached; - search.incumbent_cost = price; // Once no plan below the incumbent is left, the bound is the incumbent. - search.model_bound = bound.map(|bound| match price { + let lower_bound = bound.map(|bound| match price { Some(price) => bound.min(price), None => bound, }); - search.gap = match (price, search.model_bound) { - (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), - (Some(_), Some(_)) => Some(0.0), - _ => None, - }; + selection.search.incumbent_cost = price; + selection.search.milp = Some(crate::MilpReport { + lower_bound, + first_lower_bound: first_bound, + gap: match (price, lower_bound) { + (Some(price), Some(bound)) if price > 0.0 => Some(((price - bound) / price).max(0.0)), + (Some(_), Some(_)) => Some(0.0), + _ => None, + }, + rounds: iterations, + variable_count: size.0, + constraint_count: size.1, + solver_time_ms: solver_ms as u64, + time_limit_reached, + }); let roots = logical .remove(&selected.from_logical) .expect("the selection's logical roots are kept"); From 43afd7a3819123ab7cce81a624299126d005af94 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:49:44 +0000 Subject: [PATCH 53/67] plan-selection: memo MILP stays a lower bound with shared slots and widened scans Slots whose builds are equal over equal inputs may be one node (a site's build at one threshold, or identical expressions): each slot carries its build's structural key, and the model charges the dearest chosen one once. Scan choices of one source may be one widened scan: the model charges the dearest scan choice of each source. The model's input and the per-group-choice lower-bound check are functions (model_input_of, lower_bound_check), and a test asserts the model is a lower bound of every group choice on #840's and #839's shapes. --- .../src/pass/subexpr/coefficients.rs | 111 ++++++ .../plan-selection/src/pass/subexpr/milp.rs | 358 ++++++++++++------ 2 files changed, 352 insertions(+), 117 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index ad3481e18..dfa3daf67 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -345,6 +345,13 @@ pub(super) fn choice_parts( .map(|((i, _, _), v)| (child_groups[*i], *v)) .collect(), scan_horizon_ms: *own_horizon, + scan_source: match node.non_asap() { + Some(NonASAPOp::Scan { source, .. }) => { + Some(super::scans::raw_horizon_key(source)) + } + _ => None, + }, + scan_project_cost_per_unit_rate: parts.scan_project_cost, }, lazy: LazyTerms { thresholds: thresholds.clone(), @@ -557,6 +564,8 @@ fn is_build(payload: &Operator) -> bool { struct Coefficients { coefficients: ChoiceCoefficients, + /// For a scan: what a pass over its output costs per computation. + scan_project_cost: f64, lazy: LazyTerms, root_key: u32, root_producer_keys: Vec, @@ -839,6 +848,9 @@ fn coefficients_of( false => rate_cost, }, work_ms: work(&part), + // Equal builds over equal inputs are one node when shared (a + // site's build at one threshold, or identical expressions). + shared_site: Some((u64::from(keys[position[&node.id]]), 0)), }); slot_keys.push(shallow[position[&node.id]]); part_keys.push(part.iter().map(|id| keys[position[id]]).collect()); @@ -912,6 +924,12 @@ fn coefficients_of( built_once_state_bytes, thresholds: Vec::new(), }, + scan_project_cost: match dag.nodes[position[&root]].payload { + Operator::NonASAP(NonASAPOp::Scan { .. }) => { + prices.cost_per_cpu_op * priced[&root].output.rows as f64 + } + _ => 0.0, + }, root_key: shallow[position[&root]], shared_node_keys, root_producer_keys: producers_of @@ -1652,6 +1670,99 @@ mod tests { .is_empty()); } + /// The model restricted to each group choice ([`super::super::milp::lower_bound_check`]) + /// is a lower bound of the group choice's exact price. + fn assert_model_is_a_lower_bound( + queries: &[&str], + intervals_ms: &[u32], + targets: &[AccuracyTarget], + ) { + let roots: Vec<(usize, QueryRoot)> = queries + .iter() + .zip(targets) + .map(|(q, target)| { + let root = lower_promql(q, target.clone()); + QueryRoot::Operator(with_promql_series_identity(&root).unwrap()) + }) + .enumerate() + .collect(); + let per_computation = generate_per_computation_candidates(roots).unwrap(); + let entries = entries(intervals_ms, targets); + let data = data(); + let models = PlanningModels::builtin(); + let memo = memo(&per_computation, &entries); + let input = super::super::milp::model_input_of( + &memo, + &entries, + &data, + models, + std::time::Instant::now(), + ) + .unwrap(); + let checked = super::super::milp::lower_bound_check( + &memo, + &per_computation, + &input, + &entries, + &data, + models, + ) + .unwrap(); + assert!(!checked.is_empty()); + for (group_choice, exact, value) in checked { + let value = value.unwrap_or_else(|| panic!("{queries:?} {group_choice:?}: infeasible")); + assert!( + value <= exact * (1.0 + 1e-9), + "{queries:?} {group_choice:?}: model {value} over exact {exact}" + ); + } + } + + /// Three quantiles of one metric over a 1-minute, a 5-minute and an + /// instant window (#840's shape), every 10 s; and a shared group with + /// a rollup. + #[test] + fn the_model_is_a_lower_bound_of_every_group_choice() { + let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); + assert_model_is_a_lower_bound( + &[ + "quantile_over_time(0.5, m[1m])", + "quantile_over_time(0.9, m[5m])", + "quantile by (job) (0.5, m)", + ], + &[10_000, 10_000, 10_000], + &[target.clone(), target.clone(), target], + ); + let exact = AccuracyTarget::exact(); + assert_model_is_a_lower_bound( + &[ + "sum by (job) (rate(m[5m]))", + "sum(rate(m[5m]))", + "sum_over_time(m[10m])", + ], + &[60_000, 30_000, 60_000], + &[exact.clone(), exact.clone(), exact], + ); + } + + /// #839's shape: sums over 10 minutes by job and in total (a rollup), + /// and two 100-second quantiles over scans of one metric projecting + /// different labels (a widened scan shares them). + #[test] + fn the_model_is_a_lower_bound_with_rollups_and_widened_scans() { + let target = AccuracyTarget::max_error(0.05).failure_probability(0.3); + assert_model_is_a_lower_bound( + &[ + "sum by (job) (sum_over_time(m[10m]))", + "sum(sum_over_time(m[10m]))", + "avg by (job) (quantile_over_time(0.99, m[100s]))", + "avg(quantile_over_time(0.95, m[100s]))", + ], + &[20_000, 20_000, 1_000, 1_000], + &[target.clone(), target.clone(), target.clone(), target], + ); + } + /// Two quantiles within 1% and 2% sharing one KLL build at a size /// threshold. #[test] diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 7ec53b54f..6319dd3b1 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -57,6 +57,11 @@ pub(crate) struct SlotCoefficients { pub cost_per_unit_rate: f64, /// Work in ms per computation at query time. pub work_ms: f64, + /// The structural key of the slot's build over its inputs (and a + /// threshold index, 0): the slots of chosen choices with the same key + /// and materialization are one node when shared (a site's build at one + /// threshold, identical expressions), so the model charges them once. + pub shared_site: Option<(u64, usize)>, } /// One choice's coefficients. @@ -107,6 +112,18 @@ pub(crate) struct Choice { /// For a scan group's choice: the horizon it reads, the longest window /// among its readers. pub scan_horizon_ms: Option, + /// For a scan group's choice: its source's raw-horizon key. Scans of + /// one source may be one widened scan (the identical-expression rule + /// over scans widened to one label set), which costs at least the + /// dearest of them: the model charges the most expensive scan choice + /// of each source, not their sum. + pub scan_source: Option, + /// For a scan group's choice: what a pass over its output costs per + /// unit of rate, at most its own price. A widened scan shared by a + /// source's scan groups costs at least the dearest scan, and each other + /// reader at least such a pass: the model charges each scan choice's + /// pass, plus the largest excess over it. + pub scan_project_cost_per_unit_rate: f64, } /// A shared group's own node, built once per evaluation when chosen. @@ -318,6 +335,7 @@ fn scaled(input: &ModelInput, scale: f64) -> ModelInput { for node in &mut k.shared_nodes { node.cost_per_unit_rate *= scale; } + choice.scan_project_cost_per_unit_rate *= scale; } for scan in &mut scaled.raw_scans { scan.cost *= scale; @@ -486,13 +504,33 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { products_pending.push((z, a, b)); z }; + // Each scan choice's direct charge, per source (`Choice::scan_source`). + let mut scan_charges: BTreeMap> = BTreeMap::new(); for (&(k, c), &v) in &y { let total = input.choices[c].coefficients.direct.cost_per_unit_rate; - let shared = shared_cost[c].min(total); - let cost = total - shared; if total == 0.0 { continue; } + if let Some(source) = input.choices[c].scan_source { + let rate = match rate(k) { + Some(r) => r, + None => slowest_rate(match k { + Consumer::Group(g) => g.0 as usize, + Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, + Consumer::Query(_) => unreachable!("a query's rate is constant"), + }), + }; + let pass = input.choices[c].scan_project_cost_per_unit_rate.min(total); + objective += rate * pass * v; + *scan_charges + .entry(source) + .or_default() + .entry(c) + .or_insert_with(|| Expression::from(0.0)) += rate * (total - pass) * v; + continue; + } + let shared = shared_cost[c].min(total); + let cost = total - shared; match rate(k) { Some(r) => { objective += r * cost * v; @@ -524,6 +562,12 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } } } + let mut scan_maxima: Vec<(Variable, Vec)> = Vec::new(); + for charges in scan_charges.into_values() { + let m = vars.add(variable().min(0)); + objective += m; + scan_maxima.push((m, charges.into_values().collect())); + } for (c, choice) in input.choices.iter().enumerate() { if shared_cost[c] == 0.0 { continue; @@ -551,8 +595,40 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } cheaper_of.push((a, b, u, c, big_m)); } + // Slots that may be one shared node (equal builds over equal inputs): + // charged once, at the dearest of the chosen ones' charges with the + // slowest class's rate. One node shared by them costs at least each + // (it reads the longest of their windows), and apart they cost their + // sum. + let mut shared_slots: BTreeMap<(u64, usize, String), Vec<(usize, f64)>> = BTreeMap::new(); for (c, choice) in input.choices.iter().enumerate() { for slot in &choice.coefficients.slots { + let Some((site, threshold)) = slot.shared_site else { + continue; + }; + let at_query = !matches!(slot.materialization, Materialization::MatAtIngestion(_)); + let charge = slot.fixed_cost + + match at_query { + true => slowest_rate(choice.group.0 as usize) * slot.cost_per_unit_rate, + false => 0.0, + }; + shared_slots + .entry((site, threshold, format!("{:?}", slot.materialization))) + .or_default() + .push((c, charge)); + } + } + let mut shared_slot_vars: Vec<(Variable, Vec<(usize, f64)>)> = Vec::new(); + for members in shared_slots.into_values() { + let m = vars.add(variable().min(0)); + objective += m; + shared_slot_vars.push((m, members)); + } + for (c, choice) in input.choices.iter().enumerate() { + for slot in &choice.coefficients.slots { + if slot.shared_site.is_some() { + continue; + } objective += slot.fixed_cost * x[c]; if slot.cost_per_unit_rate != 0.0 && !matches!(slot.materialization, Materialization::MatAtIngestion(_)) @@ -976,6 +1052,16 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { )); } } + for (m, charges) in scan_maxima { + for charge in charges { + constraints.push(constraint!(m >= charge)); + } + } + for (m, members) in &shared_slot_vars { + for &(c, charge) in members { + constraints.push(constraint!(*m >= charge * x[c])); + } + } // Cuts. for (set, outside, bound) in &input.excluded { let mut chosen = Expression::from(0.0); @@ -1256,92 +1342,8 @@ pub(crate) fn search( let started = std::time::Instant::now(); let memo = super::search::memo(per_computation, entries); let memo_stats = super::search::stats(&memo, started.elapsed().as_secs_f64() * 1e3); - let scale = super::pricing::frozen_calibration(&memo, entries, data, models)?; - let mut cache = super::pricing::PriceCache::new(scale); - let mut choices: Vec = Vec::new(); - let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); - let mut held: Vec = Vec::new(); - let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { - unreachable!("the caller routes cost models 2 and 3 to the legacy search"); - }; - for g in 0..memo.groups.len() { - let (group_choices, lazy) = super::coefficients::choices( - &memo, - GroupId(g as u32), - entries, - data, - models, - &mut cache, - )?; - // A group's copies (apart) cost at least its shared choice with its - // node not built once: each copy is charged per consumer as the - // shared node is, its own scans read at least the shortest horizon, - // and its builds are sized for fewer readers but built per copy. So - // the model keeps the shared choices (a lower bound), and the exact - // pricing of a group choice prices both. - // Widened scans are reached through the identical-expression - // variants of the exact pricing, never chosen directly. - for (choice, lazy) in group_choices.into_iter().zip(lazy) { - if choice.copy_of.is_some() - || memo.candidates[g][choice.candidate].origin - == super::CandidateOrigin::WidenedScans - { - continue; - } - held.push(held_at_query_time(&choice, &lazy, prices)); - raw_scans.extend( - lazy.raw_scans - .iter() - .map(|(source, bytes, via)| (choices.len(), source.clone(), *bytes, *via)), - ); - choices.push(choice); - } - } + let mut input = model_input_of(&memo, entries, data, models, started)?; let trace = std::env::var_os("JP_MILP_TRACE").is_some(); - if trace { - for (c, choice) in choices.iter().enumerate() { - let k = &choice.coefficients; - let negative = k.direct.cost_per_unit_rate < 0.0 - || k.direct.work_ms < 0.0 - || k.slots - .iter() - .any(|s| s.fixed_cost < 0.0 || s.cost_per_unit_rate < 0.0 || s.work_ms < 0.0) - || k.shared_nodes.iter().any(|n| n.cost_per_unit_rate < 0.0); - if negative { - eprintln!( - "milp: NEGATIVE choice {c}: g{} cand{} {:?} assignment {:?}", - choice.group.0, choice.candidate, k, choice.assignment - ); - } - } - } - if trace { - eprintln!( - "milp: memo {} groups, {} choices in {:.1}s", - memo.groups.len(), - choices.len(), - started.elapsed().as_secs_f64() - ); - } - let mut input = model_input(&memo, choices, entries, models)?; - let mut sources: Vec = Vec::new(); - for (c, source, bytes, via) in raw_scans { - let index = match sources.iter().position(|s| *s == source) { - Some(i) => i, - None => { - sources.push(source); - sources.len() - 1 - } - }; - input.raw_scans.push(RawScan { - choice: c, - source: index, - cost: prices.cost_per_disk_byte_second * bytes as f64, - via, - }); - } - input.raw_data_retained = models.deployment_capabilities.raw_data_retained; - input.held = held; let mut selector = PlanSelector::new(entries, data, models, false)?; let mut logical: HashMap> = HashMap::new(); let mut number = 0usize; @@ -1501,36 +1503,9 @@ pub(crate) fn search( // Debugging aid: every group choice's exact price against the // model's value restricted to it. let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); - for group_choice in super::search::group_choices(&memo).unwrap_or_default() { - let mut one = PlanSelector::new(entries, data, models, false)?; - let mut scratch = HashMap::new(); - let mut k = 0usize; - super::search::offer_group_choice( - &memo, - per_computation, - &group_choice, - entries, - data, - models, - &mut one, - &mut scratch, - &mut k, - )?; - let Some(exact) = one.cheapest_price() else { - continue; - }; - let mut restricted = model_input_clone(&input); - restricted.excluded.clear(); - let outside: Vec = (0..restricted.choices.len()) - .filter(|&c| { - group_choice[restricted.choices[c].group.0 as usize] - != Some(restricted.choices[c].candidate) - }) - .collect(); - restricted.excluded.push((outside, Vec::new(), 0.0)); - let value = solve(&restricted) - .map_err(SelectionError::Solver)? - .map(|s| s.objective); + for (group_choice, exact, value) in + lower_bound_check(&memo, per_computation, &input, entries, data, models)? + { let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); if bad || exact < incumbent * (1.0 - 1e-9) { eprintln!( @@ -1613,6 +1588,153 @@ pub(crate) fn search( Ok((selection, selected, roots)) } +/// The model's input for `memo`: each group's shared choices (copies and +/// widened scans are left to the exact pricing) with their lazy lower +/// bounds, and the raw scans. +pub(super) fn model_input_of( + memo: &Memo, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + started: std::time::Instant, +) -> Result { + let scale = super::pricing::frozen_calibration(memo, entries, data, models)?; + let mut cache = super::pricing::PriceCache::new(scale); + let mut choices: Vec = Vec::new(); + let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); + let mut held: Vec = Vec::new(); + let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { + unreachable!("the caller routes cost models 2 and 3 to the legacy search"); + }; + for g in 0..memo.groups.len() { + let (group_choices, lazy) = super::coefficients::choices( + memo, + GroupId(g as u32), + entries, + data, + models, + &mut cache, + )?; + // A group's copies (apart) cost at least its shared choice with its + // node not built once: each copy is charged per consumer as the + // shared node is, its own scans read at least the shortest horizon, + // and its builds are sized for fewer readers but built per copy. So + // the model keeps the shared choices (a lower bound), and the exact + // pricing of a group choice prices both. + // Widened scans are reached through the identical-expression + // variants of the exact pricing, never chosen directly. + for (choice, lazy) in group_choices.into_iter().zip(lazy) { + if choice.copy_of.is_some() + || memo.candidates[g][choice.candidate].origin + == super::CandidateOrigin::WidenedScans + { + continue; + } + held.push(held_at_query_time(&choice, &lazy, prices)); + raw_scans.extend( + lazy.raw_scans + .iter() + .map(|(source, bytes, via)| (choices.len(), source.clone(), *bytes, *via)), + ); + choices.push(choice); + } + } + let trace = std::env::var_os("JP_MILP_TRACE").is_some(); + if trace { + for (c, choice) in choices.iter().enumerate() { + let k = &choice.coefficients; + let negative = k.direct.cost_per_unit_rate < 0.0 + || k.direct.work_ms < 0.0 + || k.slots + .iter() + .any(|s| s.fixed_cost < 0.0 || s.cost_per_unit_rate < 0.0 || s.work_ms < 0.0) + || k.shared_nodes.iter().any(|n| n.cost_per_unit_rate < 0.0); + if negative { + eprintln!( + "milp: NEGATIVE choice {c}: g{} cand{} {:?} assignment {:?}", + choice.group.0, choice.candidate, k, choice.assignment + ); + } + } + } + if trace { + eprintln!( + "milp: memo {} groups, {} choices in {:.1}s", + memo.groups.len(), + choices.len(), + started.elapsed().as_secs_f64() + ); + } + let mut input = model_input(memo, choices, entries, models)?; + let mut sources: Vec = Vec::new(); + for (c, source, bytes, via) in raw_scans { + let index = match sources.iter().position(|s| *s == source) { + Some(i) => i, + None => { + sources.push(source); + sources.len() - 1 + } + }; + input.raw_scans.push(RawScan { + choice: c, + source: index, + cost: prices.cost_per_disk_byte_second * bytes as f64, + via, + }); + } + input.raw_data_retained = models.deployment_capabilities.raw_data_retained; + input.held = held; + Ok(input) +} + +/// Every group choice of `memo` with its exact price (the cheapest plan +/// [`super::search::offer_group_choice`] finds) and the model's value +/// restricted to it, which should be a lower bound of the price; group +/// choices without a valid plan are left out. +pub(crate) fn lower_bound_check( + memo: &Memo, + per_computation: &PerComputationCandidates, + input: &ModelInput, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result)>, SelectionError> { + let mut out = Vec::new(); + for group_choice in super::search::group_choices(memo).unwrap_or_default() { + let mut one = PlanSelector::new(entries, data, models, false)?; + let mut scratch = HashMap::new(); + let mut k = 0usize; + super::search::offer_group_choice( + memo, + per_computation, + &group_choice, + entries, + data, + models, + &mut one, + &mut scratch, + &mut k, + )?; + let Some(exact) = one.cheapest_price() else { + continue; + }; + let mut restricted = model_input_clone(input); + restricted.excluded.clear(); + let outside: Vec = (0..restricted.choices.len()) + .filter(|&c| { + group_choice[restricted.choices[c].group.0 as usize] + != Some(restricted.choices[c].candidate) + }) + .collect(); + restricted.excluded.push((outside, Vec::new(), 0.0)); + let value = solve(&restricted) + .map_err(SelectionError::Solver)? + .map(|s| s.objective); + out.push((group_choice, exact, value)); + } + Ok(out) +} + fn model_input_clone(input: &ModelInput) -> ModelInput { ModelInput { group_count: input.group_count, @@ -1833,6 +1955,8 @@ mod tests { copy_of: None, scan_children: Vec::new(), scan_horizon_ms: None, + scan_source: None, + scan_project_cost_per_unit_rate: 0.0, } } From 7dc270e47da2bccbca4552ef9c3d2f9d5f7f67d5 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:10:11 +0000 Subject: [PATCH 54/67] plan-selection: memo MILP lower bound across groups and rollups Shared slots and shareable nodes several groups' choices may have are keyed loosely (kind, family, panes, raw sources) and charged once per key, at the dearest chosen one; their held memory is left out of the per-choice held term. A scan choice is charged a pass over its output per reader plus, once per source, the largest excess over it (a widened scan). rollup_rules reads the finer site the rollup itself requires. Tests: the model is a lower bound of every group choice on #840's shape, its sketch rollup, and #839's shape. --- .../src/pass/subexpr/coefficients.rs | 81 +++++++++++++++++- .../plan-selection/src/pass/subexpr/milp.rs | 82 +++++++++++++++++-- .../src/pass/subexpr/windows.rs | 12 ++- 3 files changed, 161 insertions(+), 14 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index dfa3daf67..0c76b3d9d 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -517,6 +517,43 @@ fn raw_windows(memo: &Memo) -> HashMap> { windows } +/// A key equal for builds that may be one node once shared: `node`'s +/// operator kind, state family and panes, and the raw sources below it. +fn build_key( + dag: &PhysicalASAPDAG, + node: &asap_types::ir::physical_export::PhysicalASAPDAGNode, + producers_of: &HashMap>, +) -> u64 { + use std::hash::{Hash, Hasher}; + let mut sources: BTreeSet = BTreeSet::new(); + let mut stack = vec![node.id]; + let mut seen = HashSet::new(); + while let Some(id) = stack.pop() { + if !seen.insert(id) { + continue; + } + if let Some(n) = dag.nodes.iter().find(|n| n.id == id) { + if let Operator::NonASAP(NonASAPOp::Scan { source, .. }) = &n.payload { + sources.insert(super::scans::raw_horizon_key(source)); + } + } + stack.extend(producers_of.get(&id).into_iter().flatten().copied()); + } + let shape = match &node.payload { + Operator::ASAP(ASAPOp::PaneChain { + family, + pane_width_ms, + pane_count, + .. + }) => format!("chain {family:?} {pane_width_ms} {pane_count}"), + Operator::ASAP(ASAPOp::SummaryAgg { family, .. }) => format!("build {family:?}"), + other => other.kind_name().to_string(), + }; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + (shape, sources).hash(&mut hasher); + hasher.finish() +} + /// The child groups of `candidate`: its own, and every other group whose /// node equals a node its node reaches and that no other child already /// reaches, which some Pass 2 candidates read without listing (a pane @@ -815,6 +852,7 @@ fn coefficients_of( cost_per_unit_rate: below.iter().map(|&n| per_rate(n)).sum(), work_ms: work(&below), state_bytes, + key: build_key(dag, &dag.nodes[position[&id]], &producers_of), }); shared_node_keys.push( below @@ -848,9 +886,11 @@ fn coefficients_of( false => rate_cost, }, work_ms: work(&part), - // Equal builds over equal inputs are one node when shared (a - // site's build at one threshold, or identical expressions). - shared_site: Some((u64::from(keys[position[&node.id]]), 0)), + // Builds that may be one node when shared (a site's build at + // one threshold, identical expressions, scans widened to one + // label set): the same kind, family and panes over the same + // sources. Keyed loosely (a wider key only lowers the bound). + shared_site: Some((build_key(dag, node, &producers_of), 0)), }); slot_keys.push(shallow[position[&node.id]]); part_keys.push(part.iter().map(|id| keys[position[id]]).collect()); @@ -916,6 +956,7 @@ fn coefficients_of( slots, child_via, shared_nodes, + cross_shared: Vec::new(), }, lazy: LazyTerms { held, @@ -1676,6 +1717,17 @@ mod tests { queries: &[&str], intervals_ms: &[u32], targets: &[AccuracyTarget], + ) { + assert_model_is_a_lower_bound_in(queries, intervals_ms, targets, true); + } + + /// [`assert_model_is_a_lower_bound`] in a deployment that keeps raw + /// data or not. + fn assert_model_is_a_lower_bound_in( + queries: &[&str], + intervals_ms: &[u32], + targets: &[AccuracyTarget], + raw_data_retained: bool, ) { let roots: Vec<(usize, QueryRoot)> = queries .iter() @@ -1689,7 +1741,13 @@ mod tests { let per_computation = generate_per_computation_candidates(roots).unwrap(); let entries = entries(intervals_ms, targets); let data = data(); - let models = PlanningModels::builtin(); + let deployment = + asap_logical_asap_optimization::deployment_capabilities::DeploymentCapabilities { + raw_data_retained, + query_time_results_retained: false, + ..Default::default() + }; + let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); let memo = memo(&per_computation, &entries); let input = super::super::milp::model_input_of( &memo, @@ -1763,6 +1821,21 @@ mod tests { ); } + /// #840's rollup: a quantile by host and job and its rollup by job, at + /// the rollup's state-size thresholds. + #[test] + fn the_model_is_a_lower_bound_with_a_sketch_rollup() { + let target = AccuracyTarget::max_error(0.01).failure_probability(0.001); + assert_model_is_a_lower_bound( + &[ + "quantile by (host, job) (0.5, m)", + "quantile by (job) (0.5, m)", + ], + &[10_000, 10_000], + &[target.clone(), target], + ); + } + /// Two quantiles within 1% and 2% sharing one KLL build at a size /// threshold. #[test] diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 6319dd3b1..ccea8bcdb 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -76,6 +76,11 @@ pub(crate) struct ChoiceCoefficients { /// than its own node: each may be one node for every copy of the group /// (the identical-expression rule), and then built once at query time. pub shared_nodes: Vec, + /// Shareable nodes another group's choices may also have (equal keys): + /// one node once shared, so charged once per key, at the dearest of + /// the chosen ones (key, price per unit of rate). Moved out of + /// `direct` and `shared_nodes` by the model's input. + pub cross_shared: Vec<(u64, f64)>, } /// A shareable node of a choice's direct part ([`ChoiceCoefficients::shared_nodes`]). @@ -88,6 +93,9 @@ pub(crate) struct SharedNodeCoefficients { pub cost_per_unit_rate: f64, pub work_ms: f64, pub state_bytes: u64, + /// A loose key of the node (kind, family, panes, raw sources): equal + /// for nodes of different groups that may be one node once shared. + pub key: u64, } /// One choice of a group: a candidate and an option per private slot. @@ -230,14 +238,17 @@ fn held_at_query_time( choice: &Choice, lazy: &LazyTerms, prices: &ResourceUnitCost, + shared_keys: &BTreeSet<(u64, usize)>, ) -> HeldAtQueryTime { let mut held = HeldAtQueryTime::default(); for term in &lazy.held { if let Via::Slot(i) = term.via { - if matches!( - choice.coefficients.slots[i].materialization, - Materialization::MatAtIngestion(_) - ) { + let slot = &choice.coefficients.slots[i]; + if matches!(slot.materialization, Materialization::MatAtIngestion(_)) + || slot + .shared_site + .is_some_and(|key| shared_keys.contains(&key)) + { continue; } } @@ -336,6 +347,9 @@ fn scaled(input: &ModelInput, scale: f64) -> ModelInput { node.cost_per_unit_rate *= scale; } choice.scan_project_cost_per_unit_rate *= scale; + for (_, cost) in &mut choice.coefficients.cross_shared { + *cost *= scale; + } } for scan in &mut scaled.raw_scans { scan.cost *= scale; @@ -618,6 +632,14 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { .push((c, charge)); } } + for (c, choice) in input.choices.iter().enumerate() { + for &(key, cost) in &choice.coefficients.cross_shared { + shared_slots + .entry((key, usize::MAX, "node".to_string())) + .or_default() + .push((c, slowest_rate(choice.group.0 as usize) * cost)); + } + } let mut shared_slot_vars: Vec<(Variable, Vec<(usize, f64)>)> = Vec::new(); for members in shared_slots.into_values() { let m = vars.add(variable().min(0)); @@ -1602,7 +1624,7 @@ pub(super) fn model_input_of( let mut cache = super::pricing::PriceCache::new(scale); let mut choices: Vec = Vec::new(); let mut raw_scans: Vec<(usize, String, u64, Via)> = Vec::new(); - let mut held: Vec = Vec::new(); + let mut lazies: Vec = Vec::new(); let CostModel::UsageOverTime(prices) = &models.plan_selection.cost_model else { unreachable!("the caller routes cost models 2 and 3 to the legacy search"); }; @@ -1630,7 +1652,7 @@ pub(super) fn model_input_of( { continue; } - held.push(held_at_query_time(&choice, &lazy, prices)); + lazies.push(lazy.clone()); raw_scans.extend( lazy.raw_scans .iter() @@ -1683,7 +1705,52 @@ pub(super) fn model_input_of( }); } input.raw_data_retained = models.deployment_capabilities.raw_data_retained; - input.held = held; + // Shareable nodes several groups' choices may have: one node once + // shared, charged once per key (`ChoiceCoefficients::cross_shared`). + let mut groups_of_node: HashMap> = HashMap::new(); + for choice in &input.choices { + for node in &choice.coefficients.shared_nodes { + groups_of_node + .entry(node.key) + .or_default() + .insert(choice.group); + } + } + for choice in &mut input.choices { + let k = &mut choice.coefficients; + let (cross, own): (Vec<_>, Vec<_>) = k + .shared_nodes + .drain(..) + .partition(|node| groups_of_node[&node.key].len() > 1); + k.shared_nodes = own; + for node in cross { + k.direct.cost_per_unit_rate -= node.cost_per_unit_rate; + k.cross_shared.push((node.key, node.cost_per_unit_rate)); + } + k.direct.cost_per_unit_rate = k.direct.cost_per_unit_rate.max(0.0); + } + // A slot whose build several groups' choices may share is one node + // charged once (`SlotCoefficients::shared_site`): its held memory is + // left out rather than charged per choice. + let mut groups_of_key: HashMap<(u64, usize), BTreeSet> = HashMap::new(); + for choice in &input.choices { + for slot in &choice.coefficients.slots { + if let Some(key) = slot.shared_site { + groups_of_key.entry(key).or_default().insert(choice.group); + } + } + } + let shared_keys: BTreeSet<(u64, usize)> = groups_of_key + .into_iter() + .filter(|(_, groups)| groups.len() > 1) + .map(|(key, _)| key) + .collect(); + input.held = input + .choices + .iter() + .zip(&lazies) + .map(|(choice, lazy)| held_at_query_time(choice, lazy, prices, &shared_keys)) + .collect(); Ok(input) } @@ -1951,6 +2018,7 @@ mod tests { slots: Vec::new(), child_via: children.iter().map(|_| Via::Direct).collect(), shared_nodes: Vec::new(), + cross_shared: Vec::new(), }, copy_of: None, scan_children: Vec::new(), diff --git a/crates/plan-selection/src/pass/subexpr/windows.rs b/crates/plan-selection/src/pass/subexpr/windows.rs index f1ddf6f8d..0a0a1b80d 100644 --- a/crates/plan-selection/src/pass/subexpr/windows.rs +++ b/crates/plan-selection/src/pass/subexpr/windows.rs @@ -1023,9 +1023,15 @@ pub(crate) fn rollup_rules(memo: &Memo, present: &BTreeSet<(usize, usize)>) -> V }; let (coarse_group, finer_group) = (rollup.0, read.finer_group.0 as usize); let coarse = Rc::as_ptr(&read.coarse_build); - let site = read.finer_state.as_ref().map(|state| CouplingSite { - key: state_key(state), - kind: CouplingKind::SharedBuildSize, + // The site the rollup requires of its finer group (its own + // `requires_sites`): the finer build's state as the finer group's + // candidates key it. + let site = read.finer_state.as_ref().and_then(|_| { + candidate(rollup) + .requires_sites + .iter() + .find(|(g, _)| g.0 as usize == finer_group) + .map(|(_, site)| *site) }); let r_node = (coarse_group, plan_node(rollup)); let mut allowed = Vec::new(); From 7381e32292823f7a37cd14aaf2c822eee7f80021 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:03:05 +0000 Subject: [PATCH 55/67] plan-selection: memo shared slots and nodes keyed exactly, up to scan widening A slot's build and a shareable node are keyed by their payload over their producers' keys, each scan by its source alone (the identical-expression rule may widen a source's scans to one label set), schemas left out. So the model charges the dearest of several groups' choices only for nodes that are one node once shared, and the per-choice sum otherwise. --- .../src/pass/subexpr/coefficients.rs | 76 +++++++++---------- 1 file changed, 35 insertions(+), 41 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 0c76b3d9d..174142915 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -291,6 +291,7 @@ pub(super) fn choice_parts( } let keys = keys_of(&physical.dag, cache.keys_mut()); let shallow = shallow_keys(&physical.dag, cache.keys_mut()); + let widened = widened_keys_of(&physical.dag, cache.keys_mut()); let private: Vec = keys.iter().map(|k| !children.contains(k)).collect(); let builds: Vec = (0..physical.dag.nodes.len()) @@ -319,6 +320,7 @@ pub(super) fn choice_parts( physical, &keys, &shallow, + &widened, &private, &child_roots, &horizons, @@ -517,41 +519,33 @@ fn raw_windows(memo: &Memo) -> HashMap> { windows } -/// A key equal for builds that may be one node once shared: `node`'s -/// operator kind, state family and panes, and the raw sources below it. -fn build_key( - dag: &PhysicalASAPDAG, - node: &asap_types::ir::physical_export::PhysicalASAPDAGNode, - producers_of: &HashMap>, -) -> u64 { - use std::hash::{Hash, Hasher}; - let mut sources: BTreeSet = BTreeSet::new(); - let mut stack = vec![node.id]; - let mut seen = HashSet::new(); - while let Some(id) = stack.pop() { - if !seen.insert(id) { - continue; - } - if let Some(n) = dag.nodes.iter().find(|n| n.id == id) { - if let Operator::NonASAP(NonASAPOp::Scan { source, .. }) = &n.payload { - sources.insert(super::scans::raw_horizon_key(source)); - } - } - stack.extend(producers_of.get(&id).into_iter().flatten().copied()); - } - let shape = match &node.payload { - Operator::ASAP(ASAPOp::PaneChain { - family, - pane_width_ms, - pane_count, - .. - }) => format!("chain {family:?} {pane_width_ms} {pane_count}"), - Operator::ASAP(ASAPOp::SummaryAgg { family, .. }) => format!("build {family:?}"), - other => other.kind_name().to_string(), - }; - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - (shape, sources).hash(&mut hasher); - hasher.finish() +/// Each node's key up to scan widening, by position: its payload over its +/// producers' keys, each scan by its source alone (the identical-expression +/// rule may widen a source's scans to one label set), schemas left out. +/// Equal keys of two groups' nodes mean the nodes are one node once shared. +fn widened_keys_of(dag: &PhysicalASAPDAG, interned: &mut HashMap) -> Vec { + let mut key_of: HashMap = HashMap::new(); + dag.nodes + .iter() + .map(|node| { + let text = match &node.payload { + Operator::NonASAP(NonASAPOp::Scan { source, .. }) => { + format!("widened scan {source:?}") + } + payload => { + let payload = payload.map_children(|id| key_of.get(id).copied()); + format!( + "widened {}", + serde_json::to_string(&payload).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() } /// The child groups of `candidate`: its own, and every other group whose @@ -703,6 +697,7 @@ fn coefficients_of( physical: &PhysicalCandidate, keys: &[u32], shallow: &[u32], + widened: &[u32], private: &[bool], child_roots: &[Option], horizons: &HashMap, @@ -852,7 +847,7 @@ fn coefficients_of( cost_per_unit_rate: below.iter().map(|&n| per_rate(n)).sum(), work_ms: work(&below), state_bytes, - key: build_key(dag, &dag.nodes[position[&id]], &producers_of), + key: u64::from(widened[position[&id]]), }); shared_node_keys.push( below @@ -886,11 +881,10 @@ fn coefficients_of( false => rate_cost, }, work_ms: work(&part), - // Builds that may be one node when shared (a site's build at - // one threshold, identical expressions, scans widened to one - // label set): the same kind, family and panes over the same - // sources. Keyed loosely (a wider key only lowers the bound). - shared_site: Some((build_key(dag, node, &producers_of), 0)), + // Builds that are one node when shared (a site's build at one + // threshold, identical expressions, scans widened to one label + // set): equal over equal inputs, scans' projections aside. + shared_site: Some((u64::from(widened[position[&node.id]]), 0)), }); slot_keys.push(shallow[position[&node.id]]); part_keys.push(part.iter().map(|id| keys[position[id]]).collect()); From 96ad7efd0713901f96f143f203ee38c1155c3618 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:47:31 +0000 Subject: [PATCH 56/67] plan-selection: memo shareable nodes leave scans to the scan choices A scan apart is its reader's own, and a shared one is its scan group's choice: counting absorbed scans as shareable nodes of several groups let the model charge only the dearest of every group's scan. --- crates/plan-selection/src/pass/subexpr/coefficients.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 174142915..1435ed297 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -812,7 +812,10 @@ fn coefficients_of( let node = &dag.nodes[position[&id]]; let shareable = matches!(node.payload, Operator::ASAP(_)) || !node.output_schema.unique_keys.is_empty(); - id != root && shareable && !is_build(&node.payload) && query_time(id) + // A scan apart is its reader's own; a shared one is its scan + // group's choice (`Choice::scan_children`). + let scan = matches!(node.payload, Operator::NonASAP(NonASAPOp::Scan { .. })); + id != root && shareable && !scan && !is_build(&node.payload) && query_time(id) }) .map(|id| { let below = walk(id) From fe387fcfd939185d3febd3fa10360215e273b33b Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:55:03 +0000 Subject: [PATCH 57/67] plan-selection: memo nodes equal up to scan widening charged once across groups Each query-time node of a choice's direct part (scans aside) carries its key up to scan widening; nodes several groups' choices have with one key are one node once shared (identical expressions over widened scans), and the model charges the dearest chosen one once. A choice's own scans are charged like scan choices: a pass per reader plus the largest excess once per source. A choice lists the children its plan reads first, with their child_via; children its candidate lists without reading (a rollup's own window) stay needed but no consumer passes into them. --- .../src/pass/subexpr/coefficients.rs | 57 +++++++++++- .../plan-selection/src/pass/subexpr/milp.rs | 91 +++++++++++++------ 2 files changed, 118 insertions(+), 30 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 1435ed297..00f029ba0 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -290,6 +290,15 @@ pub(super) fn choice_parts( } } let keys = keys_of(&physical.dag, cache.keys_mut()); + // The children its plan reads (a rollup reads its + // finer group's state, not its own window's group). + let (read_children, read_roots): (Vec, Vec>) = + kept_children + .iter() + .zip(&child_roots) + .filter(|(_, root)| root.is_some_and(|r| keys.contains(&r))) + .map(|(c, r)| (*c, *r)) + .unzip(); let shallow = shallow_keys(&physical.dag, cache.keys_mut()); let widened = widened_keys_of(&physical.dag, cache.keys_mut()); let private: Vec = @@ -322,7 +331,7 @@ pub(super) fn choice_parts( &shallow, &widened, &private, - &child_roots, + &read_roots, &horizons, own_horizon.or_else(|| scan_lookback_ms(memo, group)), &part_entries, @@ -337,7 +346,17 @@ pub(super) fn choice_parts( group, candidate: index, implements: candidate.implements.clone(), - children: kept_children.clone(), + // The children it reads first (`child_via` + // is theirs), then the ones its candidate + // lists without reading: needed by the + // plan, but no consumer passes into them. + children: read_children + .iter() + .chain( + kept_children.iter().filter(|c| !read_children.contains(c)), + ) + .copied() + .collect(), assignment, coefficients: parts.coefficients, copy_of, @@ -798,6 +817,18 @@ fn coefficients_of( cost_per_unit_rate: direct.iter().map(|&id| per_rate(id) + kept_state(id)).sum(), work_ms: work(&direct), }; + let absorbed_scans: Vec<(u64, f64, f64)> = direct + .iter() + .filter(|&&id| query_time(id)) + .filter_map(|&id| match &dag.nodes[position[&id]].payload { + Operator::NonASAP(NonASAPOp::Scan { source, .. }) => Some(( + super::scans::raw_horizon_key(source), + per_rate(id) + kept_state(id), + prices.cost_per_cpu_op * priced[&id].output.rows as f64, + )), + _ => None, + }) + .collect(); // Shareable nodes of the direct part (another copy's or group's // identical node may be one with it), each with what it computes when // built once: itself and the direct nodes only it reaches. @@ -850,7 +881,10 @@ fn coefficients_of( cost_per_unit_rate: below.iter().map(|&n| per_rate(n)).sum(), work_ms: work(&below), state_bytes, - key: u64::from(widened[position[&id]]), + members: below + .iter() + .map(|n| u64::from(widened[position[n]])) + .collect(), }); shared_node_keys.push( below @@ -954,6 +988,23 @@ fn coefficients_of( child_via, shared_nodes, cross_shared: Vec::new(), + absorbed_scans, + direct_nodes: direct + .iter() + .filter(|&&id| query_time(id)) + .filter(|&&id| { + !matches!( + dag.nodes[position[&id]].payload, + Operator::NonASAP(NonASAPOp::Scan { .. }) + ) + }) + .map(|&id| { + ( + u64::from(widened[position[&id]]), + per_rate(id) + kept_state(id), + ) + }) + .collect(), }, lazy: LazyTerms { held, diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index ccea8bcdb..7d9c41812 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -81,6 +81,19 @@ pub(crate) struct ChoiceCoefficients { /// the chosen ones (key, price per unit of rate). Moved out of /// `direct` and `shared_nodes` by the model's input. pub cross_shared: Vec<(u64, f64)>, + /// The query-time scans of the direct part (apart: this choice's own + /// copies), as (source key, price, a pass over its output) per unit of + /// rate. Scans of one source may be one widened scan: like the scan + /// choices (`Choice::scan_source`), each is charged its pass, and the + /// largest excess once per source. Moved out of `direct` by the + /// model's input. + pub absorbed_scans: Vec<(u64, f64, f64)>, + /// Each query-time node of the direct part but scans, as (key up to + /// scan widening, price per unit of rate). Nodes of several groups' + /// choices with one key are one node once shared (identical + /// expressions, scans widened to one label set): the model's input + /// moves them to `cross_shared`. + pub direct_nodes: Vec<(u64, f64)>, } /// A shareable node of a choice's direct part ([`ChoiceCoefficients::shared_nodes`]). @@ -88,14 +101,14 @@ pub(crate) struct ChoiceCoefficients { /// once at query time it costs `cost_per_unit_rate` per unit of the union /// rate of the queries reaching it instead, and holds `state_bytes` while /// its longest reader runs. -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct SharedNodeCoefficients { pub cost_per_unit_rate: f64, pub work_ms: f64, pub state_bytes: u64, - /// A loose key of the node (kind, family, panes, raw sources): equal - /// for nodes of different groups that may be one node once shared. - pub key: u64, + /// The widened keys of the nodes it covers ([`ChoiceCoefficients::direct_nodes`]), + /// itself first. + pub members: Vec, } /// One choice of a group: a candidate and an option per private slot. @@ -350,6 +363,10 @@ fn scaled(input: &ModelInput, scale: f64) -> ModelInput { for (_, cost) in &mut choice.coefficients.cross_shared { *cost *= scale; } + for (_, cost, pass) in &mut choice.coefficients.absorbed_scans { + *cost *= scale; + *pass *= scale; + } } for scan in &mut scaled.raw_scans { scan.cost *= scale; @@ -522,18 +539,28 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { let mut scan_charges: BTreeMap> = BTreeMap::new(); for (&(k, c), &v) in &y { let total = input.choices[c].coefficients.direct.cost_per_unit_rate; + let rate_or_slowest = match rate(k) { + Some(r) => r, + None => slowest_rate(match k { + Consumer::Group(g) => g.0 as usize, + Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, + Consumer::Query(_) => unreachable!("a query's rate is constant"), + }), + }; + for &(source, cost, pass) in &input.choices[c].coefficients.absorbed_scans { + let pass = pass.min(cost); + objective += rate_or_slowest * pass * v; + *scan_charges + .entry(source) + .or_default() + .entry(c) + .or_insert_with(|| Expression::from(0.0)) += rate_or_slowest * (cost - pass) * v; + } if total == 0.0 { continue; } if let Some(source) = input.choices[c].scan_source { - let rate = match rate(k) { - Some(r) => r, - None => slowest_rate(match k { - Consumer::Group(g) => g.0 as usize, - Consumer::Slot { choice, .. } => input.choices[choice].group.0 as usize, - Consumer::Query(_) => unreachable!("a query's rate is constant"), - }), - }; + let rate = rate_or_slowest; let pass = input.choices[c].scan_project_cost_per_unit_rate.min(total); objective += rate * pass * v; *scan_charges @@ -1705,28 +1732,36 @@ pub(super) fn model_input_of( }); } input.raw_data_retained = models.deployment_capabilities.raw_data_retained; - // Shareable nodes several groups' choices may have: one node once - // shared, charged once per key (`ChoiceCoefficients::cross_shared`). + // Nodes several groups' choices have, equal up to scan widening: one + // node once shared, charged once per key at the dearest chosen one + // (`ChoiceCoefficients::cross_shared`). let mut groups_of_node: HashMap> = HashMap::new(); for choice in &input.choices { - for node in &choice.coefficients.shared_nodes { - groups_of_node - .entry(node.key) - .or_default() - .insert(choice.group); + for &(key, _) in &choice.coefficients.direct_nodes { + groups_of_node.entry(key).or_default().insert(choice.group); } } for choice in &mut input.choices { let k = &mut choice.coefficients; - let (cross, own): (Vec<_>, Vec<_>) = k - .shared_nodes - .drain(..) - .partition(|node| groups_of_node[&node.key].len() > 1); - k.shared_nodes = own; - for node in cross { - k.direct.cost_per_unit_rate -= node.cost_per_unit_rate; - k.cross_shared.push((node.key, node.cost_per_unit_rate)); + for &(_, cost, _) in &k.absorbed_scans { + k.direct.cost_per_unit_rate -= cost; + } + let cross: Vec<(u64, f64)> = k + .direct_nodes + .iter() + .copied() + .filter(|(key, _)| groups_of_node[key].len() > 1) + .collect(); + for &(key, cost) in &cross { + k.direct.cost_per_unit_rate -= cost; + k.cross_shared.push((key, cost)); + for node in &mut k.shared_nodes { + if node.members.contains(&key) { + node.cost_per_unit_rate -= cost; + } + } } + k.shared_nodes.retain(|node| node.cost_per_unit_rate > 0.0); k.direct.cost_per_unit_rate = k.direct.cost_per_unit_rate.max(0.0); } // A slot whose build several groups' choices may share is one node @@ -2019,6 +2054,8 @@ mod tests { child_via: children.iter().map(|_| Via::Direct).collect(), shared_nodes: Vec::new(), cross_shared: Vec::new(), + absorbed_scans: Vec::new(), + direct_nodes: Vec::new(), }, copy_of: None, scan_children: Vec::new(), From 30aed58897c059e95d30df6e7434a95d1cd2821b Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:19:40 +0000 Subject: [PATCH 58/67] plan-selection: memo cross-group nodes at min(per consumer, once); slot consumers only of chosen choices A node several groups' choices have (equal up to scan widening) costs its charge per consumer, or built once at least the dearest chosen one's at its group's slowest rate: the model charges the smaller. A slot's consumer reaches and passes into nothing unless its choice is chosen. A shareable node's part leaves its scans to the scan charges. A debug test dumps a seeded random workload as stage_pipeline JSON. --- .../tests/joint_search_equivalence.rs | 11 ++++ .../src/pass/subexpr/coefficients.rs | 7 +++ .../plan-selection/src/pass/subexpr/milp.rs | 55 ++++++++++++++++++- 3 files changed, 71 insertions(+), 2 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 8f0b07795..d6adcbd76 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1807,3 +1807,14 @@ fn a_sketch_rollup_beside_other_quantiles_selects_the_exhaustive_price() { 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])); } + +/// Debugging aid: the random workload of seed `JP_SEED` (over the gauge +/// `m`) as stage_pipeline's `--workload` JSON, at `JP_SEED_OUT`. +#[test] +#[ignore = "debugging aid"] +fn dump_random_workload() { + let seed: u64 = std::env::var("JP_SEED").map_or(43, |s| s.parse().unwrap()); + let out = std::env::var("JP_SEED_OUT").unwrap_or_else(|_| format!("seed{seed}.json")); + let workload = random_workload(seed); + std::fs::write(&out, serde_json::to_string_pretty(&workload).unwrap()).unwrap(); +} diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 00f029ba0..a0c854866 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -849,10 +849,17 @@ fn coefficients_of( id != root && shareable && !scan && !is_build(&node.payload) && query_time(id) }) .map(|id| { + // Its scans are charged as scans (`absorbed_scans`). let below = walk(id) .0 .into_iter() .filter(|n| direct_set.contains(n)) + .filter(|n| { + !matches!( + dag.nodes[position[n]].payload, + Operator::NonASAP(NonASAPOp::Scan { .. }) + ) + }) .collect(); (id, below) }) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 7d9c41812..c364c3707 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -537,6 +537,8 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { }; // Each scan choice's direct charge, per source (`Choice::scan_source`). let mut scan_charges: BTreeMap> = BTreeMap::new(); + // Each cross-group node's charges per consumer, and their largest sum. + let mut cross_per_consumer: BTreeMap = BTreeMap::new(); for (&(k, c), &v) in &y { let total = input.choices[c].coefficients.direct.cost_per_unit_rate; let rate_or_slowest = match rate(k) { @@ -547,6 +549,13 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { Consumer::Query(_) => unreachable!("a query's rate is constant"), }), }; + for &(key, cost) in &input.choices[c].coefficients.cross_shared { + let entry = cross_per_consumer + .entry(key) + .or_insert_with(|| (Expression::from(0.0), 0.0)); + entry.0 += rate_or_slowest * cost * v; + entry.1 += rate_or_slowest * cost; + } for &(source, cost, pass) in &input.choices[c].coefficients.absorbed_scans { let pass = pass.min(cost); objective += rate_or_slowest * pass * v; @@ -659,14 +668,38 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { .push((c, charge)); } } + // A node several groups' choices have (`cross_shared`) costs, in the + // plan, its charge per consumer (one node computed per consumer, or a + // copy per group) or, built once, at least the dearest chosen one's + // charge at its group's slowest rate: the model charges the smaller. + let mut cross_once: BTreeMap> = BTreeMap::new(); for (c, choice) in input.choices.iter().enumerate() { for &(key, cost) in &choice.coefficients.cross_shared { - shared_slots - .entry((key, usize::MAX, "node".to_string())) + cross_once + .entry(key) .or_default() .push((c, slowest_rate(choice.group.0 as usize) * cost)); } } + let mut cross_vars: Vec<( + Variable, + Variable, + Variable, + Vec<(usize, f64)>, + Expression, + f64, + )> = Vec::new(); + for (key, members) in cross_once { + let (per_consumer, per_consumer_max) = cross_per_consumer + .remove(&key) + .unwrap_or_else(|| (Expression::from(0.0), 0.0)); + let big_m = per_consumer_max + members.iter().map(|(_, w)| *w).fold(0.0, f64::max); + let once = vars.add(variable().min(0)); + let charged = vars.add(variable().min(0)); + let u = vars.add(binary()); + objective += charged; + cross_vars.push((once, charged, u, members, per_consumer, big_m)); + } let mut shared_slot_vars: Vec<(Variable, Vec<(usize, f64)>)> = Vec::new(); for members in shared_slots.into_values() { let m = vars.add(variable().min(0)); @@ -924,6 +957,17 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } constraints.push(constraint!(into >= p)); } + // A slot's consumer exists only when its choice is chosen. + for (&(k, _), &r) in &reach { + if let Consumer::Slot { choice, .. } = k { + constraints.push(constraint!(r <= x[choice])); + } + } + for (&(k, _), &v) in &y { + if let Consumer::Slot { choice, .. } = k { + constraints.push(constraint!(v <= x[choice])); + } + } for (c, choice) in input.choices.iter().enumerate() { for (child, via) in choice.children.iter().zip(&choice.coefficients.child_via) { let h = child.0 as usize; @@ -1106,6 +1150,13 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { constraints.push(constraint!(m >= charge)); } } + for (once, charged, u, members, per_consumer, big_m) in cross_vars { + for (c, charge) in members { + constraints.push(constraint!(once >= charge * x[c])); + } + constraints.push(constraint!(charged >= per_consumer - big_m * u)); + constraints.push(constraint!(charged >= once - big_m * (1 - u))); + } for (m, members) in &shared_slot_vars { for &(c, charge) in members { constraints.push(constraint!(*m >= charge * x[c])); From df2427da08ecdb26eb09f2d9b492af4be94bf2a4 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:25:00 +0000 Subject: [PATCH 59/67] plan-selection: memo consumers exist only with their choice or built-once group A slot's consumer reaches, passes into and is charged by nothing unless its choice is chosen; a group's own consumer unless the group is built once. Coefficient tests charge a built-once node's scans once. --- .../src/pass/subexpr/coefficients.rs | 53 +++++++++++++++++++ .../plan-selection/src/pass/subexpr/milp.rs | 44 ++++++++++----- 2 files changed, 84 insertions(+), 13 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index a0c854866..aa208d51a 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -1614,6 +1614,59 @@ mod tests { .copied() .collect(); lazy.held.retain(|term| !sub.contains(&term.node_key)); + // Its scans below are computed as it is (once when it + // is built once, per consumer otherwise). + let scans_below: Vec = below(node) + .into_iter() + .filter(|id| { + matches!(nodes[id].payload, Operator::NonASAP(NonASAPOp::Scan { .. })) + }) + .collect(); + let materialized_here = + nodes[&node].materialization != Materialization::NotMaterialized; + for (&scan, &(_, cost, _)) in + scans_below.iter().zip(&coefficients.absorbed_scans) + { + if !materialized_here { + break; + } + actual -= cost * rates[&p.root].1; + // Its held memory, once at the scan's rate. + let scan_key = keys[position_of[&scan]]; + let scan_held: Vec = lazy + .held + .iter() + .filter(|t| t.node_key == scan_key) + .copied() + .collect(); + lazy.held.retain(|t| t.node_key != scan_key); + if built_once_nodes.insert(scan) { + actual += cost * rates[&scan].1; + let readers: Vec<(usize, f64)> = reach + .iter() + .enumerate() + .filter(|(_, r)| r.contains(&scan)) + .map(|(q, _)| (q, query_rate(q))) + .collect(); + let rate: f64 = readers.iter().map(|(_, r)| r).sum(); + for term in scan_held { + let evaluation_ms = match term.per_reader && rate > 0.0 { + true => { + readers.iter().map(|&(q, r)| r * latency[q]).sum::() + / rate + } + false => { + readers.iter().map(|&(q, _)| latency[q]).fold(0.0, f64::max) + } + }; + held_extra += prices.cost_per_memory_byte_second + * term.bytes as f64 + * rates[&scan].1 + * evaluation_ms + / 1_000.0; + } + } + } if built_once_nodes.insert(node) { actual += shared.cost_per_unit_rate * rates[&node].1; let readers: Vec<(usize, f64)> = reach diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index c364c3707..022b27bfc 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -64,6 +64,18 @@ pub(crate) struct SlotCoefficients { pub shared_site: Option<(u64, usize)>, } +/// A cross-group node's variables: its once charge, its charge, the +/// binary choosing between per consumer and once, its members' once +/// charges, its per-consumer charge and their big M. +type CrossNodeVars = ( + Variable, + Variable, + Variable, + Vec<(usize, f64)>, + Expression, + f64, +); + /// One choice's coefficients. #[derive(Debug, Clone, PartialEq)] pub(crate) struct ChoiceCoefficients { @@ -681,14 +693,7 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { .push((c, slowest_rate(choice.group.0 as usize) * cost)); } } - let mut cross_vars: Vec<( - Variable, - Variable, - Variable, - Vec<(usize, f64)>, - Expression, - f64, - )> = Vec::new(); + let mut cross_vars: Vec = Vec::new(); for (key, members) in cross_once { let (per_consumer, per_consumer_max) = cross_per_consumer .remove(&key) @@ -957,15 +962,28 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } constraints.push(constraint!(into >= p)); } - // A slot's consumer exists only when its choice is chosen. + // A slot's consumer exists only when its choice is chosen, and a + // group's own consumer only when the group is built once. + let exists = |k: Consumer| -> Option { + match k { + Consumer::Slot { choice, .. } => Some(x[choice]), + Consumer::Group(g) => once.get(&g).copied(), + Consumer::Query(_) => None, + } + }; for (&(k, _), &r) in &reach { - if let Consumer::Slot { choice, .. } = k { - constraints.push(constraint!(r <= x[choice])); + if let Some(e) = exists(k) { + constraints.push(constraint!(r <= e)); + } + } + for (&(k, _), &p) in &passes { + if let Some(e) = exists(k) { + constraints.push(constraint!(p <= e)); } } for (&(k, _), &v) in &y { - if let Consumer::Slot { choice, .. } = k { - constraints.push(constraint!(v <= x[choice])); + if let Some(e) = exists(k) { + constraints.push(constraint!(v <= e)); } } for (c, choice) in input.choices.iter().enumerate() { From 5e4c084a21ef85a6d708abc301a9c18ddb60a6b8 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:25:06 +0000 Subject: [PATCH 60/67] Regenerate the viewer fixtures for the search report's incumbent_cost --- tools/dag-viewer/examples/planner-layering-example1.json | 1 + .../examples/random-mixed-seed-862-long-workload-01.json | 1 + tools/dag-viewer/examples/random-mixed-seed-862-workload-01.json | 1 + .../examples/random-mixed-seed-862-x5-workload-01.json | 1 + 4 files changed, 4 insertions(+) diff --git a/tools/dag-viewer/examples/planner-layering-example1.json b/tools/dag-viewer/examples/planner-layering-example1.json index 0dcb640fa..75554e00d 100644 --- a/tools/dag-viewer/examples/planner-layering-example1.json +++ b/tools/dag-viewer/examples/planner-layering-example1.json @@ -128159,6 +128159,7 @@ "search": { "candidates_checked": 72, "components": [], + "incumbent_cost": 4.377000084179999, "queries": [], "raw_scan_assignments_skipped": 0, "raw_scan_candidates_skipped": 0, diff --git a/tools/dag-viewer/examples/random-mixed-seed-862-long-workload-01.json b/tools/dag-viewer/examples/random-mixed-seed-862-long-workload-01.json index 3c75a84b6..a8fcea4d0 100644 --- a/tools/dag-viewer/examples/random-mixed-seed-862-long-workload-01.json +++ b/tools/dag-viewer/examples/random-mixed-seed-862-long-workload-01.json @@ -20202,6 +20202,7 @@ "search": { "candidates_checked": 933, "components": [], + "incumbent_cost": 4.5581790928001605e-6, "interchangeable_queries": [ [ 2, diff --git a/tools/dag-viewer/examples/random-mixed-seed-862-workload-01.json b/tools/dag-viewer/examples/random-mixed-seed-862-workload-01.json index 870ec1441..ff3c246b5 100644 --- a/tools/dag-viewer/examples/random-mixed-seed-862-workload-01.json +++ b/tools/dag-viewer/examples/random-mixed-seed-862-workload-01.json @@ -20202,6 +20202,7 @@ "search": { "candidates_checked": 933, "components": [], + "incumbent_cost": 1.5764190368555072e-6, "interchangeable_queries": [ [ 2, diff --git a/tools/dag-viewer/examples/random-mixed-seed-862-x5-workload-01.json b/tools/dag-viewer/examples/random-mixed-seed-862-x5-workload-01.json index 55661c497..646f37bee 100644 --- a/tools/dag-viewer/examples/random-mixed-seed-862-x5-workload-01.json +++ b/tools/dag-viewer/examples/random-mixed-seed-862-x5-workload-01.json @@ -20202,6 +20202,7 @@ "search": { "candidates_checked": 933, "components": [], + "incumbent_cost": 7.882484714275631e-6, "interchangeable_queries": [ [ 2, From 48e30572906a8ff4c434a6e1dd53f34231db8ae1 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:38:45 +0000 Subject: [PATCH 61/67] plan-selection: memo MILP prices scaled per evaluation; held memory at the union rate The scale came from the largest charge including slots' fixed costs: a slot keeping tens of GB of state shrank every other charge below the solver's dual tolerance, so the solver stopped above the optimum with costly variables left at 1 (seed 43: 0.3% over the exact price). It is now the largest per-evaluation charge at the fastest query's rate. A part several queries of one class read runs at their union rate, so each reader's held memory is charged at union/sum of its rate. --- .../plan-selection/src/pass/subexpr/milp.rs | 44 +++++++++++++++---- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 022b27bfc..5af56992d 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -334,19 +334,20 @@ pub(crate) fn solve(input: &ModelInput) -> Result, String> /// tolerance (a relative gap of about 1e-6 at most) may have its optimum. const GAP_TOLERANCE: f64 = 1e-4; -/// The factor that brings `input`'s largest charge to about 1. +/// The factor that brings `input`'s largest charge per evaluation to about +/// 1: each choice's charges at the fastest query's rate, and raw scans. +/// A slot's fixed cost (its kept state) is left out: a huge one would +/// shrink every other charge below the solver's tolerances (about 1e-7), +/// where it no longer tells charges apart and stops above the optimum. fn price_scale(input: &ModelInput) -> f64 { - let fastest = input.queries.iter().map(|q| q.rate).fold(1.0, f64::max); + let fastest = input.queries.iter().map(|q| q.rate).fold(0.0, f64::max); let largest = input .choices .iter() .flat_map(|choice| { let k = &choice.coefficients; - std::iter::once(k.direct.cost_per_unit_rate * fastest).chain( - k.slots - .iter() - .map(move |s| s.fixed_cost + s.cost_per_unit_rate * fastest), - ) + std::iter::once(k.direct.cost_per_unit_rate * fastest) + .chain(k.slots.iter().map(move |s| s.cost_per_unit_rate * fastest)) }) .chain(input.raw_scans.iter().map(|s| s.cost)) .filter(|v| v.is_finite()) @@ -1141,11 +1142,38 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { .sum() }) .collect(); + // A part read by several queries runs at their classes' union rate, + // not their rates' sum: per reader, its held memory is charged at + // `union / sum` of that reader's rate (Stage 3's rate-weighted mean + // latency at the part's rate). The least ratio over the group's class + // sets bounds it from below. + let per_reader_share: Vec = input + .group_queries + .iter() + .map(|queries_of| { + input + .union_rate + .iter() + .filter_map(|(set, &union)| { + let rate_sum: f64 = queries_of + .iter() + .filter(|&&q| set.contains(&input.schedule[q])) + .map(|&q| input.queries[q].rate) + .sum(); + let covered = set + .iter() + .all(|class| queries_of.iter().any(|&q| input.schedule[q] == *class)); + (covered && rate_sum > 0.0).then_some((union / rate_sum).min(1.0)) + }) + .fold(1.0, f64::min) + }) + .collect(); for (&(q, c), &w) in &work_reached { let held = held_of(c); let longest = longest_latency[q]; if let Some(&h) = held_per_reader.get(&(q, c)) { - let coefficient = held.per_reader_per_ms * input.queries[q].rate; + let share = per_reader_share[input.choices[c].group.0 as usize]; + let coefficient = held.per_reader_per_ms * input.queries[q].rate * share; constraints.push(constraint!( h >= coefficient * (latency_of(q) - longest * (1 - w)) )); From 123974cfb352772e5a54e59e8c2437cfbc0b949f Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:26:43 +0000 Subject: [PATCH 62/67] Memo MILP MVP: drop the debugging environment hooks; a quick three-way test over the 20 smallest random workloads --- .../tests/joint_search_equivalence.rs | 65 +++++-- crates/plan-selection/src/lib.rs | 6 - .../src/pass/subexpr/coefficients.rs | 9 +- .../plan-selection/src/pass/subexpr/milp.rs | 179 +----------------- 4 files changed, 54 insertions(+), 205 deletions(-) diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index d6adcbd76..2baf63a3f 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1496,23 +1496,49 @@ fn three_way( /// enough to search exhaustively, over one or two gauges, with raw data /// kept or not, under each of `settings` in turn (every one with `every`). fn assert_three_way_on_random_workloads(settings: &[PlanSelectionSettings], every: bool) { + assert_three_way_on_seeds(settings, every, 60, 20); +} + +/// The random workload of `seed`, over one gauge or (every third seed) +/// two. +fn seeded_workload(seed: u64) -> PlanningWorkload { + let gauges: &[&str] = match seed % 3 { + 0 => &["m", "n"], + _ => &["m"], + }; + random_workload_over(seed, gauges) +} + +/// The number of whole-workload combinations of `workload`'s candidates. +fn combinations_of(workload: &PlanningWorkload) -> usize { + let roots = lower_promql_roots(workload).unwrap(); + asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates( + roots.iter().cloned().enumerate().collect(), + ) + .unwrap() + .combinations() +} + +/// [`assert_three_way_on_random_workloads`] on the `count` seeds of 1 to +/// 60 with the fewest combinations (all of them small enough to search +/// exhaustively), at least `at_least` of them. +fn assert_three_way_on_seeds( + settings: &[PlanSelectionSettings], + every: bool, + count: usize, + at_least: usize, +) { + let mut seeds: Vec<(usize, u64)> = (1..=60u64) + .map(|seed| (combinations_of(&seeded_workload(seed)), seed)) + .filter(|&(combinations, _)| combinations <= MAX_EXHAUSTIVE_COMBINATIONS) + .collect(); + seeds.sort_unstable(); + seeds.truncate(count); 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); + for (_, seed) in seeds { + let workload = seeded_workload(seed); 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(); @@ -1541,7 +1567,7 @@ fn assert_three_way_on_random_workloads(settings: &[PlanSelectionSettings], ever } } assert!( - searched >= 20, + searched >= at_least, "only {searched} workloads were small enough" ); assert!(mismatches.is_empty(), "{}", mismatches.join("\n")); @@ -1554,6 +1580,15 @@ fn three_way_random_workloads_under_cost_model_1() { assert_three_way_on_random_workloads(&cost_models(1)[..1], true); } +/// The three searches agree under cost model 1 on the 20 random workloads +/// with the fewest combinations: the quick subset of +/// [`three_way_random_workloads_under_cost_model_1`]. +#[test] +#[ignore = "the exhaustive search over 20 seeds takes a minute or more in a debug build; run with --release --ignored"] +fn three_way_twenty_smallest_random_workloads_under_cost_model_1() { + assert_three_way_on_seeds(&cost_models(1)[..1], true, 20, 20); +} + /// The three searches agree on the random workloads under cost model 2. #[test] #[ignore = "the exhaustive search and the brute force over every seed under cost model 2 take minutes in a debug build; run with --release --ignored"] diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 4a0bef45c..676717d47 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1120,7 +1120,6 @@ impl SearchAlgorithm { } } - /// Why the legacy search ran instead of the requested memo search. #[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "snake_case")] @@ -2126,11 +2125,6 @@ impl<'a> PlanSelector<'a> { Ok((selection, built)) } - /// The rejections recorded so far. - pub(crate) fn rejections(&self) -> &[Rejection] { - &self.rejected - } - /// Under a price ranking, the cheapest price of the valid candidates /// offered so far. pub(crate) fn cheapest_price(&self) -> Option { diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index aa208d51a..2e7d120e4 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -1857,14 +1857,7 @@ mod tests { }; let models = PlanningModels::builtin().with_deployment_capabilities(&deployment); let memo = memo(&per_computation, &entries); - let input = super::super::milp::model_input_of( - &memo, - &entries, - &data, - models, - std::time::Instant::now(), - ) - .unwrap(); + let input = super::super::milp::model_input_of(&memo, &entries, &data, models).unwrap(); let checked = super::super::milp::lower_bound_check( &memo, &per_computation, diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 5af56992d..338465212 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -11,7 +11,7 @@ //! query runs, composed accuracy) are added back by the caller's lazy cuts: //! the model's optimum is a lower bound on the plan's Stage 3 price. -use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use good_lp::solvers::highs::highs; use good_lp::{ @@ -221,10 +221,6 @@ pub(crate) struct ModelInput { /// Per choice, the memory its parts hold while a query's evaluation /// runs ([`HeldAtQueryTime`]). pub held: Vec, - /// Only plans the model prices below this: the incumbent's price. The - /// model is then infeasible exactly when no group choice left can beat - /// the incumbent. - pub cutoff: Option, /// The solver's time limit in seconds, if any. pub time_limit_s: Option, /// Solve the linear relaxation (every binary in [0, 1]): a lower bound @@ -391,14 +387,13 @@ fn scaled(input: &ModelInput, scale: f64) -> ModelInput { for (_, bound) in &mut scaled.extra_bounds { *bound *= scale; } - scaled.cutoff = input.cutoff.map(|c| c * scale); scaled } fn solve_scaled(input: &ModelInput) -> Result, String> { let mut vars = ProblemVariables::new(); let n = input.group_count; - let relaxed = input.relaxed || std::env::var_os("JP_MILP_LP").is_some(); + let relaxed = input.relaxed; let binary = || { if relaxed { variable().min(0).max(1) @@ -653,9 +648,6 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { let u = vars.add(binary()); objective += a + b; objective += shared_cost[c] * slowest_rate(input.choices[c].group.0 as usize) * x[c]; - if std::env::var_os("JP_MILP_NOMIN").is_some() { - continue; - } cheaper_of.push((a, b, u, c, big_m)); } // Slots that may be one shared node (equal builds over equal inputs): @@ -810,22 +802,6 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { for &v in &kept_on_disk { objective += v; } - if std::env::var_os("JP_MILP_TRACE").is_some() { - use good_lp::IntoAffineExpression; - let negative: Vec<(Variable, f64)> = (&objective) - .linear_coefficients() - .filter(|(_, c)| *c < 0.0 || c.is_nan()) - .collect(); - if !negative.is_empty() { - let x_set: HashSet = x.iter().copied().collect(); - eprintln!( - "milp: {} negative objective coefficients, {} on choices: {:?}", - negative.len(), - negative.iter().filter(|(v, _)| x_set.contains(v)).count(), - &negative[..negative.len().min(5)] - ); - } - } let variable_count = vars.len() as u64; let model = vars.minimise(objective.clone()).using(highs); let mut constraints = Vec::new(); @@ -1294,100 +1270,27 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } } } - if let Some(cutoff) = input - .cutoff - .filter(|_| std::env::var_os("JP_MILP_CUTOFF").is_some()) - { - constraints.push(constraint!(objective.clone() <= cutoff)); - } let constraint_count = constraints.len() as u64; let mut model = model .with_all(constraints) .with_mip_gap(0.0) .map_err(|e| format!("MIP gap: {e:?}"))?; - if std::env::var_os("JP_MILP_VERBOSE").is_some() { - model.set_verbose(true); - } // No absolute gap either: prices are scaled to about 1. model = model.set_option("mip_abs_gap", 0.0); if let Some(limit) = input.time_limit_s { model = model.set_time_limit(limit.max(0.001)); } - if let Ok(options) = std::env::var("JP_HIGHS_OPTIONS") { - for option in options.split(',').filter(|o| !o.is_empty()) { - let (key, value) = option.split_once('=').unwrap_or((option, "")); - if let Ok(v) = value.parse::() { - model = model.set_option(key, v); - } else if let Ok(v) = value.parse::() { - model = model.set_option(key, v); - } else if let Ok(v) = value.parse::() { - model = model.set_option(key, v); - } else { - model = model.set_option(key, value); - } - } - } let solution = match model.solve() { Ok(solution) => solution, Err(good_lp::ResolutionError::Infeasible) => return Ok(None), Err(error) => return Err(error.to_string()), }; - if std::env::var_os("JP_MILP_TRACE").is_some() { - use good_lp::IntoAffineExpression; - let low: Vec<(f64, f64)> = (&objective) - .linear_coefficients() - .map(|(v, c)| (c, solution.value(v))) - .filter(|(_, value)| *value < -1e-9) - .collect(); - if !low.is_empty() { - eprintln!( - "milp: {} objective variables below 0: {:?}", - low.len(), - &low[..low.len().min(8)] - ); - } - } let chosen: Vec = x .iter() .enumerate() .filter(|(_, v)| solution.value(**v) > 0.5) .map(|(c, _)| c) .collect(); - if std::env::var_os("JP_MILP_TRACE").is_some() { - if relaxed { - for (g, v) in once.iter() { - eprintln!(" once g{} = {:.3}", g.0, solution.value(*v)); - } - for (g, v) in need.iter().enumerate() { - eprintln!(" need g{} = {:.3}", g, solution.value(*v)); - } - } - let shown: Vec = if relaxed { - (0..x.len()) - .filter(|&c| solution.value(x[c]) > 1e-6) - .collect() - } else { - chosen.clone() - }; - for &c in &shown { - if relaxed { - eprintln!(" x = {:.3}", solution.value(x[c])); - } - let choice: &Choice = &input.choices[c]; - let reach: Vec = y - .iter() - .filter(|((_, d), v)| *d == c && solution.value(**v) > 1e-9) - .map(|((k, _), v)| format!("{k:?}={:.2}", solution.value(*v))) - .collect(); - eprintln!( - " g{} cand{} direct {:.4e}/rate work {:.2}ms shared {:?} slots {:?} scans {:?} horizon {:?} assignment {:?} y {:?}", - choice.group.0, choice.candidate, choice.coefficients.direct.cost_per_unit_rate, - choice.coefficients.direct.work_ms, choice.coefficients.shared_nodes.len(), - choice.coefficients.slots.iter().map(|s| (s.fixed_cost, s.cost_per_unit_rate)).collect::>(), - choice.scan_children, choice.scan_horizon_ms, choice.assignment, reach - ); - } - } let built_once = once .iter() .filter(|(_, v)| solution.value(**v) > 0.5) @@ -1488,8 +1391,7 @@ pub(crate) fn search( let started = std::time::Instant::now(); let memo = super::search::memo(per_computation, entries); let memo_stats = super::search::stats(&memo, started.elapsed().as_secs_f64() * 1e3); - let mut input = model_input_of(&memo, entries, data, models, started)?; - let trace = std::env::var_os("JP_MILP_TRACE").is_some(); + let mut input = model_input_of(&memo, entries, data, models)?; let mut selector = PlanSelector::new(entries, data, models, false)?; let mut logical: HashMap> = HashMap::new(); let mut number = 0usize; @@ -1512,9 +1414,6 @@ pub(crate) fn search( input.time_limit_s = remaining_s; iterations += 1; let solve_started = std::time::Instant::now(); - input.cutoff = selector - .cheapest_price() - .map(|incumbent| incumbent - 1e-9 * incumbent.abs()); let solved = solve(&input).map_err(SelectionError::Solver)?; solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; // The solver stopped at the time limit: its plan is no bound, but @@ -1533,9 +1432,6 @@ pub(crate) fn search( bound = Some(solution.objective); first_bound.get_or_insert(solution.objective); } - if trace { - eprintln!("milp: round {iterations} bound {:.6e} incumbent {:?} solve {:.0}ms vars {} cons {} at {:.1}s", solution.objective, selector.cheapest_price(), solve_started.elapsed().as_secs_f64()*1e3, solution.variables, solution.constraints, started.elapsed().as_secs_f64()); - } size = (solution.variables, solution.constraints); if let Some(incumbent) = selector.cheapest_price() { if !last_round && solution.objective >= incumbent - 1e-9 * incumbent.abs() { @@ -1548,9 +1444,6 @@ pub(crate) fn search( for &c in &solution.chosen { group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); } - if trace { - eprintln!("milp: round {iterations} group choice {group_choice:?}"); - } let chosen_candidates: Vec<(GroupId, &super::ImplementationCandidate)> = group_choice .iter() .enumerate() @@ -1592,15 +1485,10 @@ pub(crate) fn search( cut = true; } } - if trace { - eprintln!("milp: round {iterations} requirements not met (pair cut {cut})"); - } if cut { continue; } } - let offer_started = std::time::Instant::now(); - let offered_before = number; super::search::offer_group_choice( &memo, per_computation, @@ -1612,17 +1500,6 @@ pub(crate) fn search( &mut logical, &mut number, )?; - if trace { - eprintln!( - "milp: round {iterations} priced {} plans in {:.1}s; last rejection {:?}", - number - offered_before, - offer_started.elapsed().as_secs_f64(), - selector - .rejections() - .last() - .map(|r| (r.cause, r.reason.chars().take(200).collect::())) - ); - } // Exclude it: every choice of each chosen (group, candidate). let groups = group_choice.iter().flatten().count(); let excluded: Vec = input @@ -1645,26 +1522,6 @@ pub(crate) fn search( .excluded .push((excluded, outside, groups as f64 - 1.0)); } - if std::env::var_os("JP_MILP_CHECK").is_some() { - // Debugging aid: every group choice's exact price against the - // model's value restricted to it. - let incumbent = selector.cheapest_price().unwrap_or(f64::INFINITY); - for (group_choice, exact, value) in - lower_bound_check(&memo, per_computation, &input, entries, data, models)? - { - let bad = value.is_none_or(|v| v > exact * (1.0 + 1e-9)); - if bad || exact < incumbent * (1.0 - 1e-9) { - eprintln!( - "check: {group_choice:?} exact {exact:.7e} model {value:?}{}", - if bad { - " NOT A LOWER BOUND" - } else { - " BELOW INCUMBENT" - } - ); - } - } - } if time_limit_reached && bound.is_none() { // No round solved to optimality: the relaxation's optimum (all cuts // kept) is still a lower bound for the gap. @@ -1742,7 +1599,6 @@ pub(super) fn model_input_of( entries: &[QueryWorkloadEntry], data: &DataWorkload, models: PlanningModels<'_>, - started: std::time::Instant, ) -> Result { let scale = super::pricing::frozen_calibration(memo, entries, data, models)?; let mut cache = super::pricing::PriceCache::new(scale); @@ -1785,32 +1641,6 @@ pub(super) fn model_input_of( choices.push(choice); } } - let trace = std::env::var_os("JP_MILP_TRACE").is_some(); - if trace { - for (c, choice) in choices.iter().enumerate() { - let k = &choice.coefficients; - let negative = k.direct.cost_per_unit_rate < 0.0 - || k.direct.work_ms < 0.0 - || k.slots - .iter() - .any(|s| s.fixed_cost < 0.0 || s.cost_per_unit_rate < 0.0 || s.work_ms < 0.0) - || k.shared_nodes.iter().any(|n| n.cost_per_unit_rate < 0.0); - if negative { - eprintln!( - "milp: NEGATIVE choice {c}: g{} cand{} {:?} assignment {:?}", - choice.group.0, choice.candidate, k, choice.assignment - ); - } - } - } - if trace { - eprintln!( - "milp: memo {} groups, {} choices in {:.1}s", - memo.groups.len(), - choices.len(), - started.elapsed().as_secs_f64() - ); - } let mut input = model_input(memo, choices, entries, models)?; let mut sources: Vec = Vec::new(); for (c, source, bytes, via) in raw_scans { @@ -1950,7 +1780,6 @@ fn model_input_clone(input: &ModelInput) -> ModelInput { raw_scans: input.raw_scans.clone(), raw_data_retained: input.raw_data_retained, held: input.held.clone(), - cutoff: input.cutoff, time_limit_s: input.time_limit_s, relaxed: input.relaxed, rollup_rules: input.rollup_rules.clone(), @@ -2124,7 +1953,6 @@ fn model_input( raw_scans: Vec::new(), raw_data_retained: true, held: Vec::new(), - cutoff: None, time_limit_s: None, relaxed: false, }) @@ -2201,7 +2029,6 @@ mod tests { raw_scans: Vec::new(), raw_data_retained: true, held: Vec::new(), - cutoff: None, time_limit_s: None, relaxed: false, }; From 2b87f3dd5566ba7ceee5060a7d789bcf9a1daf00 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:35:22 +0000 Subject: [PATCH 63/67] stage_pipeline tests: per-call names for the mixed workload files, which two parallel tests shared --- crates/devtools/tests/stage_pipeline.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/devtools/tests/stage_pipeline.rs b/crates/devtools/tests/stage_pipeline.rs index 95e81d811..1e125df96 100644 --- a/crates/devtools/tests/stage_pipeline.rs +++ b/crates/devtools/tests/stage_pipeline.rs @@ -1981,9 +1981,13 @@ fn mixed_language_workload_files(sql_arrival: Value, sql_facts: Value) -> (PathB }, "data_workload": {"arrival": sql_arrival, "sources": {"flows": sql_facts}}, }); + // Tests run as threads of one process, so the names need a per-call + // counter. + static CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + let call = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); let path = |name: &str| { std::env::temp_dir().join(format!( - "stage_pipeline_{}_mixed_{name}.json", + "stage_pipeline_{}_{call}_mixed_{name}.json", std::process::id() )) }; From 3ef44df6bfc2828f3843498952dd02c3787cd4c6 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:57:18 +0000 Subject: [PATCH 64/67] Memo MILP: no safety-net cuts outside the reference evaluation; a solver time-limit stop with no solution keeps the incumbent; query latency leaves out nodes another group may share; the lower-bound tests also run without raw data, and over a query reading two groups of one metric --- crates/plan-selection/src/lib.rs | 4 +- .../src/pass/subexpr/coefficients.rs | 19 ++++- .../plan-selection/src/pass/subexpr/milp.rs | 78 +++++++------------ 3 files changed, 48 insertions(+), 53 deletions(-) diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 676717d47..b12f6305d 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1057,7 +1057,9 @@ pub struct MilpReport { /// model is against the incumbent cost. pub first_lower_bound: Option, /// (incumbent cost − lower bound) / incumbent cost, at least 0: 0 when - /// the incumbent is proven optimal. + /// the incumbent is proven optimal over the memo, within the solver's + /// tolerances (a relative 1e-9 in the stopping rule, about 1e-4 when + /// HiGHS stops at its own gap tolerance). pub gap: Option, /// Its rounds (each: a solve, the decoded plan's incumbent evaluation, /// a cut). diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index 2e7d120e4..dc1ed748f 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -1827,6 +1827,7 @@ mod tests { targets: &[AccuracyTarget], ) { assert_model_is_a_lower_bound_in(queries, intervals_ms, targets, true); + assert_model_is_a_lower_bound_in(queries, intervals_ms, targets, false); } /// [`assert_model_is_a_lower_bound`] in a deployment that keeps raw @@ -1867,7 +1868,8 @@ mod tests { models, ) .unwrap(); - assert!(!checked.is_empty()); + // Without raw data a workload may have no valid plan at all. + assert!(!raw_data_retained || !checked.is_empty()); for (group_choice, exact, value) in checked { let value = value.unwrap_or_else(|| panic!("{queries:?} {group_choice:?}: infeasible")); assert!( @@ -1922,6 +1924,21 @@ mod tests { ); } + /// One query reading two groups over one metric, whose scans may be one + /// node: its latency and held memory stay lower bounds. + #[test] + fn the_model_is_a_lower_bound_for_a_query_over_two_groups_of_one_metric() { + let exact = AccuracyTarget::exact(); + assert_model_is_a_lower_bound( + &[ + "max_over_time(m[5m]) - min_over_time(m[5m])", + "max_over_time(m[5m])", + ], + &[60_000, 60_000], + &[exact.clone(), exact], + ); + } + /// #840's rollup: a quantile by host and job and its rollup by job, at /// the rollup's state-size thresholds. #[test] diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 338465212..ab6cbee2d 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -750,13 +750,21 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { at_query_time.insert((q, g), vars.add(variable().min(0).max(1))); } } + // A choice's query-time work, without its nodes another group's + // choice may share (shareable nodes and shared slots): Stage 3 counts a + // shared node once per query, so they are left out to keep latency a + // lower bound. let query_time_work = |c: usize| -> f64 { let coefficients = &input.choices[c].coefficients; - coefficients.direct.work_ms + let shared: f64 = coefficients.shared_nodes.iter().map(|n| n.work_ms).sum(); + (coefficients.direct.work_ms - shared).max(0.0) + coefficients .slots .iter() - .filter(|s| !matches!(s.materialization, Materialization::MatAtIngestion(_))) + .filter(|s| { + !matches!(s.materialization, Materialization::MatAtIngestion(_)) + && s.shared_site.is_none() + }) .map(|s| s.work_ms) .sum::() }; @@ -1414,7 +1422,17 @@ pub(crate) fn search( input.time_limit_s = remaining_s; iterations += 1; let solve_started = std::time::Instant::now(); - let solved = solve(&input).map_err(SelectionError::Solver)?; + let solved = match solve(&input) { + Ok(solved) => solved, + // The solver stopped at the time limit before finding any + // solution: the incumbent so far stands. + Err(error) if remaining_s.is_some() && error.contains("NoSolutionFound") => { + time_limit_reached = true; + solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; + break; + } + Err(error) => return Err(SelectionError::Solver(error)), + }; solver_ms += solve_started.elapsed().as_secs_f64() * 1e3; // The solver stopped at the time limit: its plan is no bound, but // it is still priced, as a last round, and may become the @@ -1444,51 +1462,9 @@ pub(crate) fn search( for &c in &solution.chosen { group_choice[input.choices[c].group.0 as usize] = Some(input.choices[c].candidate); } - let chosen_candidates: Vec<(GroupId, &super::ImplementationCandidate)> = group_choice - .iter() - .enumerate() - .filter_map(|(g, c)| Some((GroupId(g as u32), &memo.candidates[g][(*c)?]))) - .collect(); - if !super::windows::requirements_met(&chosen_candidates) { - // The model encodes every requirement (`rollup_rules`), so this - // is a safety net: cut the smallest failing pairs of (group, - // candidate), else the whole group choice below. - let mut cut = false; - for (i, a) in chosen_candidates.iter().enumerate() { - if a.1.rollup_of.is_none() { - continue; - } - for (j, b) in chosen_candidates.iter().enumerate() { - if i != j && !super::windows::requirements_met(&[*a, *b]) { - let members: Vec = (0..input.choices.len()) - .filter(|&c| { - let choice = &input.choices[c]; - (choice.group == a.0 - && choice.candidate == group_choice[a.0 .0 as usize].unwrap()) - || (choice.group == b.0 - && choice.candidate - == group_choice[b.0 .0 as usize].unwrap()) - }) - .collect(); - input.excluded.push((members, Vec::new(), 1.0)); - cut = true; - } - } - if !super::windows::requirements_met(&[*a]) { - let members: Vec = (0..input.choices.len()) - .filter(|&c| { - input.choices[c].group == a.0 - && Some(input.choices[c].candidate) == group_choice[a.0 .0 as usize] - }) - .collect(); - input.excluded.push((members, Vec::new(), 0.0)); - cut = true; - } - } - if cut { - continue; - } - } + // A group choice that fails a requirement the model does not + // encode is offered as any other: the reference evaluation finds + // no valid plan in it, and the cut below removes exactly it. super::search::offer_group_choice( &memo, per_computation, @@ -1548,11 +1524,11 @@ pub(crate) fn search( } if !without_plan.is_empty() { return Err(SelectionError::Solver(format!( - "no valid plan for {}{}", + "no valid plan for {} (raw data {})", without_plan.join(", "), match models.deployment_capabilities.raw_data_retained { - true => "", - false => " in a deployment that keeps no raw data", + true => "retained", + false => "not retained", } ))); } From e45c42c1518a4b207310672170765f496d98f22d Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:58:16 +0000 Subject: [PATCH 65/67] Restore runtime_install.rs, which an earlier merge truncated --- crates/planner/tests/runtime_install.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/planner/tests/runtime_install.rs b/crates/planner/tests/runtime_install.rs index 2732b426f..aa940bf38 100644 --- a/crates/planner/tests/runtime_install.rs +++ b/crates/planner/tests/runtime_install.rs @@ -3300,6 +3300,7 @@ async fn a_classic_histogram_quantile_keeps_labels_and_handles_inf_and_empty_inp ); assert!((value - 0.98).abs() < 1e-9, "{value}"); } +} // ── PromQL absent() (issue #895) ───────────────────────────────────────────── From a1f1a62a2ab7bda76d707a732f7f41c1a3502fad Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:11:29 +0000 Subject: [PATCH 66/67] Memo MILP: identical scans and same-class query consumers charged once A query over two groups of one metric reads one scan node and one node per group for its schedule class, where the model charged each reader: - a scan's pass is charged once per identical scan (source and price), the largest of its readers' charges; - a query-time slot's raw scans come out of its slot charge and are charged as scans of their source, so two slots over one scan share it; - a choice's direct charge for queries of one class is their largest, not their sum. --- .../src/pass/subexpr/coefficients.rs | 24 +++++- .../plan-selection/src/pass/subexpr/milp.rs | 79 ++++++++++++++++--- 2 files changed, 91 insertions(+), 12 deletions(-) diff --git a/crates/plan-selection/src/pass/subexpr/coefficients.rs b/crates/plan-selection/src/pass/subexpr/coefficients.rs index dc1ed748f..91efa9ebe 100644 --- a/crates/plan-selection/src/pass/subexpr/coefficients.rs +++ b/crates/plan-selection/src/pass/subexpr/coefficients.rs @@ -912,7 +912,24 @@ fn coefficients_of( let via = Via::Slot(slots.len()); lazy_of(&part, via); let at_ingestion = matches!(node.materialization, Materialization::MatAtIngestion(_)); - let rate_cost: f64 = part.iter().map(|&id| per_rate(id)).sum(); + // A query-time slot's raw scans are charged as scans. + let scans: Vec<(u64, f64, f64)> = match at_ingestion { + true => Vec::new(), + false => part + .iter() + .filter(|&&id| query_time(id)) + .filter_map(|&id| match &dag.nodes[position[&id]].payload { + Operator::NonASAP(NonASAPOp::Scan { source, .. }) => Some(( + super::scans::raw_horizon_key(source), + per_rate(id), + prices.cost_per_cpu_op * priced[&id].output.rows as f64, + )), + _ => None, + }) + .collect(), + }; + let rate_cost: f64 = part.iter().map(|&id| per_rate(id)).sum::() + - scans.iter().map(|&(_, cost, _)| cost).sum::(); let state: f64 = part.iter().map(|&id| kept_state(id)).sum(); slots.push(SlotCoefficients { materialization: node.materialization, @@ -929,6 +946,7 @@ fn coefficients_of( // threshold, identical expressions, scans widened to one label // set): equal over equal inputs, scans' projections aside. shared_site: Some((u64::from(widened[position[&node.id]]), 0)), + scans, }); slot_keys.push(shallow[position[&node.id]]); part_keys.push(part.iter().map(|id| keys[position[id]]).collect()); @@ -1708,7 +1726,9 @@ mod tests { // memory too. for (i, slot) in coefficients.slots.iter().enumerate() { if charged_slots.insert(p.slots[i]) { - actual += slot.fixed_cost + slot.cost_per_unit_rate * rates[&p.slots[i]].1; + let scans: f64 = slot.scans.iter().map(|&(_, cost, _)| cost).sum(); + actual += + slot.fixed_cost + (slot.cost_per_unit_rate + scans) * rates[&p.slots[i]].1; } else { lazy.held.retain(|term| term.via != Via::Slot(i)); lazy.raw_scans.retain(|(_, _, via)| *via != Via::Slot(i)); diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index ab6cbee2d..7244668df 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -46,7 +46,7 @@ pub(crate) struct PartCoefficients { } /// One materialized private slot of a choice. -#[derive(Debug, Clone, Copy, PartialEq)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct SlotCoefficients { pub materialization: Materialization, /// Price per second that does not depend on a rate (ingestion-time @@ -62,6 +62,10 @@ pub(crate) struct SlotCoefficients { /// and materialization are one node when shared (a site's build at one /// threshold, identical expressions), so the model charges them once. pub shared_site: Option<(u64, usize)>, + /// A query-time slot's raw scans, out of `cost_per_unit_rate`: like + /// [`ChoiceCoefficients::absorbed_scans`], charged as scans of their + /// source (another slot's or part's identical scan is one node). + pub scans: Vec<(u64, f64, f64)>, } /// A cross-group node's variables: its once charge, its charge, the @@ -154,8 +158,9 @@ pub(crate) struct Choice { /// For a scan group's choice: what a pass over its output costs per /// unit of rate, at most its own price. A widened scan shared by a /// source's scan groups costs at least the dearest scan, and each other - /// reader at least such a pass: the model charges each scan choice's - /// pass, plus the largest excess over it. + /// reader at least such a pass: the model charges each distinct scan's + /// pass once (readers of one scan share its node), plus the largest + /// excess over it. pub scan_project_cost_per_unit_rate: f64, } @@ -364,6 +369,10 @@ fn scaled(input: &ModelInput, scale: f64) -> ModelInput { for slot in &mut k.slots { slot.fixed_cost *= scale; slot.cost_per_unit_rate *= scale; + for (_, cost, pass) in &mut slot.scans { + *cost *= scale; + *pass *= scale; + } } for node in &mut k.shared_nodes { node.cost_per_unit_rate *= scale; @@ -544,7 +553,14 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { z }; // Each scan choice's direct charge, per source (`Choice::scan_source`). - let mut scan_charges: BTreeMap> = BTreeMap::new(); + let mut scan_charges: BTreeMap> = BTreeMap::new(); + // Each scan choice's pass, per identical scan (its source and price): + // readers of one scan share its node and pay its pass once. + let mut pass_charges: BTreeMap<(u64, u64), BTreeMap<(usize, usize), Expression>> = + BTreeMap::new(); + // Each choice's direct charge per schedule class of its query + // consumers: one node serves the class's queries. + let mut class_charges: BTreeMap<(usize, usize), Vec> = BTreeMap::new(); // Each cross-group node's charges per consumer, and their largest sum. let mut cross_per_consumer: BTreeMap = BTreeMap::new(); for (&(k, c), &v) in &y { @@ -566,11 +582,15 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } for &(source, cost, pass) in &input.choices[c].coefficients.absorbed_scans { let pass = pass.min(cost); - objective += rate_or_slowest * pass * v; + *pass_charges + .entry((source, cost.to_bits())) + .or_default() + .entry((c, 0)) + .or_insert_with(|| Expression::from(0.0)) += rate_or_slowest * pass * v; *scan_charges .entry(source) .or_default() - .entry(c) + .entry((c, 0)) .or_insert_with(|| Expression::from(0.0)) += rate_or_slowest * (cost - pass) * v; } if total == 0.0 { @@ -579,11 +599,15 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { if let Some(source) = input.choices[c].scan_source { let rate = rate_or_slowest; let pass = input.choices[c].scan_project_cost_per_unit_rate.min(total); - objective += rate * pass * v; + *pass_charges + .entry((source, total.to_bits())) + .or_default() + .entry((c, 0)) + .or_insert_with(|| Expression::from(0.0)) += rate * pass * v; *scan_charges .entry(source) .or_default() - .entry(c) + .entry((c, 0)) .or_insert_with(|| Expression::from(0.0)) += rate * (total - pass) * v; continue; } @@ -591,7 +615,14 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { let cost = total - shared; match rate(k) { Some(r) => { - objective += r * cost * v; + match k { + // Queries of one class read one node of the choice. + Consumer::Query(q) => class_charges + .entry((c, input.schedule[q])) + .or_default() + .push(r * cost * v), + _ => objective += r * cost * v, + } if shared != 0.0 { per_consumer_shared[c].push((v, r * shared)); } @@ -620,8 +651,36 @@ fn solve_scaled(input: &ModelInput) -> Result, String> { } } } + // Query-time slots' scans, at the slot's group's slowest rate, each + // slot's apart from its choice's other parts. + for (c, choice) in input.choices.iter().enumerate() { + let slowest = slowest_rate(choice.group.0 as usize); + for (i, slot) in choice.coefficients.slots.iter().enumerate() { + if matches!(slot.materialization, Materialization::MatAtIngestion(_)) { + continue; + } + for &(source, cost, pass) in &slot.scans { + let pass = pass.min(cost); + *pass_charges + .entry((source, cost.to_bits())) + .or_default() + .entry((c, i + 1)) + .or_insert_with(|| Expression::from(0.0)) += slowest * pass * x[c]; + *scan_charges + .entry(source) + .or_default() + .entry((c, i + 1)) + .or_insert_with(|| Expression::from(0.0)) += slowest * (cost - pass) * x[c]; + } + } + } let mut scan_maxima: Vec<(Variable, Vec)> = Vec::new(); - for charges in scan_charges.into_values() { + for charges in class_charges.into_values() { + let m = vars.add(variable().min(0)); + objective += m; + scan_maxima.push((m, charges)); + } + for charges in scan_charges.into_values().chain(pass_charges.into_values()) { let m = vars.add(variable().min(0)); objective += m; scan_maxima.push((m, charges.into_values().collect())); From 73be8a3afc6865250555c53a658c001a01a72ade Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:18:07 +0000 Subject: [PATCH 67/67] Memo MILP: a failed relaxation after the time limit leaves the gap unknown instead of failing; the report says any solver limit --- crates/plan-selection/src/lib.rs | 5 +++-- crates/plan-selection/src/pass/subexpr/milp.rs | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index b12f6305d..76dc9af44 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -1071,8 +1071,9 @@ pub struct MilpReport { /// together. pub solver_time_ms: u64, /// Whether it stopped at its time limit - /// ([`crate::Stage1ToStage3Search::milp_time_limit_ms`]): the plan is - /// then proven optimal only if `gap` is 0. + /// ([`crate::Stage1ToStage3Search::milp_time_limit_ms`]) or another + /// solver limit (iterations, memory, interrupt): the plan is then + /// proven optimal only if `gap` is 0. pub time_limit_reached: bool, } diff --git a/crates/plan-selection/src/pass/subexpr/milp.rs b/crates/plan-selection/src/pass/subexpr/milp.rs index 7244668df..4cc900d35 100644 --- a/crates/plan-selection/src/pass/subexpr/milp.rs +++ b/crates/plan-selection/src/pass/subexpr/milp.rs @@ -1562,7 +1562,8 @@ pub(crate) fn search( // kept) is still a lower bound for the gap. input.relaxed = true; input.time_limit_s = None; - if let Some(relaxation) = solve(&input).map_err(SelectionError::Solver)? { + // A failed relaxation leaves the gap unknown, never the incumbent lost. + if let Ok(Some(relaxation)) = solve(&input) { bound = Some(relaxation.objective); } }