Skip to content

Commit fccecf4

Browse files
zzylolclaude
andcommitted
fix(asap-aware-mapping): SketchApplicabilityRule must dedup shared nodes like SharedSubexpressionRule does
collect_sketch_findings (SketchApplicabilityRule) walked the workload as a tree, with no Rc-identity tracking - unlike SharedSubexpressionRule's register_site, which explicitly dedupes by Rc::as_ptr. A bindable Aggregate reachable via two paths after share_common_subtrees (two different roots, or two branches of one root, e.g. median(x) / median(x)) was reported as two separate "sketch-approximation is applicable" findings, even though it's bound exactly once - one CostModel::cse_share_decision per shared Rc, per bind::implement_workload_with, not one per path that reaches it. Overstated the number of independent opportunities. for_each_operator_child now passes an Option<usize> Rc pointer alongside each child (None only for Concat's by-value branches, which have no Rc identity of their own) so collect_sketch_findings can thread a `visited` HashSet across the whole evaluate() call - shared across all roots, not reset per root - and skip re-descending into an already-visited pointer, mirroring register_site's exact dedup rationale. New test: a_shared_sketchable_aggregate_is_reported_only_once, the same median(x)/median(x) shape pre_asap::cse's own single_query_shares_its_own_repeated_subtree test uses, pinning that a single shared Aggregate produces exactly one finding. Found while auditing other open PRs for the same tree-vs-DAG issue fixed in PR #243 (default_cse_recompute_cost's naive serde_json-length proxy over-counting internally-shared subtrees). Verified: cargo build --workspace --all-targets, cargo test --workspace (444 passed, 0 failed), cargo fmt --all -- --check, cargo clippy --workspace --all-targets --all-features -- -D warnings. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent fb8bfd5 commit fccecf4

1 file changed

Lines changed: 102 additions & 14 deletions

File tree

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

Lines changed: 102 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@
7373
//! [`boundary`]: crate::boundary
7474
//! [`bind`]: crate::bind
7575
76-
use std::collections::HashMap;
76+
use std::collections::{HashMap, HashSet};
7777
use std::fmt::Display;
7878
use std::rc::Rc;
7979

@@ -184,12 +184,26 @@ impl ApplicabilityRule for SketchApplicabilityRule<'_> {
184184

185185
fn evaluate(&self, roots: &[(String, Rc<QueryExpr>)]) -> Vec<ApplicabilityFinding> {
186186
let mut findings = Vec::new();
187+
// One `visited` set shared across every root — not reset per root —
188+
// so a bindable `Aggregate` reachable from more than one place (two
189+
// different roots after `share_common_subtrees`, or two branches of
190+
// the same root, e.g. `median(x) / median(x)`) is only ever reported
191+
// once. It really is bound exactly once (stage 2/4's memoization in
192+
// `bind::implement_workload_with` reuses or, per a `CostModel`'s
193+
// `cse_share_decision`, independently rebinds it — either way, one
194+
// decision per shared `Rc`, not one per path that reaches it), so
195+
// reporting "sketch-approximation is applicable" once per path would
196+
// overstate the number of independent opportunities. Same dedup
197+
// rationale as `SharedSubexpressionRule`'s `register_site`, applied
198+
// here too since this traversal has the identical DAG-not-tree shape.
199+
let mut visited: HashSet<usize> = HashSet::new();
187200
for (id, root) in roots {
188201
collect_sketch_findings(
189202
root,
190203
&format!("root {id:?}"),
191204
self.cost_model,
192205
&mut findings,
206+
&mut visited,
193207
);
194208
}
195209
findings
@@ -201,12 +215,16 @@ impl ApplicabilityRule for SketchApplicabilityRule<'_> {
201215
/// itself requires — a multi-intent or `HAVING`-bearing `Aggregate`, or one
202216
/// under a logical parent that would subsume it, stays logical at actual
203217
/// binding time too, so it is not reported here either), recursing through
204-
/// every operator position a nested `Aggregate` could appear in.
218+
/// every operator position a nested `Aggregate` could appear in. `visited`
219+
/// (shared across the whole call, not reset per node) skips re-descending
220+
/// into a child already reached via a different path — see the dedup
221+
/// rationale on [`SketchApplicabilityRule::evaluate`].
205222
fn collect_sketch_findings(
206223
node: &QueryExpr,
207224
location: &str,
208225
cost_model: &dyn CostModel,
209226
findings: &mut Vec<ApplicabilityFinding>,
227+
visited: &mut HashSet<usize>,
210228
) {
211229
if let QueryExpr::Aggregate {
212230
measures, having, ..
@@ -227,8 +245,17 @@ fn collect_sketch_findings(
227245
}
228246
}
229247
}
230-
for_each_operator_child(node, location, |child, child_location| {
231-
collect_sketch_findings(child, &child_location, cost_model, findings);
248+
for_each_operator_child(node, location, |child, ptr, child_location| {
249+
// A `None` ptr (a `Concat` branch, stored by value — see
250+
// `for_each_operator_child`'s own doc) has no `Rc` identity of its
251+
// own to dedup on, so it's always walked; an `Rc`-backed child is
252+
// only walked the first time its pointer is seen.
253+
if let Some(ptr) = ptr {
254+
if !visited.insert(ptr) {
255+
return;
256+
}
257+
}
258+
collect_sketch_findings(child, &child_location, cost_model, findings, visited);
232259
});
233260
}
234261

@@ -384,14 +411,23 @@ fn walk_rc_children(node: &QueryExpr, location: &str, sites: &mut HashMap<usize,
384411
fn for_each_operator_child(
385412
node: &QueryExpr,
386413
location: &str,
387-
mut visit: impl FnMut(&QueryExpr, String),
414+
mut visit: impl FnMut(&QueryExpr, Option<usize>, String),
388415
) {
389416
use QueryExpr::*;
417+
// `Some(Rc::as_ptr(rc) as usize)` for an `Rc`-backed child position, so
418+
// a caller (`collect_sketch_findings`) can dedup on it the same way
419+
// `walk_rc_children`/`register_site` dedup on `Rc::as_ptr` directly —
420+
// `None` for `Concat`'s branches, the one operator position stored by
421+
// value instead (`Vec<QueryExpr>`, never itself `Rc`-aliased — see
422+
// `pre_asap::cse`'s module doc on why), which therefore has no pointer
423+
// identity of its own to dedup on at this position.
390424
match node {
391425
Scan { .. } | PromqlScalar(_) | QueryTimestamp => {}
392-
PromqlVectorFromScalar(c) | PromqlScalarFromVector(c) => {
393-
visit(c, format!("{location} > child"))
394-
}
426+
PromqlVectorFromScalar(c) | PromqlScalarFromVector(c) => visit(
427+
c,
428+
Some(Rc::as_ptr(c) as usize),
429+
format!("{location} > child"),
430+
),
395431
PromqlRelabel { child, .. }
396432
| PromqlInfoEnrich { child, .. }
397433
| PromqlSeriesSample { child, .. }
@@ -404,19 +440,39 @@ fn for_each_operator_child(
404440
| TimeShift { child, .. }
405441
| SQLWindowFunc { child, .. }
406442
| Sort { child, .. }
407-
| Limit { child, .. } => visit(child, format!("{location} > child")),
443+
| Limit { child, .. } => visit(
444+
child,
445+
Some(Rc::as_ptr(child) as usize),
446+
format!("{location} > child"),
447+
),
408448
Concat { children } => {
409449
for (i, c) in children.iter().enumerate() {
410-
visit(c, format!("{location} > concat[{i}]"));
450+
visit(c, None, format!("{location} > concat[{i}]"));
411451
}
412452
}
413453
Join { left, right, .. } | SetOp { left, right, .. } => {
414-
visit(left, format!("{location} > left"));
415-
visit(right, format!("{location} > right"));
454+
visit(
455+
left,
456+
Some(Rc::as_ptr(left) as usize),
457+
format!("{location} > left"),
458+
);
459+
visit(
460+
right,
461+
Some(Rc::as_ptr(right) as usize),
462+
format!("{location} > right"),
463+
);
416464
}
417465
BinaryOp { lhs, rhs, .. } => {
418-
visit(lhs, format!("{location} > lhs"));
419-
visit(rhs, format!("{location} > rhs"));
466+
visit(
467+
lhs,
468+
Some(Rc::as_ptr(lhs) as usize),
469+
format!("{location} > lhs"),
470+
);
471+
visit(
472+
rhs,
473+
Some(Rc::as_ptr(rhs) as usize),
474+
format!("{location} > rhs"),
475+
);
420476
}
421477
Column(_)
422478
| Literal(_)
@@ -584,6 +640,38 @@ mod tests {
584640
.all(|f| f.optimization != OptimizationKind::SketchApproximation));
585641
}
586642

643+
/// A sketch-applicable `Aggregate` reachable via two paths that
644+
/// `share_common_subtrees` collapses onto one `Rc` — the same
645+
/// `median(x) / median(x)` shape `pre_asap::cse`'s own
646+
/// `single_query_shares_its_own_repeated_subtree` test uses — must be
647+
/// reported once, not once per path. It is bound exactly once (one
648+
/// `CostModel::cse_share_decision`, per `bind::implement_workload_with`),
649+
/// so reporting it twice would overstate the number of independent
650+
/// sketch-approximation opportunities in this workload.
651+
#[test]
652+
fn a_shared_sketchable_aggregate_is_reported_only_once() {
653+
let quantile = agg(vec![2], default_quantile(0.99), metric_scan(&["job"]));
654+
let root = QueryExpr::BinaryOp {
655+
op: asap_types::pre_asap::query_expr::BinaryOpKind::Compare(
656+
asap_types::pre_asap::expr_ir::CompareOpKind::Eq,
657+
),
658+
lhs: Rc::new(quantile.clone()),
659+
rhs: Rc::new(quantile),
660+
vector_match: None,
661+
};
662+
let findings = find_applicable_optimizations(vec![("ratio", root)]);
663+
let sketch: Vec<_> = findings
664+
.iter()
665+
.filter(|f| f.optimization == OptimizationKind::SketchApproximation)
666+
.collect();
667+
assert_eq!(
668+
sketch.len(),
669+
1,
670+
"a single shared Aggregate must produce one finding, not one \
671+
per path that reaches it: got {findings:?}"
672+
);
673+
}
674+
587675
// ── SharedSubexpressionRule ─────────────────────────────────────────
588676

589677
#[test]

0 commit comments

Comments
 (0)