From 3e40ee0a590ed61c292a5e7050300b60ca9d4fe7 Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 25 Aug 2026 19:59:36 -0600 Subject: [PATCH 01/27] feat(devtools): materialize post-ASAP replacements for dag_export --post-asap Add a --post-asap flag to the dag_export devtools binary that runs asap-aware-mapping's replacement search (default_strategies() plus the unwired AvgToSumOverCountStrategy) over the lowered workload and surfaces what it found, two ways: - NamedGraph.replacements: Vec - one small, self-contained before/after pair per independently-discovered replacement site (sketch mapping, CSE share/recompute, workload-aware rollup, Hydra grouping, avg -> sum/count rewrite), each with a strategy label, rationale, rank/cost, and an after tagged {"kind": "Summary"|"Rewrite", "graph": {...}}. - NamedGraph.post_graph: Option - one merged, whole-query graph per query with every winning candidate spliced directly into the query's own pre-ASAP shape in place, built via a new asap_types::dag_export::export_post_asap. asap_types::dag_export gains the generic, crate-agnostic shapes for this (SummaryDagNode/SummaryDagGraph/export_summary, TargetReplacement/ TargetReplacementAfter, PostAsapSubstitution/export_post_asap) following the same layering discipline DagNode::notes already established: this crate never runs the search or picks a winner, only defines shapes a higher layer (the dag_export binary) populates after the fact. DagNode's source_expr becomes Option (None for a post-ASAP-originated node with no corresponding QueryExpr). Without --post-asap, output is byte-identical to before (both new fields are empty/None and skipped from JSON). Sample outputs covering all four --post-asap replacement kinds saved to /tmp/post_asap_samples/*.json for the frontend dag-viewer track. Co-Authored-By: Claude Sonnet 5 --- crates/devtools/src/bin/dag_export.rs | 413 ++++++++++++++++++- crates/types/src/dag_export.rs | 555 ++++++++++++++++++++++++-- 2 files changed, 928 insertions(+), 40 deletions(-) diff --git a/crates/devtools/src/bin/dag_export.rs b/crates/devtools/src/bin/dag_export.rs index 591773c5..99a83237 100644 --- a/crates/devtools/src/bin/dag_export.rs +++ b/crates/devtools/src/bin/dag_export.rs @@ -17,8 +17,54 @@ // note. Pass it to actually exercise that path, e.g.: // cargo run -p asap-lower --bin dag_export -- \ // --epsilon 0.01 --sql "SELECT quantile(0.99, latency) FROM metrics" --name p99 +// +// `--post-asap` is optional and off by default. When passed, this binary +// additionally runs `asap_aware_mapping::replacement::search_workload_with` +// (with `default_strategies()` plus `AvgToSumOverCountStrategy`, which isn't +// one of that crate's own defaults) over every lowered query and ranks each +// discovered `MemoGroup` via `PlanSpace::cost_sorted`. The best-ranked +// candidate per group feeds two additive outputs: +// +// - one `asap_types::dag_export::TargetReplacement` per group on whichever +// query's `NamedGraph.replacements` contains that target node (matched +// by `DagNode::hash` + structural equality, the same collision-safe +// pattern `annotate_with_explanations` below already uses for notes) — +// a small, self-contained "before -> after" pair per replacement site; +// - one merged `NamedGraph.post_graph`: a single flattened graph per +// query with every winning candidate spliced directly into the query's +// own pre-ASAP shape in place, built via +// `asap_types::dag_export::export_post_asap`. +// +// Together these surface every one of the four concrete replacement kinds: +// the sketch family `SketchAlgorithmStrategy`/`HydraGroupingStrategy` bound, +// the CSE share/recompute choice `SharedSubtreeStrategy` found, the +// workload-aware roll-up `RollupStrategy` derived, and the `avg -> +// sum/count` rewrite `AvgToSumOverCountStrategy` proposes. Without +// `--post-asap`, every existing invocation of this binary produces +// byte-identical output to before (`NamedGraph.replacements` is empty and +// `post_graph` is `None`, both skipped from the JSON entirely in that +// case). E.g.: +// cargo run -p asap-lower --bin dag_export -- \ +// --post-asap --epsilon 0.01 \ +// --sql "SELECT quantile(0.95, latency) FROM metrics" --name q1 + +use std::collections::HashMap; +use std::rc::Rc; -use asap_types::dag_export::{self, DagGraph, DagNote, NamedGraph, WorkloadGraph}; +use asap_aware_mapping::cost_model::DefaultCostModel; +use asap_aware_mapping::replacement::{ + default_strategies, search_workload_with, Replacement, ReplacementProvenance, + ReplacementStrategy, ReplacementSubDAG, +}; +use asap_aware_mapping::rewrite::AvgToSumOverCountStrategy; +use asap_types::dag_export::{ + self, DagGraph, DagNote, NamedGraph, PostAsapSubstitution, TargetReplacement, + TargetReplacementAfter, WorkloadGraph, +}; +use asap_types::post_asap::{GroupingStrategy, SummaryExpr}; +use asap_types::pre_asap::agg_intent::AggIntent; +use asap_types::pre_asap::cse::{structural_hash, HashCache}; +use asap_types::pre_asap::query_expr::QueryExpr; use asap_types::pre_asap::schema::{Column, DataType, Schema}; use asap_types::types::AccuracyTarget; @@ -56,12 +102,13 @@ fn catalog() -> SqlCatalog { /// Parses `--sql "" --name "