Skip to content

Commit d3afe8e

Browse files
zzylolclaude
andcommitted
feat(asap-aware-mapping): implement CSE stage 4 (Volcano/Cascades cost-based sharing)
Replaces the previous doc-comment-only decision for #237 with a real, working implementation, per review feedback: - The design discussion moves out of code comments into docs/cse-cost-model-decision.md. - CostModel gains CseCandidate/ShareDecision/cse_share_decision, plus two overridable cost-estimation hooks (cse_recompute_cost, cse_shared_maintenance_cost) with concrete default heuristics. The default cse_share_decision body performs a genuine Volcano/Cascades- style cost comparison — share iff the estimated cost of maintaining one shared summary is no greater than the estimated total cost of recomputing it independently at every consumer — rather than a fixed System R-style rule. - bind::implement_workload_with now computes each shared subtree's true consumer_count via a whole-workload pre-pass, decides ShareDecision once per shared subtree (from full knowledge of the workload's sharing structure), and applies it consistently to every occurrence. This closes out #223's stage 4 and #212's original "add CSE" tracking issue. Stage 3 (structural-hash unification) landed separately in #244. New tests: cost_model.rs unit tests pinning the default cost comparison's behavior in both directions and confirming custom CostModel overrides are actually consulted; a bind.rs test confirming a CostModel that declines sharing produces genuinely independent SummaryNodes even for two roots that are the same Rc<QueryExpr>. All pre-existing CSE tests (including the DefaultCostModel-sharing integration tests) still pass unchanged — the default cost weights were calibrated against a real lowered query so today's sharing behavior for realistic subtrees is preserved. Verified: cargo build --workspace --all-targets, cargo test --workspace, cargo fmt --all -- --check, cargo clippy --workspace --all-targets --all-features -- -D warnings, cargo doc --workspace --no-deps (no new broken-link warnings). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 227bad0 commit d3afe8e

4 files changed

Lines changed: 480 additions & 167 deletions

File tree

‎crates/asap-aware-mapping/src/bind.rs‎

Lines changed: 86 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ use asap_types::pre_asap::schema::Schema;
4545
use thiserror::Error;
4646

4747
use crate::boundary::{implementation_for_with, Implementation};
48-
use crate::cost_model::{CostModel, DefaultCostModel};
48+
use crate::cost_model::{CostModel, CseCandidate, DefaultCostModel, ShareDecision};
4949

5050
/// Errors from the pre-ASAP → post-ASAP binding pass.
5151
#[derive(Debug, Error)]
@@ -93,8 +93,9 @@ pub fn implement_tree_with(
9393

9494
/// Bind a whole workload's worth of already-CSE'd roots
9595
/// ([`asap_types::pre_asap::cse::share_common_subtrees`]'s output), reusing
96-
/// one bound [`SummaryNode`] wherever two roots share the same `Rc` (issue
97-
/// #212, #222, #223 stage 2).
96+
/// one bound [`SummaryNode`] wherever two roots share the same `Rc` *and*
97+
/// `cost_model` decides it's worth it (issue #212, #222, #223 stages 2 and
98+
/// 4, #237).
9899
///
99100
/// This is a real caller for `share_common_subtrees`, wired up deliberately:
100101
/// the pass's own landing plan calls out that its predecessor
@@ -111,35 +112,69 @@ pub fn implement_tree_with(
111112
/// `asap_aware_mapping::boundary::Matcher`'s documented, deliberately-unfilled
112113
/// job, not this one's.
113114
///
115+
/// A first pass over `roots` counts each distinct `Rc<QueryExpr>` pointer's
116+
/// true `consumer_count` across the whole workload, so the
117+
/// [`CseCandidate`]/[`CostModel::cse_share_decision`] cost comparison (see
118+
/// `docs/cse-cost-model-decision.md`) sees the real total, not a running
119+
/// count that grows as roots are processed left to right. The decision is
120+
/// made once, the first time a shared pointer is bound, and cached alongside
121+
/// the bound `SummaryNode` so every later occurrence of that same `Rc`
122+
/// applies the same decision consistently — either every consumer reuses one
123+
/// shared `SummaryNode`, or every consumer (including the first) binds
124+
/// independently.
125+
///
114126
/// Only whole-root sharing is memoized (matching two workload roots that are
115127
/// themselves the same `Rc<QueryExpr>` after CSE) — [`implement_tree`] is
116-
/// called at most once per distinct root pointer, but it still walks each
117-
/// such tree's own internal structure fresh; a subtree shared only *below*
118-
/// two different roots' top level does not additionally memoize inside that
119-
/// walk. Widening this to sub-root memoization is future work.
128+
/// called at most once per distinct root pointer when the decision is
129+
/// `Share`, but it still walks each such tree's own internal structure
130+
/// fresh; a subtree shared only *below* two different roots' top level does
131+
/// not additionally memoize inside that walk. Widening this to sub-root
132+
/// memoization is future work.
120133
pub fn implement_workload<Id>(
121134
roots: Vec<(Id, Rc<QueryExpr>)>,
122135
) -> Vec<(Id, Result<Rc<SummaryNode>, ImplementError>)> {
123136
implement_workload_with(roots, &DefaultCostModel)
124137
}
125138

126-
/// Like [`implement_workload`], but ranks candidate summaries via
127-
/// `cost_model` instead of the built-in static preference order (see
139+
/// Like [`implement_workload`], but ranks candidate summaries — and decides
140+
/// CSE sharing — via `cost_model` instead of the built-in defaults (see
128141
/// [`crate::cost_model`]).
129142
pub fn implement_workload_with<Id>(
130143
roots: Vec<(Id, Rc<QueryExpr>)>,
131144
cost_model: &dyn CostModel,
132145
) -> Vec<(Id, Result<Rc<SummaryNode>, ImplementError>)> {
133-
let mut memo: std::collections::HashMap<*const QueryExpr, Rc<SummaryNode>> =
146+
let mut consumer_count: std::collections::HashMap<*const QueryExpr, usize> =
147+
std::collections::HashMap::new();
148+
for (_, expr) in &roots {
149+
*consumer_count.entry(Rc::as_ptr(expr)).or_insert(0) += 1;
150+
}
151+
152+
let mut memo: std::collections::HashMap<*const QueryExpr, (Rc<SummaryNode>, ShareDecision)> =
134153
std::collections::HashMap::new();
135154
roots
136155
.into_iter()
137156
.map(|(id, expr)| {
138157
let ptr = Rc::as_ptr(&expr);
139158
let result = match memo.get(&ptr) {
140-
Some(cached) => Ok(Rc::clone(cached)),
159+
Some((cached, ShareDecision::Share)) => Ok(Rc::clone(cached)),
160+
Some((_, ShareDecision::RecomputeIndependently)) => {
161+
implement_tree_with(&expr, cost_model)
162+
}
141163
None => implement_tree_with(&expr, cost_model).inspect(|node| {
142-
memo.insert(ptr, Rc::clone(node));
164+
let count = consumer_count[&ptr];
165+
let decision = if count > 1 {
166+
let candidate = CseCandidate {
167+
subtree: &expr,
168+
bound_summary: node,
169+
consumer_count: count,
170+
};
171+
cost_model.cse_share_decision(&candidate)
172+
} else {
173+
// Only one consumer: nothing to compare against, and
174+
// this branch is never consulted again for `ptr`.
175+
ShareDecision::Share
176+
};
177+
memo.insert(ptr, (Rc::clone(node), decision));
143178
}),
144179
};
145180
(id, result)
@@ -866,4 +901,43 @@ mod tests {
866901
};
867902
assert_eq!(col, &ColumnRef::Named("bytes".into()));
868903
}
904+
905+
/// Issue #237, #223 stage 4: a `CostModel` that declines CSE sharing
906+
/// makes `implement_workload_with` bind each occurrence independently,
907+
/// even though the two roots are the exact same `Rc<QueryExpr>` (as
908+
/// `share_common_subtrees` would hand back for two identical workload
909+
/// entries) — the opposite of `DefaultCostModel`'s unconditional-share
910+
/// behavior pinned by `crates/integration-tests/tests/cse.rs`.
911+
#[test]
912+
fn implement_workload_with_recomputes_independently_when_cost_model_declines_sharing() {
913+
struct NeverShareCse;
914+
impl CostModel for NeverShareCse {
915+
fn rank_candidates(
916+
&self,
917+
_intent: &AggIntent,
918+
candidates: &[SketchKind],
919+
) -> Vec<SketchKind> {
920+
candidates.to_vec()
921+
}
922+
fn cse_share_decision(&self, _candidate: &CseCandidate) -> ShareDecision {
923+
ShareDecision::RecomputeIndependently
924+
}
925+
}
926+
927+
let shared = Rc::new(agg(vec![2], default_quantile(0.99), metric_scan(&["job"])));
928+
let bound = implement_workload_with(
929+
vec![("a", Rc::clone(&shared)), ("b", Rc::clone(&shared))],
930+
&NeverShareCse,
931+
);
932+
let [(_, ra), (_, rb)] = bound.as_slice() else {
933+
panic!("expected 2 bound results");
934+
};
935+
let ra = ra.as_ref().expect("a failed to bind");
936+
let rb = rb.as_ref().expect("b failed to bind");
937+
assert!(
938+
!Rc::ptr_eq(ra, rb),
939+
"a CostModel that declines CSE sharing must bind each occurrence \
940+
independently, even for two roots that are the same Rc<QueryExpr>"
941+
);
942+
}
869943
}

0 commit comments

Comments
 (0)