Skip to content

Commit 63f7824

Browse files
zzylolzzylol
committed
feat(promql): retain summary leaves in arithmetic plans
1 parent dd00701 commit 63f7824

10 files changed

Lines changed: 266 additions & 4 deletions

File tree

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

Lines changed: 98 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,7 @@ use asap_types::post_asap::{
357357
use asap_types::pre_asap::agg_intent::{agg_is_mergeable, AggIntent};
358358
use asap_types::pre_asap::cse::{share_common_subtrees, structural_hash, HashCache};
359359
use asap_types::pre_asap::expr_ir::ColumnRef;
360-
use asap_types::pre_asap::query_expr::{QueryExpr, QueryExprError, Reduction};
360+
use asap_types::pre_asap::query_expr::{BinaryOpKind, QueryExpr, QueryExprError, Reduction};
361361
use asap_types::pre_asap::schema::Schema;
362362
use asap_types::types::AccuracyTarget;
363363
use asap_types::workload::{QueryRecurrence, QueryWorkload, RepeatedDemand};
@@ -1263,6 +1263,17 @@ impl<'a> SketchAlgorithmStrategy<'a> {
12631263
/// differs — see [`realize_child_with`]).
12641264
fn propose_with(&self, root: &Rc<QueryExpr>, intent_override: Option<&AggIntent>) -> Proposals {
12651265
let mut proposals = Proposals::default();
1266+
if intent_override.is_none() && is_supported_exact_binary(root) {
1267+
if let Ok(Some(node)) = realize_exact_binary(root, self.models) {
1268+
proposals.candidates.push(ReplacementSubDAG {
1269+
replacement: Replacement::Summary(node),
1270+
strategy: "SketchAlgorithmStrategy",
1271+
provenance: ReplacementProvenance::SummaryImplementation,
1272+
rationale: "preserve exact PromQL arithmetic over independently realized summary operands".into(),
1273+
});
1274+
}
1275+
return proposals;
1276+
}
12661277
let Some(declared) = bindable_intent(root) else {
12671278
return proposals;
12681279
};
@@ -1412,7 +1423,7 @@ fn aggregate_child(node: &QueryExpr) -> Option<&Rc<QueryExpr>> {
14121423

14131424
impl ReplacementStrategy for SketchAlgorithmStrategy<'_> {
14141425
fn matches(&self, target: &TargetSubDAG<'_>) -> bool {
1415-
bindable_intent(target.root).is_some()
1426+
bindable_intent(target.root).is_some() || is_supported_exact_binary(target.root)
14161427
}
14171428

14181429
fn replacements(&self, target: &TargetSubDAG<'_>) -> Vec<ReplacementSubDAG> {
@@ -1527,6 +1538,9 @@ pub(crate) fn realize_child_with(
15271538
models: Models<'_>,
15281539
end_to_end_target: Option<&AccuracyTarget>,
15291540
) -> Result<Rc<SummaryNode>, ImplementError> {
1541+
if let Some(composed) = realize_exact_binary(root, models)? {
1542+
return Ok(composed);
1543+
}
15301544
let overridden = end_to_end_target.and_then(|target| {
15311545
let declared = bindable_intent(root)?;
15321546
match accuracy_target(declared) {
@@ -1560,6 +1574,88 @@ pub(crate) fn realize_child_with(
15601574
}
15611575
}
15621576

1577+
/// Preserve an exact arithmetic root while allowing each vector operand to
1578+
/// select its own summary implementation. If either vector arm cannot be
1579+
/// accelerated, return `None` so the caller keeps the whole query exact;
1580+
/// mixed raw/summary snapshots are never constructed.
1581+
fn realize_exact_binary(
1582+
root: &Rc<QueryExpr>,
1583+
models: Models<'_>,
1584+
) -> Result<Option<Rc<SummaryNode>>, ImplementError> {
1585+
let QueryExpr::BinaryOp {
1586+
op,
1587+
lhs,
1588+
rhs,
1589+
vector_match,
1590+
} = root.as_ref()
1591+
else {
1592+
return Ok(None);
1593+
};
1594+
if !matches!(op, BinaryOpKind::Arithmetic(_) | BinaryOpKind::Pow) || vector_match.is_some() {
1595+
return Ok(None);
1596+
}
1597+
1598+
let lhs_scalar = is_promql_scalar(lhs);
1599+
let rhs_scalar = is_promql_scalar(rhs);
1600+
if lhs_scalar && rhs_scalar {
1601+
return Ok(None);
1602+
}
1603+
1604+
let lhs_node = realize_binary_operand(lhs, models)?;
1605+
let rhs_node = realize_binary_operand(rhs, models)?;
1606+
let lhs_accelerated = lhs_scalar || !matches!(lhs_node.expr, SummaryExpr::KeepPreAsap(_));
1607+
let rhs_accelerated = rhs_scalar || !matches!(rhs_node.expr, SummaryExpr::KeepPreAsap(_));
1608+
if !lhs_accelerated || !rhs_accelerated {
1609+
return Ok(None);
1610+
}
1611+
1612+
Ok(Some(Rc::new(SummaryNode {
1613+
expr: SummaryExpr::ExactBinary {
1614+
lhs: lhs_node,
1615+
rhs: rhs_node,
1616+
op: op.clone(),
1617+
vector_match: vector_match.clone(),
1618+
},
1619+
schema: lift(&root.output_schema()?),
1620+
guarantee: Some(ResultGuarantee::exact("ExactBinary")),
1621+
})))
1622+
}
1623+
1624+
fn is_supported_exact_binary(root: &QueryExpr) -> bool {
1625+
matches!(
1626+
root,
1627+
QueryExpr::BinaryOp {
1628+
op: BinaryOpKind::Arithmetic(_) | BinaryOpKind::Pow,
1629+
vector_match: None,
1630+
..
1631+
}
1632+
)
1633+
}
1634+
1635+
fn is_promql_scalar(expr: &QueryExpr) -> bool {
1636+
matches!(
1637+
expr,
1638+
QueryExpr::PromqlScalarBridge(_) | QueryExpr::Literal(_)
1639+
)
1640+
}
1641+
1642+
fn realize_binary_operand(
1643+
operand: &Rc<QueryExpr>,
1644+
models: Models<'_>,
1645+
) -> Result<Rc<SummaryNode>, ImplementError> {
1646+
if is_promql_scalar(operand) {
1647+
return Ok(Rc::new(SummaryNode {
1648+
expr: SummaryExpr::KeepPreAsap(Rc::clone(operand)),
1649+
schema: SummarySchema {
1650+
fields: Vec::new(),
1651+
time_index: None,
1652+
},
1653+
guarantee: Some(ResultGuarantee::exact("PromQL scalar")),
1654+
}));
1655+
}
1656+
realize_child_with(operand, models, None)
1657+
}
1658+
15631659
/// `intent` with its `AccuracyTarget` replaced by `target` — a no-op for an
15641660
/// intent that carries none (see [`accuracy_target`]).
15651661
fn override_accuracy(intent: &AggIntent, target: &AccuracyTarget) -> AggIntent {

‎crates/asap-aware-mapping/src/summary_maintenance_cost/estimator.rs‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,11 @@ pub(super) fn estimate_heterogeneous_summary(
3636
}
3737
}
3838
SummaryExpr::SummarySubtract { left, right }
39+
| SummaryExpr::ExactBinary {
40+
lhs: left,
41+
rhs: right,
42+
..
43+
}
3944
| SummaryExpr::SummaryJoin {
4045
outer: left,
4146
inner: right,
@@ -258,6 +263,28 @@ pub(super) fn estimate_heterogeneous_summary(
258263
return Ok(());
259264
}
260265
match &node.expr {
266+
SummaryExpr::ExactBinary { lhs, rhs, .. } => {
267+
visit_ops(
268+
lhs,
269+
seen,
270+
by_node,
271+
evidence,
272+
scope,
273+
evaluation_count,
274+
cpu_ops,
275+
io_bytes,
276+
)?;
277+
visit_ops(
278+
rhs,
279+
seen,
280+
by_node,
281+
evidence,
282+
scope,
283+
evaluation_count,
284+
cpu_ops,
285+
io_bytes,
286+
)?;
287+
}
261288
SummaryExpr::KeepPreAsap(_) => {
262289
let retained = evidence
263290
.retained_queries
@@ -365,6 +392,11 @@ pub(super) fn estimate_heterogeneous_summary(
365392
.for_each(|child| collect_aggs(child, seen, out));
366393
}
367394
SummaryExpr::SummarySubtract { left, right }
395+
| SummaryExpr::ExactBinary {
396+
lhs: left,
397+
rhs: right,
398+
..
399+
}
368400
| SummaryExpr::SummaryJoin {
369401
outer: left,
370402
inner: right,
@@ -559,6 +591,11 @@ fn validate_summary_edges_and_physical_ids(
559591
children.iter().map(|child| child.as_ref()).collect()
560592
}
561593
SummaryExpr::SummarySubtract { left, right }
594+
| SummaryExpr::ExactBinary {
595+
lhs: left,
596+
rhs: right,
597+
..
598+
}
562599
| SummaryExpr::SummaryJoin {
563600
outer: left,
564601
inner: right,
@@ -573,6 +610,9 @@ fn validate_summary_edges_and_physical_ids(
573610
evidence: &StreamingNodeEvidence,
574611
) -> Result<(String, Vec<EdgeStatistics>, EdgeStatistics), AnalyticalCostError> {
575612
match &node.expr {
613+
SummaryExpr::ExactBinary { .. } => Err(AnalyticalCostError::MissingOrStale(
614+
"exact_binary_physical_evidence",
615+
)),
576616
SummaryExpr::KeepPreAsap(_) => {
577617
let retained = evidence
578618
.retained_queries
@@ -723,6 +763,11 @@ pub(super) fn estimate_transient_liveness(
723763
children.iter().map(|child| child.as_ref()).collect()
724764
}
725765
SummaryExpr::SummarySubtract { left, right }
766+
| SummaryExpr::ExactBinary {
767+
lhs: left,
768+
rhs: right,
769+
..
770+
}
726771
| SummaryExpr::SummaryJoin {
727772
outer: left,
728773
inner: right,
@@ -756,6 +801,9 @@ pub(super) fn estimate_transient_liveness(
756801
evidence: &StreamingNodeEvidence,
757802
) -> Result<(u64, u64), AnalyticalCostError> {
758803
match &node.expr {
804+
SummaryExpr::ExactBinary { .. } => Err(AnalyticalCostError::MissingOrStale(
805+
"exact_binary_physical_evidence",
806+
)),
759807
SummaryExpr::KeepPreAsap(_) => evidence
760808
.retained_queries
761809
.get(&(node as *const _))
@@ -838,6 +886,11 @@ pub(super) fn evidence_nodes(root: &SummaryNode) -> (Vec<&SummaryNode>, Vec<&Sum
838886
}
839887
}
840888
SummaryExpr::SummarySubtract { left, right }
889+
| SummaryExpr::ExactBinary {
890+
lhs: left,
891+
rhs: right,
892+
..
893+
}
841894
| SummaryExpr::SummaryJoin {
842895
outer: left,
843896
inner: right,
@@ -1185,6 +1238,10 @@ fn count_operations(root: &SummaryNode) -> Result<SummaryOperationCounts, Analyt
11851238
visit(left, seen, counts)?;
11861239
visit(right, seen, counts)?;
11871240
}
1241+
SummaryExpr::ExactBinary { lhs, rhs, .. } => {
1242+
visit(lhs, seen, counts)?;
1243+
visit(rhs, seen, counts)?;
1244+
}
11881245
SummaryExpr::SummaryDelete { summary_input, .. } => {
11891246
counts.deletes_per_update = counts
11901247
.deletes_per_update

‎crates/asap-aware-mapping/src/summary_maintenance_cost/model.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3131,6 +3131,11 @@ mod tests {
31313131
}
31323132
}
31333133
SummaryExpr::SummarySubtract { left, right }
3134+
| SummaryExpr::ExactBinary {
3135+
lhs: left,
3136+
rhs: right,
3137+
..
3138+
}
31343139
| SummaryExpr::SummaryJoin {
31353140
outer: left,
31363141
inner: right,
@@ -3306,6 +3311,11 @@ mod tests {
33063311
}
33073312
}
33083313
SummaryExpr::SummarySubtract { left, right }
3314+
| SummaryExpr::ExactBinary {
3315+
lhs: left,
3316+
rhs: right,
3317+
..
3318+
}
33093319
| SummaryExpr::SummaryJoin {
33103320
outer: left,
33113321
inner: right,
@@ -3341,6 +3351,11 @@ mod tests {
33413351
}
33423352
}
33433353
SummaryExpr::SummarySubtract { left, right }
3354+
| SummaryExpr::ExactBinary {
3355+
lhs: left,
3356+
rhs: right,
3357+
..
3358+
}
33443359
| SummaryExpr::SummaryJoin {
33453360
outer: left,
33463361
inner: right,

‎crates/asap-aware-mapping/src/summary_maintenance_cost/window.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,11 @@ pub(super) fn summary_aggregation_identities(root: &SummaryNode) -> HashSet<*con
5050
}
5151
}
5252
SummaryExpr::SummarySubtract { left, right }
53+
| SummaryExpr::ExactBinary {
54+
lhs: left,
55+
rhs: right,
56+
..
57+
}
5358
| SummaryExpr::SummaryJoin {
5459
outer: left,
5560
inner: right,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ fn annotate_lifecycle_deployments(
197197
fn summary_children(expr: &SummaryExpr) -> Vec<&Rc<SummaryNode>> {
198198
match expr {
199199
SummaryExpr::KeepPreAsap(_) => vec![],
200+
SummaryExpr::ExactBinary { lhs, rhs, .. } => vec![lhs, rhs],
200201
SummaryExpr::SummaryAgg { child, .. } => vec![child],
201202
SummaryExpr::SummaryJoin { outer, inner, .. }
202203
| SummaryExpr::SummarySubtract {

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -961,6 +961,11 @@ fn collect_summary_aggs(
961961
collect_summary_aggs(child, seen, output);
962962
}
963963
SummaryExpr::SummaryJoin { outer, inner, .. }
964+
| SummaryExpr::ExactBinary {
965+
lhs: outer,
966+
rhs: inner,
967+
..
968+
}
964969
| SummaryExpr::SummarySubtract {
965970
left: outer,
966971
right: inner,

‎crates/integration-tests/tests/promql_to_post_asap.rs‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,63 @@ fn dtype<'a>(schema: &'a SummarySchema, name: &str) -> &'a SummaryFamilyType {
5353
.dtype
5454
}
5555

56+
fn lower_and_realize(query: &str) -> Rc<SummaryNode> {
57+
let pre = lower_promql(query, AccuracyTarget::Exact).expect("lowering failed");
58+
realize(&pre).expect("binding failed")
59+
}
60+
61+
#[test]
62+
fn promql_binary_arithmetic_retains_two_summary_leaves() {
63+
for op in ["+", "-", "*", "/", "%", "^"] {
64+
let root = lower_and_realize(&format!("rate(a[1m]) {op} rate(b[1m])"));
65+
let SummaryExpr::ExactBinary { lhs, rhs, .. } = &root.expr else {
66+
panic!("expected ExactBinary for {op}, got {:?}", root.expr);
67+
};
68+
assert!(matches!(lhs.expr, SummaryExpr::SummaryAgg { .. }));
69+
assert!(matches!(rhs.expr, SummaryExpr::SummaryAgg { .. }));
70+
}
71+
}
72+
73+
#[test]
74+
fn promql_binary_arithmetic_preserves_both_scalar_operand_orders() {
75+
for query in ["rate(a[1m]) / 2", "2 / rate(a[1m])"] {
76+
let root = lower_and_realize(query);
77+
let SummaryExpr::ExactBinary { lhs, rhs, .. } = &root.expr else {
78+
panic!("expected ExactBinary for {query}, got {:?}", root.expr);
79+
};
80+
assert!(matches!(
81+
lhs.expr,
82+
SummaryExpr::SummaryAgg { .. } | SummaryExpr::KeepPreAsap(_)
83+
));
84+
assert!(matches!(
85+
rhs.expr,
86+
SummaryExpr::SummaryAgg { .. } | SummaryExpr::KeepPreAsap(_)
87+
));
88+
assert!(
89+
matches!(lhs.expr, SummaryExpr::SummaryAgg { .. })
90+
|| matches!(rhs.expr, SummaryExpr::SummaryAgg { .. })
91+
);
92+
}
93+
}
94+
95+
#[test]
96+
fn promql_binary_arithmetic_falls_back_as_a_whole_for_unsupported_arm() {
97+
let root = lower_and_realize("rate(a[1m]) + avg_over_time(b[1m])");
98+
assert!(matches!(root.expr, SummaryExpr::KeepPreAsap(_)));
99+
}
100+
101+
#[test]
102+
fn promql_binary_arithmetic_preserves_nested_structure_and_rejects_modifiers() {
103+
let nested = lower_and_realize("(rate(a[1m]) + rate(b[1m])) / 2");
104+
let SummaryExpr::ExactBinary { lhs, .. } = &nested.expr else {
105+
panic!("expected outer ExactBinary, got {:?}", nested.expr);
106+
};
107+
assert!(matches!(lhs.expr, SummaryExpr::ExactBinary { .. }));
108+
109+
let modified = lower_and_realize("rate(a[1m]) + on(job) rate(b[1m])");
110+
assert!(matches!(modified.expr, SummaryExpr::KeepPreAsap(_)));
111+
}
112+
56113
/// `quantile(0.99, rate(http_requests_total[5m]))` at ε = 0.01:
57114
///
58115
/// ```text

0 commit comments

Comments
 (0)