From 11de144127d381c2bd6a55d7fb48976620e4cf36 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:26:04 +0000 Subject: [PATCH 01/12] refactor: establish logical memo before Pass 1 and decouple legacy infrastructure --- .../tests/joint_search_equivalence.rs | 14 +- crates/logical-asap-optimization/src/lib.rs | 3 + crates/logical-asap-optimization/src/memo.rs | 322 +++++++++++++ crates/plan-selection/src/lib.rs | 7 +- .../src/search/legacy/components.rs | 80 +--- .../plan-selection/src/search/legacy/mod.rs | 2 - .../src/search/legacy/stage1_to_stage3.rs | 431 +----------------- .../src/search/memo/brute_force.rs | 88 +++- .../plan-selection/src/search/memo/groups.rs | 191 +++----- .../src/search/memo/incumbent.rs | 82 ++++ crates/plan-selection/src/search/memo/mod.rs | 31 +- .../src/search/milp/clusters.rs | 11 +- .../plan-selection/src/search/milp/model.rs | 69 ++- crates/plan-selection/src/search/mod.rs | 6 +- .../src/search/{legacy => }/pane_sharing.rs | 0 crates/plan-selection/src/search/pipeline.rs | 411 +++++++++++++++++ .../src/search/workload_dependencies.rs | 81 ++++ 17 files changed, 1162 insertions(+), 667 deletions(-) create mode 100644 crates/logical-asap-optimization/src/memo.rs create mode 100644 crates/plan-selection/src/search/memo/incumbent.rs rename crates/plan-selection/src/search/{legacy => }/pane_sharing.rs (100%) create mode 100644 crates/plan-selection/src/search/pipeline.rs create mode 100644 crates/plan-selection/src/search/workload_dependencies.rs diff --git a/crates/integration-tests/tests/joint_search_equivalence.rs b/crates/integration-tests/tests/joint_search_equivalence.rs index 196052bc4..c1772bf04 100644 --- a/crates/integration-tests/tests/joint_search_equivalence.rs +++ b/crates/integration-tests/tests/joint_search_equivalence.rs @@ -1988,8 +1988,8 @@ fn dump_random_workload() { } /// With no time left for a single MILP round, memo-milp still returns a -/// plan: the warm-start incumbent, each query planned alone by the legacy -/// search with no sharing, offered before round 1. Under cost model 1 its +/// plan: the warm-start incumbent, each query planned alone by the memo +/// bounded baseline with no sharing, offered before round 1. Under cost model 1 its /// cost is the sum of the queries' own plans' costs, and the run reports /// the time limit and a gap to the MILP lower bound. #[test] @@ -2016,7 +2016,7 @@ fn a_time_limited_memo_milp_run_returns_the_warm_start_plan() { std::slice::from_ref(entry), &data, models, - Stage1ToStage3Search::default(), + Stage1ToStage3Search::with_algorithm(SearchAlgorithm::MemoBruteForce), ) .unwrap(); let selection = run.stage3_plan_selection.unwrap(); @@ -2039,7 +2039,13 @@ fn a_time_limited_memo_milp_run_returns_the_warm_start_plan() { let selection = run.stage3_plan_selection.unwrap(); assert_eq!(selection.selected, "P0"); let cost = selection.search.incumbent_cost.unwrap(); - assert!((cost - alone).abs() <= 1e-9 * alone, "{cost} vs {alone}"); + // The bounded baseline is feasible, not necessarily the per-query optimum. + assert!(cost + 1e-9 * alone >= alone, "{cost} vs {alone}"); + assert!(cost.is_finite()); + assert_eq!( + Some(cost), + selection.costs[&selection.selected].cost_per_second + ); let milp = selection.search.milp.unwrap(); assert!(milp.time_limit_reached, "{milp:?}"); assert!(milp.gap.is_some(), "{milp:?}"); diff --git a/crates/logical-asap-optimization/src/lib.rs b/crates/logical-asap-optimization/src/lib.rs index 9c5e441e0..c2d13af5a 100644 --- a/crates/logical-asap-optimization/src/lib.rs +++ b/crates/logical-asap-optimization/src/lib.rs @@ -70,3 +70,6 @@ pub mod per_computation_candidates; pub mod asap_aware_cse; pub mod workload_candidates; + +/// Logical group identity available before Pass 1 candidate generation. +pub mod memo; diff --git a/crates/logical-asap-optimization/src/memo.rs b/crates/logical-asap-optimization/src/memo.rs new file mode 100644 index 000000000..2d8bec3ea --- /dev/null +++ b/crates/logical-asap-optimization/src/memo.rs @@ -0,0 +1,322 @@ +//! Logical group identity established before implementation candidate generation. +//! +//! This module depends only on logical IR. It preserves the existing CSE +//! legality rules and owns the original roots so pointer-based lookup remains +//! valid while later stages populate implementation candidates. + +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; +use std::time::{Duration, Instant}; + +use asap_types::ir::cse::share_common_sub_dags; +use asap_types::ir::{OperatorNode, QueryRoot}; + +/// A group, by index in [`LogicalMemo::groups`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct MemoGroupId(pub u32); + +/// One distinct logical sub-expression of the workload. +#[derive(Debug, Clone)] +pub 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, +} + +/// Initial logical groups, without implementation choices or solver state. +#[derive(Debug, Clone)] +pub struct LogicalMemo { + pub groups: Vec, + original_roots: Vec, + group_of: GroupOf, + /// Time spent constructing logical identity, excluding later Pass 1 work. + pub build_duration: Duration, +} + +impl LogicalMemo { + pub fn from_roots(roots: &[(Id, QueryRoot)]) -> Self { + let started = Instant::now(); + let (groups, group_of) = hash_consed(roots); + Self { + groups, + original_roots: roots.iter().map(|(_, root)| root.clone()).collect(), + group_of, + build_duration: started.elapsed(), + } + } + + /// Group of a node from the original workload, not a rewritten candidate. + pub fn group_for_node(&self, node: &Rc) -> Option { + self.group_of.get(&Rc::as_ptr(node)).copied() + } + + /// Check root kind and ordered operator dependencies, not scalar values. + /// Changing scalar-only output expressions does not change logical groups. + /// Never compare operator DAGs recursively: shared diamonds and NaN literals + /// make structural equality inappropriate for pointer-indexed identities. + pub fn matches_root_dependencies(&self, roots: &[(Id, QueryRoot)]) -> bool { + self.original_roots.len() == roots.len() + && self + .original_roots + .iter() + .zip(roots) + .all(|(original, (_, root))| { + std::mem::discriminant(original) == std::mem::discriminant(root) && { + let original_nodes = operator_roots(original); + let nodes = operator_roots(root); + original_nodes.len() == nodes.len() + && original_nodes + .iter() + .zip(&nodes) + .all(|(a, b)| Rc::ptr_eq(a, b)) + } + }) + } +} + +/// Each original node (of `roots`), by address, to its group. +type GroupOf = HashMap<*const OperatorNode, MemoGroupId>; + +/// 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, MemoGroupId> = 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, MemoGroupId>, +) -> MemoGroupId { + 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 = MemoGroupId(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: MemoGroupId, 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, MemoGroupId>, + 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); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::ir::ScalarExpr; + + fn query(text: &str) -> QueryRoot { + QueryRoot::Operator(lower_promql(text, AccuracyTarget::max_error(0.01))) + } + + /// Logical interning shares equivalent results without needing Pass 1 choices. + #[test] + fn repeated_results_share_groups_but_quantile_levels_do_not() { + let roots = vec![ + (0, query("sum by (job) (m)")), + (1, query("sum by (job) (m)")), + (2, query("quantile_over_time(0.95, m[5m])")), + (3, query("quantile_over_time(0.99, m[5m])")), + ]; + let memo = LogicalMemo::from_roots(&roots); + let root_group = |q| { + memo.groups + .iter() + .position(|g| g.roots_of.contains(&q)) + .unwrap() + }; + assert_eq!(root_group(0), root_group(1)); + assert_ne!(root_group(2), root_group(3)); + for (index, group) in memo.groups.iter().enumerate() { + assert!(group + .children + .iter() + .all(|child| (child.0 as usize) < index)); + for child in &group.children { + assert!(memo.groups[child.0 as usize] + .parents + .contains(&MemoGroupId(index as u32))); + } + } + } + + /// Scalar references are roots of their query; literals need no invented node. + #[test] + fn scalar_and_literal_roots_keep_workload_positions() { + let QueryRoot::Operator(node) = query("sum by (job) (m)") else { + unreachable!() + }; + let roots = vec![ + (10, QueryRoot::Scalar(ScalarExpr::literal_f64(1.0))), + ( + 20, + QueryRoot::Scalar(ScalarExpr::PromqlScalarFromVector(node.clone())), + ), + (30, QueryRoot::Operator(node.clone())), + ]; + let memo = LogicalMemo::from_roots(&roots); + let group = &memo.groups[memo.group_for_node(&node).unwrap().0 as usize]; + assert_eq!(group.roots_of, vec![1, 2]); + assert_eq!(group.queries, vec![1, 2]); + assert!(memo.matches_root_dependencies(&roots)); + let mut reordered = roots.clone(); + reordered.swap(0, 1); + assert!(!memo.matches_root_dependencies(&reordered)); + assert!(LogicalMemo::from_roots(&roots[..1]).groups.is_empty()); + } + + /// Non-reflexive floating literals must not invalidate unchanged dependencies. + #[test] + fn unchanged_nan_root_matches_by_dependency_identity() { + let roots = vec![(0, QueryRoot::Scalar(ScalarExpr::literal_f64(f64::NAN)))]; + let memo = LogicalMemo::from_roots(&roots); + assert!(memo.matches_root_dependencies(&roots)); + } + + /// Dependency checks must not unfold shared operator DAGs once per path. + #[test] + fn checking_shared_diamond_uses_root_identity() { + use asap_types::ir::operator::operator_properties::BinaryOpKind; + use asap_types::ir::scalar::CompareOpKind; + use asap_types::ir::{BinaryOperator, NonASAPOp, Operator}; + let QueryRoot::Operator(mut node) = query("m") else { + unreachable!() + }; + for _ in 0..24 { + node = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::BinaryOp { + operator: BinaryOperator { + checked_relative_division: false, + checked_finite_division: false, + kind: BinaryOpKind::Compare(CompareOpKind::Eq), + vector_match: None, + }, + return_bool: false, + lhs: node.clone(), + rhs: node, + })) + .unwrap(); + } + // Construct the identity guard directly so this regression isolates + // root matching from the independently tested CSE/interner. + let roots = vec![(0, QueryRoot::Operator(node))]; + let memo = LogicalMemo { + groups: Vec::new(), + original_roots: roots.iter().map(|(_, r)| r.clone()).collect(), + group_of: GroupOf::new(), + build_duration: Duration::ZERO, + }; + assert!(memo.matches_root_dependencies(&roots)); + } + + /// Structural equality alone cannot validate pointer-indexed target lookups. + #[test] + fn rebuilt_equal_roots_do_not_match_the_original_workload() { + let roots = vec![(0, query("sum by (job) (m)"))]; + let memo = LogicalMemo::from_roots(&roots); + let rebuilt = vec![(0, query("sum by (job) (m)"))]; + assert!(!memo.matches_root_dependencies(&rebuilt)); + assert!(memo.matches_root_dependencies(&roots.clone())); + } +} diff --git a/crates/plan-selection/src/lib.rs b/crates/plan-selection/src/lib.rs index 80f73f516..2be94df65 100644 --- a/crates/plan-selection/src/lib.rs +++ b/crates/plan-selection/src/lib.rs @@ -288,7 +288,7 @@ pub use accuracy::{ AccuracyEvidenceProvider, AccuracyModel, BuiltinAccuracyModel, NoAccuracyEvidence, }; pub mod cost; -pub use search::legacy::pane_sharing; +pub use search::pane_sharing; /// The Stage 1 to Stage 3 driver and the optimization-pass interface. pub mod search; pub use search as pass; @@ -1172,8 +1172,9 @@ pub struct MilpReport { #[serde(default, skip_serializing_if = "Vec::is_empty")] pub clusters: Vec, /// The cost per second of the warm-start plan: each query planned - /// alone, with no sharing between queries; `None` when some query has - /// no plan alone. + /// alone by the bounded memo baseline, with no sharing between queries. + /// This is not necessarily the per-query optimum; `None` when the + /// bounded attempts do not find a validated plan for every query. #[serde(default, skip_serializing_if = "Option::is_none")] pub queries_alone_cost: Option, } diff --git a/crates/plan-selection/src/search/legacy/components.rs b/crates/plan-selection/src/search/legacy/components.rs index 5c314ccd0..f4ebd7769 100644 --- a/crates/plan-selection/src/search/legacy/components.rs +++ b/crates/plan-selection/src/search/legacy/components.rs @@ -25,7 +25,6 @@ //! search combines those rules ([`combined_pass2`]), whatever the rules' //! enumeration caps over the whole workload. -use std::collections::HashMap; use std::rc::Rc; use asap_logical_asap_optimization::asap_aware_cse::{ @@ -38,13 +37,12 @@ use asap_logical_asap_optimization::per_computation_candidates::{ }; use asap_logical_asap_optimization::workload_candidates::Pass2Candidate; use asap_physical_asap_optimization::optimize_physical_asap_assigned; -use asap_types::ir::operator::Source; -use asap_types::ir::{NonASAPOp, Operator, OperatorNode, QueryRoot}; +use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; use asap_types::physical::Materialization; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; use crate::search::legacy::query_candidates::{next_digits, query_digits}; -use crate::search::legacy::stage1_to_stage3::{ +use crate::search::pipeline::{ run_stage1_to_stage3_from_pass1, selected_logical, LogicalCandidate, Stage1ToStage3Run, Stage1ToStage3Search, }; @@ -59,58 +57,6 @@ use crate::{ /// of summing the same charges in another grouping. const UNIQUE_MARGIN: f64 = 4e-9; -/// The workload's queries grouped into components that read disjoint -/// sources, each in workload order, ordered by their first query. A query -/// that reads no source is a component of its own. -pub(crate) fn source_components(roots: &[QueryRoot]) -> Vec> { - let mut parent: Vec = (0..roots.len()).collect(); - fn find(parent: &mut [usize], q: usize) -> usize { - let mut root = q; - while parent[root] != root { - root = parent[root]; - } - let mut q = q; - while parent[q] != root { - let next = parent[q]; - parent[q] = root; - q = next; - } - root - } - // The first query read each source, by name (data facts are keyed by - // name, whatever the source's kind). - let mut first_reader: HashMap = HashMap::new(); - for (query, root) in roots.iter().enumerate() { - let starts: Vec<&Rc> = match root { - QueryRoot::Operator(node) => vec![node], - QueryRoot::Scalar(expr) => expr.operator_refs(), - }; - for node in starts.into_iter().flat_map(OperatorNode::reachable) { - let Operator::NonASAP(NonASAPOp::Scan { source, .. }) = &node.operator else { - continue; - }; - let name = match source { - Source::TimeSeries { metric } => metric.clone(), - Source::Table { table_ref } => table_ref.clone(), - }; - let other = *first_reader.entry(name).or_insert(query); - let (a, b) = (find(&mut parent, query), find(&mut parent, other)); - parent[a.max(b)] = a.min(b); - } - } - let mut components: Vec> = Vec::new(); - let mut index_of: HashMap = HashMap::new(); - for query in 0..roots.len() { - let root = find(&mut parent, query); - let index = *index_of.entry(root).or_insert_with(|| { - components.push(Vec::new()); - components.len() - 1 - }); - components[index].push(query); - } - components -} - /// The workload's queries grouped into components that no candidate can /// link: [`source_components`], each split further into the queries that /// can read a common node, in workload order, ordered by their first query. @@ -658,22 +604,6 @@ fn combined_pass2(runs: &[Stage1ToStage3Run]) -> Option` at /// the start of its id or reason, counted within the component) renamed to /// the workload's. -pub(crate) fn renumber(rejection: Rejection, component: &[usize]) -> Rejection { - let rename = |text: &str| -> String { - let Some(rest) = text.strip_prefix('q') else { - return text.to_string(); - }; - let digits: String = rest.chars().take_while(char::is_ascii_digit).collect(); - match digits.parse::() { - Ok(n) if n >= 1 && n <= component.len() => { - format!("q{}{}", component[n - 1] + 1, &rest[digits.len()..]) - } - _ => text.to_string(), - } - }; - Rejection { - id: rename(&rejection.id), - cause: rejection.cause, - reason: rename(&rejection.reason), - } -} +use crate::search::workload_dependencies::{ + remap_rejection_query_indices as renumber, source_connected_queries as source_components, +}; diff --git a/crates/plan-selection/src/search/legacy/mod.rs b/crates/plan-selection/src/search/legacy/mod.rs index f3cb5bfcf..ca12eee21 100644 --- a/crates/plan-selection/src/search/legacy/mod.rs +++ b/crates/plan-selection/src/search/legacy/mod.rs @@ -3,7 +3,5 @@ pub(crate) mod components; pub(crate) mod interchangeable; -/// Costed reuse of compatible physical pane producers. -pub mod pane_sharing; pub(crate) mod query_candidates; pub(crate) mod stage1_to_stage3; diff --git a/crates/plan-selection/src/search/legacy/stage1_to_stage3.rs b/crates/plan-selection/src/search/legacy/stage1_to_stage3.rs index 213cb90f0..32cb11210 100644 --- a/crates/plan-selection/src/search/legacy/stage1_to_stage3.rs +++ b/crates/plan-selection/src/search/legacy/stage1_to_stage3.rs @@ -1,348 +1,32 @@ -//! Stage 1 (logical ASAP-aware optimization) → Stage 2 (physical ASAP-aware -//! optimization) → Stage 3 (plan selection), starting from the query-language -//! frontends' roots; Stage 0 runs before it in the facade. -//! -//! This is the planner pipeline (#509 MVP). [`Stage1ToStage3Pass`] is the -//! [`OptimizationPass`](crate::search::OptimizationPass) that runs it. -//! -//! [`run_stage1_to_stage3`] runs, for one workload: -//! 1. Stage 1, logical ASAP-aware optimization: Pass 1 per-computation candidate -//! generation, then composition of one whole-workload candidate per choice -//! of a per-computation candidate for every target (the Cartesian product, -//! in mixed-radix order, the last target varying fastest), each followed -//! by its Pass 2 candidates: those of the identical-expression rule -//! ([`identical_expression_pass2_candidates`]: one per subset of its -//! sharing candidates), then those of the summary-capability rule -//! ([`summary_capability_pass2_candidates`]: one per choice of shared -//! summary builds), each followed by its own identical-expression -//! candidates, and, for the composed candidate and each -//! summary-capability candidate, those of the window-composition rule -//! ([`window_composition_pass2_candidates`]: one per non-empty subset of -//! the builds a repeating query could build as tumbling panes, and one -//! more for a subset whose nested windows on one pane width can read one -//! shared chain), each -//! followed by its own identical-expression candidates, and, for the -//! composed candidate, those of the reuse rules (exact top-k reuse, then -//! each rollup Pass 2 candidate followed by its own exact top-k reuse -//! Pass 2 candidates: [`topk_reuse_pass2_candidates`], -//! [`rollup_pass2_candidates`]), each followed by its own -//! identical-expression candidates; every distinct candidate once; the -//! independent candidate is kept; -//! 2. Stage 2, physical ASAP-aware optimization of every Stage 1 candidate -//! ([`optimize_physical_asap`]): one physical candidate per -//! materialization assignment of its summary builds; -//! 3. Stage 3, plan selection over the physical candidates ([`select_plan`]). -//! -//! Stage 0, the query-language frontends, runs before it (the `asap-planner` -//! facade lowers the workload). The window-composition rule offers -//! per-query tumbling panes only; sliding windows, Exponential Histograms -//! and a pane width shared across queries are not planned yet. -//! -//! The search follows `docs/design_docs/proposals/joint-planning-search.md`. -//! Each query's own candidates are checked before composition (step 2, -//! [`query_candidates`](crate::search::legacy::query_candidates)), and a whole-workload -//! choice that picks a discarded one is not composed. Every other -//! candidate is priced as it is built and only the best so far is kept -//! (step 6; under cost model 3's `pareto-front`, the front so far), unless -//! the caller asks to keep every candidate ([`Stage1ToStage3Search`]). +//! Temporary whole-workload search implementation, pending memo feature parity. -use std::rc::Rc; - -use asap_types::ir::schema::schema_support::with_promql_series_identity; -use asap_types::ir::{OperatorNode, QueryRoot}; -use asap_types::workload::{DataWorkload, QueryLanguage, QueryWorkloadEntry}; +use asap_types::ir::QueryRoot; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; use crate::search::legacy::interchangeable::Classes; use crate::search::legacy::query_candidates::QueryCandidates; -use crate::search::{OptimizationInput, OptimizationPass, OptimizeError, PlanOutput, QueryPlan}; use crate::PlanningModels; use crate::{ - CalibrationFit, CostModel, PlanSelector, QuerySearchReport, Rejection, RejectionCause, - RoutedToLegacy, SearchAlgorithm, Selection, SelectionError, -}; -use asap_logical_asap_optimization::per_computation_candidates::{ - generate_per_computation_candidates, PerComputationCandidateError, PerComputationCandidates, + CalibrationFit, PlanSelector, QuerySearchReport, Rejection, RejectionCause, Selection, + SelectionError, }; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; use asap_logical_asap_optimization::workload_candidates::{ - stage1_logical_asap_optimization, stage1_logical_asap_optimization_skipping, Pass2Candidate, - Stage1Candidate, + stage1_logical_asap_optimization, stage1_logical_asap_optimization_skipping, Stage1Candidate, }; use asap_physical_asap_optimization::{ every_assignment_scans_raw_at_query_time, optimize_physical_asap, optimize_physical_asap_without_query_time_raw_scans, PhysicalCandidate, }; -/// One whole-workload CandidateLogicalASAPDAG: `choice[t]` picks target -/// `t`'s candidate, or its fused computation's -/// (`targets[t].computation_for(choice[t])`), of the Pass 1 inventory. `roots` are the -/// composed query roots in input order, or why the choice did not compose. -#[derive(Debug)] -pub struct LogicalCandidate { - /// `L`, numbered in enumeration order. - pub id: String, - pub choice: Vec, - pub roots: Result, PerComputationCandidateError>, - /// For a Pass 2 candidate: what it changed of its Pass 1 candidate. - pub pass2: Option, -} - -/// A composed logical candidate that physical ASAP-aware optimization could -/// not implement, with the reason. -#[derive(Debug, Clone, PartialEq)] -pub struct PrunedCandidate { - pub from_logical: String, - pub reason: String, -} - -/// The output of Stage 1, Stage 2 and Stage 3 for one workload -/// ([`run_stage1_to_stage3`]). The candidate lists are filled only when the -/// search keeps every candidate ([`Stage1ToStage3Search::keep_every_candidate`]). -#[derive(Debug)] -pub struct Stage1ToStage3Run { - /// Stage 1, Pass 1: the per-computation candidates, with the roots they - /// read. - pub stage1_pass1_per_computation: PerComputationCandidates, - /// Stage 1, Pass 1: the number of composed whole-workload candidates, - /// the size of the Cartesian product of the targets' candidates, - /// saturating at `usize::MAX`. Pass 2 adds to them. - pub combinations: usize, - /// Stage 1: every whole-workload logical candidate, in enumeration - /// order: each composed Pass 1 candidate, followed by its Pass 2 - /// candidates (those with [`LogicalCandidate::pass2`]). Empty unless - /// every candidate is kept. - pub stage1_logical_asap_optimization: Vec, - /// Stage 2: the physical candidates of each Stage 1 candidate that - /// composed and was implemented, one per materialization assignment: - /// `P` (nothing materialized), then `P-m`, from `L`. Empty - /// unless every candidate is kept. - pub stage2_physical_asap_optimization: Vec, - /// Stage 2: the Stage 1 candidates Stage 2 could not implement, with the - /// reasons. Empty unless every candidate is kept. - pub stage2_pruned: Vec, - /// Stage 3: plan selection over the Stage 2 candidates. - pub stage3_plan_selection: Result, - /// Stage 3: the selected physical candidate, as Stage 2 built it - /// ([`Self::selected`]). - pub stage3_selected: Option, - /// Stage 3: the selected candidate's logical roots, with the weight - /// requirements it confirmed made proofs ([`Self::selected_logical`]). - pub stage3_selected_logical: Option>, - /// Stage 1: the logical candidate the selected physical candidate - /// implements, with its roots as Stage 1 composed them (before Stage 3 - /// sized them): what Stage 2 implemented. - pub stage1_selected: Option, -} - -impl Stage1ToStage3Run { - /// The selected physical candidate, if selection succeeded. - pub fn selected(&self) -> Option<&PhysicalCandidate> { - self.stage3_selected.as_ref() - } - - /// The composed roots of the logical candidate the selected physical - /// candidate implements: the plan before physical operator - /// implementation and materialization, with the weight requirements - /// Stage 3 confirmed made proofs, as in [`Selection::plan`]. - pub fn selected_logical(&self) -> Option<&[QueryRoot]> { - self.stage3_selected_logical.as_deref() - } -} - -/// How [`run_stage1_to_stage3_with`] searches, and what it keeps. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct Stage1ToStage3Search { - /// Keep every Stage 1 and Stage 2 candidate and every Stage 3 rejection - /// in the run, for a developer tool that shows them; default `false`: - /// only the best so far is kept. - pub keep_every_candidate: bool, - /// Compose every whole-workload choice, without checking each query's - /// own candidates first: the exhaustive search the pruned one is tested - /// against; default `false`. - pub exhaustive: 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, - /// 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::MilpReport::gap`]) and - /// [`crate::MilpReport::time_limit_reached`]: not proven - /// optimal unless the gap is 0. - pub milp_time_limit_ms: Option, -} - -impl Stage1ToStage3Search { - /// The exhaustive search keeping every candidate: what the planner did - /// before the search was pruned. - pub fn exhaustive() -> Self { - Self { - keep_every_candidate: true, - exhaustive: true, - search: SearchAlgorithm::Legacy, - milp_time_limit_ms: None, - } - } - - /// The default search run with `algorithm`. - pub fn with_algorithm(algorithm: SearchAlgorithm) -> Self { - Self { - search: algorithm, - ..Self::default() - } - } -} - -/// Run Stage 1 → Stage 2 → Stage 3 over the frontends' `roots` with the -/// default search ([`Stage1ToStage3Search::default`]: each query's -/// candidates checked first, only the best kept). `entries[i]` is the -/// workload entry of `roots[i]`: its accuracy target, latency bound, -/// recurrence and predictability. Only a Pass 1 failure is an error; a -/// failed selection is reported in the run. -pub fn run_stage1_to_stage3( - roots: Vec<(Id, QueryRoot)>, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: PlanningModels<'_>, -) -> Result, PerComputationCandidateError> { - run_stage1_to_stage3_with( - roots, - entries, - data, - models, - Stage1ToStage3Search::default(), - ) -} - -/// [`run_stage1_to_stage3`] with the search `search`. Keeping every -/// candidate, a candidate that does not compose or cannot be implemented is -/// listed with its reason. -pub fn run_stage1_to_stage3_with( - roots: Vec<(Id, QueryRoot)>, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: PlanningModels<'_>, - search: Stage1ToStage3Search, -) -> Result, PerComputationCandidateError> { - let per_computation = generate_per_computation_candidates(roots)?; - Ok(run_stage1_to_stage3_from_pass1( - per_computation, - entries, - data, - models, - search, - )) -} - -/// [`run_stage1_to_stage3_with`] after Pass 1: searches the Stage 1 -/// candidates of `per_computation` (the Pass 1 inventory of the frontends' -/// roots, [`stage1_logical_asap_optimization`]) without the summaries the -/// deployment does not support -/// ([`PerComputationCandidates::without_unsupported_summaries`]), so a -/// caller can read the inventory, e.g. its -/// [`combinations`](PerComputationCandidates::combinations), before -/// enumerating. -pub fn run_stage1_to_stage3_from_pass1( - per_computation: PerComputationCandidates, - entries: &[QueryWorkloadEntry], - data: &DataWorkload, - models: PlanningModels<'_>, - search: Stage1ToStage3Search, -) -> Stage1ToStage3Run { - 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 = - crate::search::memo::brute_force::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); - let found = crate::search::milp::model::search( - &pruned, - entries, - data, - models, - search.milp_time_limit_ms, - ); - return memo_run(pruned, found, entries, data, SearchAlgorithm::MemoMilp); - } - } - }; - 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: crate::search::milp::model::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)); - let from_logical = selected.as_ref().map(|s| s.from_logical.clone()); - 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, - // 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, - }), - } -} +use crate::search::pipeline::{ + selected_logical, LogicalCandidate, PrunedCandidate, Stage1ToStage3Run, Stage1ToStage3Search, +}; -/// The legacy search ([`SearchAlgorithm::Legacy`]) of -/// [`run_stage1_to_stage3_from_pass1`]: independent components searched +/// The legacy search ([`crate::SearchAlgorithm::Legacy`]) of +/// [`crate::run_stage1_to_stage3_from_pass1`]: independent components searched /// separately when that is proven exact, else the workload as one. -fn legacy_search( +pub(crate) fn legacy_search( per_computation: PerComputationCandidates, entries: &[QueryWorkloadEntry], data: &DataWorkload, @@ -372,7 +56,7 @@ fn legacy_search( run } -/// [`run_stage1_to_stage3_from_pass1`] over the whole workload as one +/// [`crate::run_stage1_to_stage3_from_pass1`] over the whole workload as one /// component. fn search_as_one( per_computation: PerComputationCandidates, @@ -473,7 +157,7 @@ fn search_as_one( /// `models.deployment_capabilities` /// ([`PerComputationCandidates::without_unsupported_summaries`]), with /// `stage1` enumerated from that pruned inventory, as -/// [`run_stage1_to_stage3_from_pass1`] does. Pruning here would come too +/// [`crate::run_stage1_to_stage3_from_pass1`] does. Pruning here would come too /// late: each of `stage1`'s choices indexes the inventory it was enumerated /// from, so a candidate with an unsupported summary would already be in /// `stage1`. @@ -793,88 +477,3 @@ fn search_stage2_and_stage3( stage1_selected, } } - -/// `roots`, the selected candidate's Stage 1 roots, as Stage 3 sized them -/// for their queries' failure probabilities, with the weight requirements -/// it confirmed made proofs. -pub(crate) fn selected_logical( - roots: &[QueryRoot], - entries: &[QueryWorkloadEntry], - data: &DataWorkload, -) -> Option> { - let targets = &crate::accuracy_targets(entries); - let sized = crate::resize_for_failure_probability(roots, targets, data).ok()?; - Some(crate::with_confirmed_weights(&sized, data)) -} - -/// The planner pipeline behind the optimization-pass interface: the -/// workload's roots in entry order, PromQL roots with each series' full -/// identity as a column (the row representation per-series state needs at -/// runtime; a root with no such representation is planned over its labels -/// alone), through [`run_stage1_to_stage3`] with every entry's accuracy target. The -/// output is the selected physical candidate's roots, with the open weight -/// requirements plan selection confirmed made proofs. -#[derive(Debug, Default, Clone, Copy)] -pub struct Stage1ToStage3Pass; - -impl OptimizationPass for Stage1ToStage3Pass { - fn name(&self) -> &'static str { - "stage1-to-stage3" - } - - fn optimize(&self, input: OptimizationInput<'_>) -> Result { - let workload = input.workload; - let promql = workload.query_workload().language == QueryLanguage::PromQL; - let mut roots: Vec<(usize, QueryRoot)> = workload - .operator_indices() - .iter() - .zip(workload.exprs()) - .map(|(&index, node)| { - let node = match promql { - true => with_promql_series_identity(node).unwrap_or_else(|_| Rc::clone(node)), - false => Rc::clone(node), - }; - (index, QueryRoot::Operator(node)) - }) - .chain( - workload - .scalar_roots() - .iter() - .map(|(index, expr)| (*index, QueryRoot::Scalar(expr.clone()))), - ) - .collect(); - roots.sort_by_key(|(index, _)| *index); - let entries: Vec<_> = workload.query_workload().entries().collect(); - let data = workload.data_workload().cloned().unwrap_or_default(); - let run = run_stage1_to_stage3(roots, &entries, &data, input.models)?; - let selection = run.stage3_plan_selection?; - let roots = selection.plan.clone(); - let mut output = PlanOutput::new(Vec::new()); - for (entry_index, root) in roots.iter().enumerate() { - match root { - QueryRoot::Operator(node) => output.plans.push(QueryPlan { - entry_index, - root: Rc::clone(node), - }), - QueryRoot::Scalar(expr) => { - // One plan per distinct operator the scalar reads. - let mut refs: Vec<&Rc> = Vec::new(); - for node in expr.operator_refs() { - if !refs.iter().any(|r| Rc::ptr_eq(r, node)) { - refs.push(node); - } - } - output - .scalar_plans - .extend(refs.into_iter().map(|node| QueryPlan { - entry_index, - root: Rc::clone(node), - })); - output.scalar_roots.push((entry_index, expr.clone())); - } - } - } - output.selection = Some(selection); - Ok(output) - } -} diff --git a/crates/plan-selection/src/search/memo/brute_force.rs b/crates/plan-selection/src/search/memo/brute_force.rs index c29652166..caa654faf 100644 --- a/crates/plan-selection/src/search/memo/brute_force.rs +++ b/crates/plan-selection/src/search/memo/brute_force.rs @@ -18,6 +18,8 @@ use asap_physical_asap_optimization::{optimize_physical_asap, PhysicalCandidate} use asap_types::ir::{NonASAPOp, OperatorNode, QueryRoot}; use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; +use asap_logical_asap_optimization::memo::LogicalMemo; + use crate::search::memo::{groups, CouplingKind, GroupId, ImplementationCandidate, Memo}; use crate::{MemoStats, PlanSelector, PlanningModels, Selection, SelectionError}; @@ -27,7 +29,20 @@ pub(crate) fn memo( per_computation: &PerComputationCandidates, entries: &[QueryWorkloadEntry], ) -> Memo { - let mut memo = groups::memo_of(per_computation); + memo_from_initial( + LogicalMemo::from_roots(&per_computation.roots), + per_computation, + entries, + ) +} + +/// Pass 1 fills existing logical groups; Pass 2 may append groups and candidates. +pub(crate) fn memo_from_initial( + initial: LogicalMemo, + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], +) -> Memo { + let mut memo = groups::populate_pass1(initial, per_computation); crate::search::memo::rules::capability::add_shared_builds(&mut memo); crate::search::memo::rules::windows::add_pass2_candidates(&mut memo, entries); crate::search::memo::rules::scans::add_scan_widening(&mut memo); @@ -78,6 +93,12 @@ pub(crate) type GroupChoice = Vec>; /// 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. pub(crate) fn group_choices(memo: &Memo) -> Option> { + group_choices_up_to(memo, usize::MAX) +} + +/// Stop traversal at the requested number of assignments, before building the +/// remaining Cartesian product. Used only for bounded incumbent construction. +pub(crate) fn group_choices_up_to(memo: &Memo, limit: usize) -> Option> { let order = memo.topological_order()?; let parents_first: Vec = order.into_iter().rev().collect(); let n = memo.groups.len(); @@ -98,17 +119,24 @@ pub(crate) fn group_choices(memo: &Memo) -> Option> { covered: &mut [u32], choice: &mut GroupChoice, out: &mut Vec, + limit: usize, ) { + if out.len() >= limit { + return; + } 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); + walk(memo, order, at + 1, needed, covered, choice, out, limit); return; } for (index, candidate) in memo.candidates[g].iter().enumerate() { + if out.len() >= limit { + break; + } for absorbed in &candidate.implements[1..] { covered[absorbed.0 as usize] += 1; } @@ -116,7 +144,7 @@ pub(crate) fn group_choices(memo: &Memo) -> Option> { needed[child.0 as usize] += 1; } choice[g] = Some(index); - walk(memo, order, at + 1, needed, covered, choice, out); + walk(memo, order, at + 1, needed, covered, choice, out, limit); choice[g] = None; for child in &candidate.children { needed[child.0 as usize] -= 1; @@ -134,6 +162,7 @@ pub(crate) fn group_choices(memo: &Memo) -> Option> { &mut covered, &mut choice, &mut out, + limit, ); Some(out) } @@ -587,9 +616,29 @@ pub(crate) fn brute_force( data: &DataWorkload, models: PlanningModels<'_>, ) -> Result<(Selection, PhysicalCandidate, Vec), SelectionError> { + brute_force_with_initial_memo( + LogicalMemo::from_roots(&per_computation.roots), + per_computation, + entries, + data, + models, + ) +} + +pub(crate) fn brute_force_with_initial_memo( + initial: LogicalMemo, + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result<(Selection, PhysicalCandidate, Vec), SelectionError> { + let initial_build_duration = initial.build_duration; let started = std::time::Instant::now(); - let memo = memo(per_computation, entries); - let memo_stats = stats(&memo, started.elapsed().as_secs_f64() * 1e3); + let memo = memo_from_initial(initial, per_computation, entries); + let memo_stats = stats( + &memo, + (initial_build_duration + 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(); @@ -626,6 +675,35 @@ mod tests { RepeatedDemand, RepetitionInterval, SourceFacts, }; + /// A baseline cap stops before materializing a billion complete assignments. + #[test] + fn bounded_choices_do_not_expand_the_remaining_product() { + let roots: Vec<_> = (0..30) + .map(|query| { + ( + query, + QueryRoot::Operator(lower_promql( + &format!("m{query}"), + AccuracyTarget::exact(), + )), + ) + }) + .collect(); + let inventory = generate_per_computation_candidates(roots).unwrap(); + let mut memo = groups::memo_of(&inventory); + for candidates in &mut memo.candidates { + candidates.push(candidates[0].clone()); + } + assert_eq!(memo.groups.len(), 30); + assert!(group_choices_up_to(&memo, 0).unwrap().is_empty()); + let choices = group_choices_up_to(&memo, 3).unwrap(); + assert_eq!(choices.len(), 3); + assert!(choices + .iter() + .all(|choice| choice.iter().all(Option::is_some))); + assert_ne!(choices[0], choices[1]); + } + /// 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 diff --git a/crates/plan-selection/src/search/memo/groups.rs b/crates/plan-selection/src/search/memo/groups.rs index b7cf32356..22dce7b3d 100644 --- a/crates/plan-selection/src/search/memo/groups.rs +++ b/crates/plan-selection/src/search/memo/groups.rs @@ -2,7 +2,7 @@ //! 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 +//! [`asap_types::ir::cse::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 @@ -33,10 +33,10 @@ use std::collections::{HashMap, HashSet}; use std::hash::{Hash, Hasher}; use std::rc::Rc; +use asap_logical_asap_optimization::memo::LogicalMemo; use asap_logical_asap_optimization::per_computation_candidates::{ compose_logical_candidate_with_budget_share, ComputationCandidate, PerComputationCandidates, }; -use asap_types::ir::cse::share_common_sub_dags; use asap_types::ir::schema::FieldDataType; use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot}; @@ -47,99 +47,27 @@ use crate::search::memo::{CandidateOrigin, GroupId, ImplementationCandidate, Mem /// `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 } + populate_pass1( + LogicalMemo::from_roots(&per_computation.roots), + per_computation, + ) } -/// 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); - } +/// Fill the already-interned logical groups without changing their identities. +/// Pass 1 and capability pruning must retain the original roots and query order. +pub(crate) fn populate_pass1( + initial: LogicalMemo, + per_computation: &PerComputationCandidates, +) -> Memo { + assert!( + initial.matches_root_dependencies(&per_computation.roots), + "Pass 1 changed the initial memo's workload" + ); + let candidates = pass1_candidates(per_computation, &initial.groups, &initial); + Memo { + groups: initial.groups, + candidates, } - 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. @@ -151,36 +79,19 @@ fn reach(groups: &[MemoGroup], g: GroupId, reached: &mut HashSet) { } } -/// 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, + group_of: &LogicalMemo, ) -> 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 g = group_of + .group_for_node(&target.aggregate) + .expect("target belongs to logical memo"); let targets = targets_of.entry(g).or_default(); if targets .iter() @@ -248,7 +159,7 @@ fn pass1_candidates( fn target_candidates( per_computation: &PerComputationCandidates, groups: &[MemoGroup], - group_of: &GroupOf, + group_of: &LogicalMemo, g: GroupId, targets: &[usize], max_budget_share: usize, @@ -277,7 +188,9 @@ fn target_candidates( 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)]; + let absorbed = group_of + .group_for_node(&per_computation.targets[a].aggregate) + .expect("absorbed target belongs to logical memo"); if !implements.contains(&absorbed) { implements.push(absorbed); } @@ -590,6 +503,54 @@ mod tests { use asap_types::ir::ASAPOp; + /// Capability pruning must change implementations, not logical identities. + #[test] + fn initial_groups_survive_pass1_expansion_and_pruning() { + let roots: Vec<_> = [ + "quantile_over_time(0.95, m[5m])", + "quantile_over_time(0.99, m[5m])", + ] + .iter() + .enumerate() + .map(|(query, text)| { + ( + query, + QueryRoot::Operator(lower_promql(text, AccuracyTarget::max_error(0.01))), + ) + }) + .collect(); + let initial = LogicalMemo::from_roots(&roots); + let original_groups = initial.groups.clone(); + let inventory = generate_per_computation_candidates(roots).unwrap(); + let expanded = populate_pass1(initial.clone(), &inventory); + let restricted_inventory = inventory.without_unsupported_summaries( + &asap_logical_asap_optimization::deployment_capabilities::DeploymentCapabilities { + supported_summaries: Some(Vec::new()), + ..Default::default() + }, + ); + let restricted = populate_pass1(initial, &restricted_inventory); + for memo in [&expanded, &restricted] { + assert_eq!(memo.groups.len(), original_groups.len()); + for (before, after) in original_groups.iter().zip(&memo.groups) { + assert!(Rc::ptr_eq(&before.node, &after.node)); + assert_eq!(before.children, after.children); + assert_eq!(before.parents, after.parents); + assert_eq!(before.queries, after.queries); + assert_eq!(before.roots_of, after.roots_of); + } + assert_ne!( + group_of_kind(memo, 0, "Aggregate"), + group_of_kind(memo, 1, "Aggregate") + ); + assert!(memo.candidates.iter().all(|c| !c.is_empty())); + } + assert!( + expanded.candidates.iter().map(Vec::len).sum::() + > restricted.candidates.iter().map(Vec::len).sum::() + ); + } + /// The memo of PromQL `queries`, each with accuracy target `target`. fn memo(queries: &[&str], target: AccuracyTarget) -> Memo { let roots = queries diff --git a/crates/plan-selection/src/search/memo/incumbent.rs b/crates/plan-selection/src/search/memo/incumbent.rs new file mode 100644 index 000000000..df9a52f35 --- /dev/null +++ b/crates/plan-selection/src/search/memo/incumbent.rs @@ -0,0 +1,82 @@ +//! A bounded, validated memo baseline, independent of legacy search. +//! +//! This is not an optimal per-query planner: try a small prefix of local +//! Pass 1 assignments and three legal materialization policies. No full +//! workload/physical-assignment Cartesian product is materialized. + +use asap_logical_asap_optimization::memo::LogicalMemo; +use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidates; +use asap_physical_asap_optimization::{optimize_physical_asap_chosen, PhysicalCandidate}; +use asap_types::ir::QueryRoot; +use asap_types::physical::Materialization; +use asap_types::workload::{DataWorkload, QueryWorkloadEntry}; + +use super::{brute_force, groups}; +use crate::{PlanSelector, PlanningModels, SelectionError}; + +const MAX_LOGICAL_ATTEMPTS: usize = 16; + +pub(crate) fn feasible_plan( + initial: LogicalMemo, + inventory: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result)>, SelectionError> { + // Pass 2 coupling configurations belong in the joint solve, not an + // unbounded pre-solve search. Pass 1 candidates have no coupling sites. + let memo = groups::populate_pass1(initial, inventory); + let Some(choices) = brute_force::group_choices_up_to(&memo, MAX_LOGICAL_ATTEMPTS) else { + return Ok(None); + }; + let mut selector = PlanSelector::new(entries, data, models, false)?; + let mut logical_roots = std::collections::HashMap::new(); + for (attempt, choice) in choices.iter().enumerate() { + let chosen: Vec<_> = choice + .iter() + .enumerate() + .map(|(group, candidate)| candidate.map(|c| memo.candidates[group][c].clone())) + .collect(); + let Some(roots) = brute_force::composed_roots(&memo, inventory, &chosen) else { + continue; + }; + for policy in 0..3 { + let id = format!("baseline-{attempt}-{policy}"); + let mut choose = |_: *const _, _: &[usize], options: &[Materialization]| { + options + .iter() + .copied() + .find(|m| match policy { + 1 => matches!(m, Materialization::MatAtIngestion(_)), + 2 => matches!(m, Materialization::MatAtQuery(_)), + _ => false, + }) + .or_else(|| options.first().copied()) + .unwrap_or(Materialization::NotMaterialized) + }; + let Ok(Some(candidate)) = optimize_physical_asap_chosen( + &id, + &roots, + entries, + data, + models.deployment_capabilities, + &mut choose, + ) else { + continue; + }; + if selector.offer(&candidate).is_ok() && selector.keeps_from_logical(&id) { + logical_roots.insert(id, roots.clone()); + } + } + } + match selector.finish() { + Ok((_, candidate)) => { + let roots = logical_roots + .remove(&candidate.from_logical) + .expect("the validated incumbent retains its logical roots"); + Ok(Some((candidate, roots))) + } + Err(SelectionError::NoValidCandidate(_)) => Ok(None), + Err(error) => Err(error), + } +} diff --git a/crates/plan-selection/src/search/memo/mod.rs b/crates/plan-selection/src/search/memo/mod.rs index 5f3e71bfc..caba2240e 100644 --- a/crates/plan-selection/src/search/memo/mod.rs +++ b/crates/plan-selection/src/search/memo/mod.rs @@ -3,11 +3,13 @@ //! workload's distinct sub-expressions instead of whole-workload //! candidates. //! -//! [`brute_force::memo`] builds the [`Memo`] in this order: +//! The stage driver establishes logical identity before Pass 1, using +//! [`asap_logical_asap_optimization::memo::LogicalMemo`]. +//! [`brute_force::memo_from_initial`] then builds the [`Memo`] in this order: //! -//! 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. +//! 1. [`groups::populate_pass1`]: fill the initial groups with Pass 1 +//! candidates without changing group IDs. Callers that already have a +//! Pass 1 inventory can use [`groups::memo_of`] to build identity first. //! 2. [`rules::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 @@ -47,6 +49,7 @@ pub(crate) mod brute_force; pub(crate) mod groups; +pub(crate) mod incumbent; pub(crate) mod rules; use std::rc::Rc; @@ -54,25 +57,7 @@ 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, -} +pub(crate) use asap_logical_asap_optimization::memo::{MemoGroup, MemoGroupId as GroupId}; /// The hash-consed workload: its groups, children before parents. #[derive(Debug, Clone, Default)] diff --git a/crates/plan-selection/src/search/milp/clusters.rs b/crates/plan-selection/src/search/milp/clusters.rs index bbd59f202..937d15bd1 100644 --- a/crates/plan-selection/src/search/milp/clusters.rs +++ b/crates/plan-selection/src/search/milp/clusters.rs @@ -111,7 +111,7 @@ pub(crate) fn split( _ => true, }; if disk_priced && models.deployment_capabilities.raw_data_retained { - for queries in crate::search::legacy::components::source_components(roots) { + for queries in crate::search::workload_dependencies::source_connected_queries(roots) { for pair in queries.windows(2) { union(&mut parent, pair[0], pair[1]); } @@ -202,12 +202,9 @@ pub(crate) fn combine( for (&q, report) in cluster.iter().zip(&s.search.queries) { queries[q] = report.clone(); } - rejected.extend( - s.rejected - .iter() - .cloned() - .map(|r| crate::search::legacy::components::renumber(r, cluster)), - ); + rejected.extend(s.rejected.iter().cloned().map(|r| { + crate::search::workload_dependencies::remap_rejection_query_indices(r, cluster) + })); } selection.rejected = rejected; selection.search.queries = queries; diff --git a/crates/plan-selection/src/search/milp/model.rs b/crates/plan-selection/src/search/milp/model.rs index f1323a0a4..19df9f490 100644 --- a/crates/plan-selection/src/search/milp/model.rs +++ b/crates/plan-selection/src/search/milp/model.rs @@ -18,6 +18,8 @@ use good_lp::{ constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, }; +use asap_logical_asap_optimization::memo::LogicalMemo; + use crate::search::memo::{GroupId, ImplementationCandidate, Memo}; use crate::search::milp::coefficients::LazyTerms; use crate::search::milp::cuts::{SharedBuildHeld, Term, TermCuts, WidenedScan}; @@ -2157,8 +2159,31 @@ pub(crate) fn search( models: PlanningModels<'_>, time_limit_ms: Option, ) -> MemoSearchResult { - let started = std::time::Instant::now(); - let memo = crate::search::memo::brute_force::memo(per_computation, entries); + search_with_initial_memo( + LogicalMemo::from_roots(&per_computation.roots), + per_computation, + entries, + data, + models, + time_limit_ms, + ) +} + +/// Search the memo whose identities were established before Pass 1. +pub(crate) fn search_with_initial_memo( + initial: LogicalMemo, + per_computation: &PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + time_limit_ms: Option, +) -> MemoSearchResult { + // Moving logical interning before Pass 1 must not remove its cost from + // the existing memo-search budget (which excludes Pass 1 generation). + let now = std::time::Instant::now(); + let started = now.checked_sub(initial.build_duration).unwrap_or(now); + let memo = + crate::search::memo::brute_force::memo_from_initial(initial, per_computation, entries); let memo_stats = crate::search::memo::brute_force::stats(&memo, started.elapsed().as_secs_f64() * 1e3); let roots: Vec = per_computation @@ -2232,17 +2257,22 @@ fn search_clusters( .iter() .map(|&q| per_computation.roots[q].clone()) .collect(); + let initial = LogicalMemo::from_roots(&roots); let inventory = asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates(roots) .ok()? .without_unsupported_summaries(models.deployment_capabilities); let cluster_entries: Vec = queries.iter().map(|&q| entries[q].clone()).collect(); - workloads.push((inventory, cluster_entries)); + workloads.push((inventory, cluster_entries, initial)); } let mut loops = Vec::with_capacity(clusters.len()); - for (queries, (inventory, cluster_entries)) in clusters.iter().zip(&workloads) { - let memo = crate::search::memo::brute_force::memo(inventory, cluster_entries); + for (queries, (inventory, cluster_entries, initial)) in clusters.iter().zip(&workloads) { + let memo = crate::search::memo::brute_force::memo_from_initial( + initial.clone(), + inventory, + cluster_entries, + ); match ClusterLoop::new( memo, inventory, @@ -2268,7 +2298,11 @@ fn search_clusters( return Some(Err(SelectionError::NoValidCandidate( rejected .into_iter() - .map(|r| crate::search::legacy::components::renumber(r, queries)) + .map(|r| { + crate::search::workload_dependencies::remap_rejection_query_indices( + r, queries, + ) + }) .collect(), ))) } @@ -2453,33 +2487,36 @@ impl<'a, Id: Clone> ClusterLoop<'a, Id> { } /// Offer the selector the warm-start incumbent before round 1: each - /// query planned alone by the legacy search, with no sharing between + /// query planned alone by the bounded memo baseline, with no sharing between /// queries, the plans side by side. It only sets the starting /// incumbent: certification still compares the MILP lower bound with /// the incumbent, and a time-limited loop has a plan to report. A - /// query the legacy search finds no plan for leaves no warm start. + /// query the bounded baseline finds no plan for leaves no warm start. fn offer_independent_plan(&mut self) { const ID: &str = "L0"; let mut roots = Vec::with_capacity(self.entries.len()); let mut assignments = Vec::with_capacity(self.entries.len()); for (q, (_, root)) in self.per_computation.roots.iter().enumerate() { + let query_roots = vec![(q, root.clone())]; + let initial = LogicalMemo::from_roots(&query_roots); let Ok(inventory) = asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates( - vec![(q, root.clone())], + query_roots, ) else { return; }; - let run = crate::search::legacy::stage1_to_stage3::run_stage1_to_stage3_from_pass1( - inventory, + let inventory = + inventory.without_unsupported_summaries(self.models.deployment_capabilities); + // Do not call the stage driver here: its default/fallback route can + // re-enter legacy search. The bounded baseline has no MILP recursion. + let Ok(Some((selected, alone))) = crate::search::memo::incumbent::feasible_plan( + initial, + &inventory, &self.entries[q..=q], self.data, self.models, - crate::Stage1ToStage3Search::default(), - ); - let (Some(selected), Some(Ok(alone))) = - (run.stage3_selected, run.stage1_selected.map(|l| l.roots)) - else { + ) else { return; }; roots.extend(alone); diff --git a/crates/plan-selection/src/search/mod.rs b/crates/plan-selection/src/search/mod.rs index 01a7e3178..c7f29137b 100644 --- a/crates/plan-selection/src/search/mod.rs +++ b/crates/plan-selection/src/search/mod.rs @@ -15,6 +15,10 @@ pub(crate) mod legacy; pub(crate) mod memo; pub(crate) mod milp; +/// Costed reuse of compatible physical pane producers. +pub mod pane_sharing; +pub(crate) mod pipeline; +pub(crate) mod workload_dependencies; use std::collections::BTreeMap; use std::rc::Rc; @@ -36,7 +40,7 @@ use crate::{PlanSelectionSettings, Selection, SelectionError}; use asap_logical_asap_optimization::deployment_capabilities::DeploymentCapabilities; use asap_logical_asap_optimization::per_computation_candidates::PerComputationCandidateError; -pub use legacy::stage1_to_stage3::{ +pub use pipeline::{ run_stage1_to_stage3, run_stage1_to_stage3_from_pass1, run_stage1_to_stage3_with, run_stage2_and_stage3, LogicalCandidate, PrunedCandidate, Stage1ToStage3Pass, Stage1ToStage3Run, Stage1ToStage3Search, diff --git a/crates/plan-selection/src/search/legacy/pane_sharing.rs b/crates/plan-selection/src/search/pane_sharing.rs similarity index 100% rename from crates/plan-selection/src/search/legacy/pane_sharing.rs rename to crates/plan-selection/src/search/pane_sharing.rs diff --git a/crates/plan-selection/src/search/pipeline.rs b/crates/plan-selection/src/search/pipeline.rs new file mode 100644 index 000000000..8207c26db --- /dev/null +++ b/crates/plan-selection/src/search/pipeline.rs @@ -0,0 +1,411 @@ +//! Stage orchestration shared by memo selection and the temporary legacy route. +//! +//! On memo routes, logical identity is established before Pass 1; deployment +//! pruning and candidate population retain those group IDs. This module owns +//! the public stage driver rather than an individual search implementation. + +use std::rc::Rc; + +use asap_logical_asap_optimization::memo::LogicalMemo; + +use asap_types::ir::schema::schema_support::with_promql_series_identity; +use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::workload::{DataWorkload, QueryLanguage, QueryWorkloadEntry}; + +use crate::search::{OptimizationInput, OptimizationPass, OptimizeError, PlanOutput, QueryPlan}; +use crate::PlanningModels; +use crate::{CostModel, RoutedToLegacy, SearchAlgorithm, Selection, SelectionError}; +use asap_logical_asap_optimization::per_computation_candidates::{ + generate_per_computation_candidates, PerComputationCandidateError, PerComputationCandidates, +}; +use asap_logical_asap_optimization::workload_candidates::Pass2Candidate; +use asap_physical_asap_optimization::PhysicalCandidate; + +use crate::search::legacy::stage1_to_stage3::legacy_search; +pub use crate::search::legacy::stage1_to_stage3::run_stage2_and_stage3; + +/// One whole-workload CandidateLogicalASAPDAG: `choice[t]` picks target +/// `t`'s candidate, or its fused computation's +/// (`targets[t].computation_for(choice[t])`), of the Pass 1 inventory. `roots` are the +/// composed query roots in input order, or why the choice did not compose. +#[derive(Debug)] +pub struct LogicalCandidate { + /// `L`, numbered in enumeration order. + pub id: String, + pub choice: Vec, + pub roots: Result, PerComputationCandidateError>, + /// For a Pass 2 candidate: what it changed of its Pass 1 candidate. + pub pass2: Option, +} + +/// A composed logical candidate that physical ASAP-aware optimization could +/// not implement, with the reason. +#[derive(Debug, Clone, PartialEq)] +pub struct PrunedCandidate { + pub from_logical: String, + pub reason: String, +} + +/// The output of Stage 1, Stage 2 and Stage 3 for one workload +/// ([`run_stage1_to_stage3`]). The candidate lists are filled only when the +/// search keeps every candidate ([`Stage1ToStage3Search::keep_every_candidate`]). +#[derive(Debug)] +pub struct Stage1ToStage3Run { + /// Stage 1, Pass 1: the per-computation candidates, with the roots they + /// read. + pub stage1_pass1_per_computation: PerComputationCandidates, + /// Stage 1, Pass 1: the number of composed whole-workload candidates, + /// the size of the Cartesian product of the targets' candidates, + /// saturating at `usize::MAX`. Pass 2 adds to them. + pub combinations: usize, + /// Stage 1: every whole-workload logical candidate, in enumeration + /// order: each composed Pass 1 candidate, followed by its Pass 2 + /// candidates (those with [`LogicalCandidate::pass2`]). Empty unless + /// every candidate is kept. + pub stage1_logical_asap_optimization: Vec, + /// Stage 2: the physical candidates of each Stage 1 candidate that + /// composed and was implemented, one per materialization assignment: + /// `P` (nothing materialized), then `P-m`, from `L`. Empty + /// unless every candidate is kept. + pub stage2_physical_asap_optimization: Vec, + /// Stage 2: the Stage 1 candidates Stage 2 could not implement, with the + /// reasons. Empty unless every candidate is kept. + pub stage2_pruned: Vec, + /// Stage 3: plan selection over the Stage 2 candidates. + pub stage3_plan_selection: Result, + /// Stage 3: the selected physical candidate, as Stage 2 built it + /// ([`Self::selected`]). + pub stage3_selected: Option, + /// Stage 3: the selected candidate's logical roots, with the weight + /// requirements it confirmed made proofs ([`Self::selected_logical`]). + pub stage3_selected_logical: Option>, + /// Stage 1: the logical candidate the selected physical candidate + /// implements, with its roots as Stage 1 composed them (before Stage 3 + /// sized them): what Stage 2 implemented. + pub stage1_selected: Option, +} + +impl Stage1ToStage3Run { + /// The selected physical candidate, if selection succeeded. + pub fn selected(&self) -> Option<&PhysicalCandidate> { + self.stage3_selected.as_ref() + } + + /// The composed roots of the logical candidate the selected physical + /// candidate implements: the plan before physical operator + /// implementation and materialization, with the weight requirements + /// Stage 3 confirmed made proofs, as in [`Selection::plan`]. + pub fn selected_logical(&self) -> Option<&[QueryRoot]> { + self.stage3_selected_logical.as_deref() + } +} + +/// How [`run_stage1_to_stage3_with`] searches, and what it keeps. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct Stage1ToStage3Search { + /// Keep every Stage 1 and Stage 2 candidate and every Stage 3 rejection + /// in the run, for a developer tool that shows them; default `false`: + /// only the best so far is kept. + pub keep_every_candidate: bool, + /// Compose every whole-workload choice, without checking each query's + /// own candidates first: the exhaustive search the pruned one is tested + /// against; default `false`. + pub exhaustive: 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, + /// 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::MilpReport::gap`]) and + /// [`crate::MilpReport::time_limit_reached`]: not proven + /// optimal unless the gap is 0. + pub milp_time_limit_ms: Option, +} + +impl Stage1ToStage3Search { + /// The exhaustive search keeping every candidate: what the planner did + /// before the search was pruned. + pub fn exhaustive() -> Self { + Self { + keep_every_candidate: true, + exhaustive: true, + search: SearchAlgorithm::Legacy, + milp_time_limit_ms: None, + } + } + + /// The default search run with `algorithm`. + pub fn with_algorithm(algorithm: SearchAlgorithm) -> Self { + Self { + search: algorithm, + ..Self::default() + } + } +} + +/// Run Stage 1 → Stage 2 → Stage 3 over the frontends' `roots` with the +/// default search ([`Stage1ToStage3Search::default`]: each query's +/// candidates checked first, only the best kept). `entries[i]` is the +/// workload entry of `roots[i]`: its accuracy target, latency bound, +/// recurrence and predictability. Only a Pass 1 failure is an error; a +/// failed selection is reported in the run. +pub fn run_stage1_to_stage3( + roots: Vec<(Id, QueryRoot)>, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, +) -> Result, PerComputationCandidateError> { + run_stage1_to_stage3_with( + roots, + entries, + data, + models, + Stage1ToStage3Search::default(), + ) +} + +/// [`run_stage1_to_stage3`] with the search `search`. Keeping every +/// candidate, a candidate that does not compose or cannot be implemented is +/// listed with its reason. +pub fn run_stage1_to_stage3_with( + roots: Vec<(Id, QueryRoot)>, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + search: Stage1ToStage3Search, +) -> Result, PerComputationCandidateError> { + // Legacy routes do not pay for memo identity they will never use. + let uses_memo = !search.exhaustive + && match search.search { + SearchAlgorithm::Legacy => false, + SearchAlgorithm::MemoBruteForce => true, + SearchAlgorithm::MemoMilp => matches!( + models.plan_selection.cost_model, + CostModel::UsageOverTime(_) + ), + }; + if uses_memo { + for (_, root) in &roots { + root.validate_structure()?; + } + } + let initial = uses_memo.then(|| LogicalMemo::from_roots(&roots)); + let per_computation = generate_per_computation_candidates(roots)?; + Ok(run_with_initial_memo( + per_computation, + entries, + data, + models, + search, + initial, + )) +} + +/// [`run_stage1_to_stage3_with`] after Pass 1: searches the Stage 1 +/// candidates of `per_computation` (the Pass 1 inventory of the frontends' +/// roots, [`asap_logical_asap_optimization::workload_candidates::stage1_logical_asap_optimization`]) without the summaries the +/// deployment does not support +/// ([`PerComputationCandidates::without_unsupported_summaries`]), so a +/// caller can read the inventory, e.g. its +/// [`combinations`](PerComputationCandidates::combinations), before +/// enumerating. +pub fn run_stage1_to_stage3_from_pass1( + per_computation: PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + search: Stage1ToStage3Search, +) -> Stage1ToStage3Run { + run_with_initial_memo(per_computation, entries, data, models, search, None) +} + +fn run_with_initial_memo( + per_computation: PerComputationCandidates, + entries: &[QueryWorkloadEntry], + data: &DataWorkload, + models: PlanningModels<'_>, + search: Stage1ToStage3Search, + initial: Option, +) -> Stage1ToStage3Run { + 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 = crate::search::memo::brute_force::brute_force_with_initial_memo( + initial.unwrap_or_else(|| LogicalMemo::from_roots(&pruned.roots)), + &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); + let found = crate::search::milp::model::search_with_initial_memo( + initial.unwrap_or_else(|| LogicalMemo::from_roots(&pruned.roots)), + &pruned, + entries, + data, + models, + search.milp_time_limit_ms, + ); + return memo_run(pruned, found, entries, data, SearchAlgorithm::MemoMilp); + } + } + }; + 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: crate::search::milp::model::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)); + let from_logical = selected.as_ref().map(|s| s.from_logical.clone()); + 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, + // 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, + }), + } +} + +/// `roots`, the selected candidate's Stage 1 roots, as Stage 3 sized them +/// for their queries' failure probabilities, with the weight requirements +/// it confirmed made proofs. +pub(crate) fn selected_logical( + roots: &[QueryRoot], + entries: &[QueryWorkloadEntry], + data: &DataWorkload, +) -> Option> { + let targets = &crate::accuracy_targets(entries); + let sized = crate::resize_for_failure_probability(roots, targets, data).ok()?; + Some(crate::with_confirmed_weights(&sized, data)) +} + +/// The planner pipeline behind the optimization-pass interface: the +/// workload's roots in entry order, PromQL roots with each series' full +/// identity as a column (the row representation per-series state needs at +/// runtime; a root with no such representation is planned over its labels +/// alone), through [`run_stage1_to_stage3`] with every entry's accuracy target. The +/// output is the selected physical candidate's roots, with the open weight +/// requirements plan selection confirmed made proofs. +#[derive(Debug, Default, Clone, Copy)] +pub struct Stage1ToStage3Pass; + +impl OptimizationPass for Stage1ToStage3Pass { + fn name(&self) -> &'static str { + "stage1-to-stage3" + } + + fn optimize(&self, input: OptimizationInput<'_>) -> Result { + let workload = input.workload; + let promql = workload.query_workload().language == QueryLanguage::PromQL; + let mut roots: Vec<(usize, QueryRoot)> = workload + .operator_indices() + .iter() + .zip(workload.exprs()) + .map(|(&index, node)| { + let node = match promql { + true => with_promql_series_identity(node).unwrap_or_else(|_| Rc::clone(node)), + false => Rc::clone(node), + }; + (index, QueryRoot::Operator(node)) + }) + .chain( + workload + .scalar_roots() + .iter() + .map(|(index, expr)| (*index, QueryRoot::Scalar(expr.clone()))), + ) + .collect(); + roots.sort_by_key(|(index, _)| *index); + let entries: Vec<_> = workload.query_workload().entries().collect(); + let data = workload.data_workload().cloned().unwrap_or_default(); + let run = run_stage1_to_stage3(roots, &entries, &data, input.models)?; + let selection = run.stage3_plan_selection?; + let roots = selection.plan.clone(); + let mut output = PlanOutput::new(Vec::new()); + for (entry_index, root) in roots.iter().enumerate() { + match root { + QueryRoot::Operator(node) => output.plans.push(QueryPlan { + entry_index, + root: Rc::clone(node), + }), + QueryRoot::Scalar(expr) => { + // One plan per distinct operator the scalar reads. + let mut refs: Vec<&Rc> = Vec::new(); + for node in expr.operator_refs() { + if !refs.iter().any(|r| Rc::ptr_eq(r, node)) { + refs.push(node); + } + } + output + .scalar_plans + .extend(refs.into_iter().map(|node| QueryPlan { + entry_index, + root: Rc::clone(node), + })); + output.scalar_roots.push((entry_index, expr.clone())); + } + } + } + output.selection = Some(selection); + Ok(output) + } +} diff --git a/crates/plan-selection/src/search/workload_dependencies.rs b/crates/plan-selection/src/search/workload_dependencies.rs new file mode 100644 index 000000000..5cd4a0f34 --- /dev/null +++ b/crates/plan-selection/src/search/workload_dependencies.rs @@ -0,0 +1,81 @@ +//! Source dependencies and query-index mapping shared by planning strategies. + +use crate::Rejection; +use asap_types::ir::operator::operator_properties::Source; +use asap_types::ir::{NonASAPOp, Operator, OperatorNode, QueryRoot}; +use std::collections::HashMap; +use std::rc::Rc; + +/// Queries connected by reads from the same raw source. +pub(crate) fn source_connected_queries(roots: &[QueryRoot]) -> Vec> { + let mut parent: Vec = (0..roots.len()).collect(); + fn find(parent: &mut [usize], q: usize) -> usize { + let mut root = q; + while parent[root] != root { + root = parent[root]; + } + let mut q = q; + while parent[q] != root { + let next = parent[q]; + parent[q] = root; + q = next; + } + root + } + // The first query read each source, by name (data facts are keyed by + // name, whatever the source's kind). + let mut first_reader: HashMap = HashMap::new(); + for (query, root) in roots.iter().enumerate() { + let starts: Vec<&Rc> = match root { + QueryRoot::Operator(node) => vec![node], + QueryRoot::Scalar(expr) => expr.operator_refs(), + }; + for node in starts.into_iter().flat_map(OperatorNode::reachable) { + let Operator::NonASAP(NonASAPOp::Scan { source, .. }) = &node.operator else { + continue; + }; + let name = match source { + Source::TimeSeries { metric } => metric.clone(), + Source::Table { table_ref } => table_ref.clone(), + }; + let other = *first_reader.entry(name).or_insert(query); + let (a, b) = (find(&mut parent, query), find(&mut parent, other)); + parent[a.max(b)] = a.min(b); + } + } + let mut components: Vec> = Vec::new(); + let mut index_of: HashMap = HashMap::new(); + for query in 0..roots.len() { + let root = find(&mut parent, query); + let index = *index_of.entry(root).or_insert_with(|| { + components.push(Vec::new()); + components.len() - 1 + }); + components[index].push(query); + } + components +} + +/// Map component-local rejection indices back to workload indices. +pub(crate) fn remap_rejection_query_indices( + rejection: Rejection, + component: &[usize], +) -> Rejection { + let rename = |text: &str| -> String { + let Some(rest) = text.strip_prefix('q') else { + return text.to_string(); + }; + let digits: String = rest.chars().take_while(char::is_ascii_digit).collect(); + match digits.parse::() { + Ok(n) if n >= 1 && n <= component.len() => { + format!("q{}{}", component[n - 1] + 1, &rest[digits.len()..]) + } + _ => text.to_string(), + } + }; + Rejection { + id: rename(&rejection.id), + cause: rejection.cause, + reason: rename(&rejection.reason), + } +} From 7e1c90f0363f8fd8227955dd068e2530eb907fd2 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:26:04 +0000 Subject: [PATCH 02/12] feat: add candidate connection MILP core and migration plan --- .../src/search/milp/candidate_connections.rs | 753 ++++++++++++++++++ crates/plan-selection/src/search/milp/mod.rs | 3 + docs/design_docs/proposals/README.md | 2 + .../joint-planning-implementation-plan.md | 143 ++++ 4 files changed, 901 insertions(+) create mode 100644 crates/plan-selection/src/search/milp/candidate_connections.rs create mode 100644 docs/design_docs/proposals/joint-planning-implementation-plan.md diff --git a/crates/plan-selection/src/search/milp/candidate_connections.rs b/crates/plan-selection/src/search/milp/candidate_connections.rs new file mode 100644 index 000000000..d2ffd6f64 --- /dev/null +++ b/crates/plan-selection/src/search/milp/candidate_connections.rs @@ -0,0 +1,753 @@ +//! Fixed-context selection over explicit physical candidates and connections. +//! +//! This is the selection core for the joint-planning migration, not a new +//! production search entry point. Candidate generation must supply executable +//! candidates and compatible producers. Costs cover a fixed accounting context; +//! demand-dependent costs, accuracy refinement, materialization policy, resource +//! limits, and physical IR decoding still belong to the surrounding pipeline. + +use std::collections::HashSet; + +use asap_logical_asap_optimization::memo::MemoGroupId; +use good_lp::solvers::highs::highs; +use good_lp::solvers::{ResolutionError, SolutionStatus}; +use good_lp::{ + constraint, variable, Expression, ProblemVariables, Solution, SolverModel, Variable, +}; + +/// Index into [`CandidateDependencyGraph::candidates`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub(crate) struct PhysicalCandidateId(pub usize); + +/// One required input, distinct even when another input accepts the same producers. +#[derive(Debug, Clone)] +pub(crate) struct ProducerRequirement { + pub compatible_producers: Vec, +} + +/// One executable choice with costs already expressed over the same fixed horizon. +#[derive(Debug, Clone)] +pub(crate) struct PhysicalCandidate { + /// Logical ownership does not impose one selected candidate per group. + pub group: MemoGroupId, + /// Local work and state only; producer costs must not appear here again. + pub local_cost: f64, + pub producer_requirements: Vec, +} + +/// Candidates allowed to provide one required query output. +#[derive(Debug, Clone)] +pub(crate) struct QueryRootCandidates { + pub compatible_candidates: Vec, +} + +#[derive(Debug, Clone, Default)] +pub(crate) struct CandidateDependencyGraph { + pub candidates: Vec, + pub query_roots: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ProducerConsumerConnection { + pub producer: PhysicalCandidateId, + pub consumer: PhysicalCandidateId, + /// Index in the consumer's `producer_requirements`. + pub requirement: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct QueryRootAssignment { + /// Index in the graph's `query_roots`. + pub query_root: usize, + pub candidate: PhysicalCandidateId, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct SelectedCandidateGraph { + pub selected_candidates: Vec, + pub connections: Vec, + pub query_roots: Vec, + pub modeled_cost: f64, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum CandidateSelectionError { + #[error("invalid candidate dependency graph: {0}")] + InvalidGraph(String), + #[error("candidate selection solver: {0}")] + Solver(String), +} + +// The core normalizes its largest local cost to one. Smaller positive +// coefficients must stay representable by HiGHS; rejecting a wider range is an +// explicit model restriction, not proof of infeasibility or a completeness claim. +const MINIMUM_RELATIVE_COST: f64 = 1e-9; + +impl CandidateDependencyGraph { + fn validate(&self) -> Result<(), CandidateSelectionError> { + let validate_ids = |ids: &[PhysicalCandidateId], context: &str| { + let mut seen = HashSet::new(); + for id in ids { + if id.0 >= self.candidates.len() { + return Err(CandidateSelectionError::InvalidGraph(format!( + "{context} references missing candidate {}", + id.0 + ))); + } + if !seen.insert(*id) { + return Err(CandidateSelectionError::InvalidGraph(format!( + "{context} repeats candidate {}", + id.0 + ))); + } + } + Ok(()) + }; + let mut total_cost = 0.0; + for (index, candidate) in self.candidates.iter().enumerate() { + if !candidate.local_cost.is_finite() || candidate.local_cost < 0.0 { + return Err(CandidateSelectionError::InvalidGraph(format!( + "candidate {index} needs a finite nonnegative local cost" + ))); + } + total_cost += candidate.local_cost; + for (requirement, input) in candidate.producer_requirements.iter().enumerate() { + validate_ids( + &input.compatible_producers, + &format!("candidate {index}, requirement {requirement}"), + )?; + } + } + if !total_cost.is_finite() { + return Err(CandidateSelectionError::InvalidGraph( + "total candidate cost overflows".into(), + )); + } + let maximum_cost = self + .candidates + .iter() + .map(|c| c.local_cost) + .fold(0.0, f64::max); + if maximum_cost > 0.0 + && self.candidates.iter().any(|candidate| { + candidate.local_cost > 0.0 + && candidate.local_cost / maximum_cost < MINIMUM_RELATIVE_COST + }) + { + return Err(CandidateSelectionError::InvalidGraph(format!( + "positive local costs must be at least {MINIMUM_RELATIVE_COST} of the largest cost; wider ranges require a different numerical policy" + ))); + } + for (index, root) in self.query_roots.iter().enumerate() { + validate_ids(&root.compatible_candidates, &format!("query root {index}"))?; + } + Ok(()) + } +} + +/// Return the cheapest covering DAG, or `None` when no such graph exists. +/// Numerical optimality is subject to HiGHS tolerances. No production timeout or +/// incumbent fallback is implied by this fixed-context core. +pub(crate) fn select_candidates( + graph: &CandidateDependencyGraph, +) -> Result, CandidateSelectionError> { + graph.validate()?; + if graph.query_roots.is_empty() { + return Ok(Some(SelectedCandidateGraph { + selected_candidates: Vec::new(), + connections: Vec::new(), + query_roots: Vec::new(), + modeled_cost: 0.0, + })); + } + if graph + .query_roots + .iter() + .any(|root| root.compatible_candidates.is_empty()) + { + return Ok(None); + } + + let count = graph.candidates.len(); + let mut variables = ProblemVariables::new(); + let candidate_selected: Vec = (0..count) + .map(|_| variables.add(variable().binary())) + .collect(); + let topological_rank: Vec = (0..count) + .map(|_| variables.add(variable().min(0).max((count - 1) as f64))) + .collect(); + let mut connections = Vec::new(); + let mut root_assignments = Vec::new(); + let mut constraints = Vec::new(); + let mut support = vec![Expression::from(0.0); count]; + + for (consumer, candidate) in graph.candidates.iter().enumerate() { + for (requirement, input) in candidate.producer_requirements.iter().enumerate() { + let mut producers_selected = Expression::from(0.0); + for &producer in &input.compatible_producers { + let connection_selected = variables.add(variable().binary()); + producers_selected += connection_selected; + support[producer.0] += connection_selected; + constraints.push(constraint!( + connection_selected <= candidate_selected[producer.0] + )); + constraints.push(constraint!( + topological_rank[producer.0] + 1.0 + <= topological_rank[consumer] + count as f64 * (1.0 - connection_selected) + )); + connections.push(( + ProducerConsumerConnection { + producer, + consumer: PhysicalCandidateId(consumer), + requirement, + }, + connection_selected, + )); + } + constraints.push(constraint!( + producers_selected == candidate_selected[consumer] + )); + } + } + for (query_root, root) in graph.query_roots.iter().enumerate() { + let mut root_selected = Expression::from(0.0); + for &candidate in &root.compatible_candidates { + let assignment_selected = variables.add(variable().binary()); + root_selected += assignment_selected; + support[candidate.0] += assignment_selected; + constraints.push(constraint!( + assignment_selected <= candidate_selected[candidate.0] + )); + root_assignments.push(( + QueryRootAssignment { + query_root, + candidate, + }, + assignment_selected, + )); + } + constraints.push(constraint!(root_selected == 1)); + } + for candidate in 0..count { + constraints.push(constraint!( + candidate_selected[candidate] <= support[candidate].clone() + )); + } + // Validation rejects positive coefficients outside the supported relative + // range; uniform currency-unit changes do not alter the selected graph. + let maximum_cost = graph + .candidates + .iter() + .map(|c| c.local_cost) + .fold(0.0, f64::max); + let divisor = if maximum_cost > 0.0 { + maximum_cost + } else { + 1.0 + }; + let objective: Expression = graph + .candidates + .iter() + .enumerate() + .map(|(index, candidate)| (candidate.local_cost / divisor) * candidate_selected[index]) + .sum(); + let model = variables + .minimise(objective) + .using(highs) + .set_option("mip_rel_gap", 0.0) + .set_option("mip_abs_gap", 0.0) + .with_all(constraints); + let solution = match model.solve() { + Ok(solution) => solution, + Err(ResolutionError::Infeasible) => return Ok(None), + Err(error) => return Err(CandidateSelectionError::Solver(error.to_string())), + }; + if !matches!(solution.status(), SolutionStatus::Optimal) { + return Err(CandidateSelectionError::Solver(format!( + "search stopped with {:?}", + solution.status() + ))); + } + let selected_candidates: Vec<_> = candidate_selected + .iter() + .enumerate() + .filter(|(_, variable)| solution.value(**variable) > 0.5) + .map(|(index, _)| PhysicalCandidateId(index)) + .collect(); + let modeled_cost = selected_candidates + .iter() + .map(|id| graph.candidates[id.0].local_cost) + .sum(); + let selected = SelectedCandidateGraph { + selected_candidates, + connections: connections + .into_iter() + .filter(|(_, selected)| solution.value(*selected) > 0.5) + .map(|(connection, _)| connection) + .collect(), + query_roots: root_assignments + .into_iter() + .filter(|(_, selected)| solution.value(*selected) > 0.5) + .map(|(assignment, _)| assignment) + .collect(), + modeled_cost, + }; + validate_selected_graph(graph, &selected).map_err(CandidateSelectionError::Solver)?; + Ok(Some(selected)) +} + +/// Check the decoded integral graph instead of trusting rounded solver values. +fn validate_selected_graph( + graph: &CandidateDependencyGraph, + selected: &SelectedCandidateGraph, +) -> Result<(), String> { + let invalid = + || "decoded candidate selection violates dependency or root constraints".to_string(); + let active: HashSet<_> = selected.selected_candidates.iter().copied().collect(); + if active.len() != selected.selected_candidates.len() + || active.iter().any(|c| c.0 >= graph.candidates.len()) + || selected.query_roots.len() != graph.query_roots.len() + { + return Err(invalid()); + } + let mut root_seen = vec![false; graph.query_roots.len()]; + for root in &selected.query_roots { + let Some(allowed) = graph.query_roots.get(root.query_root) else { + return Err(invalid()); + }; + if root_seen[root.query_root] + || !active.contains(&root.candidate) + || !allowed.compatible_candidates.contains(&root.candidate) + { + return Err(invalid()); + } + root_seen[root.query_root] = true; + } + let mut producers: Vec>> = graph + .candidates + .iter() + .map(|c| vec![None; c.producer_requirements.len()]) + .collect(); + for connection in &selected.connections { + if !active.contains(&connection.producer) || !active.contains(&connection.consumer) { + return Err(invalid()); + } + let Some(requirement) = graph.candidates[connection.consumer.0] + .producer_requirements + .get(connection.requirement) + else { + return Err(invalid()); + }; + let target = &mut producers[connection.consumer.0][connection.requirement]; + if target.is_some() + || !requirement + .compatible_producers + .contains(&connection.producer) + { + return Err(invalid()); + } + *target = Some(connection.producer); + } + if active + .iter() + .any(|c| producers[c.0].iter().any(Option::is_none)) + { + return Err(invalid()); + } + // Kahn's algorithm checks cycles without recursion on a potentially deep DAG. + let mut outstanding: Vec<_> = producers.iter().map(Vec::len).collect(); + let mut consumers = vec![Vec::new(); graph.candidates.len()]; + for connection in &selected.connections { + consumers[connection.producer.0].push(connection.consumer.0); + } + let mut ready: Vec<_> = active + .iter() + .filter(|c| outstanding[c.0] == 0) + .map(|c| c.0) + .collect(); + let mut visited = 0; + while let Some(producer) = ready.pop() { + visited += 1; + for &consumer in &consumers[producer] { + outstanding[consumer] -= 1; + if outstanding[consumer] == 0 { + ready.push(consumer); + } + } + } + if visited != active.len() { + return Err(invalid()); + } + let mut reachable = HashSet::new(); + let mut pending: Vec<_> = selected + .query_roots + .iter() + .map(|root| root.candidate) + .collect(); + while let Some(candidate) = pending.pop() { + if reachable.insert(candidate) { + pending.extend(producers[candidate.0].iter().flatten().copied()); + } + } + if reachable != active { + return Err(invalid()); + } + let expected_cost: f64 = active + .iter() + .map(|c| graph.candidates[c.0].local_cost) + .sum(); + if !selected.modeled_cost.is_finite() + || (selected.modeled_cost - expected_cost).abs() > 1e-9 * expected_cost.abs() + { + return Err("decoded candidate cost disagrees with selected candidates".into()); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn candidate(group: u32, cost: f64, inputs: &[&[usize]]) -> PhysicalCandidate { + PhysicalCandidate { + group: MemoGroupId(group), + local_cost: cost, + producer_requirements: inputs + .iter() + .map(|sources| ProducerRequirement { + compatible_producers: sources + .iter() + .copied() + .map(PhysicalCandidateId) + .collect(), + }) + .collect(), + } + } + + fn dependency_graph( + candidates: Vec, + roots: &[&[usize]], + ) -> CandidateDependencyGraph { + CandidateDependencyGraph { + candidates, + query_roots: roots + .iter() + .map(|choices| QueryRootCandidates { + compatible_candidates: choices + .iter() + .copied() + .map(PhysicalCandidateId) + .collect(), + }) + .collect(), + } + } + + /// Independent small-input oracle: enumerate roots and edges, then traverse + /// their reachable graph. It uses DFS cycle detection, not rank constraints. + fn exhaustive_cost(graph: &CandidateDependencyGraph) -> Option { + let edges: Vec<_> = graph + .candidates + .iter() + .enumerate() + .flat_map(|(consumer, c)| { + c.producer_requirements + .iter() + .enumerate() + .flat_map(move |(requirement, input)| { + input + .compatible_producers + .iter() + .map(move |producer| (consumer, requirement, producer.0)) + }) + }) + .collect(); + let roots: Vec<_> = graph + .query_roots + .iter() + .enumerate() + .flat_map(|(q, root)| root.compatible_candidates.iter().map(move |c| (q, c.0))) + .collect(); + let bits = edges.len() + roots.len(); + assert!( + bits <= 20, + "oracle is deliberately restricted to tiny graphs" + ); + let mut best: Option = None; + for assignment in 0..(1usize << bits) { + let chosen = |bit: usize| assignment & (1usize << bit) != 0; + let selected_roots: Vec<_> = roots + .iter() + .enumerate() + .filter(|(i, _)| chosen(edges.len() + i)) + .map(|(_, root)| *root) + .collect(); + if (0..graph.query_roots.len()) + .any(|q| selected_roots.iter().filter(|(root, _)| *root == q).count() != 1) + { + continue; + } + let mut dependencies = vec![Vec::new(); graph.candidates.len()]; + for (index, &(consumer, requirement, producer)) in edges.iter().enumerate() { + if chosen(index) { + dependencies[consumer].push((requirement, producer)); + } + } + fn visit( + candidate: usize, + graph: &CandidateDependencyGraph, + dependencies: &[Vec<(usize, usize)>], + marks: &mut [u8], + ) -> bool { + if marks[candidate] == 1 { + return false; + } + if marks[candidate] == 2 { + return true; + } + marks[candidate] = 1; + for requirement in 0..graph.candidates[candidate].producer_requirements.len() { + let producers: Vec<_> = dependencies[candidate] + .iter() + .filter(|(r, _)| *r == requirement) + .collect(); + if producers.len() != 1 || !visit(producers[0].1, graph, dependencies, marks) { + return false; + } + } + marks[candidate] = 2; + true + } + let mut marks = vec![0; graph.candidates.len()]; + if !selected_roots + .iter() + .all(|(_, root)| visit(*root, graph, &dependencies, &mut marks)) + { + continue; + } + if dependencies + .iter() + .enumerate() + .any(|(c, edges)| marks[c] == 0 && !edges.is_empty()) + { + continue; + } + let cost = graph + .candidates + .iter() + .zip(&marks) + .filter(|(_, mark)| **mark == 2) + .map(|(c, _)| c.local_cost) + .sum::(); + best = Some(best.map_or(cost, |previous| previous.min(cost))); + } + best + } + + fn compare_with_oracle(graph: &CandidateDependencyGraph) -> Option { + let expected = exhaustive_cost(graph); + let actual = select_candidates(graph).unwrap(); + match (&actual, expected) { + (Some(plan), Some(expected)) => { + assert!((plan.modeled_cost - expected).abs() <= 1e-8 * expected.abs().max(1.0)) + } + (None, None) => {} + _ => panic!("MILP {actual:?}, exhaustive {expected:?}"), + } + actual + } + + /// One activated build serves two query roots without paying per connection. + #[test] + fn shared_build_cost_is_owned_once() { + let graph = dependency_graph( + vec![ + candidate(0, 8.0, &[]), + candidate(1, 1.0, &[&[0]]), + candidate(2, 1.0, &[&[0]]), + candidate(1, 9.0, &[]), + candidate(2, 9.0, &[]), + ], + &[&[1, 3], &[2, 4]], + ); + let selected = compare_with_oracle(&graph).unwrap(); + assert_eq!(selected.modeled_cost, 10.0); + assert_eq!(selected.connections.len(), 2); + assert!(selected + .connections + .iter() + .all(|c| c.producer == PhysicalCandidateId(0))); + } + + /// Distinct consumer contracts can require two physical choices of one group. + #[test] + fn one_group_can_have_multiple_selected_candidates() { + let graph = dependency_graph( + vec![ + candidate(0, 3.0, &[]), + candidate(0, 5.0, &[]), + candidate(1, 1.0, &[&[0]]), + candidate(2, 1.0, &[&[1]]), + ], + &[&[2], &[3]], + ); + let selected = compare_with_oracle(&graph).unwrap(); + assert_eq!(selected.modeled_cost, 10.0); + assert_eq!( + selected + .selected_candidates + .iter() + .filter(|c| graph.candidates[c.0].group == MemoGroupId(0)) + .count(), + 2 + ); + } + + /// A join needs both requirements, even when one producer supplies both. + #[test] + fn every_requirement_has_a_producer() { + let graph = dependency_graph( + vec![candidate(0, 2.0, &[]), candidate(1, 1.0, &[&[0], &[0]])], + &[&[1]], + ); + let selected = compare_with_oracle(&graph).unwrap(); + assert_eq!(selected.modeled_cost, 3.0); + assert_eq!( + selected + .connections + .iter() + .map(|c| c.requirement) + .collect::>(), + vec![0, 1] + ); + let impossible = dependency_graph(vec![candidate(0, 1.0, &[&[]])], &[&[0]]); + assert!(compare_with_oracle(&impossible).is_none()); + } + + /// Root-reachable cycles are infeasible; unused zero-cost cycles cannot activate. + #[test] + fn cycles_and_unused_nodes_are_excluded() { + let graph = dependency_graph( + vec![ + candidate(0, 0.0, &[]), + candidate(1, 0.0, &[&[2]]), + candidate(2, 0.0, &[&[1]]), + candidate(3, 0.0, &[]), + ], + &[&[0]], + ); + let selected = compare_with_oracle(&graph).unwrap(); + assert_eq!(selected.selected_candidates, vec![PhysicalCandidateId(0)]); + let cyclic = dependency_graph(graph.candidates, &[&[1]]); + assert!(compare_with_oracle(&cyclic).is_none()); + let self_cycle = dependency_graph(vec![candidate(0, 0.0, &[&[0]])], &[&[0]]); + assert!(compare_with_oracle(&self_cycle).is_none()); + } + + /// Dense compatibility choices agree with exhaustive graph selection across costs. + #[test] + fn competing_producers_match_exhaustive_enumeration() { + for build_cost in [0.0, 1.0, 4.0, 12.0] { + for local_cost in [0.5, 3.0, 8.0] { + let graph = dependency_graph( + vec![ + candidate(0, build_cost, &[]), + candidate(0, 2.0, &[]), + candidate(1, 0.25, &[&[0, 1]]), + candidate(1, local_cost, &[]), + candidate(2, 0.25, &[&[0, 1]]), + candidate(2, local_cost, &[]), + ], + &[&[2, 3], &[4, 5]], + ); + compare_with_oracle(&graph); + } + } + } + + /// Currency-unit scaling must not erase the cheaper candidate's advantage. + #[test] + fn small_uniform_cost_units_preserve_selection() { + let graph = dependency_graph( + vec![candidate(0, 3e-15, &[]), candidate(0, 1e-15, &[])], + &[&[0, 1]], + ); + let selected = select_candidates(&graph).unwrap().unwrap(); + assert_eq!(selected.selected_candidates, vec![PhysicalCandidateId(1)]); + assert_eq!(selected.modeled_cost, 1e-15); + } + + /// An unrelated large candidate must not erase relevant root-choice costs. + #[test] + fn unsupported_cost_dynamic_range_is_rejected() { + let graph = dependency_graph( + vec![ + candidate(0, 1e300, &[]), + candidate(1, 1e-300, &[]), + candidate(1, 2e-300, &[]), + ], + &[&[1, 2]], + ); + assert!(matches!( + select_candidates(&graph), + Err(CandidateSelectionError::InvalidGraph(_)) + )); + } + + /// Decoder validation catches missing requirements and unused selected nodes. + #[test] + fn decoded_graph_is_checked_independently_of_solver_values() { + let graph = dependency_graph( + vec![ + candidate(0, 2.0, &[]), + candidate(1, 1.0, &[&[0]]), + candidate(2, 0.0, &[]), + ], + &[&[1]], + ); + let selected = select_candidates(&graph).unwrap().unwrap(); + assert!(validate_selected_graph(&graph, &selected).is_ok()); + let mut missing_input = selected.clone(); + missing_input.connections.clear(); + assert!(validate_selected_graph(&graph, &missing_input).is_err()); + let mut unused = selected.clone(); + unused.selected_candidates.push(PhysicalCandidateId(2)); + assert!(validate_selected_graph(&graph, &unused).is_err()); + let mut bad_root = selected; + bad_root.query_roots[0].query_root = 9; + assert!(validate_selected_graph(&graph, &bad_root).is_err()); + } + + /// Invalid indices and numeric costs fail before entering the native solver. + #[test] + fn malformed_graphs_return_errors() { + for cost in [f64::NAN, f64::INFINITY, -1.0] { + assert!(matches!( + select_candidates(&dependency_graph(vec![candidate(0, cost, &[])], &[&[0]])), + Err(CandidateSelectionError::InvalidGraph(_)) + )); + } + for invalid in [ + dependency_graph(vec![candidate(0, 1.0, &[&[1]])], &[&[0]]), + dependency_graph(vec![candidate(0, 1.0, &[])], &[&[1]]), + dependency_graph(vec![candidate(0, 1.0, &[&[0, 0]])], &[&[0]]), + dependency_graph(vec![candidate(0, 1.0, &[])], &[&[0, 0]]), + ] { + assert!(matches!( + select_candidates(&invalid), + Err(CandidateSelectionError::InvalidGraph(_)) + )); + } + assert_eq!( + select_candidates(&CandidateDependencyGraph::default()) + .unwrap() + .unwrap() + .modeled_cost, + 0.0 + ); + assert!(select_candidates(&dependency_graph(Vec::new(), &[&[]])) + .unwrap() + .is_none()); + } +} diff --git a/crates/plan-selection/src/search/milp/mod.rs b/crates/plan-selection/src/search/milp/mod.rs index 8897ce86d..4a5eea02a 100644 --- a/crates/plan-selection/src/search/milp/mod.rs +++ b/crates/plan-selection/src/search/milp/mod.rs @@ -9,3 +9,6 @@ pub(crate) mod cuts; pub(crate) mod decode; pub(crate) mod model; pub(crate) mod report; + +// Candidate-connection core being integrated into the staged memo pipeline. +pub(crate) mod candidate_connections; diff --git a/docs/design_docs/proposals/README.md b/docs/design_docs/proposals/README.md index 6c4ed49f8..9c64d556f 100644 --- a/docs/design_docs/proposals/README.md +++ b/docs/design_docs/proposals/README.md @@ -12,3 +12,5 @@ extensions. A design document is not a promise of downstream runtime support. - [Decoupling operators from scalar expressions](decoupling_op_and_expr.md) - [ASAPPlanner layering](planner-layering.md) - [Schema and physical data for ASAP primitives](asap-primitive-schema.md) + +- [Joint planning implementation plan](joint-planning-implementation-plan.md) diff --git a/docs/design_docs/proposals/joint-planning-implementation-plan.md b/docs/design_docs/proposals/joint-planning-implementation-plan.md new file mode 100644 index 000000000..51d4edcd0 --- /dev/null +++ b/docs/design_docs/proposals/joint-planning-implementation-plan.md @@ -0,0 +1,143 @@ +# Joint Planning Implementation Plan + +Status: implementation in progress. The first draft implements the logical memo/stage boundary and extracts shared infrastructure; it does not yet replace the selection model or remove legacy search. + +## Outcome + +Replace the legacy whole-workload search with one memo-based planning pipeline. Build logical group identity before Stage 1, generate candidates into those groups, expose reusable state through state groups, enumerate executable physical candidates, and select a connected execution DAG jointly. The finished production pipeline must not call or fall back to legacy search. + +This is a migration of supported behavior, not just a solver replacement. Keep existing accuracy semantics, physical eligibility checks, cost objectives, scalar outputs, and execution export valid. New runtime algorithms and scheduling are outside scope. No latency or speedup target is claimed before benchmarks. + +## Current gaps + +Current `run_stage1_to_stage3_with` generates `PerComputationCandidates` before memo construction. `groups::memo_of` constructs identity and fills candidates together. Shared builds and panes use `CouplingSite` choices; MILP `Choice.copy_of` largely represents all-shared or per-query copies. The proposed producer-state groups and arbitrary partial-sharing connections are not implemented. + +Legacy is currently more than a selectable algorithm. The stage driver and public exports live under `search/legacy`; memo clustering uses its source-component and renumbering helpers; the MILP baseline calls legacy planning. Unsupported cost objectives and exhaustive requests also route there. These dependencies must be removed before deleting the directory. + +## Naming and ownership + +Names should identify domain concepts without requiring readers to decode mathematical notation. Do not introduce synonyms for the same concept or rename unrelated identifiers in bulk. The following are proposed responsibilities, not a requirement to create a new type for every row: + +| Concept | Human-readable name | Responsibility | +| --- | --- | --- | +| Logical identity | `MemoGroup`, `MemoGroupId` | Requested result or reusable state contract. | +| Logical computation choice | `LogicalCandidate` | Rule-generated computation with logical dependencies. | +| Executable choice | `PhysicalCandidate`, `PhysicalCandidateId` | Supported operators, configuration, materialization, and execution scope. | +| Consumer dependency | `ProducerRequirement` | One required value/state, including distinct left/right dependencies when applicable. | +| Compatible supplier set | `compatible_producers` | Physical candidates satisfying one consumer requirement. | +| Selected dependency | `ProducerConsumerConnection` | `producer`, `consumer`, and `requirement`. | +| Query coverage | `QueryRootAssignment` | Which selected candidate supplies a query output. | +| Evaluation demand | `QueryEvaluationSchedule`, `ExecutionDemand` | Frequency/alignment, reuse context, and derived execution demand. | +| Selected graph | `SelectedExecutionPlan` | Selected candidates, connections, root assignments, and validated output. | +| Solver activation | `candidate_selected` | Solver variable associated with a physical candidate. | +| Solver connection | `connection_selected` | Solver variable associated with a compatible connection. | +| Cycle exclusion | `topological_rank` | Dependency rank, never execution time. | +| Event-based demand, if needed | `executes_at_evaluation`, `connection_required_at_evaluation` | Clearly named auxiliaries for query evaluation demand. | +| Independent reference | `enumerate_feasible_plans` | Small-input correctness oracle over the same candidate graph. | + +The current `PhysicalCandidate` means a whole-workload DAG. Resolve that collision explicitly: rename that wrapper to `PhysicalWorkloadPlan` (or reuse an existing equivalent type) before using `PhysicalCandidate` for a local choice. Migration must update callers and documentation together. Prefer preserving a clear existing name to adding aliases indefinitely. + +Logical memo data and rule expansion belong in `logical-asap-optimization`, which Stage 2 already depends on. Physical candidates and materialization legality belong in `physical-asap-optimization`; prices and selection belong in `plan-selection`. Earlier stages must not depend on the solver crate. Move common driver/report types out of the legacy namespace without adding a new framework crate merely to share names. + +## Implementation sequence + +### 1. Capture behavior and extract shared infrastructure + +Inventory supported frontends, scalar/operator roots, sharing rules, materializations, and all configured objectives. Convert representative current outputs into acceptance fixtures containing legality and expected costs, allowing equivalent plan shapes and numerical ties. Record generation/solve/runtime-validation baselines before changing representation. + +Move the stage driver, source-component analysis, query-index remapping, and reusable pane functionality out of `search/legacy`. Preserve public behavior while removing memo code's dependency on legacy utility modules. Extract behavior, not legacy search control. + +Acceptance: existing tests pass; every remaining legacy dependency is an explicit migration item. No circular crate dependencies or duplicate authoritative implementations. + +### 2. Establish logical groups before Pass 1 + +Split root interning and group identity from candidate population. Construct the initial memo immediately after frontend logical validation. Preserve filters, windows, accuracy requirements, keyless sharing restrictions, scalar references, and query identity. Pass 1 expands those groups directly, with cache keys including relevant consumer requirements. + +Group references remain stable throughout expansion, either directly or through canonical-ID resolution after a proven-equivalence merge. Stage 0 creates only the initial groups: Pass 1 and Pass 2 can create intermediate groups, reuse existing groups, and add candidates. Approximate candidates target a result contract and do not assert exact-value equivalence. Compatibility requirements remain explicit when consumer contracts differ. + +Acceptance: group identity exists before candidate generation; identical eligible subexpressions share generation work; P95 and P99 remain distinct; pruning changes candidates without invalidating references. Test literal/scalar queries, multiple consumers, fusion boundaries, and canonical-ID resolution after group merging. + +### 3. Expose reusable state and generate physical candidates + +Pass 2 expands the memo rather than operating on a fixed set of Stage-0 groups: + +| Rule effect | Required behavior | +| --- | --- | +| Discover a shared build, pane chain, or rollup state. | Create or reuse a state group and add consumer candidates referencing it. | +| Prove two existing groups logically equivalent. | Merge their canonical identities, reconcile candidates and references, and invalidate affected caches. Sharing a physical component is not proof of logical equivalence. | +| Fuse or bypass a computation in one candidate. | Change that candidate's dependencies; preserve the group if other candidates or query roots still need it. | +| Remove unreachable memo content. | Garbage-collect only when no retained candidate or query root can reach it. A group absent from one selected plan is not automatically removable from the search space. | + +For example: + +```text +Before Pass 2: + G95 = {ExactP95, LocalKLLQuantileP95} + G99 = {ExactP99, LocalKLLQuantileP99} + +After discovering reusable KLL state: + G_KLL = {BuildKLL(...)} + G95 += EstimateP95(G_KLL) + G99 += EstimateP99(G_KLL) +``` + +`G_KLL` is new or reused; `G95` and `G99` remain distinct, and their local candidates remain available. New groups and candidates feed the rule worklist so later matches can use them. A group merge must reconcile producer requirements, parent references, query roots, and rule-application caches; it must not silently equate different accuracy or state contracts. Stage 3 selects the reachable execution subgraph rather than deleting every unselected group from the memo. + +Move shared KLL builds, pane chains, and raw retention into explicit state groups. Consumers acquire producer requirements instead of carrying implicit duplicate producer costs. Preserve scan-widening, rollup, top-k, and fusion semantics; only exposed outputs may serve external consumers. + +Stage 2 generates legal local physical candidates with concrete configuration, internal materialization assignments, and finite execution scopes. Generate sparse compatible producer connections. Specify scope/binding distinctions needed for multiple candidates from one group to coexist. Do not enumerate every consumer subset or claim completeness without a bounded generation rule. Any restricted scope inventory must be reported. + +Acceptance: P95/P99 can share one compatible KLL candidate, select distinct capacities, or stay local; timing and accuracy incompatibilities prevent connections. A fused result does not accidentally expose absorbed intermediate outputs. Eligibility agrees with the current physical evaluator. + +### 4. Implement joint selection and reconstruction + +Add candidate selection, root assignments, producer requirements, connection activation, reverse support, acyclicity, and resource-specific configuration exclusion. Permit several selected physical candidates per logical group. Reconstruct the selected graph directly into the supported physical IR, preserving distinct execution scopes and deduplicating only the selected shared candidate. + +Build an independent tiny-input enumerator over candidate subsets and connections. It is a test oracle, not the old whole-workload search retained as a fallback. Exercise the actual MILP with the same finite candidates and costs. + +Acceptance: MILP and enumeration agree on optimal cost for small graphs; shared producer cost is charged once per owned workload, each required input is present, cycles and unsupported outputs are rejected, and zero-cost candidates cannot remain disconnected. + +### 5. Complete workload costs and feasibility + +Migrate existing per-node pricing, union-rate logic, and sound term cuts into the new representation. Derive demand from selected consumers. Prefer compact schedule classes or proven rate formulas to materializing every query evaluation event; validate event equivalence before compression. Keep a tiny explicit-event reference for aligned/offset schedule tests. + +Implement separate demand and lifetime handling for query-time recomputation, ingestion maintenance, and supported retained state. Own source retention in shared resource candidates. Keep fixed request overhead out of the selection objective but in reporting and applicable budgets. + +Preserve all existing objective modes: usage-over-time, provisioned capacity, and cluster capacity, including its configured latency/memory priorities and Pareto behavior. Pareto output requires appropriate repeated solves or frontier exploration, not one arbitrary weighted objective. Encode valid accuracy/latency bounds directly where available; retain evaluator refinement with sound local cuts and full-assignment rejection only as a fallback. An unresolved objective cannot silently route to legacy or disappear from the API. + +Acceptance: objective and resource accounting match the configured evaluator; shared/repeated execution costs match explicit-event examples; all existing supported objective modes and materializations have acceptance coverage. No unsupported retained-state semantics are inferred from query-time equations. + +### 6. Replace production entry points and delete legacy + +Build a validated feasible baseline from the same candidate graph, using bounded constructive or feasibility search; do not call legacy to obtain it. Return the best validated incumbent under the planning budget. If none exists, return an explicit no-plan/timeout status rather than relaxing requirements. Budget generation, model construction, solve, and validation, with distinct reporting of solver timeout and proven infeasibility. + +Switch the production driver to the new memo pipeline. Remove `SearchAlgorithm::Legacy`, `RoutedToLegacy`, legacy CLI modes, implicit cost-model/exhaustive fallback, and legacy baseline calls. Replace any retained diagnostic enumeration with the new candidate-graph enumerator and explicit size limits. Delete `search/legacy` after its still-needed functionality has moved; remove temporary adapters and old exhaustive selection machinery. + +Update stage-pipeline tooling, fuzz entry points, reports, serialized configuration handling, examples, and tests. Obsolete options should receive a clear migration error, not silently select different behavior. No permanent legacy feature flag is part of the final design. + +Acceptance: repository search finds no active legacy algorithm, fallback, imports, or user-facing selection mode. Production tests cover all migrated objectives. New exhaustive oracle and recorded acceptance fixtures replace legacy-based differential tests. + +### 7. Measure scalability and tune + +Benchmark little-sharing workloads, many consumers of few states, dense compatibility graphs, many window/configuration combinations, and evaluator-heavy constraints. Report candidate generation time, model construction time, cumulative MILP solve time, validation time/refinement rounds, peak planner memory, validated execution cost, and total planning wall time. Include solver status/gap and inventory restrictions. + +Compare against recorded pre-migration measurements and independent per-query plans generated by the new pipeline. Investigate candidate count, connection density, demand representation, and weak cuts before adding more solver machinery. Keep runtime execution measurements separate from modeled execution cost. + +Acceptance: correctness and feasibility hold under time limits; no unbounded event expansion or consumer-partition enumeration is hidden in the default path. Publish measured tradeoffs without assuming the richer model always plans faster. + +## Delivery policy + +Use reviewable commits following the sequence above. Temporary comparison against legacy during migration is permitted only before the final replacement; the delivery is not complete until legacy production code and fallback have been removed. Do not delete an objective or supported behavior merely to make the migration smaller. Keep the scope and any deferred capability explicit at each checkpoint. + +The first implementation change should establish the stage/memo ownership boundary and its regression tests, not introduce a second permanent standalone solver. Subsequent changes should attach to that pipeline. Code names, errors, reports, and tests must use the same producer/consumer vocabulary as the design. + +## First draft status + +- Implemented: `LogicalMemo`, `MemoGroup`, and `MemoGroupId` in the logical optimization crate, independent of solver types. +- Implemented: initialize logical groups before Pass 1 on memo routes; populate those groups after capability pruning; apply existing Pass 2 rules afterward. The from-Pass-1 API remains supported. +- Implemented: shared stage orchestration, source-dependency analysis, diagnostic index mapping, and pane reuse moved out of the legacy namespace. +- Implemented: an internal fixed-context candidate-connection MILP core, including same-group multi-selection and an independent tiny-graph enumerator. Production IR decoding and workload-dependent cost/constraint adapters are pending; this core is not yet a replacement production search mode. +- Verified: initial group identity survives Pass 1 population and summary pruning; scalar/literal roots retain query positions; distinct quantile results remain distinct groups. +- Pending: Pass 1 rule generation directly into memo groups, shared-state contracts, arbitrary candidate connections, complete cost-model migration, bounded baseline generation, and deletion of legacy search/fallback. The baseline now tries at most 16 local memo assignments and three materialization policies per query, without calling the legacy route or materializing all combinations. A strict wall-clock budget across generation and evaluators remains pending. + +The current first draft preserves legacy runtime behavior while preparing its removal. It must not be described as completion of the full implementation plan. From 3d9e7725e4133de4b89ffa952a12da6755a3d324 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:28:12 +0000 Subject: [PATCH 03/12] test: account for frontend nodes in bounded memo enumeration regression --- crates/plan-selection/src/search/memo/brute_force.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/plan-selection/src/search/memo/brute_force.rs b/crates/plan-selection/src/search/memo/brute_force.rs index caa654faf..550750949 100644 --- a/crates/plan-selection/src/search/memo/brute_force.rs +++ b/crates/plan-selection/src/search/memo/brute_force.rs @@ -694,7 +694,8 @@ mod tests { for candidates in &mut memo.candidates { candidates.push(candidates[0].clone()); } - assert_eq!(memo.groups.len(), 30); + // PromQL lowering may introduce additional logical nodes per root. + assert!(memo.groups.len() >= 30); assert!(group_choices_up_to(&memo, 0).unwrap().is_empty()); let choices = group_choices_up_to(&memo, 3).unwrap(); assert_eq!(choices.len(), 3); From 59740e72f836e016711e58df0f3e57bc271b508d Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:33:57 +0000 Subject: [PATCH 04/12] refactor: own memo candidate population in logical optimization --- crates/logical-asap-optimization/src/memo.rs | 4 + .../src/memo/candidates.rs | 129 ++ .../src/memo/groups.rs | 1232 +++++++++++++++++ .../src/search/memo/brute_force.rs | 8 +- .../plan-selection/src/search/memo/groups.rs | 1160 +--------------- crates/plan-selection/src/search/memo/mod.rs | 111 +- .../joint-planning-implementation-plan.md | 5 +- 7 files changed, 1388 insertions(+), 1261 deletions(-) create mode 100644 crates/logical-asap-optimization/src/memo/candidates.rs create mode 100644 crates/logical-asap-optimization/src/memo/groups.rs diff --git a/crates/logical-asap-optimization/src/memo.rs b/crates/logical-asap-optimization/src/memo.rs index 2d8bec3ea..f3b6051aa 100644 --- a/crates/logical-asap-optimization/src/memo.rs +++ b/crates/logical-asap-optimization/src/memo.rs @@ -4,6 +4,10 @@ //! legality rules and owns the original roots so pointer-based lookup remains //! valid while later stages populate implementation candidates. +/// Logical implementation candidates populated after initial identity construction. +pub mod candidates; +mod groups; + use std::collections::{HashMap, HashSet}; use std::rc::Rc; use std::time::{Duration, Instant}; diff --git a/crates/logical-asap-optimization/src/memo/candidates.rs b/crates/logical-asap-optimization/src/memo/candidates.rs new file mode 100644 index 000000000..eb4afad13 --- /dev/null +++ b/crates/logical-asap-optimization/src/memo/candidates.rs @@ -0,0 +1,129 @@ +//! Logical implementation candidates attached to stable memo groups. +//! +//! Pass 1 population owns no pricing or solver state. Existing Pass 2 rule +//! metadata remains here during migration; state groups will replace the +//! implicit shared-state decisions in a subsequent change. + +use std::rc::Rc; + +use asap_types::ir::OperatorNode; +use asap_types::physical::Materialization; + +pub use super::groups::{try_populate_pass1, NodeKeys}; +use super::{MemoGroup, MemoGroupId}; + +/// The hash-consed workload: its groups, children before parents. +#[derive(Debug, Clone, Default)] +pub struct CandidateMemo { + 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 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, + /// 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, + /// 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, + /// 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`) + /// This is an eligibility restriction, not a selected assignment. + pub materializations: Vec, + /// 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<(MemoGroupId, CouplingSite)>, +} + +/// 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 struct RollupRead { + /// The finer group it reads. + pub finer_group: MemoGroupId, + /// 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. +#[derive(Debug, Clone, PartialEq, Eq)] +pub 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, + /// 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. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub 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 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, +} + +/// A Pass 1 inventory does not belong to the initial logical memo. +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum CandidatePopulationError { + #[error("Pass 1 root dependencies differ from the initial logical memo")] + RootDependenciesChanged, + #[error("Pass 1 target {target} is not a valid measure of an initial aggregate group")] + InvalidTarget { target: usize }, + #[error("Pass 1 target {target} references missing absorbed target {absorbed}")] + InvalidAbsorption { target: usize, absorbed: usize }, +} diff --git a/crates/logical-asap-optimization/src/memo/groups.rs b/crates/logical-asap-optimization/src/memo/groups.rs new file mode 100644 index 000000000..cba7af7dc --- /dev/null +++ b/crates/logical-asap-optimization/src/memo/groups.rs @@ -0,0 +1,1232 @@ +//! 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 +//! [`asap_types::ir::cse::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 +//! [`crate::per_computation_candidates::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. +//! 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'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 +//! 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}; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use super::LogicalMemo; +use crate::per_computation_candidates::{ + compose_logical_candidate_with_budget_share, ComputationCandidate, PerComputationCandidates, +}; +use asap_types::ir::schema::FieldDataType; +use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot}; + +use super::candidates::{ + CandidateMemo, CandidateOrigin, CandidatePopulationError, ImplementationCandidate, +}; +use super::{MemoGroup, MemoGroupId}; + +/// Populate implementation candidates without changing initial group identities. +/// The inventory must retain the original ordered root dependencies. This phase +/// consumes existing Pass 1 choices; it does not yet generate rules per group. +pub fn try_populate_pass1( + initial: LogicalMemo, + per_computation: &PerComputationCandidates, +) -> Result { + if !initial.matches_root_dependencies(&per_computation.roots) { + return Err(CandidatePopulationError::RootDependenciesChanged); + } + for (index, target) in per_computation.targets.iter().enumerate() { + if initial.group_for_node(&target.aggregate).is_none() + || !matches!(target.aggregate.non_asap(), Some(asap_types::ir::NonASAPOp::Aggregate { measures, .. }) if target.measure < measures.len()) + { + return Err(CandidatePopulationError::InvalidTarget { target: index }); + } + if let Some(fused) = &target.fused { + if let Some(&absorbed) = fused + .absorbs + .iter() + .find(|&&absorbed| absorbed >= per_computation.targets.len()) + { + return Err(CandidatePopulationError::InvalidAbsorption { + target: index, + absorbed, + }); + } + } + } + let candidates = pass1_candidates(per_computation, &initial.groups, &initial); + Ok(CandidateMemo { + groups: initial.groups, + candidates, + }) +} + +/// `g` and every group below it. +fn reach(groups: &[MemoGroup], g: MemoGroupId, reached: &mut HashSet) { + if reached.insert(g) { + for &child in &groups[g.0 as usize].children { + reach(groups, child, reached); + } + } +} + +/// Every group's Pass 1 candidates, by group index. +fn pass1_candidates( + per_computation: &PerComputationCandidates, + groups: &[MemoGroup], + group_of: &LogicalMemo, +) -> 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 + .group_for_node(&target.aggregate) + .expect("target belongs to logical memo"); + 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![MemoGroupId(g as u32)], + node: Rc::clone(&groups[g].node), + 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, + requires_sites: Vec::new(), + }; + // 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(&MemoGroupId(g as u32)) { + None => vec![own(g)], + Some(targets) => target_candidates( + per_computation, + groups, + group_of, + MemoGroupId(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)]), + }) + .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: &LogicalMemo, + g: MemoGroupId, + 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. + 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 { + 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 + .group_for_node(&per_computation.targets[a].aggregate) + .expect("absorbed target belongs to logical memo"); + if !implements.contains(&absorbed) { + implements.push(absorbed); + } + } + } + } + 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) + 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 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; + 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 { + schema: group_node.schema.clone(), + ..(*node).clone() + }) + }; + (node, children) + }; + candidates.push(ImplementationCandidate { + implements: implements.clone(), + node, + children, + origin: CandidateOrigin::Pass1, + coupling_sites: Vec::new(), + materializations: Vec::new(), + budget_share, + requires_exact_children: Vec::new(), + rollup_of: None, + requires_sites: Vec::new(), + }); + } + // 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) +} + +/// 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 +/// computed then). +fn compose_alone( + per_computation: &PerComputationCandidates, + aggregate: &Rc, + targets: &[usize], + choice: &[usize], + budget_share: usize, +) -> Result, crate::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_with_budget_share(&alone, &alone_choice, budget_share) +} + +/// `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, + rebuilt: &mut HashMap<*const OperatorNode, Rc>, +) -> Rc { + let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); + for child in node.children() { + let key = Rc::as_ptr(child); + if replaced.contains_key(&key) { + continue; + } + 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 => 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, +/// 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)] +pub struct NodeKeys { + memo: HashMap<*const OperatorNode, (Rc, u64)>, +} + +impl NodeKeys { + pub 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 + } +} + +impl CandidateMemo { + /// 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 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(MemoGroupId(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::*; + + fn populate_pass1( + initial: LogicalMemo, + inventory: &PerComputationCandidates, + ) -> CandidateMemo { + try_populate_pass1(initial, inventory).unwrap() + } + + fn memo_of(inventory: &PerComputationCandidates) -> CandidateMemo { + populate_pass1(LogicalMemo::from_roots(&inventory.roots), inventory) + } + use crate::per_computation_candidates::{ + compose_logical_candidate, generate_per_computation_candidates, + }; + use crate::test_support::lower_promql; + use asap_types::ir::properties::AccuracyTarget; + + use asap_types::ir::ASAPOp; + + /// The cross-stage API rejects mismatched root identities without panicking. + #[test] + fn population_rejects_an_inventory_from_other_roots() { + let roots = vec![( + 0, + QueryRoot::Operator(lower_promql( + "quantile_over_time(0.95, m[5m])", + AccuracyTarget::max_error(0.01), + )), + )]; + let initial = LogicalMemo::from_roots(&roots); + let other_roots = vec![( + 0, + QueryRoot::Operator(lower_promql( + "quantile_over_time(0.95, m[5m])", + AccuracyTarget::max_error(0.01), + )), + )]; + let inventory = generate_per_computation_candidates(other_roots).unwrap(); + assert!(matches!( + try_populate_pass1(initial, &inventory), + Err(CandidatePopulationError::RootDependenciesChanged) + )); + } + + /// Invalid target and fusion indices fail at the stage boundary, not in composition. + #[test] + fn population_checks_target_and_absorption_indices() { + let roots = vec![( + 0, + QueryRoot::Operator(lower_promql( + "quantile_over_time(0.95, m[5m])", + AccuracyTarget::max_error(0.01), + )), + )]; + let initial = LogicalMemo::from_roots(&roots); + let inventory = generate_per_computation_candidates(roots).unwrap(); + let mut invalid_measure = inventory.clone(); + invalid_measure.targets[0].measure = usize::MAX; + assert!(matches!( + try_populate_pass1(initial.clone(), &invalid_measure), + Err(CandidatePopulationError::InvalidTarget { target: 0 }) + )); + let mut invalid_absorption = inventory; + let missing = invalid_absorption.targets.len(); + invalid_absorption.targets[0].fused = + Some(crate::per_computation_candidates::FusedComputation { + absorbs: vec![missing], + candidates: Vec::new(), + pruned: Vec::new(), + }); + assert!(matches!(try_populate_pass1(initial, &invalid_absorption), + Err(CandidatePopulationError::InvalidAbsorption { target: 0, absorbed }) if absorbed == missing)); + } + + /// Capability pruning must change implementations, not logical identities. + #[test] + fn initial_groups_survive_pass1_expansion_and_pruning() { + let roots: Vec<_> = [ + "quantile_over_time(0.95, m[5m])", + "quantile_over_time(0.99, m[5m])", + ] + .iter() + .enumerate() + .map(|(query, text)| { + ( + query, + QueryRoot::Operator(lower_promql(text, AccuracyTarget::max_error(0.01))), + ) + }) + .collect(); + let initial = LogicalMemo::from_roots(&roots); + let original_groups = initial.groups.clone(); + let inventory = generate_per_computation_candidates(roots).unwrap(); + let expanded = populate_pass1(initial.clone(), &inventory); + let restricted_inventory = inventory.without_unsupported_summaries( + &crate::deployment_capabilities::DeploymentCapabilities { + supported_summaries: Some(Vec::new()), + ..Default::default() + }, + ); + let restricted = populate_pass1(initial, &restricted_inventory); + for memo in [&expanded, &restricted] { + assert_eq!(memo.groups.len(), original_groups.len()); + for (before, after) in original_groups.iter().zip(&memo.groups) { + assert!(Rc::ptr_eq(&before.node, &after.node)); + assert_eq!(before.children, after.children); + assert_eq!(before.parents, after.parents); + assert_eq!(before.queries, after.queries); + assert_eq!(before.roots_of, after.roots_of); + } + assert_ne!( + group_of_kind(memo, 0, "Aggregate"), + group_of_kind(memo, 1, "Aggregate") + ); + assert!(memo.candidates.iter().all(|c| !c.is_empty())); + } + assert!( + expanded.candidates.iter().map(Vec::len).sum::() + > restricted.candidates.iter().map(Vec::len).sum::() + ); + } + + /// The memo of PromQL `queries`, each with accuracy target `target`. + fn memo(queries: &[&str], target: AccuracyTarget) -> CandidateMemo { + 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: &CandidateMemo, 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![MemoGroupId(product as u32)]); + assert_eq!(memo.groups[product].children, vec![MemoGroupId(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![MemoGroupId(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![MemoGroupId(topk as u32), sums]); + assert_eq!(candidate.children, sums_input); + 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 [`crate::per_computation_candidates::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: &CandidateMemo, 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(MemoGroupId) + .collect::>() + ); + let quantile = group_of_kind(&memo, 0, "Aggregate"); + let appended = MemoGroupId(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: MemoGroupId| order.iter().position(|&x| x == g).unwrap(); + assert!(position(appended) < position(MemoGroupId(quantile as u32))); + // A cycle has no order. + memo.candidates[1][0] + .children + .push(MemoGroupId(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 CandidateMemo, + 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: &[MemoGroupId], + ) -> 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 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; + 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); + // 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 { + 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 crate::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![MemoGroupId(1), MemoGroupId(2)]); + } +} diff --git a/crates/plan-selection/src/search/memo/brute_force.rs b/crates/plan-selection/src/search/memo/brute_force.rs index 550750949..eefd23290 100644 --- a/crates/plan-selection/src/search/memo/brute_force.rs +++ b/crates/plan-selection/src/search/memo/brute_force.rs @@ -114,7 +114,6 @@ pub(crate) fn group_choices_up_to(memo: &Memo, limit: usize) -> Option Option= limit { return; } - let Some(&GroupId(g)) = order.get(at) else { + let Some(&GroupId(g)) = order.first() 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, limit); + walk(memo, &order[1..], needed, covered, choice, out, limit); return; } for (index, candidate) in memo.candidates[g].iter().enumerate() { @@ -144,7 +143,7 @@ pub(crate) fn group_choices_up_to(memo: &Memo, limit: usize) -> Option Option(per_computation: &PerComputationCandidates) -> Memo { - populate_pass1( - LogicalMemo::from_roots(&per_computation.roots), - per_computation, - ) +pub(crate) fn memo_of(inventory: &PerComputationCandidates) -> Memo { + populate_pass1(LogicalMemo::from_roots(&inventory.roots), inventory) } -/// Fill the already-interned logical groups without changing their identities. -/// Pass 1 and capability pruning must retain the original roots and query order. +/// Internal callers generate and prune this inventory from the retained roots. pub(crate) fn populate_pass1( initial: LogicalMemo, - per_computation: &PerComputationCandidates, + inventory: &PerComputationCandidates, ) -> Memo { - assert!( - initial.matches_root_dependencies(&per_computation.roots), - "Pass 1 changed the initial memo's workload" - ); - let candidates = pass1_candidates(per_computation, &initial.groups, &initial); - Memo { - groups: initial.groups, - candidates, - } -} - -/// `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); - } - } -} - -/// Every group's Pass 1 candidates, by group index. -fn pass1_candidates( - per_computation: &PerComputationCandidates, - groups: &[MemoGroup], - group_of: &LogicalMemo, -) -> 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 - .group_for_node(&target.aggregate) - .expect("target belongs to logical memo"); - 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(), - materializations: Vec::new(), - 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. - 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)], - Some(targets) => target_candidates( - per_computation, - groups, - 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)]), - }) - .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: &LogicalMemo, - 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. - 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 { - 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 - .group_for_node(&per_computation.targets[a].aggregate) - .expect("absorbed target belongs to logical memo"); - if !implements.contains(&absorbed) { - implements.push(absorbed); - } - } - } - } - 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) - 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 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; - 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 { - schema: group_node.schema.clone(), - ..(*node).clone() - }) - }; - (node, children) - }; - candidates.push(ImplementationCandidate { - implements: implements.clone(), - node, - children, - origin: CandidateOrigin::Pass1, - coupling_sites: Vec::new(), - materializations: Vec::new(), - budget_share, - requires_exact_children: Vec::new(), - rollup_of: None, - requires_sites: Vec::new(), - }); - } - // 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) -} - -/// 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 -/// computed then). -fn compose_alone( - per_computation: &PerComputationCandidates, - aggregate: &Rc, - targets: &[usize], - choice: &[usize], - budget_share: 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_with_budget_share(&alone, &alone_choice, budget_share) -} - -/// `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, - rebuilt: &mut HashMap<*const OperatorNode, Rc>, -) -> Rc { - let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); - for child in node.children() { - let key = Rc::as_ptr(child); - if replaced.contains_key(&key) { - continue; - } - 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 => 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, -/// 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)] -pub(crate) struct NodeKeys { - memo: HashMap<*const OperatorNode, (Rc, u64)>, -} - -impl NodeKeys { - pub(crate) 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 - } -} - -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::{ - compose_logical_candidate, generate_per_computation_candidates, - }; - use asap_logical_asap_optimization::test_support::lower_promql; - use asap_types::ir::properties::AccuracyTarget; - - use asap_types::ir::ASAPOp; - - /// Capability pruning must change implementations, not logical identities. - #[test] - fn initial_groups_survive_pass1_expansion_and_pruning() { - let roots: Vec<_> = [ - "quantile_over_time(0.95, m[5m])", - "quantile_over_time(0.99, m[5m])", - ] - .iter() - .enumerate() - .map(|(query, text)| { - ( - query, - QueryRoot::Operator(lower_promql(text, AccuracyTarget::max_error(0.01))), - ) - }) - .collect(); - let initial = LogicalMemo::from_roots(&roots); - let original_groups = initial.groups.clone(); - let inventory = generate_per_computation_candidates(roots).unwrap(); - let expanded = populate_pass1(initial.clone(), &inventory); - let restricted_inventory = inventory.without_unsupported_summaries( - &asap_logical_asap_optimization::deployment_capabilities::DeploymentCapabilities { - supported_summaries: Some(Vec::new()), - ..Default::default() - }, - ); - let restricted = populate_pass1(initial, &restricted_inventory); - for memo in [&expanded, &restricted] { - assert_eq!(memo.groups.len(), original_groups.len()); - for (before, after) in original_groups.iter().zip(&memo.groups) { - assert!(Rc::ptr_eq(&before.node, &after.node)); - assert_eq!(before.children, after.children); - assert_eq!(before.parents, after.parents); - assert_eq!(before.queries, after.queries); - assert_eq!(before.roots_of, after.roots_of); - } - assert_ne!( - group_of_kind(memo, 0, "Aggregate"), - group_of_kind(memo, 1, "Aggregate") - ); - assert!(memo.candidates.iter().all(|c| !c.is_empty())); - } - assert!( - expanded.candidates.iter().map(Vec::len).sum::() - > restricted.candidates.iter().map(Vec::len).sum::() - ); - } - - /// 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")); - } - } - - /// 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); - } - - /// 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 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; - 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); - // 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 { - 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)]); - } + try_populate_pass1(initial, inventory) + .expect("internally generated Pass 1 inventory must match its logical memo") } diff --git a/crates/plan-selection/src/search/memo/mod.rs b/crates/plan-selection/src/search/memo/mod.rs index caba2240e..77d540739 100644 --- a/crates/plan-selection/src/search/memo/mod.rs +++ b/crates/plan-selection/src/search/memo/mod.rs @@ -52,117 +52,14 @@ pub(crate) mod groups; pub(crate) mod incumbent; pub(crate) mod rules; -use std::rc::Rc; - -use asap_types::ir::OperatorNode; use asap_types::physical::Materialization; pub(crate) use asap_logical_asap_optimization::memo::{MemoGroup, MemoGroupId as GroupId}; -/// 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, - /// 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, - /// 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, - /// 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`) - /// ([`rules::scans::scan_materializations`]). - pub materializations: Vec, - /// 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 -/// 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. -#[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, - /// 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. -#[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, -} +pub(crate) use asap_logical_asap_optimization::memo::candidates::{ + CandidateMemo as Memo, CandidateOrigin, CouplingKind, CouplingSite, ImplementationCandidate, + RollupRead, +}; /// The thresholds chosen for a cluster's coupling sites, by site key: the /// context a candidate is priced in. diff --git a/docs/design_docs/proposals/joint-planning-implementation-plan.md b/docs/design_docs/proposals/joint-planning-implementation-plan.md index 51d4edcd0..baff637f5 100644 --- a/docs/design_docs/proposals/joint-planning-implementation-plan.md +++ b/docs/design_docs/proposals/joint-planning-implementation-plan.md @@ -8,7 +8,7 @@ Replace the legacy whole-workload search with one memo-based planning pipeline. This is a migration of supported behavior, not just a solver replacement. Keep existing accuracy semantics, physical eligibility checks, cost objectives, scalar outputs, and execution export valid. New runtime algorithms and scheduling are outside scope. No latency or speedup target is claimed before benchmarks. -## Current gaps +## Starting gaps (before this migration) Current `run_stage1_to_stage3_with` generates `PerComputationCandidates` before memo construction. `groups::memo_of` constructs identity and fills candidates together. Shared builds and panes use `CouplingSite` choices; MILP `Choice.copy_of` largely represents all-shared or per-query copies. The proposed producer-state groups and arbitrary partial-sharing connections are not implemented. @@ -134,10 +134,11 @@ The first implementation change should establish the stage/memo ownership bounda ## First draft status - Implemented: `LogicalMemo`, `MemoGroup`, and `MemoGroupId` in the logical optimization crate, independent of solver types. +- Implemented: `CandidateMemo` and Pass 1 candidate population now belong to logical optimization; population rejects inventories with incompatible roots or invalid target references. Selection retains only a temporary forwarding adapter and pricing context. - Implemented: initialize logical groups before Pass 1 on memo routes; populate those groups after capability pruning; apply existing Pass 2 rules afterward. The from-Pass-1 API remains supported. - Implemented: shared stage orchestration, source-dependency analysis, diagnostic index mapping, and pane reuse moved out of the legacy namespace. - Implemented: an internal fixed-context candidate-connection MILP core, including same-group multi-selection and an independent tiny-graph enumerator. Production IR decoding and workload-dependent cost/constraint adapters are pending; this core is not yet a replacement production search mode. - Verified: initial group identity survives Pass 1 population and summary pruning; scalar/literal roots retain query positions; distinct quantile results remain distinct groups. -- Pending: Pass 1 rule generation directly into memo groups, shared-state contracts, arbitrary candidate connections, complete cost-model migration, bounded baseline generation, and deletion of legacy search/fallback. The baseline now tries at most 16 local memo assignments and three materialization policies per query, without calling the legacy route or materializing all combinations. A strict wall-clock budget across generation and evaluators remains pending. +- Pending: Pass 1 rule generation directly into memo groups, shared-state contracts, arbitrary candidate connections, complete cost-model migration, a strict planning-time budget, and deletion of legacy search/fallback. The baseline now tries at most 16 local memo assignments and three materialization policies per query, without calling the legacy route or materializing all combinations. A strict wall-clock budget across generation and evaluators remains pending. The current first draft preserves legacy runtime behavior while preparing its removal. It must not be described as completion of the full implementation plan. From 557ad5ae710d0da0d4f82eebfbfcea0062a1f6d2 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:35:30 +0000 Subject: [PATCH 05/12] perf: reuse local Pass 1 rule generation across memo groups --- .../src/per_computation_candidates.rs | 83 ++++++++++++++++++- .../tests/per_computation_candidates.rs | 48 +++++++++++ .../plan-selection/src/search/milp/model.rs | 5 +- crates/plan-selection/src/search/pipeline.rs | 8 +- 4 files changed, 138 insertions(+), 6 deletions(-) diff --git a/crates/logical-asap-optimization/src/per_computation_candidates.rs b/crates/logical-asap-optimization/src/per_computation_candidates.rs index ed7b3d098..40cbd8ce8 100644 --- a/crates/logical-asap-optimization/src/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/per_computation_candidates.rs @@ -449,6 +449,8 @@ fn absorbed_targets(targets: &[ComputationTarget]) -> impl Iterator HavingPlacement { /// consumer of the aggregate applies the same predicate. pub fn generate_per_computation_candidates( roots: Vec<(Id, QueryRoot)>, +) -> Result, PerComputationCandidateError> { + generate_with_local_candidate_cache(roots, None, &mut HashMap::new()) +} + +/// Expand consumer-independent local rules once per logical group and measure. +/// Consumer predicates, nullable-input checks, fusion, and absorption remain +/// occurrence-specific. Cached choices are cloned before any contextual pruning. +pub fn generate_per_computation_candidates_with_memo( + roots: Vec<(Id, QueryRoot)>, + memo: &crate::memo::LogicalMemo, +) -> Result, PerComputationCandidateError> { + if !memo.matches_root_dependencies(&roots) { + return Err(PerComputationCandidateError::MemoRootsChanged); + } + generate_with_local_candidate_cache(roots, Some(memo), &mut HashMap::new()) +} + +type LocalCandidateCache = HashMap<(crate::memo::MemoGroupId, usize), ComputationChoices>; + +fn generate_with_local_candidate_cache( + roots: Vec<(Id, QueryRoot)>, + memo: Option<&crate::memo::LogicalMemo>, + local_candidates: &mut LocalCandidateCache, ) -> Result, PerComputationCandidateError> { let mut seen = HashSet::new(); // Each aggregate with the first root that reaches it. @@ -676,8 +701,15 @@ pub fn generate_per_computation_candidates( }); } for (measure, intent) in measures.iter().enumerate() { - let mut choices = - computation_candidates_for_intent(intent).map_err(|error| match error { + let group_measure = memo + .and_then(|memo| memo.group_for_node(&node)) + .map(|group| (group, measure)); + let cached = group_measure + .and_then(|key| local_candidates.get(&key)) + .cloned(); + let mut choices = cached + .map_or_else(|| computation_candidates_for_intent(intent), Ok) + .map_err(|error| match error { PerComputationCandidateError::InvalidAccuracy { .. } => { PerComputationCandidateError::InvalidAccuracy { root: Some(root.clone()), @@ -686,6 +718,11 @@ pub fn generate_per_computation_candidates( } error => error, })?; + if let Some(key) = group_measure { + local_candidates + .entry(key) + .or_insert_with(|| choices.clone()); + } if choices.candidates == [ComputationCandidate::Exact] && choices.pruned.is_empty() { continue; } @@ -4331,6 +4368,48 @@ mod tests { } } + // Repeated groups reuse local rules without changing contextual candidates. + #[test] + fn memo_local_rules_cache_repeated_groups_without_changing_fusion() { + let queries = [ + "sum by (job) (requests)", + "sum by (job) (requests)", + "topk(3, sum by (job) (requests))", + ]; + let roots: Vec<_> = queries + .iter() + .enumerate() + .map(|(id, query)| { + let inventory = inventory(query); + (id, inventory.roots[0].1.clone()) + }) + .collect(); + let memo = crate::memo::LogicalMemo::from_roots(&roots); + let baseline = generate_per_computation_candidates(roots.clone()).unwrap(); + let mut cache = LocalCandidateCache::new(); + let cached = generate_with_local_candidate_cache(roots, Some(&memo), &mut cache).unwrap(); + assert!( + cache.len() < cached.targets.len(), + "repeated eligible groups reuse local rules" + ); + assert_eq!(baseline.targets.len(), cached.targets.len()); + for (before, after) in baseline.targets.iter().zip(&cached.targets) { + assert_eq!(before.choices, after.choices); + assert_eq!(before.having, after.having); + assert_eq!(before.fused, after.fused); + } + for target in 0..baseline.targets.len() { + for choice in 0..baseline.targets[target].choice_count() { + let mut selected = vec![0; baseline.targets.len()]; + selected[target] = choice; + assert_eq!( + compose_logical_candidate(&baseline, &selected).unwrap(), + compose_logical_candidate(&cached, &selected).unwrap() + ); + } + } + } + fn inventory(query: &str) -> PerComputationCandidates { let root = crate::test_support::lower_promql( query, diff --git a/crates/logical-asap-optimization/tests/per_computation_candidates.rs b/crates/logical-asap-optimization/tests/per_computation_candidates.rs index 15a8c7804..f1c11a864 100644 --- a/crates/logical-asap-optimization/tests/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/tests/per_computation_candidates.rs @@ -1290,3 +1290,51 @@ fn grouped_avg_composition_preserves_the_whole_schema() { .any(|node| matches!(node.asap(), Some(ASAPOp::FinalizeExactAccumulator { .. })))); assert_eq!(out.schema, root.schema); } + +// Cached local rules preserve filtered and unfiltered consumers independently. +#[test] +fn memo_cached_rules_preserve_occurrence_specific_having() { + use asap_logical_asap_optimization::memo::LogicalMemo; + use asap_logical_asap_optimization::per_computation_candidates::{ + compose_logical_candidate, generate_per_computation_candidates_with_memo, + }; + let shared = grouped_count(None); + let other = grouped_count(None); + let key = equals(0, ScalarValue::Utf8("a".into())); + let roots = vec![ + ( + 0, + QueryRoot::Operator(having_filter(shared.clone(), key.clone())), + ), + (1, QueryRoot::Operator(shared)), + (2, QueryRoot::Operator(having_filter(other, key))), + ]; + let memo = LogicalMemo::from_roots(&roots); + let baseline = generate_per_computation_candidates(roots.clone()).unwrap(); + let cached = generate_per_computation_candidates_with_memo(roots, &memo).unwrap(); + assert_eq!(cached.targets.len(), 2); + assert_ne!( + cached.targets[0].having[0].placement, + cached.targets[1].having[0].placement + ); + for (before, after) in baseline.targets.iter().zip(&cached.targets) { + assert_eq!(before.choices, after.choices); + assert_eq!(before.having, after.having); + assert_eq!(before.fused, after.fused); + } + for target in 0..baseline.targets.len() { + for choice in 0..baseline.targets[target].choice_count() { + let mut selected = vec![0; baseline.targets.len()]; + selected[target] = choice; + assert_eq!( + compose_logical_candidate(&baseline, &selected).unwrap(), + compose_logical_candidate(&cached, &selected).unwrap() + ); + } + } + let other_roots = vec![(0, QueryRoot::Operator(grouped_count(None)))]; + assert!(matches!( + generate_per_computation_candidates_with_memo(other_roots, &memo), + Err(PerComputationCandidateError::MemoRootsChanged) + )); +} diff --git a/crates/plan-selection/src/search/milp/model.rs b/crates/plan-selection/src/search/milp/model.rs index 19df9f490..af3c2ec32 100644 --- a/crates/plan-selection/src/search/milp/model.rs +++ b/crates/plan-selection/src/search/milp/model.rs @@ -2259,7 +2259,7 @@ fn search_clusters( .collect(); let initial = LogicalMemo::from_roots(&roots); let inventory = - asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates(roots) + asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates_with_memo(roots, &initial) .ok()? .without_unsupported_summaries(models.deployment_capabilities); let cluster_entries: Vec = @@ -2500,8 +2500,9 @@ impl<'a, Id: Clone> ClusterLoop<'a, Id> { let query_roots = vec![(q, root.clone())]; let initial = LogicalMemo::from_roots(&query_roots); let Ok(inventory) = - asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates( + asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates_with_memo( query_roots, + &initial, ) else { return; diff --git a/crates/plan-selection/src/search/pipeline.rs b/crates/plan-selection/src/search/pipeline.rs index 8207c26db..3ed879706 100644 --- a/crates/plan-selection/src/search/pipeline.rs +++ b/crates/plan-selection/src/search/pipeline.rs @@ -16,7 +16,8 @@ use crate::search::{OptimizationInput, OptimizationPass, OptimizeError, PlanOutp use crate::PlanningModels; use crate::{CostModel, RoutedToLegacy, SearchAlgorithm, Selection, SelectionError}; use asap_logical_asap_optimization::per_computation_candidates::{ - generate_per_computation_candidates, PerComputationCandidateError, PerComputationCandidates, + generate_per_computation_candidates, generate_per_computation_candidates_with_memo, + PerComputationCandidateError, PerComputationCandidates, }; use asap_logical_asap_optimization::workload_candidates::Pass2Candidate; use asap_physical_asap_optimization::PhysicalCandidate; @@ -193,7 +194,10 @@ pub fn run_stage1_to_stage3_with( } } let initial = uses_memo.then(|| LogicalMemo::from_roots(&roots)); - let per_computation = generate_per_computation_candidates(roots)?; + let per_computation = match &initial { + Some(memo) => generate_per_computation_candidates_with_memo(roots, memo)?, + None => generate_per_computation_candidates(roots)?, + }; Ok(run_with_initial_memo( per_computation, entries, From d7fc5f61225a53414216a98108e96ac824a127fb Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:35:35 +0000 Subject: [PATCH 06/12] refactor: move shared summary expansion into logical optimization --- crates/logical-asap-optimization/src/memo.rs | 2 + .../src/memo/groups.rs | 6 +- .../src/memo/shared_builds.rs | 728 ++++++++++++++++++ .../src/search/memo/rules/capability.rs | 726 +---------------- .../joint-planning-implementation-plan.md | 1 + 5 files changed, 739 insertions(+), 724 deletions(-) create mode 100644 crates/logical-asap-optimization/src/memo/shared_builds.rs diff --git a/crates/logical-asap-optimization/src/memo.rs b/crates/logical-asap-optimization/src/memo.rs index f3b6051aa..b71752f57 100644 --- a/crates/logical-asap-optimization/src/memo.rs +++ b/crates/logical-asap-optimization/src/memo.rs @@ -6,6 +6,8 @@ /// Logical implementation candidates populated after initial identity construction. pub mod candidates; +/// Pass 2 summary-capability expansion and consistent producer/reader resizing. +pub mod shared_builds; mod groups; use std::collections::{HashMap, HashSet}; diff --git a/crates/logical-asap-optimization/src/memo/groups.rs b/crates/logical-asap-optimization/src/memo/groups.rs index cba7af7dc..35ecf5b09 100644 --- a/crates/logical-asap-optimization/src/memo/groups.rs +++ b/crates/logical-asap-optimization/src/memo/groups.rs @@ -4,9 +4,9 @@ //! Groups are the workload's logical DAG interned by //! [`asap_types::ir::cse::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. +//! legal (an ASAP node, a scan, or one whose schema has a unique key). +//! Structurally equal keyless aggregates remain separate; eligible scans +//! and keyed aggregates can share a group. //! //! A Pass 1 target's aggregate group gets one candidate per combination of //! its measures' Pass 1 candidates: the aggregate alone composed by diff --git a/crates/logical-asap-optimization/src/memo/shared_builds.rs b/crates/logical-asap-optimization/src/memo/shared_builds.rs new file mode 100644 index 000000000..ac02389ef --- /dev/null +++ b/crates/logical-asap-optimization/src/memo/shared_builds.rs @@ -0,0 +1,728 @@ +//! 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, HashSet}; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use crate::accuracy::estimators::size_for_all_readers; +use crate::asap_aware_cse::with_resized_build; +use crate::per_computation_candidates::PerComputationCandidateError; +use asap_types::ir::operator::Reduction; +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::candidates::NodeKeys; +use super::candidates::{ + CandidateMemo as Memo, CandidateOrigin, CouplingKind, CouplingSite, ImplementationCandidate, +}; +use super::MemoGroupId as GroupId; + +/// 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 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 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 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. +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, + } +} + +/// `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 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 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 fn share_equal_builds(roots: &[QueryRoot], resized: &[Rc]) -> Vec { + // Equal by node key: guarantees, composed again over a plan, may differ + // from the resized candidate's. + let mut keys = 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 { .. })) + { + continue; + } + 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((key, 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)) +} + +/// 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'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)?; + let Some(ASAPOp::SummaryAgg { + child, + family, + input, + reduction, + grouping, + filter, + }) = build.asap() + else { + return None; + }; + // 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:?}"), + }; + let mut hasher = DefaultHasher::new(); + ( + family_key, + input.0, + format!("{reduction:?}/{grouping:?}/{filter:?}/{:?}", input.1), + ) + .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 crate::memo::MemoGroup; + use crate::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; + + fn memo_of( + inventory: &crate::per_computation_candidates::PerComputationCandidates, + ) -> Memo { + super::super::candidates::try_populate_pass1( + super::super::LogicalMemo::from_roots(&inventory.roots), + inventory, + ) + .expect("generated inventory belongs to its logical memo") + } + + /// 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| crate::memo::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(crate::memo::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(), + materializations: Vec::new(), + budget_share: 1, + requires_exact_children: Vec::new(), + rollup_of: None, + requires_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))); + } + + /// 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::per_computation_candidates::generate_per_computation_candidates; + use crate::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]]); + } + + /// 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::per_computation_candidates::generate_per_computation_candidates; + use crate::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/search/memo/rules/capability.rs b/crates/plan-selection/src/search/memo/rules/capability.rs index 80c56f72a..d825a0288 100644 --- a/crates/plan-selection/src/search/memo/rules/capability.rs +++ b/crates/plan-selection/src/search/memo/rules/capability.rs @@ -1,723 +1,7 @@ -//! 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`]. +//! Compatibility imports during the shared-state rule migration. +//! Logical Pass 2 expansion and build resizing live with the logical memo; +//! selection only chooses configurations and prices their physical execution. -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::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, SketchKind, SketchParams}; -use asap_types::ir::schema::FieldDataType; -use asap_types::ir::QueryRoot; -use asap_types::ir::{ASAPOp, OperatorNode}; - -use crate::search::memo::groups::NodeKeys; -use crate::search::memo::{ - CandidateOrigin, CouplingKind, CouplingSite, GroupId, ImplementationCandidate, Memo, +pub(crate) use asap_logical_asap_optimization::memo::shared_builds::{ + add_shared_builds, resized_at_threshold, share_equal_builds, shared_build_of, size_thresholds, }; - -/// 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, - } -} - -/// `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 { - // Equal by node key: guarantees, composed again over a plan, may differ - // from the resized candidate's. - let mut keys = crate::search::memo::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 { .. })) - { - continue; - } - 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((key, 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)) -} - -/// 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'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)?; - let Some(ASAPOp::SummaryAgg { - child, - family, - input, - reduction, - grouping, - filter, - }) = build.asap() - else { - return None; - }; - // 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:?}"), - }; - let mut hasher = DefaultHasher::new(); - ( - family_key, - input.0, - format!("{reduction:?}/{grouping:?}/{filter:?}/{:?}", input.1), - ) - .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 crate::search::memo::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}; - 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| crate::search::memo::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(crate::search::memo::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(), - materializations: Vec::new(), - budget_share: 1, - requires_exact_children: Vec::new(), - rollup_of: None, - requires_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))); - } - - /// 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::search::memo::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]]); - } - - /// 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::search::memo::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/docs/design_docs/proposals/joint-planning-implementation-plan.md b/docs/design_docs/proposals/joint-planning-implementation-plan.md index baff637f5..a0c5906bb 100644 --- a/docs/design_docs/proposals/joint-planning-implementation-plan.md +++ b/docs/design_docs/proposals/joint-planning-implementation-plan.md @@ -136,6 +136,7 @@ The first implementation change should establish the stage/memo ownership bounda - Implemented: `LogicalMemo`, `MemoGroup`, and `MemoGroupId` in the logical optimization crate, independent of solver types. - Implemented: `CandidateMemo` and Pass 1 candidate population now belong to logical optimization; population rejects inventories with incompatible roots or invalid target references. Selection retains only a temporary forwarding adapter and pricing context. - Implemented: initialize logical groups before Pass 1 on memo routes; populate those groups after capability pruning; apply existing Pass 2 rules afterward. The from-Pass-1 API remains supported. +- Implemented: consumer-independent Pass 1 rules are cached by logical group and measure; per-occurrence HAVING, nullable pruning, fusion, and absorption remain contextual. Shared-summary Pass 2 rules now live in logical optimization, with temporary selection reexports. - Implemented: shared stage orchestration, source-dependency analysis, diagnostic index mapping, and pane reuse moved out of the legacy namespace. - Implemented: an internal fixed-context candidate-connection MILP core, including same-group multi-selection and an independent tiny-graph enumerator. Production IR decoding and workload-dependent cost/constraint adapters are pending; this core is not yet a replacement production search mode. - Verified: initial group identity survives Pass 1 population and summary pruning; scalar/literal roots retain query positions; distinct quantile results remain distinct groups. From 25fa398e35e87411aa5cf394a4b886e46c70adbd Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:50:53 +0000 Subject: [PATCH 07/12] feat: constrain additive resources in candidate connection selection --- .../src/search/milp/candidate_connections.rs | 105 +++++++++++++++++- 1 file changed, 101 insertions(+), 4 deletions(-) diff --git a/crates/plan-selection/src/search/milp/candidate_connections.rs b/crates/plan-selection/src/search/milp/candidate_connections.rs index d2ffd6f64..e08808d8b 100644 --- a/crates/plan-selection/src/search/milp/candidate_connections.rs +++ b/crates/plan-selection/src/search/milp/candidate_connections.rs @@ -3,8 +3,8 @@ //! This is the selection core for the joint-planning migration, not a new //! production search entry point. Candidate generation must supply executable //! candidates and compatible producers. Costs cover a fixed accounting context; -//! demand-dependent costs, accuracy refinement, materialization policy, resource -//! limits, and physical IR decoding still belong to the surrounding pipeline. +//! demand-dependent costs, accuracy refinement, materialization policy, +//! nonadditive resource limits, and physical IR decoding still belong to the surrounding pipeline. use std::collections::HashSet; @@ -41,10 +41,21 @@ pub(crate) struct QueryRootCandidates { pub compatible_candidates: Vec, } +/// Additive usage over a fixed accounting context, owned once per candidate. +/// Peak temporaries and connection-dependent demand require a different model. +#[derive(Debug, Clone)] +pub(crate) struct ResourceBudget { + pub name: String, + pub capacity: f64, + /// Same index order as the graph's candidates. + pub candidate_usage: Vec, +} + #[derive(Debug, Clone, Default)] pub(crate) struct CandidateDependencyGraph { pub candidates: Vec, pub query_roots: Vec, + pub resource_budgets: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -84,7 +95,7 @@ pub(crate) enum CandidateSelectionError { const MINIMUM_RELATIVE_COST: f64 = 1e-9; impl CandidateDependencyGraph { - fn validate(&self) -> Result<(), CandidateSelectionError> { + pub(super) fn validate(&self) -> Result<(), CandidateSelectionError> { let validate_ids = |ids: &[PhysicalCandidateId], context: &str| { let mut seen = HashSet::new(); for id in ids { @@ -103,6 +114,22 @@ impl CandidateDependencyGraph { } Ok(()) }; + for budget in &self.resource_budgets { + if budget.name.is_empty() + || !budget.capacity.is_finite() + || budget.capacity < 0.0 + || budget.candidate_usage.len() != self.candidates.len() + || budget + .candidate_usage + .iter() + .any(|usage| !usage.is_finite() || *usage < 0.0) + || !budget.candidate_usage.iter().sum::().is_finite() + { + return Err(CandidateSelectionError::InvalidGraph(format!( + "resource budget {:?} needs finite nonnegative capacity and one finite nonnegative usage per candidate", budget.name + ))); + } + } let mut total_cost = 0.0; for (index, candidate) in self.candidates.iter().enumerate() { if !candidate.local_cost.is_finite() || candidate.local_cost < 0.0 { @@ -233,6 +260,22 @@ pub(crate) fn select_candidates( candidate_selected[candidate] <= support[candidate].clone() )); } + for budget in &graph.resource_budgets { + // Normalize each resource independently of currency and other resources. + let scale = budget + .candidate_usage + .iter() + .copied() + .fold(budget.capacity, f64::max) + .max(f64::MIN_POSITIVE); + let usage: Expression = budget + .candidate_usage + .iter() + .zip(&candidate_selected) + .map(|(usage, selected)| (*usage / scale) * *selected) + .sum(); + constraints.push(constraint!(usage <= budget.capacity / scale)); + } // Validation rejects positive coefficients outside the supported relative // range; uniform currency-unit changes do not alter the selected graph. let maximum_cost = graph @@ -297,7 +340,7 @@ pub(crate) fn select_candidates( } /// Check the decoded integral graph instead of trusting rounded solver values. -fn validate_selected_graph( +pub(super) fn validate_selected_graph( graph: &CandidateDependencyGraph, selected: &SelectedCandidateGraph, ) -> Result<(), String> { @@ -310,6 +353,15 @@ fn validate_selected_graph( { return Err(invalid()); } + for budget in &graph.resource_budgets { + let usage: f64 = active.iter().map(|c| budget.candidate_usage[c.0]).sum(); + if usage > budget.capacity { + return Err(format!( + "decoded selection exceeds resource budget {:?}", + budget.name + )); + } + } let mut root_seen = vec![false; graph.query_roots.len()]; for root in &selected.query_roots { let Some(allowed) = graph.query_roots.get(root.query_root) else { @@ -408,6 +460,38 @@ fn validate_selected_graph( mod tests { use super::*; + // A shared producer owns resident state once, regardless of its reader count. + #[test] + fn resource_limits_price_shared_state_once_and_change_the_selected_plan() { + let mut graph = dependency_graph( + vec![ + candidate(0, 8.0, &[]), + candidate(1, 1.0, &[&[0]]), + candidate(2, 1.0, &[&[0]]), + candidate(1, 12.0, &[]), + candidate(2, 12.0, &[]), + ], + &[&[1, 3], &[2, 4]], + ); + graph.resource_budgets.push(ResourceBudget { + name: "resident memory".into(), + capacity: 6.0, + candidate_usage: vec![6.0, 0.0, 0.0, 1.0, 1.0], + }); + assert_eq!(compare_with_oracle(&graph).unwrap().modeled_cost, 10.0); + let shared = compare_with_oracle(&graph).unwrap(); + graph.resource_budgets[0].capacity = 5.0; + assert!(validate_selected_graph(&graph, &shared).is_err()); + assert_eq!(compare_with_oracle(&graph).unwrap().modeled_cost, 24.0); + graph.resource_budgets[0].capacity = 1.0; + assert!(compare_with_oracle(&graph).is_none()); + graph.resource_budgets[0].candidate_usage.pop(); + assert!(matches!( + select_candidates(&graph), + Err(CandidateSelectionError::InvalidGraph(_)) + )); + } + fn candidate(group: u32, cost: f64, inputs: &[&[usize]]) -> PhysicalCandidate { PhysicalCandidate { group: MemoGroupId(group), @@ -431,6 +515,7 @@ mod tests { ) -> CandidateDependencyGraph { CandidateDependencyGraph { candidates, + resource_budgets: Vec::new(), query_roots: roots .iter() .map(|choices| QueryRootCandidates { @@ -533,6 +618,18 @@ mod tests { { continue; } + if graph.resource_budgets.iter().any(|budget| { + budget + .candidate_usage + .iter() + .zip(&marks) + .filter(|(_, mark)| **mark == 2) + .map(|(usage, _)| usage) + .sum::() + > budget.capacity + }) { + continue; + } let cost = graph .candidates .iter() From 11308a11f33b78ee372a62e71d038c6043ab4c45 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:53:51 +0000 Subject: [PATCH 08/12] feat: expose concrete shared summary state as memo groups --- crates/logical-asap-optimization/src/memo.rs | 2 + .../src/memo/state_groups.rs | 506 ++++++++++++++++++ 2 files changed, 508 insertions(+) create mode 100644 crates/logical-asap-optimization/src/memo/state_groups.rs diff --git a/crates/logical-asap-optimization/src/memo.rs b/crates/logical-asap-optimization/src/memo.rs index b71752f57..e2ddfc914 100644 --- a/crates/logical-asap-optimization/src/memo.rs +++ b/crates/logical-asap-optimization/src/memo.rs @@ -8,6 +8,8 @@ pub mod candidates; /// Pass 2 summary-capability expansion and consistent producer/reader resizing. pub mod shared_builds; +/// Concrete summary-state groups for the connection-based planner migration. +pub mod state_groups; mod groups; use std::collections::{HashMap, HashSet}; diff --git a/crates/logical-asap-optimization/src/memo/state_groups.rs b/crates/logical-asap-optimization/src/memo/state_groups.rs new file mode 100644 index 000000000..89bd23a26 --- /dev/null +++ b/crates/logical-asap-optimization/src/memo/state_groups.rs @@ -0,0 +1,506 @@ +//! Opt-in concrete summary-state groups for the candidate-connection planner. +//! +//! This does not change the production selection path. Only simple Pass 1 builds +//! directly over canonical input groups are expanded. Capacity, guarantee, +//! update/filter semantics, and materialization annotations must match exactly; +//! larger-capacity substitution and private preprocessing are not inferred. + +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; + +use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode}; + +use super::candidates::{CandidateMemo, CandidateOrigin, ImplementationCandidate}; +use super::{MemoGroup, MemoGroupId}; + +#[derive(Debug)] +pub struct StateGroupExpansion { + pub memo: CandidateMemo, + pub state_groups: Vec, + pub added_reader_candidates: usize, +} + +#[derive(Debug, thiserror::Error)] +pub enum StateGroupExpansionError { + #[error("memo group and candidate vectors have different lengths")] + MismatchedGroups, + #[error("memo contains an out-of-range group reference")] + InvalidGroupReference, + #[error("memo group capacity exceeds the group identifier range")] + TooManyGroups, + #[error("memo dependencies contain a cycle")] + CyclicDependencies, + #[error("explicit summary-state groups have already been introduced")] + AlreadyExpanded, + #[error("a previously discovered reader no longer exposes its summary boundary")] + InvalidReaderBoundary, +} + +struct BuildReaders { + build: Rc, + children: Vec, + readers: Vec<(usize, usize)>, +} + +/// Preserve local candidates and add explicit shared-state reader candidates. +/// The input must precede explicit state-group expansion. Unsupported candidate +/// shapes are retained unchanged, rather than guessed into a shared interface. +pub fn expand_concrete_summary_state_groups( + mut memo: CandidateMemo, +) -> Result { + validate_references(&memo)?; + if memo + .groups + .iter() + .any(|group| matches!(group.node.asap(), Some(ASAPOp::SummaryAgg { .. }))) + { + return Err(StateGroupExpansionError::AlreadyExpanded); + } + if memo.topological_order().is_none() { + return Err(StateGroupExpansionError::CyclicDependencies); + } + let mut builds: Vec = Vec::new(); + for (group, candidates) in memo.candidates.iter().enumerate() { + for (index, candidate) in candidates.iter().enumerate() { + let Some((build, children)) = simple_build(&memo, candidate) else { + continue; + }; + if let Some(existing) = builds + .iter_mut() + .find(|entry| same_concrete_build(&entry.build, &build)) + { + existing.readers.push((group, index)); + } else { + builds.push(BuildReaders { + build, + children, + readers: vec![(group, index)], + }); + } + } + } + let mut state_groups = Vec::new(); + let mut added_reader_candidates = 0; + for build in builds { + if !build + .readers + .iter() + .any(|&(group, _)| group != build.readers[0].0) + { + continue; + } + let state_group = MemoGroupId( + u32::try_from(memo.groups.len()) + .map_err(|_| StateGroupExpansionError::TooManyGroups)?, + ); + let mut queries: Vec<_> = build + .readers + .iter() + .flat_map(|&(group, _)| memo.groups[group].queries.iter().copied()) + .collect(); + queries.sort_unstable(); + queries.dedup(); + memo.groups.push(MemoGroup { + node: Rc::clone(&build.build), + children: build.children.clone(), + parents: Vec::new(), + queries, + roots_of: Vec::new(), + }); + memo.candidates.push(vec![ImplementationCandidate { + implements: vec![state_group], + node: Rc::clone(&build.build), + children: build.children, + origin: CandidateOrigin::SharedBuild, + coupling_sites: Vec::new(), + budget_share: 1, + requires_exact_children: Vec::new(), + rollup_of: None, + materializations: Vec::new(), + requires_sites: Vec::new(), + }]); + for (group, index) in build.readers { + let mut reader = memo.candidates[group][index].clone(); + // Each reader may have an equal but separately allocated build. + let (own_build, _) = simple_build(&memo, &reader) + .ok_or(StateGroupExpansionError::InvalidReaderBoundary)?; + reader.node = replace_node(&reader.node, &own_build, &build.build); + let boundaries: HashMap<_, _> = reader + .children + .iter() + .map(|&id| (Rc::as_ptr(&memo.groups[id.0 as usize].node), id)) + .chain([(Rc::as_ptr(&build.build), state_group)]) + .collect(); + reader.children = boundary_groups(&reader.node, &boundaries); + reader.origin = CandidateOrigin::SharedBuild; + memo.candidates[group].push(reader); + added_reader_candidates += 1; + } + state_groups.push(state_group); + } + // Parents describe every possible group/candidate dependency, not a selected plan. + for group in &mut memo.groups { + group.parents.clear(); + } + for group in 0..memo.groups.len() { + let dependencies: HashSet<_> = memo.groups[group] + .children + .iter() + .chain( + memo.candidates[group] + .iter() + .flat_map(|candidate| &candidate.children), + ) + .copied() + .collect(); + for child in dependencies { + memo.groups[child.0 as usize] + .parents + .push(MemoGroupId(group as u32)); + } + } + if memo.topological_order().is_none() { + return Err(StateGroupExpansionError::CyclicDependencies); + } + // The chosen representative input may now serve additional queries through + // state reuse; propagate potential reach along the expanded dependency DAG. + for group in memo + .topological_order() + .ok_or(StateGroupExpansionError::CyclicDependencies)? + .into_iter() + .rev() + { + let queries = memo.groups[group.0 as usize].queries.clone(); + let children: HashSet<_> = memo.groups[group.0 as usize] + .children + .iter() + .chain( + memo.candidates[group.0 as usize] + .iter() + .flat_map(|candidate| &candidate.children), + ) + .copied() + .collect(); + for child in children { + memo.groups[child.0 as usize].queries.extend(&queries); + memo.groups[child.0 as usize].queries.sort_unstable(); + memo.groups[child.0 as usize].queries.dedup(); + } + } + Ok(StateGroupExpansion { + memo, + state_groups, + added_reader_candidates, + }) +} + +fn validate_references(memo: &CandidateMemo) -> Result<(), StateGroupExpansionError> { + if memo.groups.len() != memo.candidates.len() { + return Err(StateGroupExpansionError::MismatchedGroups); + } + if memo.groups.len() > u32::MAX as usize { + return Err(StateGroupExpansionError::TooManyGroups); + } + for group in &memo.groups { + if group + .children + .iter() + .chain(&group.parents) + .any(|id| id.0 as usize >= memo.groups.len()) + { + return Err(StateGroupExpansionError::InvalidGroupReference); + } + } + for candidate in memo.candidates.iter().flatten() { + if candidate + .children + .iter() + .chain(&candidate.implements) + .chain(&candidate.requires_exact_children) + .chain(candidate.requires_sites.iter().map(|(group, _)| group)) + .any(|id| id.0 as usize >= memo.groups.len()) + || candidate + .rollup_of + .as_ref() + .is_some_and(|rollup| rollup.finer_group.0 as usize >= memo.groups.len()) + { + return Err(StateGroupExpansionError::InvalidGroupReference); + } + } + Ok(()) +} + +fn simple_build( + memo: &CandidateMemo, + candidate: &ImplementationCandidate, +) -> Option<(Rc, Vec)> { + if candidate.origin != CandidateOrigin::Pass1 + || candidate.implements.len() != 1 + || !candidate.coupling_sites.is_empty() + || !candidate.requires_sites.is_empty() + || !candidate.requires_exact_children.is_empty() + || candidate.rollup_of.is_some() + || !candidate.materializations.is_empty() + || candidate.budget_share != 1 + { + return None; + } + let boundaries: HashMap<_, _> = candidate + .children + .iter() + .map(|&group| (Rc::as_ptr(&memo.groups[group.0 as usize].node), group)) + .collect(); + let mut stack = vec![Rc::clone(&candidate.node)]; + let mut seen = HashSet::new(); + let mut build = None; + while let Some(node) = stack.pop() { + if boundaries.contains_key(&Rc::as_ptr(&node)) || !seen.insert(Rc::as_ptr(&node)) { + continue; + } + if matches!(node.asap(), Some(ASAPOp::SummaryAgg { .. })) { + if build.is_some() { + return None; + } + build = Some(Rc::clone(&node)); + } + stack.extend(node.children().into_iter().cloned()); + } + let build = build?; + if Rc::ptr_eq(&build, &candidate.node) { + return None; + } + let mut children = Vec::new(); + for child in build.children() { + let group = *boundaries.get(&Rc::as_ptr(child))?; + if !children.contains(&group) { + children.push(group); + } + } + Some((build, children)) +} + +fn boundary_groups( + node: &Rc, + boundaries: &HashMap<*const OperatorNode, MemoGroupId>, +) -> Vec { + let mut found = Vec::new(); + let mut stack = vec![Rc::clone(node)]; + let mut seen = HashSet::new(); + while let Some(node) = stack.pop() { + if !seen.insert(Rc::as_ptr(&node)) { + continue; + } + if let Some(&group) = boundaries.get(&Rc::as_ptr(&node)) { + if !found.contains(&group) { + found.push(group); + } + } else { + stack.extend(node.children().into_iter().rev().cloned()); + } + } + found +} + +// Immediate inputs are canonical memo nodes. Compare local fields without +// recursively walking their potentially shared DAGs; serialization additionally +// distinguishes signed-zero parameters, matching conservative CSE semantics. +fn same_concrete_build(left: &OperatorNode, right: &OperatorNode) -> bool { + let left_inputs = left.children(); + let right_inputs = right.children(); + left_inputs.len() == right_inputs.len() + && left_inputs + .iter() + .zip(right_inputs) + .all(|(left, right)| Rc::ptr_eq(left, right) || same_range_input(left, right)) + && same_local_fields(left, right) +} + +// TimeRange lacks a unique row key and therefore is not generally memo-interned. +// Equal ranges over the same canonical input nevertheless feed identical data +// into a summary build, as the existing ASAP-aware summary reuse rule allows. +fn same_range_input(left: &OperatorNode, right: &OperatorNode) -> bool { + let ( + Some(NonASAPOp::TimeRange { + child: left_child, .. + }), + Some(NonASAPOp::TimeRange { + child: right_child, .. + }), + ) = (left.non_asap(), right.non_asap()) + else { + return false; + }; + Rc::ptr_eq(left_child, right_child) && same_local_fields(left, right) +} + +fn same_local_fields(left: &OperatorNode, right: &OperatorNode) -> bool { + let signature = |node: &OperatorNode| { + ( + node.operator.map_children(|_| ()), + node.result_kind, + node.schema.clone(), + node.guarantee.clone(), + node.timing, + node.materialization, + ) + }; + let left = signature(left); + let right = signature(right); + left == right + && matches!((serde_json::to_string(&left), serde_json::to_string(&right)), (Ok(left), Ok(right)) if left == right) +} + +fn replace_node( + node: &Rc, + old: &Rc, + new: &Rc, +) -> Rc { + replace_node_cached(node, old, new, &mut HashMap::new()) +} + +fn replace_node_cached( + node: &Rc, + old: &Rc, + new: &Rc, + rebuilt: &mut HashMap<*const OperatorNode, Rc>, +) -> Rc { + if Rc::ptr_eq(node, old) { + return Rc::clone(new); + } + if let Some(existing) = rebuilt.get(&Rc::as_ptr(node)) { + return Rc::clone(existing); + } + let mut changed = false; + let operator = node.operator.map_children(|child| { + let replacement = replace_node_cached(child, old, new, rebuilt); + changed |= !Rc::ptr_eq(child, &replacement); + replacement + }); + let result = if changed { + Rc::new(OperatorNode { + operator, + ..(**node).clone() + }) + } else { + Rc::clone(node) + }; + rebuilt.insert(Rc::as_ptr(node), Rc::clone(&result)); + result +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memo::candidates::try_populate_pass1; + use crate::memo::LogicalMemo; + use crate::per_computation_candidates::generate_per_computation_candidates; + use asap_types::ir::properties::AccuracyTarget; + use asap_types::ir::QueryRoot; + + fn workload(queries: &[(&str, f64)]) -> CandidateMemo { + let roots: Vec<_> = queries + .iter() + .enumerate() + .map(|(id, (query, error))| { + ( + id, + QueryRoot::Operator(crate::test_support::lower_promql( + query, + AccuracyTarget::max_error(*error).failure_probability(0.001), + )), + ) + }) + .collect(); + let initial = LogicalMemo::from_roots(&roots); + let inventory = generate_per_computation_candidates(roots).unwrap(); + try_populate_pass1(initial, &inventory).unwrap() + } + + // Equal-sized quantiles share one canonical state while local choices remain. + #[test] + fn quantile_readers_share_concrete_state_and_retain_local_candidates() { + let original = workload(&[ + ("quantile_over_time(0.95, latency[5m])", 0.01), + ("quantile_over_time(0.99, latency[5m])", 0.01), + ]); + let original_count = original.groups.len(); + let expanded = expand_concrete_summary_state_groups(original.clone()).unwrap(); + assert!(!expanded.state_groups.is_empty()); + assert!(expanded.added_reader_candidates >= 2); + let order = expanded.memo.topological_order().unwrap(); + for state in &expanded.state_groups { + let group = &expanded.memo.groups[state.0 as usize]; + assert_eq!(group.queries, vec![0, 1]); + for input in &group.children { + assert_eq!(expanded.memo.groups[input.0 as usize].queries, vec![0, 1]); + assert!(group + .node + .children() + .iter() + .any(|node| Rc::ptr_eq(node, &expanded.memo.groups[input.0 as usize].node))); + } + assert_eq!(group.parents.len(), 2); + assert!(group.roots_of.is_empty()); + for &parent in &group.parents { + assert!( + order.iter().position(|id| id == state).unwrap() + < order.iter().position(|id| id == &parent).unwrap() + ); + let reader = expanded.memo.candidates[parent.0 as usize] + .iter() + .find(|candidate| candidate.children.contains(state)) + .unwrap(); + assert!(OperatorNode::reachable(&reader.node) + .iter() + .any(|node| Rc::ptr_eq(node, &group.node))); + assert_eq!(reader.children, vec![*state]); + } + } + for group in 0..original_count { + for (before, after) in original.candidates[group] + .iter() + .zip(&expanded.memo.candidates[group]) + { + assert!(Rc::ptr_eq(&before.node, &after.node)); + assert_eq!(before.children, after.children); + } + } + assert!(matches!( + expand_concrete_summary_state_groups(expanded.memo), + Err(StateGroupExpansionError::AlreadyExpanded) + )); + } + + // Capacity and source predicates are part of the concrete producer identity. + #[test] + fn distinct_sizes_and_filters_do_not_share_concrete_builds() { + for queries in [ + [ + ("quantile_over_time(0.95, latency[5m])", 0.01), + ("quantile_over_time(0.99, latency[5m])", 0.02), + ], + [ + ("quantile_over_time(0.95, latency{job=\"a\"}[5m])", 0.01), + ("quantile_over_time(0.99, latency{job=\"b\"}[5m])", 0.01), + ], + ] { + let original = workload(&queries); + let count = original.groups.len(); + let expanded = expand_concrete_summary_state_groups(original).unwrap(); + assert!(expanded.state_groups.is_empty()); + assert_eq!(expanded.memo.groups.len(), count); + assert_eq!(expanded.added_reader_candidates, 0); + } + } + + // Public expansion rejects malformed IDs before indexing dependency arrays. + #[test] + fn malformed_references_fail_before_dependency_traversal() { + let mut memo = workload(&[("quantile_over_time(0.95, latency[5m])", 0.01)]); + memo.groups[0].children.push(MemoGroupId(u32::MAX)); + assert!(matches!( + expand_concrete_summary_state_groups(memo), + Err(StateGroupExpansionError::InvalidGroupReference) + )); + } +} From 46a9516389baaacf76430141c2d5184b72d4fd0f Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:53:51 +0000 Subject: [PATCH 09/12] feat: reconstruct selected candidate connections with typed inputs --- .../search/milp/candidate_reconstruction.rs | 325 ++++++++++++++++++ crates/plan-selection/src/search/milp/mod.rs | 2 + 2 files changed, 327 insertions(+) create mode 100644 crates/plan-selection/src/search/milp/candidate_reconstruction.rs diff --git a/crates/plan-selection/src/search/milp/candidate_reconstruction.rs b/crates/plan-selection/src/search/milp/candidate_reconstruction.rs new file mode 100644 index 000000000..488164968 --- /dev/null +++ b/crates/plan-selection/src/search/milp/candidate_reconstruction.rs @@ -0,0 +1,325 @@ +//! Reconstruct selected immediate operators, preserving candidate identity. +//! +//! This decoder checks typed connections, not Stage 2 execution legality or +//! query accuracy targets. The caller must run Stage 2 and the plan validator. +use std::rc::Rc; + +use asap_logical_asap_optimization::per_computation_candidates::recompose_guarantees; +use asap_types::ir::{OperatorNode, QueryRoot}; + +use super::candidate_connections::{ + validate_selected_graph, CandidateDependencyGraph, PhysicalCandidateId, SelectedCandidateGraph, +}; + +/// Templates indexed by globally unique candidate ID. Each template describes +/// one immediate operator; its ordered `Operator::map_children` positions correspond one-for-one to the +/// candidate's producer requirements. Children serve as typed input contracts, +/// never as implicit selected dependencies. +pub(super) type CandidateOperatorTemplates = Vec>>; + +/// Enumerate in exactly the traversal order used when rebinding. `children()` +/// puts direct children before scalar references; `map_children` follows fields. +fn ordered_inputs(node: &OperatorNode) -> Vec> { + let mut inputs = Vec::new(); + node.operator.map_children(|child| { + inputs.push(Rc::clone(child)); + Rc::clone(child) + }); + inputs +} + +/// Operator output references in `graph.query_roots` order. A caller with scalar +/// roots must map these references back into its original scalar expressions +/// (and preserve literal-only roots); an output reference is not a whole query. +/// Templates' timing/materialization are deliberately discarded: Stage 2 must +/// apply the selected materialization policy to the reconstructed connections. +pub(super) fn reconstruct_candidate_outputs( + graph: &CandidateDependencyGraph, + selected: &SelectedCandidateGraph, + templates: &CandidateOperatorTemplates, +) -> Result>, String> { + graph.validate().map_err(|e| e.to_string())?; + validate_selected_graph(graph, selected)?; + let count = graph.candidates.len(); + let mut inputs = graph + .candidates + .iter() + .map(|candidate| vec![None; candidate.producer_requirements.len()]) + .collect::>(); + let mut consumers = vec![Vec::new(); count]; + for connection in &selected.connections { + inputs[connection.consumer.0][connection.requirement] = Some(connection.producer); + consumers[connection.producer.0].push(connection.consumer.0); + } + let mut remaining = inputs.iter().map(Vec::len).collect::>(); + let mut ready = Vec::new(); + for &PhysicalCandidateId(id) in &selected.selected_candidates { + let template = templates + .get(id) + .and_then(Option::as_ref) + .ok_or_else(|| format!("candidate {id} has no operator template"))?; + if template.children().len() != inputs[id].len() { + return Err(format!( + "candidate {id} template inputs do not match producer requirements" + )); + } + if remaining[id] == 0 { + ready.push(id); + } + } + let mut rebuilt: Vec>> = vec![None; count]; + while let Some(id) = ready.pop() { + let template = templates[id].as_ref().unwrap(); + let expected = ordered_inputs(template); + let mut producers = Vec::with_capacity(inputs[id].len()); + for (position, producer) in inputs[id].iter().enumerate() { + let producer = rebuilt[producer.unwrap().0].as_ref().unwrap(); + // Serialization includes sketch error/failure parameters whose + // PartialEq intentionally omits some sizing metadata. + if expected[position].result_kind != producer.result_kind + || expected[position].schema != producer.schema + || serde_json::to_value(&expected[position].schema).map_err(|e| e.to_string())? + != serde_json::to_value(&producer.schema).map_err(|e| e.to_string())? + { + return Err(format!("candidate {id} input {position} has incompatible schema or state configuration")); + } + producers.push(Rc::clone(producer)); + } + let mut position = 0; + let operator = template.operator.map_children(|_| { + let producer = Rc::clone(&producers[position]); + position += 1; + producer + }); + let node = Rc::new( + OperatorNode::with_schema(operator, template.schema.clone()) + .with_guarantee(template.guarantee.clone()), + ); + node.validate_structure() + .map_err(|e| format!("candidate {id}: {e}"))?; + rebuilt[id] = Some(node); + for &consumer in &consumers[id] { + remaining[consumer] -= 1; + if remaining[consumer] == 0 { + ready.push(consumer); + } + } + } + let mut roots = vec![None; graph.query_roots.len()]; + for root in &selected.query_roots { + roots[root.query_root] = Some(QueryRoot::Operator(Rc::clone( + rebuilt[root.candidate.0].as_ref().unwrap(), + ))); + } + let roots = roots.into_iter().map(Option::unwrap).collect::>(); + recompose_guarantees(&roots) + .map_err(|e| e.to_string())? + .into_iter() + .map(|root| match root { + QueryRoot::Operator(node) => Ok(node), + QueryRoot::Scalar(_) => Err("unexpected scalar during operator reconstruction".into()), + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::super::candidate_connections::{ + select_candidates, PhysicalCandidate, ProducerRequirement, QueryRootCandidates, + }; + use super::*; + use asap_logical_asap_optimization::memo::MemoGroupId; + use asap_logical_asap_optimization::test_support::lower_promql; + + fn catalog(root: Rc) -> (CandidateDependencyGraph, CandidateOperatorTemplates) { + fn add( + node: &Rc, + graph: &mut CandidateDependencyGraph, + templates: &mut CandidateOperatorTemplates, + ) -> PhysicalCandidateId { + if let Some(id) = templates + .iter() + .position(|t| t.as_ref().is_some_and(|t| Rc::ptr_eq(t, node))) + { + return PhysicalCandidateId(id); + } + let producer_requirements = ordered_inputs(node) + .iter() + .map(|child| ProducerRequirement { + compatible_producers: vec![add(child, graph, templates)], + }) + .collect(); + let id = templates.len(); + templates.push(Some(Rc::clone(node))); + graph.candidates.push(PhysicalCandidate { + group: MemoGroupId(id as u32), + local_cost: 1.0, + producer_requirements, + }); + PhysicalCandidateId(id) + } + let mut graph = CandidateDependencyGraph::default(); + let mut templates = Vec::new(); + let root = add(&root, &mut graph, &mut templates); + graph.query_roots.push(QueryRootCandidates { + compatible_candidates: vec![root], + }); + (graph, templates) + } + + // Shared producers retain one identity; separately selected readers do not merge. + #[test] + fn two_readers_share_one_build_and_distinct_candidates_keep_identity() { + let (mut graph, mut templates) = catalog(lower_promql( + "sum(sum(metric))", + asap_types::ir::properties::AccuracyTarget::exact(), + )); + let root = templates.len() - 1; + assert_eq!(graph.candidates[root].producer_requirements.len(), 1); + // Same group and same template, deliberately distinct physical candidates. + graph.candidates.push(graph.candidates[root].clone()); + templates.push(templates[root].clone()); + graph.query_roots.push(QueryRootCandidates { + compatible_candidates: vec![PhysicalCandidateId(root + 1)], + }); + let selected = select_candidates(&graph).unwrap().unwrap(); + let outputs = reconstruct_candidate_outputs(&graph, &selected, &templates).unwrap(); + assert!(!Rc::ptr_eq(&outputs[0], &outputs[1])); + assert!(Rc::ptr_eq( + outputs[0].children()[0], + outputs[1].children()[0] + )); + assert!(outputs.iter().all(|node| node.timing.is_none())); + } + + // Input positions remain independent even when template references are identical. + #[test] + fn repeated_input_positions_can_select_distinct_suppliers() { + let root = lower_promql( + "metric + metric", + asap_types::ir::properties::AccuracyTarget::exact(), + ); + let (mut graph, mut templates) = catalog(root); + let root = templates.len() - 1; + assert_eq!(graph.candidates[root].producer_requirements.len(), 2); + let first = graph.candidates[root].producer_requirements[0].compatible_producers[0]; + let duplicate = PhysicalCandidateId(templates.len()); + graph.candidates.push(graph.candidates[first.0].clone()); + templates.push(templates[first.0].clone()); + // Both template positions name the same expected child. Connections may differ. + let template = templates[root].as_ref().unwrap(); + let child = Rc::clone(template.children()[0]); + templates[root] = Some(Rc::new(OperatorNode::with_schema( + template.operator.map_children(|_| Rc::clone(&child)), + template.schema.clone(), + ))); + graph.candidates[root].producer_requirements[0].compatible_producers = vec![first]; + graph.candidates[root].producer_requirements[1].compatible_producers = vec![duplicate]; + let selected = select_candidates(&graph).unwrap().unwrap(); + let outputs = reconstruct_candidate_outputs(&graph, &selected, &templates).unwrap(); + let children = outputs[0].children(); + assert!(!Rc::ptr_eq(children[0], children[1])); + } + + // Invalid template metadata must fail before returning reconstructed roots. + #[test] + fn missing_template_and_incompatible_input_schema_are_rejected() { + let (graph, mut templates) = catalog(lower_promql( + "sum(metric)", + asap_types::ir::properties::AccuracyTarget::exact(), + )); + let selected = select_candidates(&graph).unwrap().unwrap(); + let saved = templates[0].take(); + assert!(reconstruct_candidate_outputs(&graph, &selected, &templates) + .unwrap_err() + .contains("no operator template")); + templates[0] = saved; + let root = templates.len() - 1; + let template = templates[root].as_ref().unwrap(); + let wrong = lower_promql( + "sum(other)", + asap_types::ir::properties::AccuracyTarget::exact(), + ); + templates[root] = Some(Rc::new(OperatorNode::with_schema( + template.operator.map_children(|_| Rc::clone(&wrong)), + template.schema.clone(), + ))); + assert!(reconstruct_candidate_outputs(&graph, &selected, &templates) + .unwrap_err() + .contains("incompatible schema")); + } + // Scalar references precede direct inputs during rebinding, unlike children(). + // Verify identities and scalar contents, so equal schemas cannot hide swaps. + #[test] + fn scalar_reference_and_direct_input_preserve_their_distinct_positions() { + use asap_types::ir::scalar::expr_ir::ScalarValue; + use asap_types::ir::scalar::ProjectItem; + use asap_types::ir::{NonASAPOp, Operator, ScalarExpr}; + + let direct = lower_promql( + "sum(direct_metric)", + asap_types::ir::properties::AccuracyTarget::exact(), + ); + let scalar = lower_promql( + "sum(scalar_metric)", + asap_types::ir::properties::AccuracyTarget::exact(), + ); + assert_eq!(direct.schema, scalar.schema); + let root = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Project { + cols: vec![ + ProjectItem { + alias: Some("scalar_result".into()), + expr: ScalarExpr::PromqlScalarFromVector(Rc::clone(&scalar)), + }, + ProjectItem { + alias: Some("constant".into()), + expr: ScalarExpr::Literal(ScalarValue::Float64(7.0)), + }, + ], + qualifier: None, + child: Rc::clone(&direct), + })) + .unwrap(); + assert!(Rc::ptr_eq(root.children()[0], &direct)); + assert!(Rc::ptr_eq(&ordered_inputs(&root)[0], &scalar)); + let (mut graph, templates) = catalog(Rc::clone(&root)); + let find = |node: &Rc| { + PhysicalCandidateId( + templates + .iter() + .position(|template| { + template + .as_ref() + .is_some_and(|template| Rc::ptr_eq(template, node)) + }) + .unwrap(), + ) + }; + graph.query_roots.push(QueryRootCandidates { + compatible_candidates: vec![find(&direct)], + }); + graph.query_roots.push(QueryRootCandidates { + compatible_candidates: vec![find(&scalar)], + }); + let selected = select_candidates(&graph).unwrap().unwrap(); + let outputs = reconstruct_candidate_outputs(&graph, &selected, &templates).unwrap(); + let Some(NonASAPOp::Project { + cols, + qualifier, + child, + }) = outputs[0].non_asap() + else { + panic!("expected reconstructed projection"); + }; + assert!(Rc::ptr_eq(child, &outputs[1])); + let ScalarExpr::PromqlScalarFromVector(scalar) = &cols[0].expr else { + panic!("scalar conversion must be preserved"); + }; + assert!(Rc::ptr_eq(scalar, &outputs[2])); + assert!(!Rc::ptr_eq(child, scalar)); + assert_eq!(cols[0].alias.as_deref(), Some("scalar_result")); + assert_eq!(cols[1].alias.as_deref(), Some("constant")); + assert_eq!(cols[1].expr, ScalarExpr::Literal(ScalarValue::Float64(7.0))); + assert!(qualifier.is_none()); + } +} diff --git a/crates/plan-selection/src/search/milp/mod.rs b/crates/plan-selection/src/search/milp/mod.rs index 4a5eea02a..b94bbcfbb 100644 --- a/crates/plan-selection/src/search/milp/mod.rs +++ b/crates/plan-selection/src/search/milp/mod.rs @@ -12,3 +12,5 @@ pub(crate) mod report; // Candidate-connection core being integrated into the staged memo pipeline. pub(crate) mod candidate_connections; + +pub(crate) mod candidate_reconstruction; From 80b942eb70c50ab52406a61995f2c993ea631be9 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:56:03 +0000 Subject: [PATCH 10/12] fix: reject unsupported resource coefficient ranges before solving --- .../src/search/milp/candidate_connections.rs | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/crates/plan-selection/src/search/milp/candidate_connections.rs b/crates/plan-selection/src/search/milp/candidate_connections.rs index e08808d8b..33294518d 100644 --- a/crates/plan-selection/src/search/milp/candidate_connections.rs +++ b/crates/plan-selection/src/search/milp/candidate_connections.rs @@ -92,6 +92,7 @@ pub(crate) enum CandidateSelectionError { // The core normalizes its largest local cost to one. Smaller positive // coefficients must stay representable by HiGHS; rejecting a wider range is an // explicit model restriction, not proof of infeasibility or a completeness claim. +// Resource rows use the same conservative positive coefficient range. const MINIMUM_RELATIVE_COST: f64 = 1e-9; impl CandidateDependencyGraph { @@ -130,6 +131,24 @@ impl CandidateDependencyGraph { ))); } } + for budget in &self.resource_budgets { + let scale = budget + .candidate_usage + .iter() + .copied() + .fold(budget.capacity, f64::max); + if scale > 0.0 + && budget + .candidate_usage + .iter() + .chain(std::iter::once(&budget.capacity)) + .any(|value| *value > 0.0 && *value / scale < MINIMUM_RELATIVE_COST) + { + return Err(CandidateSelectionError::InvalidGraph(format!( + "resource budget {:?} exceeds the supported positive coefficient range ({MINIMUM_RELATIVE_COST}); rescale or reformulate the resource model", budget.name + ))); + } + } let mut total_cost = 0.0; for (index, candidate) in self.candidates.iter().enumerate() { if !candidate.local_cost.is_finite() || candidate.local_cost < 0.0 { @@ -492,6 +511,28 @@ mod tests { )); } + // Unreachable large states must not silently erase a live resource constraint. + #[test] + fn unsupported_resource_dynamic_range_is_explicitly_rejected() { + let mut graph = dependency_graph( + vec![ + candidate(0, 2.0, &[]), + candidate(0, 1.0, &[]), + candidate(1, 1.0, &[]), + ], + &[&[0, 1]], + ); + graph.resource_budgets.push(ResourceBudget { + name: "memory".into(), + capacity: 1.0, + candidate_usage: vec![1.0, 2.0, 1e20], + }); + assert!(matches!( + select_candidates(&graph), + Err(CandidateSelectionError::InvalidGraph(_)) + )); + } + fn candidate(group: u32, cost: f64, inputs: &[&[usize]]) -> PhysicalCandidate { PhysicalCandidate { group: MemoGroupId(group), From 3f52c9ffec57001bb56b4af12e1390d589526a83 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 02:59:11 +0000 Subject: [PATCH 11/12] feat: bridge memo candidates to typed connection selection and Stage 2 --- .../src/search/milp/candidate_connections.rs | 10 +- .../src/search/milp/candidate_inventory.rs | 475 ++++++++++++++++++ .../search/milp/candidate_reconstruction.rs | 16 +- crates/plan-selection/src/search/milp/mod.rs | 1 + .../joint-planning-implementation-plan.md | 5 +- 5 files changed, 495 insertions(+), 12 deletions(-) create mode 100644 crates/plan-selection/src/search/milp/candidate_inventory.rs diff --git a/crates/plan-selection/src/search/milp/candidate_connections.rs b/crates/plan-selection/src/search/milp/candidate_connections.rs index 33294518d..697664a0e 100644 --- a/crates/plan-selection/src/search/milp/candidate_connections.rs +++ b/crates/plan-selection/src/search/milp/candidate_connections.rs @@ -93,7 +93,7 @@ pub(crate) enum CandidateSelectionError { // coefficients must stay representable by HiGHS; rejecting a wider range is an // explicit model restriction, not proof of infeasibility or a completeness claim. // Resource rows use the same conservative positive coefficient range. -const MINIMUM_RELATIVE_COST: f64 = 1e-9; +const MINIMUM_RELATIVE_COEFFICIENT: f64 = 1e-9; impl CandidateDependencyGraph { pub(super) fn validate(&self) -> Result<(), CandidateSelectionError> { @@ -142,10 +142,10 @@ impl CandidateDependencyGraph { .candidate_usage .iter() .chain(std::iter::once(&budget.capacity)) - .any(|value| *value > 0.0 && *value / scale < MINIMUM_RELATIVE_COST) + .any(|value| *value > 0.0 && *value / scale < MINIMUM_RELATIVE_COEFFICIENT) { return Err(CandidateSelectionError::InvalidGraph(format!( - "resource budget {:?} exceeds the supported positive coefficient range ({MINIMUM_RELATIVE_COST}); rescale or reformulate the resource model", budget.name + "resource budget {:?} exceeds the supported positive coefficient range ({MINIMUM_RELATIVE_COEFFICIENT}); restrict or reformulate the resource model", budget.name ))); } } @@ -177,11 +177,11 @@ impl CandidateDependencyGraph { if maximum_cost > 0.0 && self.candidates.iter().any(|candidate| { candidate.local_cost > 0.0 - && candidate.local_cost / maximum_cost < MINIMUM_RELATIVE_COST + && candidate.local_cost / maximum_cost < MINIMUM_RELATIVE_COEFFICIENT }) { return Err(CandidateSelectionError::InvalidGraph(format!( - "positive local costs must be at least {MINIMUM_RELATIVE_COST} of the largest cost; wider ranges require a different numerical policy" + "positive local costs must be at least {MINIMUM_RELATIVE_COEFFICIENT} of the largest cost; wider ranges require a different numerical policy" ))); } for (index, root) in self.query_roots.iter().enumerate() { diff --git a/crates/plan-selection/src/search/milp/candidate_inventory.rs b/crates/plan-selection/src/search/milp/candidate_inventory.rs new file mode 100644 index 000000000..f2651d1be --- /dev/null +++ b/crates/plan-selection/src/search/milp/candidate_inventory.rs @@ -0,0 +1,475 @@ +//! Opt-in bridge from logical candidate memo to immediate candidate connections. +//! +//! The caller supplies node-local costs for one fixed accounting context. These +//! are not runtime cost estimates. Consumer-dependent demand, accuracy budgets, +//! physical materialization assignments, and legacy coupling constraints remain +//! unsupported and are never silently erased by this adapter. This is a logical +//! inventory bridge for solver experiments, not Stage 2 physical enumeration. +//! A cost callback cannot certify execution legality: callers must run Stage 2 +//! and validate the reconstructed plan before execution. +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; + +use asap_logical_asap_optimization::memo::candidates::CandidateMemo; +use asap_logical_asap_optimization::memo::MemoGroupId; +use asap_types::ir::OperatorNode; +use asap_types::physical::Materialization; + +use super::candidate_connections::{ + CandidateDependencyGraph, PhysicalCandidate, PhysicalCandidateId, ProducerRequirement, + QueryRootCandidates, +}; +use super::candidate_reconstruction::CandidateOperatorTemplates; + +#[derive(Debug)] +pub(super) struct CandidateInventory { + pub graph: CandidateDependencyGraph, + pub templates: CandidateOperatorTemplates, + /// Root implementation IDs, excluding each implementation's private nodes. + pub group_candidates: Vec>, +} + +#[derive(Debug, thiserror::Error)] +pub(super) enum CandidateInventoryError { + #[error("invalid candidate memo: {0}")] + InvalidMemo(String), + #[error("unsupported candidate {candidate} in group {group}: {reason}")] + Unsupported { + group: usize, + candidate: usize, + reason: &'static str, + }, + #[error("node-local fixed-context cost: {0}")] + Cost(String), +} + +/// Each output reference is an operator group, not necessarily an entire query. +/// Scalar query assembly and Stage 2/accuracy validation remain with the caller. +pub(super) fn build_candidate_inventory( + memo: &CandidateMemo, + output_groups: &[MemoGroupId], + mut node_local_cost: impl FnMut(&OperatorNode) -> Result, +) -> Result { + let invalid = |message: &str| CandidateInventoryError::InvalidMemo(message.into()); + if memo.groups.len() != memo.candidates.len() { + return Err(invalid("group and candidate vector lengths differ")); + } + if output_groups + .iter() + .any(|group| group.0 as usize >= memo.groups.len()) + { + return Err(invalid("query output references missing group")); + } + let mut inventory = CandidateInventory { + graph: CandidateDependencyGraph::default(), + templates: Vec::new(), + group_candidates: vec![Vec::new(); memo.groups.len()], + }; + for (group, candidates) in memo.candidates.iter().enumerate() { + for (index, candidate) in candidates.iter().enumerate() { + let unsupported = |reason| CandidateInventoryError::Unsupported { + group, + candidate: index, + reason, + }; + if candidate.implements != [MemoGroupId(group as u32)] { + return Err(unsupported("fused or differently owned computation")); + } + if !candidate.coupling_sites.is_empty() + || !candidate.requires_sites.is_empty() + || !candidate.requires_exact_children.is_empty() + || candidate.rollup_of.is_some() + { + return Err(unsupported( + "legacy shared-decision or implementation constraints", + )); + } + if candidate.budget_share != 1 { + return Err(unsupported("composed accuracy budget shares")); + } + if !candidate.materializations.is_empty() { + return Err(unsupported("materialization eligibility restrictions")); + } + if candidate + .children + .iter() + .any(|child| child.0 as usize >= memo.groups.len()) + { + return Err(invalid("candidate references missing input group")); + } + let id = PhysicalCandidateId(inventory.graph.candidates.len()); + inventory.group_candidates[group].push(id); + inventory.templates.push(Some(Rc::clone(&candidate.node))); + inventory.graph.candidates.push(PhysicalCandidate { + group: MemoGroupId(group as u32), + local_cost: 0.0, + producer_requirements: Vec::new(), + }); + } + } + for (group, candidates) in memo.candidates.iter().enumerate() { + for (index, candidate) in candidates.iter().enumerate() { + let mut boundaries = HashMap::new(); + for &child in &candidate.children { + let pointer = Rc::as_ptr(&memo.groups[child.0 as usize].node); + if boundaries + .insert(pointer, child) + .is_some_and(|previous| previous != child) + { + return Err(invalid("one boundary node names different input groups")); + } + } + if boundaries.contains_key(&Rc::as_ptr(&candidate.node)) { + return Err(CandidateInventoryError::Unsupported { + group, + candidate: index, + reason: "root aliases an input group without an operator", + }); + } + let mut builder = ImmediateCandidateBuilder { + inventory: &mut inventory, + boundaries, + private_nodes: HashMap::new(), + used_boundaries: HashSet::new(), + owner: MemoGroupId(group as u32), + candidate_index: index, + node_local_cost: &mut node_local_cost, + }; + let root_id = builder.inventory.group_candidates[group][index]; + builder.add(&candidate.node, Some(root_id))?; + if candidate + .children + .iter() + .any(|group| !builder.used_boundaries.contains(group)) + { + return Err(invalid( + "declared input group has no canonical node boundary", + )); + } + } + } + for group in output_groups { + let expected = &memo.groups[group.0 as usize].node; + let mut compatible_candidates = Vec::new(); + for &id in &inventory.group_candidates[group.0 as usize] { + let template = inventory.templates[id.0] + .as_ref() + .ok_or_else(|| invalid("missing root template"))?; + if compatible_output(expected, template)? { + compatible_candidates.push(id); + } + } + inventory.graph.query_roots.push(QueryRootCandidates { + compatible_candidates, + }); + } + inventory + .graph + .validate() + .map_err(|error| invalid(&error.to_string()))?; + Ok(inventory) +} + +struct ImmediateCandidateBuilder<'a, F> { + inventory: &'a mut CandidateInventory, + boundaries: HashMap<*const OperatorNode, MemoGroupId>, + private_nodes: HashMap<*const OperatorNode, PhysicalCandidateId>, + used_boundaries: HashSet, + owner: MemoGroupId, + candidate_index: usize, + node_local_cost: &'a mut F, +} + +impl Result> ImmediateCandidateBuilder<'_, F> { + fn add( + &mut self, + node: &Rc, + reserved: Option, + ) -> Result { + if let Some(&existing) = self.private_nodes.get(&Rc::as_ptr(node)) { + return Ok(existing); + } + if node.timing.is_some() || node.materialization != Materialization::NotMaterialized { + return Err(CandidateInventoryError::Unsupported { + group: self.owner.0 as usize, + candidate: self.candidate_index, + reason: "assigned execution timing or materialization", + }); + } + let mut inputs = Vec::new(); + node.operator.map_children(|child| { + inputs.push(Rc::clone(child)); + Rc::clone(child) + }); + let mut requirements = Vec::new(); + for input in inputs { + let producers = if let Some(&group) = self.boundaries.get(&Rc::as_ptr(&input)) { + self.used_boundaries.insert(group); + let mut compatible = Vec::new(); + for &candidate in &self.inventory.group_candidates[group.0 as usize] { + let template = + self.inventory.templates[candidate.0] + .as_ref() + .ok_or_else(|| { + CandidateInventoryError::InvalidMemo("missing root template".into()) + })?; + if compatible_output(&input, template)? { + compatible.push(candidate); + } + } + compatible + } else { + vec![self.add(&input, None)?] + }; + requirements.push(ProducerRequirement { + compatible_producers: producers, + }); + } + let cost = (self.node_local_cost)(node).map_err(CandidateInventoryError::Cost)?; + if !cost.is_finite() || cost < 0.0 { + return Err(CandidateInventoryError::Cost( + "cost must be finite and nonnegative".into(), + )); + } + let candidate = PhysicalCandidate { + group: self.owner, + local_cost: cost, + producer_requirements: requirements, + }; + let id = if let Some(id) = reserved { + self.inventory.graph.candidates[id.0] = candidate; + id + } else { + let id = PhysicalCandidateId(self.inventory.graph.candidates.len()); + self.inventory.graph.candidates.push(candidate); + self.inventory.templates.push(Some(Rc::clone(node))); + id + }; + self.private_nodes.insert(Rc::as_ptr(node), id); + Ok(id) + } +} + +fn compatible_output( + expected: &OperatorNode, + producer: &OperatorNode, +) -> Result { + let serialize = |node: &OperatorNode| { + serde_json::to_string(&node.schema) + .map_err(|error| CandidateInventoryError::InvalidMemo(error.to_string())) + }; + Ok(expected.result_kind == producer.result_kind + && expected.schema == producer.schema + && serialize(expected)? == serialize(producer)?) +} + +#[cfg(test)] +mod tests { + use super::super::candidate_connections::select_candidates; + use super::super::candidate_reconstruction::reconstruct_candidate_outputs; + use super::*; + use asap_logical_asap_optimization::memo::candidates::try_populate_pass1; + use asap_logical_asap_optimization::memo::state_groups::expand_concrete_summary_state_groups; + use asap_logical_asap_optimization::memo::LogicalMemo; + 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::{FieldDataType, SketchAlgorithm}; + use asap_types::ir::{ASAPOp, NonASAPOp, QueryRoot}; + + fn quantile_memo() -> (CandidateMemo, Vec) { + let roots = [0.95, 0.99] + .into_iter() + .enumerate() + .map(|(query, quantile)| { + ( + query, + QueryRoot::Operator(lower_promql( + &format!("quantile_over_time({quantile}, latency[5m])"), + AccuracyTarget::max_error(0.01).failure_probability(0.001), + )), + ) + }) + .collect::>(); + let initial = LogicalMemo::from_roots(&roots); + let output_groups = (0..2) + .map(|query| { + MemoGroupId( + initial + .groups + .iter() + .position(|group| group.roots_of.contains(&query)) + .unwrap() as u32, + ) + }) + .collect(); + let candidates = generate_per_computation_candidates(roots).unwrap(); + let memo = try_populate_pass1(initial, &candidates).unwrap(); + ( + expand_concrete_summary_state_groups(memo).unwrap().memo, + output_groups, + ) + } + + fn is_kll(node: &OperatorNode) -> bool { + matches!(node.asap(), Some(ASAPOp::SummaryAgg { family: FieldDataType::Sketch(kind, _), .. }) if *kind.algorithm() == SketchAlgorithm::Kll) + } + + // Synthetic fixed node costs favor one KLL producer; no runtime cost claim. + #[test] + fn expanded_quantiles_select_and_reconstruct_one_shared_kll() { + let (memo, outputs) = quantile_memo(); + let inventory = build_candidate_inventory(&memo, &outputs, |node| { + Ok(if is_kll(node) { + 8.0 + } else if matches!(node.asap(), Some(ASAPOp::SummaryAgg { .. })) { + 80.0 + } else if matches!(node.asap(), Some(ASAPOp::SummaryEstimate { .. })) { + 1.0 + } else if matches!(node.non_asap(), Some(NonASAPOp::Aggregate { .. })) { + 100.0 + } else { + 0.0 + }) + }) + .unwrap(); + for (group, candidates) in memo.candidates.iter().enumerate() { + assert_eq!(inventory.group_candidates[group].len(), candidates.len()); + } + let selected = select_candidates(&inventory.graph).unwrap().unwrap(); + assert!((selected.modeled_cost - 10.0).abs() < 1e-8); + let reconstructed = + reconstruct_candidate_outputs(&inventory.graph, &selected, &inventory.templates) + .unwrap(); + let builds: Vec<_> = reconstructed + .iter() + .map(|root| { + OperatorNode::reachable(root) + .into_iter() + .find(|node| is_kll(node)) + .unwrap() + }) + .collect(); + assert!(Rc::ptr_eq(&builds[0], &builds[1])); + // Legality is independently checked by the real physical stage; the + // synthetic objective above is deliberately not an evaluator estimate. + stage2_smoke(&reconstructed); + assert_eq!( + selected + .selected_candidates + .iter() + .filter(|id| is_kll(inventory.templates[id.0].as_ref().unwrap())) + .count(), + 1 + ); + } + + fn stage2_smoke(outputs: &[Rc]) { + use asap_physical_asap_optimization::optimize_physical_asap_chosen; + use asap_types::physical::Medium; + use asap_types::workload::{ + DataArrival, DataWorkload, DurationMs, MetricType, Predictability, Query, + QueryRecurrence, QueryRequirements, QueryWorkloadEntry, RepeatedDemand, + RepetitionInterval, SourceFacts, + }; + let entries: Vec<_> = [0.95, 0.99] + .into_iter() + .map(|quantile| QueryWorkloadEntry { + query: Query(format!("quantile_over_time({quantile}, latency[5m])")), + requirements: QueryRequirements { + accuracy: AccuracyTarget::max_error(0.01).failure_probability(0.001), + 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: [( + "latency".into(), + SourceFacts { + series_count: Some(1_000), + metric_type: Some(MetricType::Gauge), + sample_interval: Some(DurationMs(15_000)), + ..SourceFacts::default() + }, + )] + .into(), + }; + let roots: Vec<_> = outputs.iter().cloned().map(QueryRoot::Operator).collect(); + let physical = optimize_physical_asap_chosen( + "candidate-inventory-smoke", + &roots, + &entries, + &data, + crate::PlanningModels::builtin().deployment_capabilities, + &mut |_, _, options| { + let query_materialization = Materialization::MatAtQuery(Medium::Memory); + if options.contains(&query_materialization) { + query_materialization + } else { + options + .first() + .copied() + .unwrap_or(Materialization::NotMaterialized) + } + }, + ) + .unwrap() + .expect("reconstructed logical connections admit a legal Stage 2 plan"); + physical.dag.validate().unwrap(); + assert_eq!(physical.roots.len(), 2); + } + + // A boundary-only alias has no immediate operator for the decoder to rebuild. + #[test] + fn rejects_root_alias_instead_of_duplicating_boundary_work() { + let (mut memo, outputs) = quantile_memo(); + let group = outputs[0].0 as usize; + let boundary = memo.candidates[group][0].children[0]; + memo.candidates[group][0].node = Rc::clone(&memo.groups[boundary.0 as usize].node); + assert!(matches!( + build_candidate_inventory(&memo, &outputs, |_| Ok(1.0)), + Err(CandidateInventoryError::Unsupported { + reason: "root aliases an input group without an operator", + .. + }) + )); + } + + // Unsupported decisions must fail adaptation, rather than disappear. + #[test] + fn rejects_accuracy_shares_materialization_restrictions_and_invalid_costs() { + let (mut memo, outputs) = quantile_memo(); + memo.candidates[0][0].budget_share = 2; + assert!(matches!( + build_candidate_inventory(&memo, &outputs, |_| Ok(1.0)), + Err(CandidateInventoryError::Unsupported { + reason: "composed accuracy budget shares", + .. + }) + )); + memo.candidates[0][0].budget_share = 1; + memo.candidates[0][0] + .materializations + .push(Materialization::NotMaterialized); + assert!(matches!( + build_candidate_inventory(&memo, &outputs, |_| Ok(1.0)), + Err(CandidateInventoryError::Unsupported { + reason: "materialization eligibility restrictions", + .. + }) + )); + memo.candidates[0][0].materializations.clear(); + assert!(matches!( + build_candidate_inventory(&memo, &outputs, |_| Ok(f64::NAN)), + Err(CandidateInventoryError::Cost(_)) + )); + } +} diff --git a/crates/plan-selection/src/search/milp/candidate_reconstruction.rs b/crates/plan-selection/src/search/milp/candidate_reconstruction.rs index 488164968..12d0035be 100644 --- a/crates/plan-selection/src/search/milp/candidate_reconstruction.rs +++ b/crates/plan-selection/src/search/milp/candidate_reconstruction.rs @@ -86,13 +86,17 @@ pub(super) fn reconstruct_candidate_outputs( producers.push(Rc::clone(producer)); } let mut position = 0; - let operator = template.operator.map_children(|_| { - let producer = Rc::clone(&producers[position]); - position += 1; - producer - }); + // Re-derive structural metadata over the selected inputs, preserving + // only explicit naming overrides. Logical rewrite templates can retain + // pre-rewrite keys/nullability; copying their entire schema is invalid. let node = Rc::new( - OperatorNode::with_schema(operator, template.schema.clone()) + template + .with_new_children(|_| { + let producer = Rc::clone(&producers[position]); + position += 1; + producer + }) + .map_err(|e| format!("candidate {id}: {e}"))? .with_guarantee(template.guarantee.clone()), ); node.validate_structure() diff --git a/crates/plan-selection/src/search/milp/mod.rs b/crates/plan-selection/src/search/milp/mod.rs index b94bbcfbb..acba506ee 100644 --- a/crates/plan-selection/src/search/milp/mod.rs +++ b/crates/plan-selection/src/search/milp/mod.rs @@ -14,3 +14,4 @@ pub(crate) mod report; pub(crate) mod candidate_connections; pub(crate) mod candidate_reconstruction; +pub(crate) mod candidate_inventory; diff --git a/docs/design_docs/proposals/joint-planning-implementation-plan.md b/docs/design_docs/proposals/joint-planning-implementation-plan.md index a0c5906bb..8c21c2f0c 100644 --- a/docs/design_docs/proposals/joint-planning-implementation-plan.md +++ b/docs/design_docs/proposals/joint-planning-implementation-plan.md @@ -138,7 +138,10 @@ The first implementation change should establish the stage/memo ownership bounda - Implemented: initialize logical groups before Pass 1 on memo routes; populate those groups after capability pruning; apply existing Pass 2 rules afterward. The from-Pass-1 API remains supported. - Implemented: consumer-independent Pass 1 rules are cached by logical group and measure; per-occurrence HAVING, nullable pruning, fusion, and absorption remain contextual. Shared-summary Pass 2 rules now live in logical optimization, with temporary selection reexports. - Implemented: shared stage orchestration, source-dependency analysis, diagnostic index mapping, and pane reuse moved out of the legacy namespace. -- Implemented: an internal fixed-context candidate-connection MILP core, including same-group multi-selection and an independent tiny-graph enumerator. Production IR decoding and workload-dependent cost/constraint adapters are pending; this core is not yet a replacement production search mode. +- Implemented: an internal fixed-context candidate-connection MILP core, including same-group multi-selection, additive resource budgets charged once per selected candidate, and an independent tiny-graph enumerator. +- Implemented: opt-in concrete summary-state groups retain local candidates and expose shared builds with explicit reader dependencies. Equal TimeRange data over the same canonical source can feed one concrete state without merging the result groups. Cross-capacity binding and complex producer shapes remain restricted. +- Implemented: typed selected-connection reconstruction preserves candidate identity and ordered scalar/operator dependencies, then recomposes guarantees. Stage 2 must still assign legal materialization and timing on the selected graph. +- Implemented: a restricted memo-to-connection inventory bridge with explicit rejection of unsupported constraints and typed producer bindings. A P95/P99 test selects one shared KLL using synthetic fixed costs, reconstructs the graph, and passes Stage 2 physicalization. This is not a runtime-cost optimality result. The production driver still uses the existing selector; workload-dependent demand, costs, and full feasibility adapters are required before replacement. - Verified: initial group identity survives Pass 1 population and summary pruning; scalar/literal roots retain query positions; distinct quantile results remain distinct groups. - Pending: Pass 1 rule generation directly into memo groups, shared-state contracts, arbitrary candidate connections, complete cost-model migration, a strict planning-time budget, and deletion of legacy search/fallback. The baseline now tries at most 16 local memo assignments and three materialization policies per query, without calling the legacy route or materializing all combinations. A strict wall-clock budget across generation and evaluators remains pending. From aa0ca89fbfc470825fc9b787fc8c27d6207dda6f Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Sat, 24 Oct 2026 03:05:57 +0000 Subject: [PATCH 12/12] style: format final memo and selection module declarations --- crates/logical-asap-optimization/src/memo.rs | 2 +- crates/plan-selection/src/search/milp/mod.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/logical-asap-optimization/src/memo.rs b/crates/logical-asap-optimization/src/memo.rs index e2ddfc914..9d095040f 100644 --- a/crates/logical-asap-optimization/src/memo.rs +++ b/crates/logical-asap-optimization/src/memo.rs @@ -6,11 +6,11 @@ /// Logical implementation candidates populated after initial identity construction. pub mod candidates; +mod groups; /// Pass 2 summary-capability expansion and consistent producer/reader resizing. pub mod shared_builds; /// Concrete summary-state groups for the connection-based planner migration. pub mod state_groups; -mod groups; use std::collections::{HashMap, HashSet}; use std::rc::Rc; diff --git a/crates/plan-selection/src/search/milp/mod.rs b/crates/plan-selection/src/search/milp/mod.rs index acba506ee..4e02d8129 100644 --- a/crates/plan-selection/src/search/milp/mod.rs +++ b/crates/plan-selection/src/search/milp/mod.rs @@ -13,5 +13,5 @@ pub(crate) mod report; // Candidate-connection core being integrated into the staged memo pipeline. pub(crate) mod candidate_connections; -pub(crate) mod candidate_reconstruction; pub(crate) mod candidate_inventory; +pub(crate) mod candidate_reconstruction;