diff --git a/crates/asap-aware-mapping/src/cost_model.rs b/crates/asap-aware-mapping/src/cost_model.rs index 9e92cf60..82a3e97e 100644 --- a/crates/asap-aware-mapping/src/cost_model.rs +++ b/crates/asap-aware-mapping/src/cost_model.rs @@ -676,7 +676,7 @@ mod tests { SummaryNode { expr: SummaryExpr::SummaryAgg { child: std::rc::Rc::new(SummaryNode { - expr: SummaryExpr::KeepPreAsap(Box::new(scan())), + expr: SummaryExpr::KeepPreAsap(Rc::new(scan())), schema: SummarySchema { fields: vec![], time_index: None, @@ -868,6 +868,7 @@ mod tests { let root = Rc::new(scan()); let target = TargetSubDAG::new(&root); let candidate = ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Summary(Rc::new(summary_node(SummaryFamilyType::Plain( asap_types::pre_asap::DataType::Float64, )))), @@ -891,6 +892,7 @@ mod tests { let target = TargetSubDAG::new(&root); let cheap = ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Summary(Rc::new(summary_node( SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum), ))), @@ -898,6 +900,7 @@ mod tests { rationale: "exact accumulator".into(), }; let pricey = ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Summary(Rc::new(summary_node(SummaryFamilyType::StatModel( asap_types::post_asap::StatModelKind::Parametric, asap_types::post_asap::StatModelParams::Parametric { @@ -936,11 +939,13 @@ mod tests { let target = TargetSubDAG::with_consumer_count(&target_root, 20); let share = ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Rewrite(Rc::clone(&target_root)), provenance: crate::replacement::ReplacementProvenance::CseShare, rationale: "build once and share".into(), }; let recompute = ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Rewrite(Rc::new((*target_root).clone())), provenance: crate::replacement::ReplacementProvenance::CseRecompute, rationale: "build independently".into(), diff --git a/crates/asap-aware-mapping/src/explanation.rs b/crates/asap-aware-mapping/src/explanation.rs index 84a5212e..bcfee70a 100644 --- a/crates/asap-aware-mapping/src/explanation.rs +++ b/crates/asap-aware-mapping/src/explanation.rs @@ -534,7 +534,8 @@ mod tests { .find(|f| f.kind == ExplanationKind::SketchApproximation) .expect("expected a sketch finding"); assert_eq!( - sketch.node_hash, expected_hash, + Some(sketch.node_hash), + expected_hash, "ReplacementExplanation::node_hash must match dag_export's DagNode::hash \ for the same QueryExpr subtree" ); diff --git a/crates/asap-aware-mapping/src/grouping.rs b/crates/asap-aware-mapping/src/grouping.rs index b8aec8c5..c5485f40 100644 --- a/crates/asap-aware-mapping/src/grouping.rs +++ b/crates/asap-aware-mapping/src/grouping.rs @@ -200,6 +200,7 @@ impl<'a> HydraGroupingStrategy<'a> { let patched = with_grouping(node, grouping); Some(ReplacementSubDAG { + strategy: "HydraGroupingStrategy", replacement: Replacement::Summary(patched), provenance: crate::replacement::ReplacementProvenance::SummaryImplementation, rationale: format!( diff --git a/crates/asap-aware-mapping/src/lib.rs b/crates/asap-aware-mapping/src/lib.rs index ad98fdc7..cd52612a 100644 --- a/crates/asap-aware-mapping/src/lib.rs +++ b/crates/asap-aware-mapping/src/lib.rs @@ -175,6 +175,7 @@ pub mod grouping; pub mod replacement; pub mod rewrite; pub mod rollup; +pub mod topk_reuse; pub use cost_model::{CostModel, DefaultCostModel}; pub use explanation::{ @@ -189,3 +190,4 @@ pub use replacement::{ MAX_SEARCH_ITERATIONS, }; pub use rewrite::AvgToSumOverCountStrategy; +pub use topk_reuse::TopKLimitReuseStrategy; diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index 257455c7..e18a9a50 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -365,6 +365,7 @@ use thiserror::Error; use crate::cost_model::{CostModel, CseCandidate, DefaultCostModel, ShareDecision}; use crate::grouping::HydraGroupingStrategy; use crate::rollup::RollupStrategy; +use crate::topk_reuse::TopKLimitReuseStrategy; /// Errors from the pre-ASAP → post-ASAP replacement/construction path /// ([`realize_child`] and [`keep_pre_asap`]). Moved here from the former @@ -446,6 +447,10 @@ pub enum Replacement { #[derive(Debug, Clone)] pub struct ReplacementSubDAG { pub replacement: Replacement, + /// Name of the [`ReplacementStrategy`] that proposed this candidate. + /// Search fills this from `ReplacementStrategy::name`; consumers must not + /// infer it from the replacement's shape or provenance. + pub strategy: &'static str, /// Machine-readable origin/role of this alternative. Selection uses this /// instead of inferring strategy semantics from replacement shape or /// pointer identity when several strategies contribute to one memo group. @@ -476,6 +481,16 @@ pub enum ReplacementProvenance { /// don't match, so a caller that skips the `matches` check first still gets a /// safe (merely uninformative) answer instead of a crash. pub trait ReplacementStrategy { + /// Stable, human-readable strategy name carried into every proposed + /// candidate and ultimately into planner diagnostics/visualizations. + fn name(&self) -> &'static str { + let short = std::any::type_name::() + .rsplit("::") + .next() + .expect("a Rust type name always has a final segment"); + short.split_once('<').map_or(short, |(base, _)| base) + } + /// Does this strategy have any replacement to offer for `target`? fn matches(&self, target: &TargetSubDAG<'_>) -> bool; @@ -1093,6 +1108,7 @@ impl ReplacementStrategy for SketchAlgorithmStrategy<'_> { let rationale = describe_implementation(intent, &implementation); let node = construct_summary(target.root, implementation, self.cost_model).ok()?; Some(ReplacementSubDAG { + strategy: "SketchAlgorithmStrategy", replacement: Replacement::Summary(node), provenance: ReplacementProvenance::SummaryImplementation, rationale, @@ -1219,10 +1235,14 @@ pub(crate) fn realize_child( /// candidate for a target, or a deployment wants to force a node its own /// runtime can't actually implement — through the same fallback this /// crate's own dispatch uses, without duplicating the schema-lift logic. -pub fn keep_pre_asap(expr: &QueryExpr) -> Result, ImplementError> { +pub fn keep_pre_asap(expr: &Rc) -> Result, ImplementError> { + keep_pre_asap_rc(Rc::clone(expr)) +} + +fn keep_pre_asap_rc(expr: Rc) -> Result, ImplementError> { let schema = expr.output_schema()?; Ok(Rc::new(SummaryNode { - expr: SummaryExpr::KeepPreAsap(Box::new(expr.clone())), + expr: SummaryExpr::KeepPreAsap(expr), schema: lift(&schema), })) } @@ -1292,7 +1312,7 @@ pub(crate) fn construct_summary( } } } - keep_pre_asap(expr) + keep_pre_asap_rc(Rc::new(expr.clone())) } /// Translate an [`Implementation`] into the `(family, needs a @@ -1493,6 +1513,7 @@ impl ReplacementStrategy for SharedSubtreeStrategy { let count = target.consumer_count; vec![ ReplacementSubDAG { + strategy: "SharedSubtreeStrategy", // The already-interned `Rc` itself: reusing it verbatim *is* // "build once and share" — no new node to construct. replacement: Replacement::Rewrite(Rc::clone(target.root)), @@ -1504,6 +1525,7 @@ impl ReplacementStrategy for SharedSubtreeStrategy { ), }, ReplacementSubDAG { + strategy: "SharedSubtreeStrategy", // A structurally-identical but freshly-allocated `Rc`: same // value (`PartialEq`), deliberately *not* the same pointer, // representing "undo the sharing and recompute independently". @@ -2528,6 +2550,14 @@ fn search_cse_workload_with<'s, Id>( }) .collect(); let rollup_strategy = RollupStrategy::new(&siblings); + let limits: Vec> = order + .iter() + .filter_map(|ptr| { + let node = &nodes[ptr]; + matches!(node.as_ref(), QueryExpr::Limit { .. }).then(|| Rc::clone(node)) + }) + .collect(); + let topk_reuse_strategy = TopKLimitReuseStrategy::new(&limits); let mut groups: HashMap<*const QueryExpr, MemoGroup> = HashMap::new(); for ptr in &order { @@ -2565,11 +2595,32 @@ fn search_cse_workload_with<'s, Id>( let mut proposed = Vec::new(); for strategy in strategies { if strategy.matches(&target) { - proposed.extend(strategy.replacements(&target)); + let name = strategy.name(); + proposed.extend(strategy.replacements(&target).into_iter().map( + |mut candidate| { + candidate.strategy = name; + candidate + }, + )); } } if rollup_strategy.matches(&target) { - proposed.extend(rollup_strategy.replacements(&target)); + let name = rollup_strategy.name(); + proposed.extend(rollup_strategy.replacements(&target).into_iter().map( + |mut candidate| { + candidate.strategy = name; + candidate + }, + )); + } + if topk_reuse_strategy.matches(&target) { + let name = topk_reuse_strategy.name(); + proposed.extend(topk_reuse_strategy.replacements(&target).into_iter().map( + |mut candidate| { + candidate.strategy = name; + candidate + }, + )); } for candidate in &proposed { @@ -4314,16 +4365,19 @@ mod tests { let mut group = MemoGroup::new(Rc::clone(&target), 2); group.candidates = vec![ ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Rewrite(Rc::clone(&target)), provenance: ReplacementProvenance::CseShare, rationale: "share".into(), }, ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Rewrite(Rc::new(target.as_ref().clone())), provenance: ReplacementProvenance::CseRecompute, rationale: "recompute".into(), }, ReplacementSubDAG { + strategy: "TestStrategy", replacement: Replacement::Rewrite(Rc::new(QueryExpr::CurrentTimestamp)), provenance: ReplacementProvenance::LogicalRewrite, rationale: "different rewrite strategy".into(), @@ -4578,6 +4632,7 @@ mod tests { fn replacements(&self, _target: &TargetSubDAG<'_>) -> Vec { vec![ReplacementSubDAG { + strategy: "ReplaceFilterChild", replacement: Replacement::Rewrite(Rc::new(QueryExpr::Dedup { cols: vec![0], child: Rc::new(metric_scan(&["replacement"])), @@ -4799,6 +4854,7 @@ mod tests { child: Rc::new(fresh_inner_layer), }; vec![ReplacementSubDAG { + strategy: "AlwaysGrowingStrategy", replacement: Replacement::Rewrite(Rc::new(outer_wrapper)), provenance: ReplacementProvenance::LogicalRewrite, rationale: format!("pathological candidate #{n}"), diff --git a/crates/asap-aware-mapping/src/rewrite.rs b/crates/asap-aware-mapping/src/rewrite.rs index 92e39d6e..0f8a5bbb 100644 --- a/crates/asap-aware-mapping/src/rewrite.rs +++ b/crates/asap-aware-mapping/src/rewrite.rs @@ -229,6 +229,7 @@ impl ReplacementStrategy for AvgToSumOverCountStrategy { return Vec::new(); }; vec![ReplacementSubDAG { + strategy: "AvgToSumOverCountStrategy", replacement: Replacement::Rewrite(rewritten), provenance: crate::replacement::ReplacementProvenance::LogicalRewrite, rationale: diff --git a/crates/asap-aware-mapping/src/rollup.rs b/crates/asap-aware-mapping/src/rollup.rs index cc43e36f..5dc3c81e 100644 --- a/crates/asap-aware-mapping/src/rollup.rs +++ b/crates/asap-aware-mapping/src/rollup.rs @@ -372,6 +372,7 @@ fn build_rollup( }; Some(ReplacementSubDAG { + strategy: "RollupStrategy", replacement: Replacement::Rewrite(Rc::new(rewritten)), provenance: crate::replacement::ReplacementProvenance::LogicalRewrite, rationale: format!( diff --git a/crates/asap-aware-mapping/src/topk_reuse.rs b/crates/asap-aware-mapping/src/topk_reuse.rs new file mode 100644 index 00000000..8329569d --- /dev/null +++ b/crates/asap-aware-mapping/src/topk_reuse.rs @@ -0,0 +1,168 @@ +//! Workload-aware reuse between compatible ordered limits. +//! +//! If two queries rank the same input and request top-k results with +//! `small_k < large_k`, the smaller result is exactly the first `small_k` +//! rows of the larger result. This strategy therefore replaces +//! `Limit(small_k, Sort(X))` with `Limit(small_k, Limit(large_k, Sort(X)))`. + +use std::rc::Rc; + +use asap_types::pre_asap::QueryExpr; + +use crate::replacement::{ + Replacement, ReplacementProvenance, ReplacementStrategy, ReplacementSubDAG, TargetSubDAG, +}; + +/// Derives a smaller top-k result from a compatible larger top-k sibling. +pub struct TopKLimitReuseStrategy { + limits: Vec>, +} + +impl TopKLimitReuseStrategy { + pub fn new(limits: &[Rc]) -> Self { + Self { + limits: limits.to_vec(), + } + } + + fn larger_sources<'a>(&'a self, target: &TargetSubDAG<'_>) -> Vec<&'a Rc> { + let QueryExpr::Limit { + n: target_n, + offset: 0, + child: target_child, + } = target.root.as_ref() + else { + return Vec::new(); + }; + + let mut sources: Vec<_> = self + .limits + .iter() + .filter(|candidate| { + if Rc::ptr_eq(candidate, target.root) { + return false; + } + let QueryExpr::Limit { + n, + offset: 0, + child, + } = candidate.as_ref() + else { + return false; + }; + n > target_n + && (Rc::ptr_eq(child, target_child) || child.as_ref() == target_child.as_ref()) + }) + .collect(); + // Prefer the smallest sufficient materialized top-k when several + // larger siblings are available. + sources.sort_by_key(|source| match source.as_ref() { + QueryExpr::Limit { n, .. } => *n, + _ => unreachable!(), + }); + sources + } +} + +impl ReplacementStrategy for TopKLimitReuseStrategy { + fn matches(&self, target: &TargetSubDAG<'_>) -> bool { + !self.larger_sources(target).is_empty() + } + + fn replacements(&self, target: &TargetSubDAG<'_>) -> Vec { + let QueryExpr::Limit { + n: target_n, + offset: 0, + .. + } = target.root.as_ref() + else { + return Vec::new(); + }; + + self.larger_sources(target) + .into_iter() + .map(|source| { + let source_n = match source.as_ref() { + QueryExpr::Limit { n, .. } => *n, + _ => unreachable!(), + }; + ReplacementSubDAG { + strategy: "TopKLimitReuseStrategy", + replacement: Replacement::Rewrite(Rc::new(QueryExpr::Limit { + n: *target_n, + offset: 0, + child: Rc::clone(source), + })), + provenance: ReplacementProvenance::LogicalRewrite, + rationale: format!( + "derives top-{target_n} from the compatible shared top-{source_n} result; both rank the identical input with the same ordering" + ), + } + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_types::pre_asap::{Schema, Source}; + + fn scan_named(metric: &str) -> Rc { + Rc::new(QueryExpr::Scan { + source: Source::TimeSeries { + metric: metric.into(), + }, + predicates: vec![], + schema: Schema::with_time_index(vec![], 0, vec![]), + }) + } + + #[test] + fn smaller_limit_reuses_larger_compatible_limit() { + let child = scan_named("m"); + let small = Rc::new(QueryExpr::Limit { + n: 5, + offset: 0, + child: Rc::clone(&child), + }); + let large = Rc::new(QueryExpr::Limit { + n: 10, + offset: 0, + child, + }); + let strategy = TopKLimitReuseStrategy::new(&[Rc::clone(&small), Rc::clone(&large)]); + let replacements = strategy.replacements(&TargetSubDAG::new(&small)); + assert_eq!(replacements.len(), 1); + let Replacement::Rewrite(rewrite) = &replacements[0].replacement else { + panic!() + }; + let QueryExpr::Limit { n: 5, child, .. } = rewrite.as_ref() else { + panic!() + }; + assert!(Rc::ptr_eq(child, &large)); + } + + #[test] + fn offset_or_different_input_is_not_reused() { + let a = scan_named("a"); + let b = scan_named("b"); + let small = Rc::new(QueryExpr::Limit { + n: 5, + offset: 0, + child: a, + }); + let large = Rc::new(QueryExpr::Limit { + n: 10, + offset: 0, + child: b, + }); + let offset = Rc::new(QueryExpr::Limit { + n: 20, + offset: 1, + child: scan_named("a"), + }); + let strategy = TopKLimitReuseStrategy::new(&[Rc::clone(&small), large, offset]); + assert!(!strategy.matches(&TargetSubDAG::new(&small))); + } +} diff --git a/crates/devtools/src/bin/dag_export.rs b/crates/devtools/src/bin/dag_export.rs index 168bc740..3dd7dde5 100644 --- a/crates/devtools/src/bin/dag_export.rs +++ b/crates/devtools/src/bin/dag_export.rs @@ -18,15 +18,50 @@ // cargo run -p asap-lower --bin dag_export -- \ // --epsilon 0.01 --sql "SELECT quantile(0.99, latency) FROM metrics" --name p99 // -// `--catalog ` is optional (default `metrics`) and applies to every -// `--sql` query in the run — pick the schema matching whatever domain corpus -// you're feeding in. One of `metrics` (the built-in `metrics`/`hosts` demo -// schema this binary has always used), `packets` (matches -// `frontend-sql/tests/data_quality_check/data/synthetic_packet_trace_queries.sql`), -// or `netflow` (matches `frontend-sql/tests/netflow/data/netflow.sql`). -// `--promql` queries ignore `--catalog` entirely (PromQL has no table schema). - -use asap_types::dag_export::{self, DagGraph, DagNote, NamedGraph, WorkloadGraph}; +// `--post-asap` is optional and off by default. When passed, this binary +// additionally runs `asap_aware_mapping::replacement::search_workload` (this +// binary took no strategies of its own — `default_strategies()` already +// includes `AvgToSumOverCountStrategy` as of #282) 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 std::time::Instant; + +use asap_aware_mapping::cost_model::DefaultCostModel; +use asap_aware_mapping::replacement::{search_workload, Replacement, ReplacementSubDAG}; +use asap_types::dag_export::{ + self, DagDecision, DagGraph, DagNote, NamedGraph, PostAsapSubstitution, TargetReplacement, + TargetReplacementAfter, WorkloadGraph, +}; +use asap_types::post_asap::SummaryExpr; +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; @@ -37,21 +72,17 @@ enum Lang { PromQl, } -fn col(name: &str, dtype: DataType) -> Column { - Column::new(name, dtype, false) -} - -fn metrics_catalog() -> SqlCatalog { +fn default_catalog() -> SqlCatalog { SqlCatalog::new() .with_table( "metrics", Schema::with_time_index( vec![ - col("ts", DataType::Timestamp), - col("service", DataType::Utf8), - col("region", DataType::Utf8), - col("latency", DataType::Float64), - col("bytes", DataType::Int64), + Column::new("ts", DataType::Timestamp, false), + Column::new("service", DataType::Utf8, false), + Column::new("region", DataType::Utf8, false), + Column::new("latency", DataType::Float64, false), + Column::new("bytes", DataType::Int64, false), ], 0, vec![], @@ -60,66 +91,76 @@ fn metrics_catalog() -> SqlCatalog { .with_table( "hosts", Schema::new(vec![ - col("service", DataType::Utf8), - col("region", DataType::Utf8), + Column::new("service", DataType::Utf8, false), + Column::new("region", DataType::Utf8, false), ]), ) } -/// Matches `synthetic_packet_trace_queries.sql`'s `packets` table. -fn packets_catalog() -> SqlCatalog { - SqlCatalog::new().with_table( - "packets", - Schema::new(vec![ - col("srcip", DataType::Utf8), - col("dstip", DataType::Utf8), - col("srcport", DataType::Int64), - col("dstport", DataType::Int64), - col("proto", DataType::Utf8), - col("time", DataType::Float64), - col("pkt_len", DataType::Int64), - ]), - ) -} - -/// Matches `netflow.sql`'s `netflow_table` table. -fn netflow_catalog() -> SqlCatalog { - SqlCatalog::new().with_table( - "netflow_table", - Schema::with_time_index( - vec![ - col("time", DataType::Timestamp), - col("srcip", DataType::Utf8), - col("dstip", DataType::Utf8), - col("srcport", DataType::Int64), - col("dstport", DataType::Int64), - col("proto", DataType::Utf8), - col("pkt_len", DataType::Int64), - ], - 0, - vec![], - ), - ) -} - -fn catalog_by_name(name: &str) -> SqlCatalog { - match name { - "metrics" => metrics_catalog(), - "packets" => packets_catalog(), - "netflow" => netflow_catalog(), - other => panic!("unknown --catalog {other:?} — expected one of: metrics, packets, netflow"), +fn catalog(custom: &[String]) -> SqlCatalog { + let mut catalog = default_catalog(); + for raw in custom { + let value: serde_json::Value = serde_json::from_str(raw) + .unwrap_or_else(|error| panic!("--table-schema must be valid JSON: {error}")); + let name = value["name"] + .as_str() + .expect("--table-schema.name must be a string"); + let columns = value["columns"] + .as_array() + .expect("--table-schema.columns must be an array"); + let columns: Vec = columns + .iter() + .map(|column| { + let column_name = column["name"] + .as_str() + .expect("column.name must be a string"); + let data_type = match column["type"] + .as_str() + .expect("column.type must be a string") + .to_ascii_lowercase() + .as_str() + { + "timestamp" => DataType::Timestamp, + "utf8" | "string" => DataType::Utf8, + "float64" | "double" => DataType::Float64, + "int64" | "bigint" => DataType::Int64, + other => panic!("unsupported column type {other:?}"), + }; + Column::new( + column_name, + data_type, + column["nullable"].as_bool().unwrap_or(true), + ) + }) + .collect(); + let schema = match value.get("time_index").and_then(|index| index.as_u64()) { + Some(index) => Schema::with_time_index(columns, index as usize, vec![]), + None => Schema::new(columns), + }; + catalog = catalog.with_table(name, schema); } + catalog } /// Parses `--sql "" --name "