Skip to content

Commit b83918a

Browse files
feat(query-engine): add native instant binary-expr evaluator, staged ahead of cutover (#572)
* feat(query-engine): add native instant binary-expr evaluator, staged ahead of cutover Builds a native (non-DataFusion) implementation of PromQL binary- arithmetic instant queries: a recursive arm evaluator over Vec<InstantVectorElement>, plus vector-vector and scalar combiners lifted from the existing range-binary path. Not yet wired into production dispatch (handle_query_promql still calls the DataFusion path) — exposed via handle_query_promql_native for equivalence testing against the DataFusion path ahead of the Stage 3 cutover. One accepted, deliberately loud behavior change: an arm with zero current precomputed data now falls back to Prometheus (matching "not acceleratable" semantics) instead of DataFusion's silent empty-result behavior, with a warn! so it's observable. Part of #567 Stage 2. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * fix(query-engine): reject cross-label-set joins in native binary-expr combiner combine_vector_vector_native joined two arms purely by positional KeyByLabelValues equality (a bare Vec<String> of values, no label names attached). DataFusion's build_binary_vector_plan joins on named columns instead, so it fails to resolve (-> None) whenever the two arms don't share the same label set. Native had no equivalent check, so it could silently fabricate a joined result whenever two differently-grouped arms' values happened to coincide (e.g. `sum(a) by (host) + sum(b) by (region)` with a host value equal to a region value). combine_vector_vector_native now takes both arms' label-name lists and returns None on a mismatch, matching DataFusion's failed-join behavior. Adds four regression tests proving the divergence (three were red before this fix) plus a same-label-set control case. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * fix(query-engine): enable topk limiting/formatting for binary-expr leaf arms evaluate_arm_native's leaf branch hardcoded execute_query_pipeline(&ctx, false, false), so a topk arm inside a binary expression (e.g. `topk(10, metric) + 0`) never got truncated to k or metric-name-prefixed — both flags are self-gated on statistic == Topk / a "k" kwarg being present (see execute_query_pipeline's doc comment), so passing (true, true) unconditionally is a no-op for non-topk arms, matching what the main non-binary instant-query path already does. Adds a regression test proving the divergence: topk(10, ...) + 0 returned all 15 unformatted rows before this fix, now correctly truncates to 10 with the metric-name prefix. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 7c2eb98 commit b83918a

3 files changed

Lines changed: 716 additions & 1 deletion

File tree

‎asap-query-engine/src/engines/simple_engine/promql.rs‎

Lines changed: 243 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use super::{
99
RangeQueryParams,
1010
};
1111
use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig};
12-
use crate::engines::query_result::{QueryResult, RangeVectorElement};
12+
use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement};
1313
use asap_types::query_requirements::build_query_requirements_promql;
1414
use asap_types::PromQLSchema;
1515
use promql_utilities::ast_matching::PromQLMatchResult;
@@ -36,6 +36,67 @@ fn detect_scalar_arm<'a>(
3636
}
3737
}
3838

39+
/// Native vector-vector combiner for one binary-expr level (instant query):
40+
/// joins two arms' results by label key and applies `op` per matching key,
41+
/// dropping non-matches (inner-join semantics, matching DataFusion's
42+
/// `build_binary_vector_plan`). Positional `KeyByLabelValues` equality is
43+
/// safe *once `lhs_labels == rhs_labels` is confirmed*: label names are
44+
/// canonically sorted by `KeyByLabelNames::new()`, so two arms with the same
45+
/// label set always order their values the same way. Returns `None` if the
46+
/// two arms don't share the same label set — mirroring DataFusion's
47+
/// `build_binary_vector_plan`, which fails to resolve a join column that
48+
/// only exists on one side.
49+
fn combine_vector_vector_native(
50+
lhs_results: Vec<InstantVectorElement>,
51+
lhs_labels: &[String],
52+
rhs_results: Vec<InstantVectorElement>,
53+
rhs_labels: &[String],
54+
op: &promql_parser::parser::token::TokenType,
55+
) -> Option<Vec<InstantVectorElement>> {
56+
if lhs_labels != rhs_labels {
57+
return None;
58+
}
59+
60+
let rhs_map: HashMap<KeyByLabelValues, f64> = rhs_results
61+
.into_iter()
62+
.map(|elem| (elem.labels, elem.value))
63+
.collect();
64+
65+
Some(
66+
lhs_results
67+
.into_iter()
68+
.filter_map(|lhs_elem| {
69+
rhs_map.get(&lhs_elem.labels).map(|&rhs_val| {
70+
let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val);
71+
InstantVectorElement::new(lhs_elem.labels.clone(), value)
72+
})
73+
})
74+
.collect(),
75+
)
76+
}
77+
78+
/// Native scalar combiner for one binary-expr level (instant query):
79+
/// applies `op(scalar, value)` or `op(value, scalar)` per `scalar_on_left`
80+
/// to every element of the vector arm's results.
81+
fn combine_scalar_native(
82+
vector_results: Vec<InstantVectorElement>,
83+
scalar: f64,
84+
op: &promql_parser::parser::token::TokenType,
85+
scalar_on_left: bool,
86+
) -> Vec<InstantVectorElement> {
87+
vector_results
88+
.into_iter()
89+
.map(|elem| {
90+
let value = if scalar_on_left {
91+
SimpleEngine::apply_range_binary_op(op, scalar, elem.value)
92+
} else {
93+
SimpleEngine::apply_range_binary_op(op, elem.value, scalar)
94+
};
95+
InstantVectorElement::new(elem.labels, value)
96+
})
97+
.collect()
98+
}
99+
39100
impl SimpleEngine {
40101
/// Aligns `end_timestamp` down to the nearest data-ingestion-interval
41102
/// boundary, unconditionally — mirroring SQL's `align_end_timestamp_sql`.
@@ -453,6 +514,127 @@ impl SimpleEngine {
453514
Some((output_labels, QueryResult::vector(results, query_time)))
454515
}
455516

517+
/// Recursively evaluates one arm of a binary arithmetic expression via
518+
/// the native pipeline (`execute_query_pipeline`), instead of building a
519+
/// DataFusion plan. Mirrors `build_arm_logical_plan`'s shape exactly,
520+
/// including its limitation: nested `Binary` arms are combined as
521+
/// vector-vector only — a scalar inside a nested arm (e.g. `(a+5)*b`)
522+
/// is not supported, matching today's DataFusion path.
523+
///
524+
/// - Leaf arm: resolved via `resolve_arm_leaf_context`, executed through
525+
/// `execute_query_pipeline`.
526+
/// - Binary arm: recursively evaluate both sides then combine with
527+
/// `combine_vector_vector_native`.
528+
/// - Scalar literal: returns `None` (handled by the caller separately).
529+
fn evaluate_arm_native(
530+
&self,
531+
arm_ast: &promql_parser::parser::Expr,
532+
time: f64,
533+
) -> Option<(Vec<InstantVectorElement>, Vec<String>)> {
534+
use promql_parser::parser::Expr;
535+
536+
match arm_ast {
537+
Expr::NumberLiteral(_) => None, // caller handles scalars
538+
Expr::Paren(paren) => self.evaluate_arm_native(&paren.expr, time),
539+
Expr::Binary(binary) => {
540+
// Nested binary expression — recurse on both sides
541+
let (lhs_results, lhs_labels) = self.evaluate_arm_native(&binary.lhs, time)?;
542+
let (rhs_results, rhs_labels) = self.evaluate_arm_native(&binary.rhs, time)?;
543+
let combined = combine_vector_vector_native(
544+
lhs_results,
545+
&lhs_labels,
546+
rhs_results,
547+
&rhs_labels,
548+
&binary.op,
549+
)?;
550+
Some((combined, lhs_labels))
551+
}
552+
_ => {
553+
let (ctx, label_names) = self.resolve_arm_leaf_context(arm_ast, time)?;
554+
// Unlike DataFusion's PrecomputedSummaryReadExec (which streams
555+
// whatever rows exist, including zero, so a currently-empty arm
556+
// still returns Some(empty vector)), execute_query_pipeline errors
557+
// when the store has no precomputed outputs at all for this arm —
558+
// that propagates to None here, triggering a full Prometheus
559+
// fallback for the whole expression instead of an empty result
560+
// for just this arm. Accepted behavior change (#567); warn loudly
561+
// so it's visible rather than silent.
562+
let results = self
563+
// (true, true): safe unconditionally — both flags are
564+
// self-gated on statistic == Topk / a "k" kwarg being
565+
// present, same as the main instant-query path (see
566+
// execute_query_pipeline's doc comment).
567+
.execute_query_pipeline(&ctx, true, true)
568+
.map_err(|e| {
569+
warn!(
570+
"Native binary-expr arm for metric '{}' produced no results \
571+
({}) — unlike DataFusion, this falls back to Prometheus for \
572+
the whole expression rather than returning an empty result \
573+
for just this arm",
574+
ctx.metric, e
575+
);
576+
e
577+
})
578+
.ok()?;
579+
Some((results, label_names))
580+
}
581+
}
582+
}
583+
584+
/// Native (non-DataFusion) implementation of binary-expr handling,
585+
/// staged ahead of #567's cutover so it can be equivalence-tested
586+
/// against the DataFusion path (`handle_binary_expr_promql`) before
587+
/// that path is rewired to call this one. Not yet reachable from
588+
/// `handle_query_promql` — see `handle_query_promql_native` for a
589+
/// test-facing entrypoint.
590+
pub fn handle_binary_expr_promql_native(
591+
&self,
592+
ast: &promql_parser::parser::Expr,
593+
time: f64,
594+
) -> Option<(KeyByLabelNames, QueryResult)> {
595+
use promql_parser::parser::Expr;
596+
597+
let query_time = Self::convert_query_time_to_data_time(time);
598+
599+
let binary = match ast {
600+
Expr::Binary(b) => b,
601+
_ => return None,
602+
};
603+
604+
let lhs = binary.lhs.as_ref();
605+
let rhs = binary.rhs.as_ref();
606+
let op = &binary.op;
607+
608+
if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) {
609+
let (vector_results, label_names) = self.evaluate_arm_native(vector_arm, time)?;
610+
let combined = combine_scalar_native(vector_results, scalar, op, scalar_on_left);
611+
return Some((
612+
KeyByLabelNames::new(label_names),
613+
QueryResult::vector(combined, query_time),
614+
));
615+
}
616+
617+
// Vector–vector
618+
let (lhs_results, lhs_labels) = self.evaluate_arm_native(lhs, time)?;
619+
let (rhs_results, rhs_labels) = self.evaluate_arm_native(rhs, time)?;
620+
let combined =
621+
combine_vector_vector_native(lhs_results, &lhs_labels, rhs_results, &rhs_labels, op)?;
622+
let output_labels = KeyByLabelNames::new(lhs_labels);
623+
Some((output_labels, QueryResult::vector(combined, query_time)))
624+
}
625+
626+
/// Parses `query` and dispatches to `handle_binary_expr_promql_native`.
627+
/// Test-facing convenience wrapper — mirrors the binary-expr branch of
628+
/// `handle_query_promql`, but for the native path.
629+
pub fn handle_query_promql_native(
630+
&self,
631+
query: String,
632+
time: f64,
633+
) -> Option<(KeyByLabelNames, QueryResult)> {
634+
let ast = promql_parser::parser::parse(&query).ok()?;
635+
self.handle_binary_expr_promql_native(&ast, time)
636+
}
637+
456638
/// Applies a PromQL binary arithmetic operator to two f64 values.
457639
fn apply_range_binary_op(
458640
op: &promql_parser::parser::token::TokenType,
@@ -1474,4 +1656,64 @@ mod topk_pipeline_tests {
14741656
assert!(pair[0] >= pair[1]);
14751657
}
14761658
}
1659+
1660+
/// A topk leaf wrapped in a binary expr (`topk(10, ...) + 0`) must still
1661+
/// get the same top-10 truncation and metric-name-prefixed formatting as
1662+
/// the bare `topk(10, ...)` query — evaluate_arm_native's leaf branch
1663+
/// used to hardcode (false, false) for enable_topk_limiting/formatting,
1664+
/// which would have returned all 15 unformatted (single-label) rows here
1665+
/// instead of the top 10 with the metric-name prefix.
1666+
#[test]
1667+
fn topk_wrapped_in_binary_expr_still_truncates_and_formats() {
1668+
let (engine, store) = build_topk_engine();
1669+
1670+
let context = engine
1671+
.build_query_execution_context_promql(TOPK_QUERY.to_string(), QUERY_TIME)
1672+
.expect("context should build");
1673+
let window = &context.store_plan.values_query;
1674+
1675+
let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32);
1676+
for i in 1..=15u64 {
1677+
let srcip = format!("10.0.0.{i}");
1678+
sketch.inner.update(&srcip, (i * 10) as f64);
1679+
}
1680+
1681+
let output =
1682+
PrecomputedOutput::new(window.start_timestamp, window.end_timestamp, None, AGG_ID);
1683+
store
1684+
.insert_precomputed_output(output, Box::new(sketch))
1685+
.expect("insert should succeed");
1686+
1687+
let (_, query_result) = engine
1688+
.handle_query_promql_native(format!("{TOPK_QUERY} + 0"), QUERY_TIME)
1689+
.expect("binary-expr-wrapped topk should still resolve via the native path");
1690+
1691+
let results = match query_result {
1692+
QueryResult::Vector(iv) => iv.values,
1693+
other => panic!("expected a vector result, got {other:?}"),
1694+
};
1695+
1696+
assert_eq!(
1697+
results.len(),
1698+
10,
1699+
"topk(10, ...) + 0 must still truncate to 10 rows"
1700+
);
1701+
for pair in results.windows(2) {
1702+
assert!(
1703+
pair[0].value >= pair[1].value,
1704+
"results must stay sorted by count descending"
1705+
);
1706+
}
1707+
assert_eq!(
1708+
results[0].labels.labels,
1709+
vec![METRIC.to_string(), "10.0.0.15".to_string()],
1710+
);
1711+
assert_eq!(results[0].value, 150.0);
1712+
for element in &results {
1713+
assert_eq!(
1714+
element.labels.labels[0], METRIC,
1715+
"binary-expr path must still prepend the metric name (PromQL top-k formatting)",
1716+
);
1717+
}
1718+
}
14771719
}

‎asap-query-engine/src/tests/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ pub mod clickhouse_forwarding_tests;
33
pub mod datafusion;
44
pub mod elastic_dsl_query_tests;
55
pub mod elastic_forwarding_tests;
6+
pub mod native_binary_instant_tests;
67
pub mod native_pipeline_merge_tests;
78
pub mod prometheus_forwarding_tests;
89
pub mod query_equivalence_tests;

0 commit comments

Comments
 (0)