From eba4fbe6669a8c5f1c9e8abc85569de4b3f5ac80 Mon Sep 17 00:00:00 2001 From: zz_y Date: Sat, 22 Aug 2026 14:59:53 -0600 Subject: [PATCH] feat(pre-asap): structural CSE via bottom-up hash-consing (stages 1+2) Rebased onto main now that #226 (the QueryExpr rename) and #227 (the SQL function catalog) have landed there. This PR's own design and content are unchanged; adapted the new cse.rs module's exhaustive QueryExpr match to the current post-rename variant names (Merge -> Concat, Distinct -> Dedup, Scalar -> PromqlScalar, EvalTime -> QueryTimestamp, WindowFunc -> SQLWindowFunc, Arith -> Arithmetic, VectorFromScalar/ScalarFromVector/Relabel/InfoJoin/Sample/Subquery -> their Promql*-prefixed names, ArithOp/CompareOp -> ArithmeticOpKind/ CompareOpKind) and reapplied mod.rs's doc/pub-use additions around the renamed expr_ir re-export. Co-Authored-By: Claude Sonnet 5 --- crates/asap-aware-mapping/src/bind.rs | 56 +++ crates/asap-aware-mapping/src/lib.rs | 15 +- crates/integration-tests/tests/cse.rs | 134 +++++++ crates/types/src/pre_asap/cse.rs | 550 ++++++++++++++++++++++++++ crates/types/src/pre_asap/mod.rs | 6 + 5 files changed, 757 insertions(+), 4 deletions(-) create mode 100644 crates/integration-tests/tests/cse.rs create mode 100644 crates/types/src/pre_asap/cse.rs diff --git a/crates/asap-aware-mapping/src/bind.rs b/crates/asap-aware-mapping/src/bind.rs index cad571396..e95dd05ae 100644 --- a/crates/asap-aware-mapping/src/bind.rs +++ b/crates/asap-aware-mapping/src/bind.rs @@ -91,6 +91,62 @@ pub fn implement_tree_with( logical(expr) } +/// Bind a whole workload's worth of already-CSE'd roots +/// ([`asap_types::pre_asap::cse::share_common_subtrees`]'s output), reusing +/// one bound [`SummaryNode`] wherever two roots share the same `Rc` (issue +/// #212, #222, #223 stage 2). +/// +/// This is a real caller for `share_common_subtrees`, wired up deliberately: +/// the pass's own landing plan calls out that its predecessor +/// (`asap-plan::cse::dedupe_subtrees`) was deleted in #192 for being unwired +/// dead code, and lands this memoization alongside it so that never becomes +/// true again. +/// +/// The memo key is `Rc::as_ptr` — pointer identity, not a second +/// `PartialEq`/structural-equality pass. `share_common_subtrees` already made +/// the (non-negotiable, `PartialEq`-checked) sharing decision; this only +/// needs to recognize when it already bound the exact `Rc` a later root +/// hands back, and is deliberately *not* a general "does an available +/// `Implementation` satisfy this one" lookup — that subsumption question is +/// `asap_aware_mapping::boundary::Matcher`'s documented, deliberately-unfilled +/// job, not this one's. +/// +/// Only whole-root sharing is memoized (matching two workload roots that are +/// themselves the same `Rc` after CSE) — [`implement_tree`] is +/// called at most once per distinct root pointer, but it still walks each +/// such tree's own internal structure fresh; a subtree shared only *below* +/// two different roots' top level does not additionally memoize inside that +/// walk. Widening this to sub-root memoization is future work. +pub fn implement_workload( + roots: Vec<(Id, Rc)>, +) -> Vec<(Id, Result, ImplementError>)> { + implement_workload_with(roots, &DefaultCostModel) +} + +/// Like [`implement_workload`], but ranks candidate summaries via +/// `cost_model` instead of the built-in static preference order (see +/// [`crate::cost_model`]). +pub fn implement_workload_with( + roots: Vec<(Id, Rc)>, + cost_model: &dyn CostModel, +) -> Vec<(Id, Result, ImplementError>)> { + let mut memo: std::collections::HashMap<*const QueryExpr, Rc> = + std::collections::HashMap::new(); + roots + .into_iter() + .map(|(id, expr)| { + let ptr = Rc::as_ptr(&expr); + let result = match memo.get(&ptr) { + Some(cached) => Ok(Rc::clone(cached)), + None => implement_tree_with(&expr, cost_model).inspect(|node| { + memo.insert(ptr, Rc::clone(node)); + }), + }; + (id, result) + }) + .collect() +} + /// Translate an [`Implementation`] into the `(family, needs a /// SummaryEstimate readout)` pair [`bind_summary_agg`] needs, or `None` for /// `PassThrough` (the caller falls back to [`logical`]). diff --git a/crates/asap-aware-mapping/src/lib.rs b/crates/asap-aware-mapping/src/lib.rs index 31599363d..b428c2edc 100644 --- a/crates/asap-aware-mapping/src/lib.rs +++ b/crates/asap-aware-mapping/src/lib.rs @@ -36,9 +36,13 @@ //! `AggIntent → Implementation` (a summary family's own `(Kind, Params)`, //! or an exact accumulator) sized to the `AccuracyTarget` (issue #98). //! [`boundary::implementation_for`] is the per-node decision; -//! [`bind::implement_tree`] drives it over a whole tree — see the -//! terminology section below for why these two are named around -//! "implementation" rather than "bind". +//! [`bind::implement_tree`] drives it over a whole tree, and +//! [`bind::implement_workload`] drives it over a whole workload's roots — +//! memoized on `Rc` identity so two roots that +//! `asap_types::pre_asap::cse::share_common_subtrees` already collapsed +//! onto one shared subtree bind to one shared `SummaryNode` too (issue +//! #212, #222, #223) — see the terminology section below for why these are +//! named around "implementation" rather than "bind". //! - [`cost_model`] — the [`CostModel`](cost_model::CostModel) trait every //! deployment's cost-based sketch selection plugs into (issues #6, #33). //! `asap-plan` itself only ships [`DefaultCostModel`](cost_model::DefaultCostModel), @@ -87,7 +91,10 @@ pub mod bind; pub mod boundary; pub mod cost_model; -pub use bind::{implement_tree, implement_tree_with, ImplementError}; +pub use bind::{ + implement_tree, implement_tree_with, implement_workload, implement_workload_with, + ImplementError, +}; pub use boundary::{ implementation_for, implementation_for_with, summary_candidates, Implementation, Matcher, }; diff --git a/crates/integration-tests/tests/cse.rs b/crates/integration-tests/tests/cse.rs new file mode 100644 index 000000000..6d28e1a70 --- /dev/null +++ b/crates/integration-tests/tests/cse.rs @@ -0,0 +1,134 @@ +//! End-to-end pre-ASAP CSE → implementation pin (issue #212, #222, #223). +//! +//! Drives the full staged pipeline this issue lands: two independently +//! lowered `QueryExpr` trees → `share_common_subtrees` (stage 1, +//! `asap-types::pre_asap::cse`) → `implement_workload` (stage 2, +//! `asap-aware-mapping`) — and asserts the sharing that stage 1 decides +//! survives into stage 2's bound output as one genuinely shared +//! `Rc`, not just one shared `Rc`. This is the "real +//! caller" the issue's landing plan requires before `share_common_subtrees` +//! is allowed to exist at all (its predecessor, +//! `asap-plan::cse::dedupe_subtrees`, was deleted in #192 for being unwired +//! dead code). + +use std::rc::Rc; + +use asap_aware_mapping::implement_workload; +use asap_frontend_promql::lower_promql; +use asap_types::pre_asap::cse::share_common_subtrees; +use asap_types::pre_asap::query_expr::QueryExpr; +use asap_types::types::AccuracyTarget; + +/// Two workload entries that happen to submit the exact same query (a +/// realistic case — two dashboards, or a query fired both standalone and as +/// part of a larger batch) collapse onto one shared `Rc` after +/// `share_common_subtrees`, and then onto one genuinely-shared, single build +/// `Rc` after `implement_workload` — no second structural- +/// equality pass at the post-ASAP layer needed for this kind of sharing. +#[test] +fn duplicate_workload_queries_share_one_bound_summary() { + // Grouped (`by (job)`), so the shared `Aggregate`'s output schema carries + // a provable unique key — the legality gate `share_common_subtrees` + // enforces (see `asap-types::pre_asap::cse`'s module doc) — and its + // `ExactAggregate(Sum)` realization is deterministic regardless of the + // accuracy target, so this pins the sharing mechanism itself rather than + // any one particular summary-family choice. + let query = "sum by (job) (http_requests_total)"; + let a = lower_promql(query, AccuracyTarget::Exact).expect("query a failed to lower"); + let b = lower_promql(query, AccuracyTarget::Exact).expect("query b failed to lower"); + + // Independently lowered: not yet sharing any `Rc`, even though they are + // structurally identical (`resolve_root` gives each call its own fresh + // tree). + assert_eq!( + a, b, + "fixture sanity: identical query text lowers identically" + ); + + let shared = share_common_subtrees(vec![("a", a), ("b", b)]); + let [(_, qa), (_, qb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!( + Rc::ptr_eq(qa, qb), + "share_common_subtrees must collapse the two identical roots onto one Rc" + ); + + let bound = implement_workload(shared); + let [(_, ra), (_, rb)] = bound.as_slice() else { + panic!("expected 2 bound results"); + }; + let ra = ra.as_ref().expect("query a failed to bind"); + let rb = rb.as_ref().expect("query b failed to bind"); + assert!( + Rc::ptr_eq(ra, rb), + "implement_workload must reuse one bound SummaryNode for the shared Rc, \ + got two independently-built nodes: {ra:?} vs {rb:?}" + ); +} + +/// The negative control: two workload entries whose queries are NOT +/// structurally identical must not be conflated by `implement_workload`'s +/// `Rc::as_ptr` memoization — different `Rc` roots always bind +/// independently. +#[test] +fn distinct_workload_queries_do_not_share_a_bound_summary() { + let a = lower_promql("sum by (job) (http_requests_total)", AccuracyTarget::Exact) + .expect("query a failed to lower"); + let b = lower_promql("sum by (job) (http_response_total)", AccuracyTarget::Exact) + .expect("query b failed to lower"); + assert_ne!(a, b, "fixture sanity: the two queries differ"); + + let shared = share_common_subtrees(vec![("a", a), ("b", b)]); + let [(_, qa), (_, qb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!(!Rc::ptr_eq(qa, qb)); + + let bound = implement_workload(shared); + let [(_, ra), (_, rb)] = bound.as_slice() else { + panic!("expected 2 bound results"); + }; + let ra = ra.as_ref().expect("query a failed to bind"); + let rb = rb.as_ref().expect("query b failed to bind"); + assert!( + !Rc::ptr_eq(ra, rb), + "distinct queries must bind independently" + ); +} + +/// Single-query CSE (a repeated sub-expression within one query) also +/// survives through `implement_workload`: the two grouped-`Aggregate` +/// branches of a `BinaryOp` collapse to one shared `Rc` in +/// `share_common_subtrees`, and to one shared bound `Rc` here. +#[test] +fn single_query_repeated_subexpression_shares_one_bound_summary() { + let query = "sum by (job) (http_requests_total) / sum by (job) (http_requests_total)"; + let expr = lower_promql(query, AccuracyTarget::Exact).expect("query failed to lower"); + + let shared = share_common_subtrees(vec![("q", expr)]); + let [(_, root)] = shared.as_slice() else { + panic!("expected 1 root"); + }; + let QueryExpr::BinaryOp { lhs, rhs, .. } = root.as_ref() else { + panic!("expected a BinaryOp root, got {root:?}"); + }; + assert!( + Rc::ptr_eq(lhs, rhs), + "the two identical sum-by-job branches must collapse onto one Rc" + ); + + // Bind both branches (as if they were two roots of a tiny sub-workload) + // through the same memoized `implement_workload`, and confirm they bind + // to the identical `Rc`. + let bound = implement_workload(vec![("lhs", Rc::clone(lhs)), ("rhs", Rc::clone(rhs))]); + let [(_, ra), (_, rb)] = bound.as_slice() else { + panic!("expected 2 bound results"); + }; + let ra = ra.as_ref().expect("lhs failed to bind"); + let rb = rb.as_ref().expect("rhs failed to bind"); + assert!( + Rc::ptr_eq(ra, rb), + "the shared branch must bind to one shared SummaryNode" + ); +} diff --git a/crates/types/src/pre_asap/cse.rs b/crates/types/src/pre_asap/cse.rs new file mode 100644 index 000000000..2939ae3b7 --- /dev/null +++ b/crates/types/src/pre_asap/cse.rs @@ -0,0 +1,550 @@ +//! Pre-ASAP structural common-subexpression elimination: bottom-up +//! hash-consing over an already-`resolve_root`'d [`QueryExpr`] tree (issue +//! #212, #222, #223). +//! +//! CSE only runs on an already-bound, already-canonicalized tree — +//! structural matching is meaningless before canonicalization has converged +//! semantically-equivalent queries onto one shape (`docs/pre-asap-ir.md` +//! design principle 3; `median(latency)` and `approx_percentile_cont(latency, +//! 0.5)` already lower to an identical `AggIntent::Quantile` today, per +//! `sql_lowering.rs`'s `median_is_the_same_intent_as_an_explicit_half_percentile` +//! test). [`share_common_subtrees`] is the single entry point, run once per +//! workload batch (or once per query — see "Single-query CSE" below) *after* +//! `resolve_root`, *before* `implement_workload` +//! ([`asap_aware_mapping::implement_workload`]). +//! +//! ## Algorithm: classic hash-consing / value-numbering +//! +//! Bottom-up: every child is interned before its parent, so a parent's +//! candidacy for sharing naturally incorporates whether its own children were +//! themselves shared — two parents whose children were independently +//! deduplicated down to the same `Rc`s are structurally identical iff their +//! own fields also match, without re-walking the subtrees. +//! +//! Only the **relational skeleton** participates — the same set of "operator" +//! children [`canonicalize`](super::canonicalize)'s `children_mut` walks +//! (`Filter`/`Project`/`Aggregate`/`Concat`/`Join`/`BinaryOp`/…). A scalar +//! subexpression reachable only through a wrapper position (`Predicate`, +//! `ProjectItem.expr`, `Aggregate.having`, `SQLWindowFunc.args`, …) stays +//! embedded as opaque data on its owning operator node, compared by +//! `QueryExpr`'s derived `PartialEq` along with the rest of that node's +//! fields, rather than separately hash-consed — the same scope +//! `canonicalize.rs` settled on ("none of the rewrite rules touch a scalar +//! subtree, so there's nothing to gain by recursing into one"). Widening this +//! to scalar positions is future work, not attempted here. +//! +//! ## Correctness: hash is a filter, `PartialEq` is the decision +//! +//! This is the one non-negotiable rule. A **false positive** here — two +//! subtrees wrongly judged shareable — is a wrong query answer, not a missed +//! optimization: two different queries would read each other's data. +//! [`structural_hash`] (`DefaultHasher`/SipHash over a canonical +//! serialization, no collision-freedom guarantee) may only narrow the +//! candidate set within one bucket; [`InternTable::intern`]'s `PartialEq` +//! check on that bucket is what actually decides sharing, every time, no +//! exceptions for "the hash probably didn't collide." +//! +//! This also means the pass is safe by construction against the case #212 +//! flagged as a real historical bug (issue #115): `AggIntent::Quantile` +//! carries its input column and its `AccuracyTarget`, both `PartialEq` +//! fields, so `Quantile(x, 0.99, ε=0.01)` and `Quantile(x, 0.99, ε=0.001)` — +//! or `Quantile(x, ..)` vs `Quantile(y, ..)` — are never merged. This is +//! intentionally conservative: it only recognizes *exact* structural +//! matches, not "a stricter-accuracy summary could also answer a looser +//! request." That subsumption question already has a documented, +//! deliberately-unfilled home (`asap_aware_mapping::boundary::Matcher`) — +//! CSE here does not attempt it. +//! +//! ## Legality: gated by `Schema::unique_keys` +//! +//! Structural equality alone is necessary but not sufficient. Per +//! [`Schema::unique_keys`](super::schema::Schema::unique_keys)'s own doc: "a +//! producer's output can only be safely shared across consumers when its row +//! identity is provably stable across reads." A candidate node with no +//! provable unique key (`Schema::has_unique_key()` false, or `output_schema` +//! not even defined for that node, e.g. a `Concat`/`SetOp` branch whose union +//! drops `unique_keys`, or an ungrouped/global `Aggregate`, whose empty `by` +//! also reports no unique key today) is **never** hoisted, even when it is +//! structurally identical to something already interned — it is always +//! inserted fresh, matching the rule the (now-deleted) prior CSE attempt +//! already encoded and the doc comment on `Aggregate`'s `child` field +//! ("`unique_keys` feeds CSE's producer-sharing legality check"). +//! +//! ## Single-query CSE falls out for free +//! +//! A repeated sub-expression within *one* query (e.g. the same grouped +//! `Aggregate` referenced twice on two `BinaryOp` branches) is deduplicated +//! by the exact same bottom-up interning — a workload of size one still +//! interns bottom-up within that one tree. No separate mechanism is needed; +//! see the `single_query_shares_its_own_repeated_subtree` test below. +//! +//! ## Landing plan (issue #223) +//! +//! This module is stage 1 of a 4-stage plan. Stage 2 +//! ([`asap_aware_mapping::implement_workload`]) is a real caller, wired at +//! the same time so this never becomes unwired dead code again (the original +//! `asap-plan::cse::dedupe_subtrees` was deleted in #192 for exactly that). +//! Stages 3 (`dag_export::structural_hash` unification) and 4 (`CostModel` +//! CSE credit) are deliberately deferred follow-ups, not attempted here. + +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use super::query_expr::QueryExpr; + +/// Bottom-up hash-consing table: structurally-equal, sharing-legal +/// [`QueryExpr`] nodes collapse onto one `Rc`. +/// +/// `buckets` is keyed by [`structural_hash`] — a coarse candidate filter +/// only (see the module-level "Correctness" section). Every entry within one +/// bucket is a full node kept around for the `PartialEq` comparison that +/// actually decides a match; a hash collision between structurally different +/// nodes just means a (harmless) linear scan of a few extra candidates. +struct InternTable { + buckets: HashMap>>, +} + +impl InternTable { + fn new() -> Self { + Self { + buckets: HashMap::new(), + } + } + + /// Intern one already-children-rebuilt node: look it up by + /// [`structural_hash`], confirm with `PartialEq`, and — only when + /// sharing is legal (see "Legality" above) — return the existing `Rc` + /// instead of allocating a new one. + fn intern(&mut self, node: QueryExpr) -> Rc { + let hash = structural_hash(&node); + // A node with no provable unique key is never *returned* as a match + // for something else — it may still go on to occupy a fresh slot in + // the bucket (harmless; it just never gets found by a later + // `PartialEq` scan that also requires `reusable`). + let reusable = node + .output_schema() + .is_ok_and(|schema| schema.has_unique_key()); + let bucket = self.buckets.entry(hash).or_default(); + if reusable { + if let Some(existing) = bucket.iter().find(|candidate| candidate.as_ref() == &node) { + return Rc::clone(existing); + } + } + let rc = Rc::new(node); + bucket.push(Rc::clone(&rc)); + rc + } +} + +/// Coarse structural hash used only to bucket [`InternTable::intern`]'s +/// candidate search — never the actual sharing decision (`PartialEq` is). +/// +/// `QueryExpr` carries `f64`s (`PromqlScalar`, `AggIntent::Quantile.q`, …), so it +/// cannot derive `std::hash::Hash`. Serializing to a canonical JSON string +/// and hashing that sidesteps the `f64` problem the same way +/// `dag_export.rs`'s own `structural_hash` does — a deliberately independent +/// implementation for now (this module's stage 1; unifying the two is stage +/// 3 of issue #223's landing plan, not done here). A NaN/infinite `f64` +/// makes JSON serialization fail; falling back to a fixed hash just puts +/// every such node in one (larger, still `PartialEq`-disambiguated) bucket. +fn structural_hash(node: &QueryExpr) -> u64 { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + let canonical = serde_json::to_string(node).unwrap_or_default(); + canonical.hash(&mut hasher); + hasher.finish() +} + +/// Recurse into `child`, then intern the result. `Rc::try_unwrap` recovers +/// the owned node without cloning in the overwhelmingly common case — a +/// tree freshly built by a front end / `resolve_root`, not yet shared by any +/// prior CSE pass, where every `Rc` is uniquely owned. Falls back to cloning +/// this node's own fields (its children stay `Rc`s, not deep-copied) only +/// when `child` is already shared — e.g. re-running CSE over a tree that +/// went through a previous `share_common_subtrees` pass; a structural +/// duplicate collapses right back onto `child` itself via `PartialEq`, an +/// already-optimal no-op. +fn intern_child(table: &mut InternTable, child: Rc) -> Rc { + match Rc::try_unwrap(child) { + Ok(owned) => intern_bottom_up(table, owned), + Err(shared) => intern_bottom_up(table, (*shared).clone()), + } +} + +/// Like [`intern_child`], for a `Concat` branch — stored by value +/// (`Vec`, not `Rc`), so this position itself can never +/// alias another parent. Interning it anyway still lets any `Rc`-typed +/// descendant of the branch participate in sharing, and registers the +/// branch's own hash/value in the table for a *different* `Concat` elsewhere +/// with a structurally identical branch (which — being in its own `Vec` +/// slot too — still can't literally share the `Rc`, but this keeps the +/// interning behavior uniform and the table's bucket contents consistent). +fn intern_owned(table: &mut InternTable, expr: QueryExpr) -> QueryExpr { + let rc = intern_bottom_up(table, expr); + Rc::try_unwrap(rc).unwrap_or_else(|shared| (*shared).clone()) +} + +/// Bottom-up: rebuild `expr`'s children (recursively interning each), then +/// intern the rebuilt node itself. +fn intern_bottom_up(table: &mut InternTable, expr: QueryExpr) -> Rc { + let rebuilt = rebuild_children(table, expr); + table.intern(rebuilt) +} + +/// Rebuild `expr` with each **operator** child (see the module doc on scope) +/// replaced by its interned `Rc`. Exhaustive over every `QueryExpr` variant, +/// matching `canonicalize.rs`'s `children_mut` exactly in which fields count +/// as an operator child — new variants fail to compile here until this match +/// is extended. +fn rebuild_children(table: &mut InternTable, expr: QueryExpr) -> QueryExpr { + use QueryExpr::*; + match expr { + Scan { .. } | PromqlScalar(_) | QueryTimestamp => expr, + PromqlVectorFromScalar(c) => PromqlVectorFromScalar(intern_child(table, c)), + PromqlScalarFromVector(c) => PromqlScalarFromVector(intern_child(table, c)), + PromqlRelabel { dst, value, child } => PromqlRelabel { + dst, + value, + child: intern_child(table, child), + }, + PromqlInfoEnrich { selector, child } => PromqlInfoEnrich { + selector, + child: intern_child(table, child), + }, + PromqlSeriesSample { by, kind, child } => PromqlSeriesSample { + by, + kind, + child: intern_child(table, child), + }, + Filter { pred, child } => Filter { + pred, + child: intern_child(table, child), + }, + Project { + cols, + qualifier, + child, + } => Project { + cols, + qualifier, + child: intern_child(table, child), + }, + Aggregate { + reduction, + measures, + output_names, + having, + child, + } => Aggregate { + reduction, + measures, + output_names, + having, + child: intern_child(table, child), + }, + Dedup { cols, child } => Dedup { + cols, + child: intern_child(table, child), + }, + Concat { children } => Concat { + children: children + .into_iter() + .map(|c| intern_owned(table, c)) + .collect(), + }, + Join { + kind, + pred, + left, + right, + } => Join { + kind, + pred, + left: intern_child(table, left), + right: intern_child(table, right), + }, + SetOp { + kind, + all, + left, + right, + } => SetOp { + kind, + all, + left: intern_child(table, left), + right: intern_child(table, right), + }, + Sort { + keys, + partition_by, + child, + } => Sort { + keys, + partition_by, + child: intern_child(table, child), + }, + Limit { n, offset, child } => Limit { + n, + offset, + child: intern_child(table, child), + }, + PromqlSubquery { + range, + resolution, + child, + } => PromqlSubquery { + range, + resolution, + child: intern_child(table, child), + }, + TimeRange { range, child } => TimeRange { + range, + child: intern_child(table, child), + }, + TimeShift { shift, child } => TimeShift { + shift, + child: intern_child(table, child), + }, + SQLWindowFunc { + func, + args, + partition_by, + order_by, + output_name, + child, + } => SQLWindowFunc { + func, + args, + partition_by, + order_by, + output_name, + child: intern_child(table, child), + }, + BinaryOp { + op, + lhs, + rhs, + vector_match, + } => BinaryOp { + op, + lhs: intern_child(table, lhs), + rhs: intern_child(table, rhs), + vector_match, + }, + // Scalar variants (issue #205) — never descended into; see the + // module doc's "Algorithm" section on scope. Left byte-for-byte + // unchanged: predicate / project-list / sort-key / window-arg + // expressions stay embedded as opaque leaf data, compared by the + // enclosing operator node's derived `PartialEq`. + Column(_) + | Literal(_) + | Compare { .. } + | BoolAnd(_) + | BoolOr(_) + | Not(_) + | IsNull(_) + | IsNotNull(_) + | Cast { .. } + | InList { .. } + | FunctionCall { .. } + | Arithmetic { .. } + | Case { .. } => expr, + } +} + +/// Share structurally-identical, sharing-legal subtrees across a workload's +/// query roots (or within one query, for `roots.len() == 1` — see the +/// module doc's "Single-query CSE" section). Every root's *value* is +/// unchanged (`PartialEq`-equal to its input) — only its internal `Rc` +/// structure may now alias another root's, or another part of its own tree. +/// +/// `roots` must already be bound + canonicalized (post-`resolve_root`). +/// `Id` is caller-chosen — a `QueryWorkload` entry's own key, an index, a +/// query name, whatever identifies one root through the pipeline; this +/// module has no opinion on its shape. +pub fn share_common_subtrees(roots: Vec<(Id, QueryExpr)>) -> Vec<(Id, Rc)> { + let mut table = InternTable::new(); + roots + .into_iter() + .map(|(id, expr)| (id, intern_bottom_up(&mut table, expr))) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::pre_asap::agg_intent::AggIntent; + use crate::pre_asap::query_expr::{BinaryOpKind, GroupKeys, Reduction, Source}; + use crate::pre_asap::schema::{Column, DataType, Schema}; + use crate::types::AccuracyTarget; + + /// `[ts, service, value, latency]`. + fn scan() -> QueryExpr { + QueryExpr::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Column::new("ts", DataType::Timestamp, false), + Column::new("service", DataType::Utf8, false), + Column::new("value", DataType::Float64, false), + Column::new("latency", DataType::Float64, false), + ], + 0, + vec![], + ), + } + } + + fn quantile_agg(by: Vec, col: Option, q: f64) -> QueryExpr { + QueryExpr::Aggregate { + reduction: Reduction::by(by), + measures: vec![AggIntent::Quantile { + col, + q, + accuracy: AccuracyTarget::Exact, + }], + output_names: vec![], + having: None, + child: Rc::new(scan()), + } + } + + #[test] + fn distinct_column_quantiles_do_not_merge() { + // Grouped (unique_keys present) so the legality gate isn't what's + // blocking the merge — only the differing `col` is. + let a = quantile_agg(vec![1], Some(2), 0.5); + let b = quantile_agg(vec![1], Some(3), 0.5); + let shared = share_common_subtrees(vec![("a", a), ("b", b)]); + let [(_, ra), (_, rb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!( + !Rc::ptr_eq(ra, rb), + "distinct-column Quantiles must not be shared" + ); + assert_ne!(ra, rb); + } + + #[test] + fn no_unique_keys_means_no_merge_even_when_structurally_identical() { + // Ungrouped (global) aggregate: `by` is empty, so + // `aggregate_output_schema` reports no unique key today — not + // hoistable even though `a` and `b` are structurally identical. + let a = quantile_agg(vec![], Some(2), 0.9); + let b = quantile_agg(vec![], Some(2), 0.9); + assert_eq!(a, b, "fixture sanity: the two trees are structurally equal"); + assert!( + !a.output_schema().unwrap().has_unique_key(), + "fixture sanity: an ungrouped aggregate has no provable unique key" + ); + let shared = share_common_subtrees(vec![("a", a), ("b", b)]); + let [(_, ra), (_, rb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!( + !Rc::ptr_eq(ra, rb), + "no unique key ⇒ never hoisted, even for an identical structural match" + ); + } + + #[test] + fn median_and_explicit_half_percentile_merge() { + // Two front-end spellings ("median" and "approx_percentile_cont(., + // 0.5)") already lower to the identical `AggIntent::Quantile { q: + // 0.5, .. }` today (see `sql_lowering.rs`'s + // `median_is_the_same_intent_as_an_explicit_half_percentile`) — here + // built directly (grouped, so a unique key is provable) as two + // independently-constructed but structurally identical trees, the + // way two different call sites in a workload would produce them. + let median = quantile_agg(vec![1], Some(2), 0.5); + let approx_percentile_cont_half = quantile_agg(vec![1], Some(2), 0.5); + let shared = share_common_subtrees(vec![ + ("median", median), + ("percentile", approx_percentile_cont_half), + ]); + let [(_, m), (_, p)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!( + Rc::ptr_eq(m, p), + "median and an explicit 0.5 percentile must merge onto one Rc" + ); + } + + #[test] + fn single_query_shares_its_own_repeated_subtree() { + // One query root referencing the same grouped Aggregate on both + // BinaryOp branches — built as two separately-allocated but + // structurally identical subtrees (`.clone()` into two distinct + // `Rc::new` calls), the shape a front end emitting a repeated + // sub-expression would actually produce (no sharing yet). A + // workload of size 1 still interns bottom-up within this one tree — + // no separate single-query mechanism needed. + let agg = quantile_agg(vec![1], Some(2), 0.5); + let root = QueryExpr::BinaryOp { + op: BinaryOpKind::Compare(crate::pre_asap::expr_ir::CompareOpKind::Eq), + lhs: Rc::new(agg.clone()), + rhs: Rc::new(agg), + vector_match: None, + }; + let shared = share_common_subtrees(vec![("q", root)]); + let [(_, root)] = shared.as_slice() else { + panic!("expected 1 root"); + }; + let QueryExpr::BinaryOp { lhs, rhs, .. } = root.as_ref() else { + panic!("expected BinaryOp root, got {root:?}"); + }; + assert!( + Rc::ptr_eq(lhs, rhs), + "the two structurally identical branches must collapse onto one Rc" + ); + } + + #[test] + fn dedup_gates_sharing_the_same_as_aggregate() { + // `Dedup { cols }` adds `cols` as a unique key — so two identical + // `Dedup` subtrees over a keyed column *do* merge, exercising the + // legality gate on a non-`Aggregate` node. + let dedup = |cols: Vec| QueryExpr::Dedup { + cols, + child: Rc::new(scan()), + }; + let a = dedup(vec![1]); + let b = dedup(vec![1]); + let shared = share_common_subtrees(vec![("a", a), ("b", b)]); + let [(_, ra), (_, rb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!( + Rc::ptr_eq(ra, rb), + "Dedup on the same cols has a provable unique key and should merge" + ); + } + + #[test] + fn group_keys_gate_still_prevented_when_partition_by_without_used() { + // Sanity on the module's advertised precedent: a `without(...)` + // grouping stays open (no unique key) even though `by` is + // non-empty-shaped structurally, so two identical `without` groups + // do not merge under the same gate that blocks the ungrouped case. + let without_agg = || QueryExpr::Aggregate { + reduction: Reduction::Reduce(GroupKeys::without(vec![0])), + measures: vec![AggIntent::Count { + accuracy: AccuracyTarget::Exact, + }], + output_names: vec![], + having: None, + child: Rc::new(scan()), + }; + let a = without_agg(); + let b = without_agg(); + assert!(!a.output_schema().unwrap().has_unique_key()); + let shared = share_common_subtrees(vec![("a", a), ("b", b)]); + let [(_, ra), (_, rb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + assert!(!Rc::ptr_eq(ra, rb)); + } +} diff --git a/crates/types/src/pre_asap/mod.rs b/crates/types/src/pre_asap/mod.rs index dacfb25ae..f7a18b215 100644 --- a/crates/types/src/pre_asap/mod.rs +++ b/crates/types/src/pre_asap/mod.rs @@ -20,6 +20,10 @@ //! relational tree or converter anymore. //! - [`canonicalize`] — post-lowering structural normalization of [`QueryExpr`] //! (issue #34), run by [`resolve_root`]. +//! - [`cse`] — workload-level structural common-subexpression elimination +//! over an already-`resolve_root`'d tree (issue #212, #222, #223), run +//! *after* `resolve_root` / `canonicalize` and *before* implementation +//! (`asap_aware_mapping::implement_workload`). //! //! Formerly the separate `asap-l2` crate; folded in here since //! `binder`/`column_resolution`/`canonicalize`/`resolve` have no @@ -30,6 +34,7 @@ pub mod agg_intent; pub mod binder; pub mod canonicalize; pub mod column_resolution; +pub mod cse; pub mod expr_ir; pub mod query_expr; pub mod resolve; @@ -45,6 +50,7 @@ pub use column_resolution::{ output_schema_for_aggregate, resolve_column_ref, resolve_column_refs, resolve_expr, ResolveError, }; +pub use cse::share_common_subtrees; pub use expr_ir::{ArithmeticOpKind, ColumnRef, CompareOpKind, ScalarValue}; pub use query_expr::{ aggregate_output_schema, AtModifier, BinaryOpKind, ColState, DataModel, GroupKeys, GroupSide,