feat(plan-selection): group selection MILP over execution contexts as Stage 3's default search (#1095)

crates/cost/src/lib.rs

@@ -41,9 +41,10 @@ pub use measured::{
 pub use operator_statistics::EdgeStatistics;
 pub use ranking::{cheaper, worst_latency_ms, PlanRanking};
 pub use usage::{
-    family_name, kept_tumbling_window_chains, node_map, node_usage, plan_cost, query_reach,
-    trailing_window, tumbling_window_retention_ms, tumbling_window_retention_ms_of_plan,
-    tumbling_windows_of, NodeMap, NodeUsage, ProducerOutput,
+    family_name, kept_tumbling_window_chains, maintained_window_ms, node_map, node_usage,
+    plan_cost, query_reach, retained_tumbling_windows, trailing_window,
+    tumbling_window_retention_ms, tumbling_window_retention_ms_of_plan, tumbling_windows_of,
+    NodeMap, NodeUsage, ProducerOutput,
 };
 
 /// Cost per second of wall time; one unit is one CPU-millisecond under

crates/cost/src/usage.rs

@@ -1487,9 +1487,9 @@ fn merged_states(
     })
 }
 
-/// The tumbling windows a chain in `retention_ms` ([`tumbling_window_retention_ms`]) keeps
-/// completed: its retention over its tumbling-window width, at least its tumbling-window count.
-fn retained_tumbling_windows(
+/// The panes a chain in `retention_ms` ([`tumbling_window_retention_ms`]) keeps
+/// completed: its retention over its pane width, at least its pane count.
+pub fn retained_tumbling_windows(
     node: &asap_types::ir::physical_export::PhysicalASAPDAGNode,
     retention_ms: &BTreeMap<PhysicalASAPNodeId, u64>,
 ) -> u64 {
@@ -1546,7 +1546,7 @@ pub fn query_reach(
 /// The window an ingestion-time grouping `node` maintains: a tumbling-window chain's
 /// tumbling windows, or the time range below it (through filters, projections and
 /// time shifts). `None` when there is none.
-fn maintained_window_ms(
+pub fn maintained_window_ms(
     node: PhysicalASAPNodeId,
     nodes: &NodeMap<'_>,
     producers_of: &HashMap<PhysicalASAPNodeId, Vec<PhysicalASAPNodeId>>,

crates/devtools/Cargo.toml

@@ -32,6 +32,11 @@ serde_json = "1"
 serde_yaml = "0.9"
 tokio = { version = "1", features = ["rt", "macros", "rt-multi-thread"] }
 
+[features]
+# Each Stage 3 selection MILP solve also reports where its LP relaxation
+# loses its bound, with `ASAP_STAGE3_LP_DIAGNOSIS=<path>` (measurement only).
+lp-diagnosis = ["asap-plan-selection/lp-diagnosis"]
+
 [dev-dependencies]
 asap-types = { path = "../types", features = ["test-support"] }
 asap-physical-asap-optimization = { path = "../physical-asap-optimization", features = ["test-support"] }

crates/devtools/src/bin/stage_pipeline.rs

@@ -72,8 +72,13 @@
 //     accepted regardless of guarantee, `accepted_regardless_of_guarantee`,
 //     `guarantee_check_rejection`, the sketches read and a one-line summary),
 //     `pruned_candidates` and `search_report` (`SearchReport`: `status`,
-//     `cost_lower_bound_per_second`, MILP size and timings); `error` when no
-//     plan was selected. `tumbling_window_retention_ms` maps each tumbling-window chain the
+//     `cost_lower_bound_per_second`, MILP size and timings, and
+//     `inexact_objective` when the new search's objective and `plan_cost`
+//     differ over terms the MILP encodes inexactly). Stage 3 runs #1095's
+//     selection MILP, which supports cost model 1 only (cost models 2 and 3
+//     are an `UnsupportedCostModel` error); `ASAP_STAGE3_SEARCH=old` runs the
+//     old engine instead; `error` when no
+//     plan was selected. `tumbling_window_retention_ms` maps each pane chain the
 //     selected plan keeps beyond one evaluation (by node id in `P1`) to how
 //     far back the runtime must keep its tumbling windows (`RealizedPlan::tumbling_window_retention_ms`).
 //     `materialized_views` are the selected plan's materialized views as

crates/devtools/tests/stage_pipeline.rs

@@ -2552,3 +2552,171 @@ fn memo_section_lists_each_step_the_selected_choice_and_the_enumeration() {
     std::fs::remove_file(out).unwrap();
     assert!(document.get("memo").is_none());
 }
+
+/// `stage_pipeline`'s document with Stage 3's new search forced on
+/// (`ASAP_STAGE3_SEARCH=milp`, #1095), set on the child process only.
+fn document_with_milp_search(args: &[&str]) -> Value {
+    static CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
+    let call = CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+    let out = std::env::temp_dir().join(format!(
+        "stage_pipeline_{}_milp_{call}.json",
+        std::process::id(),
+    ));
+    let status = Command::new(env!("CARGO_BIN_EXE_stage_pipeline"))
+        .env("ASAP_STAGE3_SEARCH", "milp")
+        .args(args)
+        .arg("--out")
+        .arg(&out)
+        .status()
+        .unwrap();
+    assert!(status.success());
+    let document = serde_json::from_str(&std::fs::read_to_string(&out).unwrap()).unwrap();
+    std::fs::remove_file(out).unwrap();
+    document
+}
+
+/// #1095 C5: with the new search on, example 1 and the RangeHeavy T5
+/// q0.5 + q0.9 workload select a plan through the real pipeline, solved to
+/// optimality or stopped at the deadline, whose `plan_cost` equals the MILP's objective exactly (the
+/// search errors on a difference and reports `inexact_objective` when the
+/// model encoded a term inexactly) and which passes every query's accuracy
+/// check.
+#[test]
+fn milp_search_objective_equals_plan_cost_on_example1_and_rangeheavy() {
+    let rangeheavy = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("../../tools/dag-viewer/memo-examples/promql/rangeheavy-t5-q05-q09-1h.json");
+    let runs: [Vec<&str>; 2] = [
+        vec![
+            "--example",
+            "planner-layering-1",
+            "--enforce-accuracy-guarantees",
+        ],
+        vec![
+            "--workload",
+            rangeheavy.to_str().unwrap(),
+            "--deployment-capabilities",
+            r#"{"raw_data_retained": true, "query_time_results_retained": false}"#,
+            "--cost-model",
+            "aws-fargate-us-east-1",
+            "--enforce-accuracy-guarantees",
+        ],
+    ];
+    for args in runs {
+        let document = document_with_milp_search(&args);
+        let selection = &document["stage3_plan_selection"];
+        assert!(selects_a_plan(selection), "{args:?}: {selection}");
+        let report = &selection["search_report"];
+        // Optimal or not by the deadline, the selected plan's `plan_cost`
+        // equals the MILP's objective: the search errors on a difference,
+        // and no term is encoded inexactly here.
+        assert!(
+            report["status"] == "optimal" || report["status"] == "deadline_reached",
+            "{args:?}: {report}"
+        );
+        assert!(
+            report.get("inexact_objective").is_none(),
+            "{args:?}: {report}"
+        );
+        assert!(report["milp_variables"].as_u64().unwrap() > 0);
+        let queries = document["workload"]["queries"].as_array().unwrap().len();
+        let accepted = selection["accepted_accuracy_per_query"].as_array().unwrap();
+        assert_eq!(accepted.len(), queries);
+        for query in accepted {
+            assert_eq!(query["accepted_regardless_of_guarantee"], false, "{query}");
+        }
+    }
+}
+
+/// #1127 regressions: with the new search on and guarantees enforced, the
+/// memo examples on which the A1 rows admitted a plan that `satisfies`
+/// rejects now select a plan that passes every query's accuracy check:
+/// ratio-of-quantiles (KLL rank errors compose no guarantee through the
+/// quotient, so no query can use KLL) and rollupheavy-t11 (HLL's
+/// failure probability is 1 over the declared 10^6 distinct users, above
+/// its contract's 2^17, so Pass 1 prunes it for both queries).
+#[test]
+fn milp_search_selects_checked_plans_on_ratio_and_rollupheavy() {
+    let root =
+        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../tools/dag-viewer/memo-examples");
+    let tables: Vec<String> = serde_json::from_str::<Vec<Value>>(
+        &std::fs::read_to_string(root.join("sql/tables.json")).unwrap(),
+    )
+    .unwrap()
+    .iter()
+    .map(Value::to_string)
+    .collect();
+    for (workload, pruned_family) in [
+        ("promql/ratio-of-quantiles-5m.json", "Kll"),
+        ("sql/rollupheavy-t11-service-endpoint-5m.json", "Hll"),
+    ] {
+        let path = root.join(workload);
+        let mut args = vec![
+            "--workload",
+            path.to_str().unwrap(),
+            "--deployment-capabilities",
+            r#"{"raw_data_retained": true, "query_time_results_retained": false}"#,
+            "--cost-model",
+            "aws-fargate-us-east-1",
+            "--enforce-accuracy-guarantees",
+        ];
+        if workload.starts_with("sql/") {
+            for table in &tables {
+                args.extend(["--table", table.as_str()]);
+            }
+        }
+        let document = document_with_milp_search(&args);
+        let selection = &document["stage3_plan_selection"];
+        assert!(selects_a_plan(selection), "{workload}: {selection}");
+        let queries = document["workload"]["queries"].as_array().unwrap().len();
+        let accepted = selection["accepted_accuracy_per_query"].as_array().unwrap();
+        assert_eq!(accepted.len(), queries);
+        for query in accepted {
+            assert_eq!(query["accepted_regardless_of_guarantee"], false, "{query}");
+            assert!(
+                !query["sketches"]
+                    .as_array()
+                    .unwrap()
+                    .iter()
+                    .any(|sketch| sketch == pruned_family),
+                "{workload}: {query}"
+            );
+        }
+        let pruned = pass1_pruned(&document);
+        assert!(
+            pruned.iter().any(|(family, reason)| family == pruned_family
+                && reason == "NoErrorBoundForQueryTarget"),
+            "{workload}: {pruned:?}"
+        );
+    }
+}
+
+/// #1095 C5: on RangeHeavy T8+T5 (`q0.9 / q0.5` with both operands also
+/// queried alone), the new search selects a DDSketch plan that passes every
+/// query's accuracy check (the A1 rows admitted KLL there, which the check
+/// rejects). Which chains it shares depends on the deadline; the weights
+/// that let one chain serve all three queries are checked in
+/// `groups::tests::one_pane_chain_serves_a_ratio_and_its_operands`.
+#[test]
+fn milp_search_selects_a_checked_ddsketch_plan_for_a_ratio_and_its_operands() {
+    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("../../tools/dag-viewer/memo-examples/promql/rangeheavy-t8-t5-1h.json");
+    let document = document_with_milp_search(&[
+        "--workload",
+        path.to_str().unwrap(),
+        "--deployment-capabilities",
+        r#"{"raw_data_retained": true, "query_time_results_retained": false}"#,
+        "--cost-model",
+        "aws-fargate-us-east-1",
+        "--enforce-accuracy-guarantees",
+    ]);
+    let selection = &document["stage3_plan_selection"];
+    assert!(selects_a_plan(selection), "{selection}");
+    for query in selection["accepted_accuracy_per_query"].as_array().unwrap() {
+        assert_eq!(query["accepted_regardless_of_guarantee"], false, "{query}");
+        assert_eq!(
+            query["sketches"],
+            serde_json::json!(["DDSketch"]),
+            "{query}"
+        );
+    }
+}

crates/devtools/tests/viewer_contract.rs

@@ -353,6 +353,7 @@ fn viewer_knows_every_view_refresh_and_rejection_cause() {
         C::LatencyBoundExceeded,
         C::LatencyOverRepeatInterval,
         C::OverClusterCapacity,
+        C::Infeasible,
         C::Costlier,
         C::Slower,
         C::MoreMemory,
@@ -366,6 +367,7 @@ fn viewer_knows_every_view_refresh_and_rejection_cause() {
             | C::LatencyBoundExceeded
             | C::LatencyOverRepeatInterval
             | C::OverClusterCapacity
+            | C::Infeasible
             | C::Costlier
             | C::Slower
             | C::MoreMemory

crates/plan-selection/Cargo.toml

@@ -25,6 +25,11 @@ highs = "=2.4.0"
 # `reference_pricer`: the whole-workload selector over enumerated plans, for
 # the developer tools' reference pricing (`asap-devtools`).
 reference-pricer = []
+# `lp-diagnosis`: with `ASAP_STAGE3_LP_DIAGNOSIS=<path>`, each selection MILP
+# solve also reports where its LP relaxation loses its bound (measurement
+# only, for the developer tools; a diagnosed run's timings are not
+# comparable with an undiagnosed one's).
+lp-diagnosis = []
 
 [dev-dependencies]
 asap-types = { path = "../types", features = ["test-support"] }