diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index a9a2d54..bdf615e 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -9,7 +9,7 @@ use super::{ RangeQueryParams, }; use crate::data_model::{AggregationIdInfo, KeyByLabelValues, QueryConfig, SchemaConfig}; -use crate::engines::query_result::{QueryResult, RangeVectorElement}; +use crate::engines::query_result::{InstantVectorElement, QueryResult, RangeVectorElement}; use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; @@ -36,6 +36,67 @@ fn detect_scalar_arm<'a>( } } +/// Native vector-vector combiner for one binary-expr level (instant query): +/// joins two arms' results by label key and applies `op` per matching key, +/// dropping non-matches (inner-join semantics, matching DataFusion's +/// `build_binary_vector_plan`). Positional `KeyByLabelValues` equality is +/// safe *once `lhs_labels == rhs_labels` is confirmed*: label names are +/// canonically sorted by `KeyByLabelNames::new()`, so two arms with the same +/// label set always order their values the same way. Returns `None` if the +/// two arms don't share the same label set — mirroring DataFusion's +/// `build_binary_vector_plan`, which fails to resolve a join column that +/// only exists on one side. +fn combine_vector_vector_native( + lhs_results: Vec, + lhs_labels: &[String], + rhs_results: Vec, + rhs_labels: &[String], + op: &promql_parser::parser::token::TokenType, +) -> Option> { + if lhs_labels != rhs_labels { + return None; + } + + let rhs_map: HashMap = rhs_results + .into_iter() + .map(|elem| (elem.labels, elem.value)) + .collect(); + + Some( + lhs_results + .into_iter() + .filter_map(|lhs_elem| { + rhs_map.get(&lhs_elem.labels).map(|&rhs_val| { + let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val); + InstantVectorElement::new(lhs_elem.labels.clone(), value) + }) + }) + .collect(), + ) +} + +/// Native scalar combiner for one binary-expr level (instant query): +/// applies `op(scalar, value)` or `op(value, scalar)` per `scalar_on_left` +/// to every element of the vector arm's results. +fn combine_scalar_native( + vector_results: Vec, + scalar: f64, + op: &promql_parser::parser::token::TokenType, + scalar_on_left: bool, +) -> Vec { + vector_results + .into_iter() + .map(|elem| { + let value = if scalar_on_left { + SimpleEngine::apply_range_binary_op(op, scalar, elem.value) + } else { + SimpleEngine::apply_range_binary_op(op, elem.value, scalar) + }; + InstantVectorElement::new(elem.labels, value) + }) + .collect() +} + impl SimpleEngine { /// Aligns `end_timestamp` down to the nearest data-ingestion-interval /// boundary, unconditionally — mirroring SQL's `align_end_timestamp_sql`. @@ -453,6 +514,127 @@ impl SimpleEngine { Some((output_labels, QueryResult::vector(results, query_time))) } + /// Recursively evaluates one arm of a binary arithmetic expression via + /// the native pipeline (`execute_query_pipeline`), instead of building a + /// DataFusion plan. Mirrors `build_arm_logical_plan`'s shape exactly, + /// including its limitation: nested `Binary` arms are combined as + /// vector-vector only — a scalar inside a nested arm (e.g. `(a+5)*b`) + /// is not supported, matching today's DataFusion path. + /// + /// - Leaf arm: resolved via `resolve_arm_leaf_context`, executed through + /// `execute_query_pipeline`. + /// - Binary arm: recursively evaluate both sides then combine with + /// `combine_vector_vector_native`. + /// - Scalar literal: returns `None` (handled by the caller separately). + fn evaluate_arm_native( + &self, + arm_ast: &promql_parser::parser::Expr, + time: f64, + ) -> Option<(Vec, Vec)> { + use promql_parser::parser::Expr; + + match arm_ast { + Expr::NumberLiteral(_) => None, // caller handles scalars + Expr::Paren(paren) => self.evaluate_arm_native(&paren.expr, time), + Expr::Binary(binary) => { + // Nested binary expression — recurse on both sides + let (lhs_results, lhs_labels) = self.evaluate_arm_native(&binary.lhs, time)?; + let (rhs_results, rhs_labels) = self.evaluate_arm_native(&binary.rhs, time)?; + let combined = combine_vector_vector_native( + lhs_results, + &lhs_labels, + rhs_results, + &rhs_labels, + &binary.op, + )?; + Some((combined, lhs_labels)) + } + _ => { + let (ctx, label_names) = self.resolve_arm_leaf_context(arm_ast, time)?; + // Unlike DataFusion's PrecomputedSummaryReadExec (which streams + // whatever rows exist, including zero, so a currently-empty arm + // still returns Some(empty vector)), execute_query_pipeline errors + // when the store has no precomputed outputs at all for this arm — + // that propagates to None here, triggering a full Prometheus + // fallback for the whole expression instead of an empty result + // for just this arm. Accepted behavior change (#567); warn loudly + // so it's visible rather than silent. + let results = self + // (true, true): safe unconditionally — both flags are + // self-gated on statistic == Topk / a "k" kwarg being + // present, same as the main instant-query path (see + // execute_query_pipeline's doc comment). + .execute_query_pipeline(&ctx, true, true) + .map_err(|e| { + warn!( + "Native binary-expr arm for metric '{}' produced no results \ + ({}) — unlike DataFusion, this falls back to Prometheus for \ + the whole expression rather than returning an empty result \ + for just this arm", + ctx.metric, e + ); + e + }) + .ok()?; + Some((results, label_names)) + } + } + } + + /// Native (non-DataFusion) implementation of binary-expr handling, + /// staged ahead of #567's cutover so it can be equivalence-tested + /// against the DataFusion path (`handle_binary_expr_promql`) before + /// that path is rewired to call this one. Not yet reachable from + /// `handle_query_promql` — see `handle_query_promql_native` for a + /// test-facing entrypoint. + pub fn handle_binary_expr_promql_native( + &self, + ast: &promql_parser::parser::Expr, + time: f64, + ) -> Option<(KeyByLabelNames, QueryResult)> { + use promql_parser::parser::Expr; + + let query_time = Self::convert_query_time_to_data_time(time); + + let binary = match ast { + Expr::Binary(b) => b, + _ => return None, + }; + + let lhs = binary.lhs.as_ref(); + let rhs = binary.rhs.as_ref(); + let op = &binary.op; + + if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) { + let (vector_results, label_names) = self.evaluate_arm_native(vector_arm, time)?; + let combined = combine_scalar_native(vector_results, scalar, op, scalar_on_left); + return Some(( + KeyByLabelNames::new(label_names), + QueryResult::vector(combined, query_time), + )); + } + + // Vector–vector + let (lhs_results, lhs_labels) = self.evaluate_arm_native(lhs, time)?; + let (rhs_results, rhs_labels) = self.evaluate_arm_native(rhs, time)?; + let combined = + combine_vector_vector_native(lhs_results, &lhs_labels, rhs_results, &rhs_labels, op)?; + let output_labels = KeyByLabelNames::new(lhs_labels); + Some((output_labels, QueryResult::vector(combined, query_time))) + } + + /// Parses `query` and dispatches to `handle_binary_expr_promql_native`. + /// Test-facing convenience wrapper — mirrors the binary-expr branch of + /// `handle_query_promql`, but for the native path. + pub fn handle_query_promql_native( + &self, + query: String, + time: f64, + ) -> Option<(KeyByLabelNames, QueryResult)> { + let ast = promql_parser::parser::parse(&query).ok()?; + self.handle_binary_expr_promql_native(&ast, time) + } + /// Applies a PromQL binary arithmetic operator to two f64 values. fn apply_range_binary_op( op: &promql_parser::parser::token::TokenType, @@ -1474,4 +1656,64 @@ mod topk_pipeline_tests { assert!(pair[0] >= pair[1]); } } + + /// A topk leaf wrapped in a binary expr (`topk(10, ...) + 0`) must still + /// get the same top-10 truncation and metric-name-prefixed formatting as + /// the bare `topk(10, ...)` query — evaluate_arm_native's leaf branch + /// used to hardcode (false, false) for enable_topk_limiting/formatting, + /// which would have returned all 15 unformatted (single-label) rows here + /// instead of the top 10 with the metric-name prefix. + #[test] + fn topk_wrapped_in_binary_expr_still_truncates_and_formats() { + let (engine, store) = build_topk_engine(); + + let context = engine + .build_query_execution_context_promql(TOPK_QUERY.to_string(), QUERY_TIME) + .expect("context should build"); + let window = &context.store_plan.values_query; + + let mut sketch = CountMinSketchWithHeapAccumulator::new(3, 1024, 32); + for i in 1..=15u64 { + let srcip = format!("10.0.0.{i}"); + sketch.inner.update(&srcip, (i * 10) as f64); + } + + let output = + PrecomputedOutput::new(window.start_timestamp, window.end_timestamp, None, AGG_ID); + store + .insert_precomputed_output(output, Box::new(sketch)) + .expect("insert should succeed"); + + let (_, query_result) = engine + .handle_query_promql_native(format!("{TOPK_QUERY} + 0"), QUERY_TIME) + .expect("binary-expr-wrapped topk should still resolve via the native path"); + + let results = match query_result { + QueryResult::Vector(iv) => iv.values, + other => panic!("expected a vector result, got {other:?}"), + }; + + assert_eq!( + results.len(), + 10, + "topk(10, ...) + 0 must still truncate to 10 rows" + ); + for pair in results.windows(2) { + assert!( + pair[0].value >= pair[1].value, + "results must stay sorted by count descending" + ); + } + assert_eq!( + results[0].labels.labels, + vec![METRIC.to_string(), "10.0.0.15".to_string()], + ); + assert_eq!(results[0].value, 150.0); + for element in &results { + assert_eq!( + element.labels.labels[0], METRIC, + "binary-expr path must still prepend the metric name (PromQL top-k formatting)", + ); + } + } } diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 563f854..d5b31ee 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -3,6 +3,7 @@ pub mod clickhouse_forwarding_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; +pub mod native_binary_instant_tests; pub mod native_pipeline_merge_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; diff --git a/asap-query-engine/src/tests/native_binary_instant_tests.rs b/asap-query-engine/src/tests/native_binary_instant_tests.rs new file mode 100644 index 0000000..1599d8c --- /dev/null +++ b/asap-query-engine/src/tests/native_binary_instant_tests.rs @@ -0,0 +1,472 @@ +//! Native instant binary-expr evaluator tests (issue #567, Stage 2). +//! +//! `handle_binary_expr_promql_native`/`evaluate_arm_native` are a new, +//! parallel implementation of PromQL binary-arithmetic instant queries, +//! built to eventually replace the DataFusion-backed +//! `handle_binary_expr_promql` (#567's Stage 3 cutover). These tests compare +//! the two paths directly — both are reachable today (`handle_query_promql` +//! for DataFusion, `handle_query_promql_native` for the new path) so +//! equivalence can be proven before anything is rewired. + +#[cfg(test)] +mod tests { + use crate::data_model::{AggregationType, KeyByLabelValues}; + use crate::engines::query_result::QueryResult; + use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::{CountMinSketchAccumulator, DeltaSetAggregatorAccumulator}; + use crate::tests::test_utilities::engine_factories::{ + create_engine_dual_input, create_engine_single_pop, create_engine_three_metrics, + create_engine_two_metrics, + }; + use crate::AggregateCore; + + const QUERY_TIME: f64 = 1000.0; + + fn vector_values(qr: QueryResult) -> Vec<(Vec, f64)> { + match qr { + QueryResult::Vector(iv) => iv + .values + .into_iter() + .map(|e| (e.labels.labels, e.value)) + .collect(), + _ => panic!("Expected vector result"), + } + } + + fn sorted(mut v: Vec<(Vec, f64)>) -> Vec<(Vec, f64)> { + v.sort_by(|a, b| a.0.cmp(&b.0)); + v + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_vector_vector_all_ops_match_datafusion() { + // `^` excluded: DataFusion's build_binary_vector_plan maps T_POW to + // Operator::BitwiseXor "as a proxy" (plan_builder.rs comment) and + // fails to even produce a result for it today — nothing to compare + // native against for that operator. See the dedicated `^` test below. + for (op, expected) in [ + ("+", 30.0), + ("-", -10.0), + ("*", 200.0), + ("/", 0.5), + ("%", 10.0), + ] { + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = format!("sum(metric_a) by (host) {op} sum(metric_b) by (host)"); + let (_, old_qr) = engine + .handle_query_promql(query.clone(), QUERY_TIME) + .unwrap_or_else(|| panic!("old path failed for op {op}")); + let (_, new_qr) = engine + .handle_query_promql_native(query.clone(), QUERY_TIME) + .unwrap_or_else(|| panic!("new path failed for op {op}")); + + let old_values = sorted(vector_values(old_qr)); + let new_values = sorted(vector_values(new_qr)); + assert_eq!(old_values.len(), 1, "op {op}"); + assert_eq!( + new_values, old_values, + "op {op}: native must match DataFusion" + ); + assert!( + (new_values[0].1 - expected).abs() < 1e-10, + "op {op}: expected {expected}, got {}", + new_values[0].1 + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_power_operator_computes_correctly_datafusion_unsupported() { + // DataFusion's `^` support is broken today (see comment above); this + // documents that native computes it correctly via f64::powf, + // independent of a DataFusion comparison. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(2.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = "sum(metric_a) by (host) ^ sum(metric_b) by (host)"; + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + let new_values = vector_values(new_qr); + assert!((new_values[0].1 - 1024.0).abs() < 1e-6, "2^10 = 1024"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_vector_scalar_both_orderings_match_datafusion() { + let engine = create_engine_single_pop( + "errors_total", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(7.0)) as Box, + )], + "sum(errors_total) by (host)", + ); + + for query in [ + "sum(errors_total) by (host) * 100", + "100 * sum(errors_total) by (host)", + ] { + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .unwrap_or_else(|| panic!("old path failed for {query}")); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .unwrap_or_else(|| panic!("new path failed for {query}")); + + let old_values = vector_values(old_qr); + let new_values = vector_values(new_qr); + assert_eq!(new_values, old_values, "{query}"); + assert!((new_values[0].1 - 700.0).abs() < 1e-10, "{query}"); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_nested_binary_matches_datafusion() { + // (metric_a + metric_b) * metric_c, all Sum, host-a: (10+20)*3 = 90 + let engine = create_engine_three_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + "metric_c", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(3.0)) as Box, + )], + "sum(metric_c) by (host)", + ); + + let query = "(sum(metric_a) by (host) + sum(metric_b) by (host)) * sum(metric_c) by (host)"; + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("old path failed"); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + + let old_values = vector_values(old_qr); + let new_values = vector_values(new_qr); + assert_eq!(new_values, old_values); + assert!((new_values[0].1 - 90.0).abs() < 1e-10, "(10+20)*3 = 90"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_binary_expr_no_data_falls_back_to_none() { + // metric_a is configured (schema + pattern match) but has zero + // precomputed data. Accepted behavior change (#567): unlike + // DataFusion (which would return an empty-but-present result), + // native falls back to Prometheus (returns None) for the whole + // expression — see evaluate_arm_native's leaf branch. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![], // no data at all for metric_a + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + assert!( + old_result.is_some(), + "DataFusion returns an empty-but-present result, not None, for a currently-empty arm" + ); + + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + assert!( + new_result.is_none(), + "native falls back to Prometheus (None) for a currently-empty arm" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_binary_expr_unsupported_arm_returns_none() { + // foo() is not a supported PromQL function -> arm lookup fails -> None, + // same as today's DataFusion path (dispatch_arithmetic_tests.rs). + let engine = create_engine_single_pop( + "requests_total", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(200.0)) as Box, + )], + "sum(requests_total) by (host)", + ); + + let query = "foo(errors_total[5m]) / sum(requests_total) by (host)"; + assert!(engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .is_none()); + assert!(engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_vector_vector_dual_population_matches_datafusion() { + // event_frequency is dual-population (CountMinSketch values + + // DeltaSetAggregator keys) -- confirms the leaf swap + // (ctx.to_logical_plan() -> execute_query_pipeline) still resolves + // keys_query correctly. Wrapped in `+ 0` to route through the + // binary-expr handler at all. + let cms = CountMinSketchAccumulator::new(2, 3); + let mut keys = DeltaSetAggregatorAccumulator::new(); + keys.add_key(KeyByLabelValues { + labels: vec!["host-a".to_string(), "evt-1".to_string()], + }); + + let engine = create_engine_dual_input( + "event_frequency", + AggregationType::CountMinSketch, + AggregationType::DeltaSetAggregator, + vec![], + vec!["host", "event"], + vec![(None, Box::new(cms))], + vec![(None, Box::new(keys))], + "count(event_frequency) by (host, event)", + ); + + let query = "count(event_frequency) by (host, event) + 0"; + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("old path failed"); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + + assert_eq!(sorted(vector_values(new_qr)), sorted(vector_values(old_qr))); + } + + // --- Regression tests: combine_vector_vector_native joins purely on + // positional KeyByLabelValues equality (raw Vec of values, no + // label names attached at all -- see key_by_label_values.rs). DataFusion's + // build_binary_vector_plan joins on named columns (lhs.{label} = + // rhs.{label}) instead, so it fails to resolve (-> None) whenever the two + // arms don't share the same label set. Native has no equivalent check, so + // it can silently join two arms grouped by *different* labels whenever + // their values happen to coincide. These tests are expected to be RED + // against the current implementation -- they exist to prove the + // divergence before combine_vector_vector_native is fixed. + + #[tokio::test(flavor = "multi_thread")] + async fn native_mismatched_label_names_with_colliding_values_diverges_from_datafusion() { + // metric_a grouped by (host), metric_b grouped by (region) -- disjoint + // label sets -- but both happen to produce the value "us-east". + // DataFusion can't join lhs.host against a plan with no "host" column + // -> None. Native joins on the bare value vector ["us-east"] == + // ["us-east"] -> a spurious combined result. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["us-east".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["region"], + vec![( + Some(vec!["us-east".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (region)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + + match old_result { + None => assert!( + new_result.is_none(), + "BUG: DataFusion refuses to join arms with different label sets (None), \ + but native silently produced {new_result:?} by matching on value alone" + ), + Some((_, old_qr)) => assert_eq!( + new_result.map(|(_, qr)| sorted(vector_values(qr))), + Some(sorted(vector_values(old_qr))), + "native and DataFusion must agree when both produce a result" + ), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_mismatched_label_names_multi_label_full_collision_diverges_from_datafusion() { + // Same bug with a 2-label grouping: (host, dc) vs (region, zone), but + // the *entire* ordered value vector coincides ("us-east", "az1" on + // both sides) -- confirms the collision isn't a single-label fluke. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host", "dc"], + vec![( + Some(vec!["us-east".to_string(), "az1".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host, dc)", + "metric_b", + AggregationType::Sum, + vec!["region", "zone"], + vec![( + Some(vec!["us-east".to_string(), "az1".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (region, zone)", + ); + + let query = "sum(metric_a) by (host, dc) + sum(metric_b) by (region, zone)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + + match old_result { + None => assert!( + new_result.is_none(), + "BUG: DataFusion refuses to join arms with different label sets (None), \ + but native silently produced {new_result:?} by matching on value alone" + ), + Some((_, old_qr)) => assert_eq!( + new_result.map(|(_, qr)| sorted(vector_values(qr))), + Some(sorted(vector_values(old_qr))), + "native and DataFusion must agree when both produce a result" + ), + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_mismatched_label_names_non_colliding_values_still_diverges_from_datafusion() { + // Same disjoint label sets (host vs region), but this time the values + // don't collide ("us-east" vs "eu-west"). DataFusion still can't + // resolve the join (column doesn't exist, independent of the data) + // -> None. Native's join correctly finds no match on the *value*, but + // that's incidental -- it still returns Some(empty) instead of None, + // because it never checked whether the label sets matched at all. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["us-east".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["region"], + vec![( + Some(vec!["eu-west".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (region)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; + let old_result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + let new_result = engine.handle_query_promql_native(query.to_string(), QUERY_TIME); + + assert_eq!( + old_result.is_none(), + new_result.is_none(), + "BUG: DataFusion returns {old_result:?} (no such column to join on), native \ + returns {new_result:?} -- native never checks that the two arms' label sets match" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn native_same_label_set_no_matching_values_both_agree_empty() { + // Control case: identical label sets (host on both sides) but + // disjoint values -- both paths should agree there's no match. This + // isolates the bug to *differing* label sets, not "no match" in + // general. + let engine = create_engine_two_metrics( + "metric_a", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + )], + "sum(metric_a) by (host)", + "metric_b", + AggregationType::Sum, + vec!["host"], + vec![( + Some(vec!["host-b".to_string()]), + Box::new(SumAccumulator::with_sum(20.0)) as Box, + )], + "sum(metric_b) by (host)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (host)"; + let (_, old_qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("old path failed"); + let (_, new_qr) = engine + .handle_query_promql_native(query.to_string(), QUERY_TIME) + .expect("new path failed"); + + let old_values = vector_values(old_qr); + assert_eq!(old_values, Vec::new()); + assert_eq!(sorted(vector_values(new_qr)), sorted(old_values)); + } +}