Skip to content

Commit 727cf89

Browse files
zzylolclaude
andcommitted
fix(devtools): price every candidate; --max-candidates only limits the plans written
stage_pipeline selected among the first --max-candidates combinations, so a cap below the combination count could select a costlier plan (Examples 3a and 4a selected P122 at 128 of 486; the cheapest is P365). Stage 3 now prices up to 4096 combinations, and the document carries the cheapest --max-candidates plans with their logical candidates, then invalid ones, and a shown_of section with the totals when not every plan is written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 9b6f89f commit 727cf89

2 files changed

Lines changed: 85 additions & 13 deletions

File tree

‎crates/devtools/src/bin/stage_pipeline.rs‎

Lines changed: 68 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,23 @@
1717
// merges something, again with identical sub-DAGs shared ("· shared
1818
// input"), and, when the summary-capability rule applies, again with
1919
// one summary sized for its strictest consumer ("· shared summary"); in
20-
// enumeration order and capped by `--max-candidates` (default 64). A
20+
// enumeration order, only those with a written physical candidate. A
2121
// repeating query's mergeable alternatives also come in tumbling panes
2222
// (Pass 2's window-composition rule), e.g. "Q1 Kll · tumbling 1m panes";
2323
// - stage2_physical_asap: per logical candidate, its physical candidates
2424
// (operator implementation; everything at query time, then one per
2525
// down-closed set of summaries maintained at ingestion time, labeled
2626
// e.g. "· ingestion time: Kll ×5 panes"), no cost;
2727
// - stage3_selection: per-candidate costs, the selected candidate, and
28-
// every other candidate as rejected (`valid: false`: inaccurate, over a
29-
// latency bound, needing a capability the deployment lacks, or could not
30-
// be built) or costlier.
28+
// every other written candidate as rejected (`valid: false`: inaccurate,
29+
// over a latency bound, needing a capability the deployment lacks, or
30+
// could not be built) or costlier.
31+
// - shown_of: present when not every plan is written, the totals.
32+
//
33+
// Stage 3 prices every candidate (up to PRICE_LIMIT combinations), so the
34+
// selection does not depend on `--max-candidates`; that flag only limits
35+
// how many plans the document carries (default 64): the cheapest first,
36+
// then invalid ones in enumeration order.
3137
//
3238
// The deployment inputs are the built-in cost and accuracy models and the
3339
// reference executor's capabilities (`asap_executor::capabilities`; without
@@ -46,6 +52,7 @@
4652
// query; without them the queries are exact. `--interval-ms` is the source
4753
// cadence PromQL needs (default 15000).
4854

55+
use std::collections::HashSet;
4956
use std::rc::Rc;
5057

5158
use asap_logical_optimizer::pass1::logical_candidates::{
@@ -134,6 +141,10 @@ fn run(args: Vec<String>) -> Result<(), String> {
134141
std::fs::write(&out, text).map_err(|e| format!("{out}: {e}"))
135142
}
136143

144+
/// Most Stage 1 combinations built and priced; beyond it the selection is
145+
/// over the first PRICE_LIMIT only, and the tool warns.
146+
const PRICE_LIMIT: usize = 4096;
147+
137148
fn stage_pipeline(
138149
workload: &PlanningWorkload,
139150
capabilities: &DeploymentCapabilities,
@@ -180,11 +191,34 @@ fn stage_pipeline(
180191
&demand,
181192
&data,
182193
models,
183-
max_candidates.max(1),
194+
PRICE_LIMIT,
184195
)
185196
.map_err(|e| format!("planning: {e}"))?;
186197
let enumeration = run.enumeration.expect("display was requested");
187198
let combinations = enumeration.combinations;
199+
if combinations > PRICE_LIMIT {
200+
eprintln!("stage_pipeline: priced the first {PRICE_LIMIT} of {combinations} combinations");
201+
}
202+
let selection = &enumeration.selection;
203+
// Every plan id, physical candidates then those that could not be
204+
// built; the cheapest `max_candidates` are written.
205+
let physical_ids: Vec<&str> = enumeration
206+
.candidates
207+
.iter()
208+
.flat_map(|c| c.physical.iter().map(|p| p.id.as_str()))
209+
.collect();
210+
let mut ranked = physical_ids.clone();
211+
let known: HashSet<&str> = ranked.iter().copied().collect();
212+
ranked.extend(
213+
selection
214+
.rejected
215+
.iter()
216+
.map(|r| r.id.as_str())
217+
.filter(|id| !known.contains(id)),
218+
);
219+
let total = |id: &str| selection.costs.get(id).map_or(f64::INFINITY, |c| c.total);
220+
ranked.sort_by(|a, b| total(a).total_cmp(&total(b)));
221+
let shown: HashSet<&str> = ranked.iter().take(max_candidates.max(1)).copied().collect();
188222
let mut candidates = Vec::new();
189223
let mut stage2 = Vec::new();
190224
for candidate in &enumeration.candidates {
@@ -209,12 +243,24 @@ fn stage_pipeline(
209243
Sharing::IdenticalExpressions => " · shared input",
210244
Sharing::SummaryCapability => " · shared summary",
211245
};
246+
let written = candidate
247+
.physical
248+
.iter()
249+
.any(|p| shown.contains(p.id.as_str()))
250+
|| shown.contains(format!("P{index}").as_str());
251+
if !written {
252+
continue;
253+
}
212254
if let Some(logical) = &candidate.logical {
213255
let roots: Vec<_> = logical.iter().map(|(_, root)| root.clone()).collect();
214256
candidates
215257
.push(json!({ "id": format!("L{index}"), "label": label, "dag": export(&roots)? }));
216258
}
217-
for p in &candidate.physical {
259+
for p in candidate
260+
.physical
261+
.iter()
262+
.filter(|p| shown.contains(p.id.as_str()))
263+
{
218264
let label = match p.materialization.as_str() {
219265
"" => label.clone(),
220266
m => format!("{label} · {m}"),
@@ -224,19 +270,27 @@ fn stage_pipeline(
224270
);
225271
}
226272
}
227-
Ok(json!({
273+
let mut document = json!({
228274
"format": "asap-stage-pipeline/v1",
229275
"workload": { "queries": workload_queries(workload) },
230276
"deployment": deployment_json(&models),
231277
"stage0_logical": { "dag": stage0 },
232278
"stage1_logical_asap": {
233279
"combinations": combinations,
234-
"capped": combinations > max_candidates,
280+
"capped": shown.len() < ranked.len(),
235281
"candidates": candidates,
236282
},
237283
"stage2_physical_asap": { "candidates": stage2 },
238-
"stage3_selection": stage3_json(&enumeration.selection),
239-
}))
284+
"stage3_selection": stage3_json(selection, &shown),
285+
});
286+
if shown.len() < ranked.len() {
287+
document["shown_of"] = json!({
288+
"logical": enumeration.candidates.len(),
289+
"physical": physical_ids.len(),
290+
"priced": selection.costs.len(),
291+
});
292+
}
293+
Ok(document)
240294
}
241295

242296
/// The deployment inputs Stage 3 used: the executor's capabilities, one
@@ -280,10 +334,12 @@ fn deployment_json(models: &PlanningModels<'_>) -> Value {
280334
})
281335
}
282336

283-
fn stage3_json(selection: &Selection) -> Value {
337+
/// Stage 3's outcome for the written plans `shown`.
338+
fn stage3_json(selection: &Selection, shown: &HashSet<&str>) -> Value {
284339
let costs: serde_json::Map<_, _> = selection
285340
.costs
286341
.iter()
342+
.filter(|(id, _)| shown.contains(id.as_str()))
287343
.map(|(id, cost)| {
288344
let per_node: serde_json::Map<_, _> = cost
289345
.per_node
@@ -304,6 +360,7 @@ fn stage3_json(selection: &Selection) -> Value {
304360
let rejected: Vec<_> = selection
305361
.rejected
306362
.iter()
363+
.filter(|r| shown.contains(r.id.as_str()))
307364
.map(|r| json!({ "id": r.id, "valid": r.valid, "reason": r.reason }))
308365
.collect();
309366
json!({ "costs": costs, "selected": selection.selected, "rejected": rejected })

‎crates/devtools/tests/stage_pipeline.rs‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,14 +105,29 @@ fn example1_document_is_valid_and_committed_fixture_is_current() {
105105
assert_eq!(accounted, all);
106106
}
107107

108-
/// The Cartesian product is cut at `--max-candidates` and says so.
108+
/// `--max-candidates` limits the plans written, cheapest first, and says
109+
/// so; Stage 3 still prices every plan, so the selection does not change.
109110
#[test]
110111
fn candidate_cap_is_recorded() {
111112
let document = generate(&["--example", "planner-layering-1", "--max-candidates", "5"]);
113+
let full = generate(&EXAMPLE1);
112114
let stage1 = &document["stage1_logical_asap"];
113115
assert_eq!(stage1["capped"], true);
114-
assert_eq!(stage1["candidates"].as_array().unwrap().len(), 5);
116+
assert!(stage1["candidates"].as_array().unwrap().len() <= 5);
115117
assert!(stage1["combinations"].as_u64().unwrap() > 5);
118+
let physical = document["stage2_physical_asap"]["candidates"]
119+
.as_array()
120+
.unwrap();
121+
assert_eq!(physical.len(), 5);
122+
let selection = &document["stage3_selection"];
123+
assert_eq!(selection["selected"], full["stage3_selection"]["selected"]);
124+
let costs = selection["costs"].as_object().unwrap();
125+
assert_eq!(costs.len(), 5, "the five cheapest plans are written");
126+
let full_physical = full["stage2_physical_asap"]["candidates"]
127+
.as_array()
128+
.unwrap();
129+
assert_eq!(document["shown_of"]["physical"], full_physical.len());
130+
assert!(full.get("shown_of").is_none());
116131
}
117132

118133
/// Example 3, Pattern B: the 5-min p99 every minute offers KLL and DDSketch

0 commit comments

Comments
 (0)