feat(plan-selection): memo over logical sub-expressions with a brute-force reference search (--search memo-brute-force)

crates/devtools/src/bin/stage_pipeline.rs

@@ -208,7 +208,7 @@ use asap_devtools::planner_layering_examples::{
     example1, example2, pattern_a, pattern_b, PatternARecurrence, SHARED_ROWS_PER_SECOND,
 };
 use asap_devtools::planning_stages::{
-    lower_promql_roots, lower_sql_roots, run_default_search_over, run_planning_stages_over,
+    lower_promql_roots, lower_sql_roots, run_planning_stages_over, run_search_over,
     PlanningStagesError, PlanningStagesRun,
 };
 use asap_devtools::SqlCatalog;
@@ -220,7 +220,7 @@ use asap_logical_asap_optimization::DeploymentCapabilities;
 use asap_plan_selection::{
     large_peak_memory_warning, ClusterLimits, ClusterObjective, CostModel, MeasuredOperationCosts,
     PlanSelectionSettings, PlanningModels, Rejection, ResourceUnitCost, ResourceUsage, Selection,
-    SelectionError,
+    SelectionError, Stage1ToStage3Search,
 };
 use asap_planner::{installable_plan_of, PlannerExecutorCompileCheck};
 use asap_types::ir::flat::{flatten, FlatDag};
@@ -239,7 +239,8 @@ const USAGE: &str =
 [--name <name>]... [--table <json>...] [--max-error <f64> --failure-probability <f64> --accuracy-metric <metric>] [--sample-interval-ms <u64>] [--series-count <u64>] [--rows-per-second <f64>] [--rows-held <u64>])) \
 [--cost-model usage-over-time|aws-fargate-us-east-1|provisioned-capacity|cluster-capacity [--cluster-memory-bytes <n> \
 --cluster-disk-bytes <n> [--cluster-cpu-cores <n>] --cluster-objective lowest-latency|least-memory|pareto-front]] \
-[--plan-selection-settings <file>] [--enforce-accuracy-guarantees] [--deployment-capabilities <json>] [--exhaustive [--max-combinations <n>]] [--max-displayed <n>] --out <file>";
+[--plan-selection-settings <file>] [--enforce-accuracy-guarantees] [--deployment-capabilities <json>] [--exhaustive [--max-combinations <n>] | --search legacy|memo-brute-force] [--max-displayed <n>] --out <file>\n\
+--search legacy is the default search (the document's \"search\" is \"default\").";
 
 /// The default `--max-combinations`: the most Stage 1 candidates (Pass 2's
 /// included) an exhaustive run enumerates.
@@ -343,6 +344,9 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
     let mut plan_selection_settings_path: Option<String> = None;
     let mut enforce_accuracy_guarantees = false;
     let mut exhaustive = false;
+    // `--search memo-brute-force`: the reference search over the memo of
+    // sub-expressions instead of the default search.
+    let mut memo_brute_force = false;
     let mut max_combinations_given = false;
     let mut args = args.into_iter();
     while let Some(flag) = args.next() {
@@ -405,6 +409,13 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
             "--plan-selection-settings" => plan_selection_settings_path = Some(value()?),
             "--enforce-accuracy-guarantees" => enforce_accuracy_guarantees = true,
             "--exhaustive" => exhaustive = true,
+            "--search" => {
+                memo_brute_force = match value()?.as_str() {
+                    "legacy" => false,
+                    "memo-brute-force" => true,
+                    other => return Err(format!("unknown search {other}").into()),
+                }
+            }
             other => return Err(format!("unknown argument {other}").into()),
         }
     }
@@ -427,6 +438,9 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
     if max_combinations_given && !exhaustive {
         return Err("--max-combinations limits the exhaustive search: it needs --exhaustive (the default search has no limit)".into());
     }
+    if exhaustive && memo_brute_force {
+        return Err("--exhaustive and --search memo-brute-force are two searches: give one".into());
+    }
     if plan_selection_settings_path.is_some()
         && (cost_model.is_some()
             || cluster_memory.is_some()
@@ -646,7 +660,16 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
             Some(max_combinations),
             models,
         ),
-        false => run_default_search_over(roots.clone(), &entries, &data, models),
+        false => run_search_over(
+            roots.clone(),
+            &entries,
+            &data,
+            models,
+            Stage1ToStage3Search {
+                memo_groups: memo_brute_force,
+                ..Stage1ToStage3Search::default()
+            },
+        ),
     };
     // A selected plan's very large peak memory, when the deployment
     // declares no memory capacity to check it against.
@@ -668,7 +691,11 @@ fn run(args: Vec<String>) -> Result<(), Failure> {
     }
     document["workload"]["data_workload"] =
         serde_json::to_value(&data).map_err(|e| e.to_string())?;
-    document["search"] = json!(if exhaustive { "exhaustive" } else { "default" });
+    document["search"] = json!(match (exhaustive, memo_brute_force) {
+        (true, _) => "exhaustive",
+        (false, true) => "memo-brute-force",
+        (false, false) => "default",
+    });
     document["accuracy_guarantees"] = json!(if enforce_accuracy_guarantees {
         "enforced"
     } else {

crates/devtools/src/planning_stages.rs

@@ -227,17 +227,29 @@ pub fn run_default_search_over(
     data: &DataWorkload,
     models: PlanningModels<'_>,
 ) -> Result<PlanningStagesRun, PlanningStagesError> {
-    validate(&roots)?;
-    let per_computation =
-        generate_per_computation_candidates(roots.iter().cloned().enumerate().collect())
-            .map_err(|e| PlanningStagesError::Pass1(e.to_string()))?;
-    let run = run_stage1_to_stage3_from_pass1(
-        per_computation,
+    run_search_over(
+        roots,
         entries,
         data,
         models,
         Stage1ToStage3Search::default(),
-    );
+    )
+}
+
+/// [`run_default_search_over`] with the search `search` (e.g. the
+/// reference search over the memo, `memo_groups`).
+pub fn run_search_over(
+    roots: Vec<QueryRoot>,
+    entries: &[QueryWorkloadEntry],
+    data: &DataWorkload,
+    models: PlanningModels<'_>,
+    search: Stage1ToStage3Search,
+) -> Result<PlanningStagesRun, PlanningStagesError> {
+    validate(&roots)?;
+    let per_computation =
+        generate_per_computation_candidates(roots.iter().cloned().enumerate().collect())
+            .map_err(|e| PlanningStagesError::Pass1(e.to_string()))?;
+    let run = run_stage1_to_stage3_from_pass1(per_computation, entries, data, models, search);
     if let Some(Ok(roots)) = run.stage1_selected.as_ref().map(|c| &c.roots) {
         validate(roots)?;
     }

crates/devtools/tests/stage_pipeline.rs

@@ -1705,3 +1705,54 @@ fn scans_projecting_different_labels_are_shared() {
         "{rolled_up} vs {without_rollup}"
     );
 }
+
+/// The memo search's document names its selected plan's logical candidate
+/// under the id its physical candidates read (`from_logical`), as the
+/// viewer's validator requires: the brute force numbers its logical plans,
+/// so the selected one need not be the first.
+#[test]
+fn the_memo_search_document_names_the_selected_logical_candidate() {
+    let document = document_of(&[
+        "--promql",
+        "quantile_over_time(0.5, data[5m])",
+        "--promql",
+        "quantile_over_time(0.9, data[5m])",
+        "--sample-interval-ms",
+        "15000",
+        "--search",
+        "memo-brute-force",
+    ]);
+    assert_eq!(document["search"], "memo-brute-force");
+    let logical: HashSet<&str> = document["stage1_logical_asap_optimization"]["candidates"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|c| c["id"].as_str().unwrap())
+        .collect();
+    let physical = document["stage2_physical_asap_optimization"]["candidates"]
+        .as_array()
+        .unwrap();
+    assert!(!physical.is_empty());
+    for candidate in physical {
+        let from = candidate["from_logical"].as_str().unwrap();
+        assert!(logical.contains(from), "{from} is not in {logical:?}");
+    }
+    let selected = document["stage3_plan_selection"]["selected"]
+        .as_str()
+        .unwrap();
+    assert!(physical.iter().any(|c| c["id"] == selected));
+}
+
+/// `--exhaustive` and `--search memo-brute-force` are two searches.
+#[test]
+fn exhaustive_and_the_memo_search_are_not_both_given() {
+    let (status, stderr) = failure(&[
+        "--promql",
+        "sum(m)",
+        "--exhaustive",
+        "--search",
+        "memo-brute-force",
+    ]);
+    assert_eq!(status, Some(2));
+    assert!(stderr.contains("give one"), "{stderr}");
+}

crates/integration-tests/tests/joint_search_equivalence.rs

@@ -503,6 +503,7 @@ fn a_known_bound_that_misses_is_not_discarded_before_composition() {
         Stage1ToStage3Search {
             keep_every_candidate: true,
             exhaustive: false,
+            memo_groups: false,
         },
     )
     .unwrap();
@@ -1157,3 +1158,183 @@ fn queries_over_other_windows_or_with_a_latency_bound_are_in_no_class() {
     let search = &run.stage3_plan_selection.as_ref().unwrap().search;
     assert!(search.interchangeable_queries.is_empty());
 }
+
+/// The memo-group search (not the default yet) on seeded random workloads
+/// over one or two gauges, with raw data kept or not, under cost model 1:
+/// it selects a plan of the exhaustive search's price, or both find none.
+#[test]
+fn memo_group_search_selects_the_exhaustive_price() {
+    let mut searched = 0;
+    let mut mismatches = Vec::new();
+    for seed in 1..=60u64 {
+        let gauges: &[&str] = match seed % 3 {
+            0 => &["m", "n"],
+            _ => &["m"],
+        };
+        let workload = random_workload_over(seed, gauges);
+        let roots = lower_promql_roots(&workload).unwrap();
+        let inventory =
+            asap_logical_asap_optimization::per_computation_candidates::generate_per_computation_candidates(
+                roots.iter().cloned().enumerate().collect(),
+            )
+            .unwrap();
+        if inventory.combinations() > MAX_EXHAUSTIVE_COMBINATIONS {
+            continue;
+        }
+        searched += 1;
+        let entries: Vec<_> = workload.query_workload.entries().collect();
+        let data = workload.data_workload.clone().unwrap();
+        let deployment = match seed % 2 {
+            0 => no_raw_data(),
+            _ => asap_logical_asap_optimization::DeploymentCapabilities::default(),
+        };
+        let models = PlanningModels::builtin().with_deployment_capabilities(&deployment);
+        let indexed = || roots.iter().cloned().enumerate().collect::<Vec<_>>();
+        let exhaustive = run_stage1_to_stage3_with(
+            indexed(),
+            &entries,
+            &data,
+            models,
+            Stage1ToStage3Search::exhaustive(),
+        )
+        .unwrap();
+        let memo = run_stage1_to_stage3_with(
+            indexed(),
+            &entries,
+            &data,
+            models,
+            Stage1ToStage3Search {
+                memo_groups: true,
+                ..Stage1ToStage3Search::default()
+            },
+        )
+        .unwrap();
+        let price = |run: &Stage1ToStage3Run<usize>| {
+            run.stage3_plan_selection
+                .as_ref()
+                .ok()
+                .map(|s| s.costs[&s.selected].cost_per_second.unwrap())
+        };
+        let (all, found) = (price(&exhaustive), price(&memo));
+        let same = match (all, found) {
+            (Some(a), Some(f)) => (a - f).abs() <= 1e-9 * a.abs(),
+            (None, None) => true,
+            _ => false,
+        };
+        if !same {
+            mismatches.push(format!(
+                "seed {seed}: exhaustive {all:?}, memo {found:?}: {:?}",
+                workload
+                    .query_workload
+                    .repeating_queries
+                    .as_ref()
+                    .unwrap()
+                    .iter()
+                    .map(|e| e.query.0.clone())
+                    .collect::<Vec<_>>()
+            ));
+        }
+    }
+    assert!(
+        searched >= 20,
+        "only {searched} workloads were small enough"
+    );
+    assert!(mismatches.is_empty(), "{}", mismatches.join("\n"));
+}
+
+/// The memo-group search selects a plan of the exhaustive search's price
+/// on `workload`, with raw data kept and without (or both find none).
+fn assert_memo_selects_the_exhaustive_price(label: &str, workload: &PlanningWorkload) {
+    let roots = lower_promql_roots(workload).unwrap();
+    let entries: Vec<_> = workload.query_workload.entries().collect();
+    let data = workload.data_workload.clone().unwrap();
+    for deployment in [
+        asap_logical_asap_optimization::DeploymentCapabilities::default(),
+        no_raw_data(),
+    ] {
+        let models = PlanningModels::builtin().with_deployment_capabilities(&deployment);
+        let run = |search| {
+            run_stage1_to_stage3_with(
+                roots.iter().cloned().enumerate().collect(),
+                &entries,
+                &data,
+                models,
+                search,
+            )
+            .unwrap()
+        };
+        let price = |run: &Stage1ToStage3Run<usize>| {
+            run.stage3_plan_selection
+                .as_ref()
+                .ok()
+                .map(|s| s.costs[&s.selected].cost_per_second.unwrap())
+        };
+        let exhaustive = price(&run(Stage1ToStage3Search::exhaustive()));
+        let memo = price(&run(Stage1ToStage3Search {
+            memo_groups: true,
+            ..Stage1ToStage3Search::default()
+        }));
+        match (exhaustive, memo) {
+            (Some(a), Some(f)) => assert!(
+                (a - f).abs() <= 1e-9 * a.abs(),
+                "{label}: exhaustive {a}, memo {f}"
+            ),
+            (a, f) => assert_eq!(a.is_some(), f.is_some(), "{label}: {a:?} vs {f:?}"),
+        }
+    }
+}
+
+/// A query's top-k over a sum, beside a query of that sum alone: the top-k's
+/// fused heap absorbs the sum, which is the other query's root, so the memo
+/// keeps the sum its own group for the second query. The memo-group search
+/// selects a plan of the exhaustive search's price, with raw data kept and
+/// without.
+#[test]
+fn a_fused_top_k_beside_its_inner_sum_selects_the_exhaustive_price() {
+    let workload = workload_of(
+        &[
+            ("topk(5, sum by (job) (sum_over_time(m[5m])))", 60_000),
+            ("sum by (job) (sum_over_time(m[5m]))", 60_000),
+        ],
+        AccuracyTarget::max_error(0.01),
+    );
+    assert_memo_selects_the_exhaustive_price("top-k beside its sum", &workload);
+}
+
+/// A per-series quantile over `data[1m]` beside `quantile by (host, job)`
+/// and `quantile by (job)` of `data` (and, in the whole workload, a 5-minute
+/// quantile), every 10 s over 10,000 series: the exhaustive plan rolls the
+/// `by (job)` quantile up from the `by (host, job)` KLL build, which the
+/// memo lists too.
+#[test]
+fn a_sketch_rollup_beside_other_quantiles_selects_the_exhaustive_price() {
+    let workload = |queries: &[&str]| -> PlanningWorkload {
+        let entries: Vec<serde_json::Value> = queries
+            .iter()
+            .map(|q| {
+                serde_json::json!({
+                    "query": q,
+                    "demand": {"fixed_interval": 10000},
+                    "requirements": {"accuracy": {"max_error": 0.01, "failure_probability": 0.001}},
+                    "predictability": {"predictable": {"known_at": null}}
+                })
+            })
+            .collect();
+        serde_json::from_value(serde_json::json!({
+            "query_workload": {"language": "promql", "query_batch": null, "repeating_queries": entries},
+            "data_workload": {
+                "arrival": {"ingesting": {"rows_per_second": 666.67}},
+                "sources": {"data": {"series_count": 10000, "metric_type": "gauge", "sample_interval": 15000}}
+            }
+        }))
+        .unwrap()
+    };
+    let (q1, q2, q3, q4) = (
+        "quantile_over_time(0.5, data[1m])",
+        "quantile_over_time(0.9, data[5m])",
+        "quantile by (host, job) (0.5, data)",
+        "quantile by (job) (0.5, data)",
+    );
+    assert_memo_selects_the_exhaustive_price("Q1, Q3, Q4", &workload(&[q1, q3, q4]));
+    assert_memo_selects_the_exhaustive_price("Q2, Q3, Q4", &workload(&[q2, q3, q4]));
+}

crates/logical-asap-optimization/src/accuracy/estimators/mod.rs

@@ -235,7 +235,7 @@ pub(crate) fn binomial_upper_tail(n: u32, p: f64, k: u32) -> f64 {
 /// shrinks, so each reader's `k` or heap is first resized to the shared
 /// (smallest) failure probability at no larger error ([`kll::shared_k`],
 /// [`univmon::shared_heap`]), and the shared build meets every reader.
-pub(crate) fn size_for_all_readers(params: &[&SketchParams]) -> Option<SketchParams> {
+pub fn size_for_all_readers(params: &[&SketchParams]) -> Option<SketchParams> {
     use SketchParams as P;
     let (first, rest) = params.split_first()?;
     let shared = rest.iter().try_fold((*first).clone(), |a, b| {