feat(dag-viewer): export and visualize planner cost/benefit annotations (#286)

crates/devtools/src/bin/dag_export.rs

@@ -53,8 +53,9 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use std::time::Instant;
 
-use asap_aware_mapping::cost_model::DefaultCostModel;
+use asap_aware_mapping::cost_model::{default_cse_recompute_cost, DefaultCostModel};
 use asap_aware_mapping::replacement::{search_workload, Replacement, ReplacementSubDAG};
+use asap_types::cost::{BaselineRef, CostAnnotation, CostInput, CostSource, CostUnit};
 use asap_types::dag_export::{
     self, DagDecision, DagGraph, DagNote, NamedGraph, PostAsapSubstitution, TargetReplacement,
     TargetReplacementAfter, WorkloadGraph,
@@ -65,6 +66,90 @@ use asap_types::pre_asap::query_expr::QueryExpr;
 use asap_types::pre_asap::schema::{Column, DataType, Schema};
 use asap_types::types::AccuracyTarget;
 
+/// `model_version` tag for every [`CostAnnotation`] this binary computes —
+/// see [`winner_cost_annotations`]'s own doc for what it's a structural
+/// proxy for and why (issue #286, issue #287 is the real rate-modeling
+/// follow-up).
+const STRUCTURAL_COST_MODEL_VERSION: &str = "dag_export-structural-cost-v1";
+
+/// Baseline/selected/benefit [`CostAnnotation`]s for one [`Winner`] — issue
+/// #286's "replacement-region baseline cost, selected cost, and benefit"
+/// granularity item, reused verbatim for [`TargetReplacement`] and for the
+/// [`DagDecision`] carried by every node the winning candidate produced or
+/// carried.
+///
+/// The baseline is always [`BaselineRef::PreAsapRecomputation`]: the cost of
+/// recomputing the winner's own target independently at every one of its
+/// `consumer_count` use sites. `per_consumer_recompute` is that one-site cost
+/// — [`default_cse_recompute_cost`] applied to the target, the exact
+/// structural-size proxy `asap_aware_mapping::cost_model`'s own
+/// `DefaultCostModel` already uses, not a second formula — passed in
+/// pre-computed rather than taking `target: &QueryExpr` and recomputing it
+/// here: this function is called once per matched *node position* a winner
+/// produced or carried (once per `target_replacement` flat entry, once per
+/// `find_winner` hit), and the same winner's target can be reached from more
+/// than one node position, so callers memoize `default_cse_recompute_cost`
+/// once per winner (see `run_post_asap_with_progress`'s
+/// `per_consumer_recompute_costs`) instead of this function re-walking the
+/// identical subtree on every call. This baseline is always computable
+/// (never `Unavailable`), unlike `selected`, which is `Unavailable` whenever
+/// `selected_cost` is `NaN` (the plugged-in cost model has no numeric
+/// estimate for that particular candidate shape — see
+/// [`RankedGroup::costs`](asap_aware_mapping::replacement::RankedGroup::costs)'s
+/// own doc).
+///
+/// Every value here is unit-tagged [`CostUnit::RelativeStructuralUnits`],
+/// not [`CostUnit::CostUnitsPerSecond`]: today's cost model has no
+/// `update_rate`/`evaluation_rate`/`query_interval` recurrence inputs at all
+/// (issue #287's job) — see `asap_types::cost`'s module doc for why this
+/// crate refuses to mislabel a structural proxy as a real rate.
+fn winner_cost_annotations(
+    per_consumer_recompute: f64,
+    consumer_count: usize,
+    selected_cost: f64,
+) -> (CostAnnotation, CostAnnotation, CostAnnotation) {
+    let baseline_value = per_consumer_recompute * consumer_count.max(1) as f64;
+    let baseline = CostAnnotation::modeled(
+        baseline_value,
+        CostUnit::RelativeStructuralUnits,
+        STRUCTURAL_COST_MODEL_VERSION,
+        vec![
+            CostInput::new("per_consumer_recompute_cost", per_consumer_recompute),
+            CostInput::new("consumer_count", consumer_count.max(1) as f64),
+        ],
+    );
+
+    if !selected_cost.is_finite() {
+        return (
+            baseline,
+            CostAnnotation::unavailable(CostUnit::RelativeStructuralUnits),
+            CostAnnotation::unavailable(CostUnit::RelativeStructuralUnits),
+        );
+    }
+
+    let selected = CostAnnotation::modeled(
+        selected_cost,
+        CostUnit::RelativeStructuralUnits,
+        STRUCTURAL_COST_MODEL_VERSION,
+        vec![],
+    )
+    .with_baseline(BaselineRef::PreAsapRecomputation, baseline_value);
+
+    let benefit = CostAnnotation {
+        value: selected.delta,
+        unit: CostUnit::RelativeStructuralUnits,
+        source: CostSource::Modeled,
+        baseline: Some(BaselineRef::PreAsapRecomputation),
+        delta: None,
+        benefit_ratio: selected.benefit_ratio,
+        model_version: Some(STRUCTURAL_COST_MODEL_VERSION.to_string()),
+        benchmark_id: None,
+        inputs: Vec::new(),
+    };
+
+    (baseline, selected, benefit)
+}
+
 use asap_devtools::{lower_promql, lower_sql, SqlCatalog};
 
 enum Lang {
@@ -251,6 +336,10 @@ struct Winner<'a> {
     target: &'a Rc<QueryExpr>,
     candidate: &'a ReplacementSubDAG,
     cost: f64,
+    /// The target's own `MemoGroup::consumer_count` — threaded through so
+    /// [`winner_cost_annotations`] can compute a baseline without a second
+    /// lookup back into `PlanSpace`.
+    consumer_count: usize,
 }
 
 /// Short explanation intended for a selected winner in node-level UI. The
@@ -316,12 +405,41 @@ fn lookup_winner(
         .find(|&i| expr == winners[i].target.as_ref())
 }
 
+/// One `(decision.id, baseline_cost, selected_cost)` triple per *distinct*
+/// [`DagDecision`] carried anywhere in `graph` — collapsing every node that
+/// shares one `decision.id` (a replacement region can span many nodes, all
+/// carrying an identical clone of the same decision) down to a single
+/// entry, so a caller summing these never counts one decision's cost once
+/// per node it happens to touch.
+fn decision_cost_entries(graph: &DagGraph) -> Vec<(u32, CostAnnotation, CostAnnotation)> {
+    let mut seen = std::collections::HashSet::new();
+    let mut entries = Vec::new();
+    for node in &graph.nodes {
+        let Some(decision) = &node.decision else {
+            continue;
+        };
+        if !seen.insert(decision.id) {
+            continue;
+        }
+        let (Some(baseline), Some(selected)) = (&decision.baseline_cost, &decision.selected_cost)
+        else {
+            continue;
+        };
+        entries.push((decision.id, baseline.clone(), selected.clone()));
+    }
+    entries
+}
+
 /// Build a [`TargetReplacement`] for `winner`, matching this file's own
-/// per-target `before`/`after` construction.
+/// per-target `before`/`after` construction. `per_consumer_recompute` is
+/// `winner`'s own memoized [`default_cse_recompute_cost`] — see
+/// `winner_cost_annotations`'s own doc for why callers pass this in instead
+/// of recomputing it here.
 fn target_replacement(
     decision_id: u32,
     target_pre_id: u32,
     winner: &Winner<'_>,
+    per_consumer_recompute: f64,
 ) -> TargetReplacement {
     let strategy = winner.candidate.strategy.to_string();
     let before = dag_export::export(winner.target);
@@ -333,15 +451,26 @@ fn target_replacement(
             TargetReplacementAfter::Rewrite(dag_export::export(rewritten))
         }
     };
+    let (baseline_cost, selected_cost, benefit) =
+        winner_cost_annotations(per_consumer_recompute, winner.consumer_count, winner.cost);
+    // Derived from `selected_cost` (not read from `winner.cost` a second
+    // time) so the legacy scalar field and the structured annotation can
+    // never drift apart at this call site — see `winner_cost_annotations`'s
+    // own doc for why `selected_cost.value` is `None` (and this is `NaN`)
+    // in exactly the same case `winner.cost` itself would already be `NaN`.
+    let cost = selected_cost.value.unwrap_or(f64::NAN);
     TargetReplacement {
         decision_id,
         target_pre_id,
         strategy,
         rationale: decision_rationale(winner),
         rank: 0,
-        cost: winner.cost,
+        cost,
         before,
         after,
+        baseline_cost: Some(baseline_cost),
+        selected_cost: Some(selected_cost),
+        benefit: Some(benefit),
     }
 }
 
@@ -461,10 +590,25 @@ fn run_post_asap_with_progress(
                 target: group.target,
                 candidate,
                 cost: group.costs[0],
+                consumer_count: group.consumer_count,
             })
         })
         .collect();
 
+    // One `default_cse_recompute_cost` walk per *winner*, not per node a
+    // winner's decision ends up cloned onto: `winner_cost_annotations` is
+    // called once per matched node position — once per `target_replacement`
+    // flat entry, and once per `find_winner` hit inside `export_post_asap`'s
+    // traversal — and a single target can be reached from more than one
+    // node position (an internally-shared subtree within one query, or the
+    // same CSE-shared target reached from several queries), so recomputing
+    // this per call would re-walk the identical subtree redundantly.
+    // Indexed in parallel with `winners`.
+    let per_consumer_recompute_costs: Vec<f64> = winners
+        .iter()
+        .map(|winner| default_cse_recompute_cost(winner.target).0)
+        .collect();
+
     // `by_hash` only narrows the search; `lookup_winner`'s own structural
     // equality check is the real decision — see that function's doc.
     let mut by_hash_cache = HashCache::new();
@@ -503,13 +647,24 @@ fn run_post_asap_with_progress(
     let mut find_winner = |expr: &QueryExpr| -> Option<PostAsapSubstitution> {
         let i = lookup_winner(&by_hash, &winners, &mut post_graph_cache, expr)?;
         let winner = &winners[i];
+        let (baseline_cost, selected_cost, benefit) = winner_cost_annotations(
+            per_consumer_recompute_costs[i],
+            winner.consumer_count,
+            winner.cost,
+        );
+        // Derived from `selected_cost`, not `winner.cost` a second time —
+        // see `target_replacement`'s identical derivation for why.
+        let cost = selected_cost.value.unwrap_or(f64::NAN);
         let decision = DagDecision {
             id: i as u32,
             strategy: winner.candidate.strategy.to_string(),
             rationale: decision_rationale(winner),
             rank: 0,
-            cost: winner.cost,
+            cost,
             role: "replacement_region",
+            baseline_cost: Some(baseline_cost),
+            selected_cost: Some(selected_cost),
+            benefit: Some(benefit),
         };
         Some(match &winners[i].candidate.replacement {
             Replacement::Rewrite(rc) => PostAsapSubstitution::Rewrite {
@@ -555,7 +710,12 @@ fn run_post_asap_with_progress(
             if let Some(i) = lookup_winner(&by_hash, &winners, &mut lookup_cache, source_expr) {
                 replacements.push((
                     name.clone(),
-                    target_replacement(i as u32, node.id, &winners[i]),
+                    target_replacement(
+                        i as u32,
+                        node.id,
+                        &winners[i],
+                        per_consumer_recompute_costs[i],
+                    ),
                 ));
                 matched[i] = true;
             }
@@ -675,6 +835,7 @@ async fn main() {
             graph,
             replacements: Vec::new(),
             post_graph: None,
+            workload_cost: None,
         });
     }
     for (explanation, matched) in explanations.iter().zip(matched) {
@@ -718,7 +879,59 @@ async fn main() {
         assign_workload_node_ids(&mut post_graphs);
     }
 
-    let workload = WorkloadGraph { queries };
+    // Whole selected-workload cost/benefit (issue #286) — per query, and
+    // for the whole selected workload. `decision.id` (== the winning
+    // `Winner`'s own index — see `run_post_asap_with_progress`) is already
+    // a collision-free dedup key for a decision shared across multiple
+    // nodes (a replacement region spans several nodes, all carrying the
+    // same `decision.id`) and across multiple queries (a CSE-shared target
+    // reachable from more than one query's root) alike, so it's reused
+    // directly as `sum_workload_costs`'s dedup key — no separate lookup
+    // needed.
+    let mut workload_entries = Vec::new();
+    for query in &mut queries {
+        let Some(post_graph) = &query.post_graph else {
+            continue;
+        };
+        let entries = decision_cost_entries(post_graph);
+        if entries.is_empty() {
+            continue;
+        }
+        match asap_types::cost::workload_cost_summary(
+            entries
+                .iter()
+                .map(|(id, baseline, selected)| (Some(*id), baseline, selected)),
+            "dag_export-workload-cost-v1",
+        ) {
+            Ok(summary) => query.workload_cost = Some(summary),
+            Err(mismatch) => eprintln!(
+                "dag_export: workload cost aggregation for {:?} skipped — {mismatch}",
+                query.name
+            ),
+        }
+        workload_entries.extend(entries);
+    }
+    let workload_cost = if workload_entries.is_empty() {
+        None
+    } else {
+        match asap_types::cost::workload_cost_summary(
+            workload_entries
+                .iter()
+                .map(|(id, baseline, selected)| (Some(*id), baseline, selected)),
+            "dag_export-workload-cost-v1",
+        ) {
+            Ok(summary) => Some(summary),
+            Err(mismatch) => {
+                eprintln!("dag_export: workload-wide cost aggregation skipped — {mismatch}");
+                None
+            }
+        }
+    };
+
+    let workload = WorkloadGraph {
+        queries,
+        workload_cost,
+    };
     if progress {
         eprintln!(
             "Total planner time: {:.2} ms",