diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 501f7218..f4e4c514 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -549,22 +549,29 @@ impl SimpleEngine { }; let merged_values = if plan.values_query.is_exact_query { - // Sliding window: no merge needed, extract buckets from timestamped data - debug!("Sliding window mode: Skipping merge (expecting 1 precompute per key)"); - values_map - .into_iter() - .map(|(key, timestamped_buckets)| { - if timestamped_buckets.len() != 1 { - warn!( - "Sliding window expected 1 precompute per key, found {}. Using first.", - timestamped_buckets.len() - ); - } - // Extract bucket from timestamped tuple - let (_, bucket) = timestamped_buckets.into_iter().next().unwrap(); - (key, bucket.as_ref().clone_boxed_core()) - }) - .collect() + // Sliding window: expected exactly 1 precompute per key today + // (ponytail: hardcoded, #554 will make >1 legitimate — don't + // block on it). The store can legitimately return more than + // expected for one exact window; merge whatever came back + // instead of arbitrarily keeping the first and dropping the + // rest (see #567). + const EXPECTED_BUCKETS_PER_KEY: usize = 1; + debug!("Sliding window mode: merging {} keys", values_map.len()); + for timestamped_buckets in values_map.values() { + if timestamped_buckets.is_empty() { + continue; + } + if timestamped_buckets.len() != EXPECTED_BUCKETS_PER_KEY { + warn!( + "Sliding window expected {} precompute(s) per key, found {}. Merging all.", + EXPECTED_BUCKETS_PER_KEY, + timestamped_buckets.len() + ); + } + } + // Sliding windows always merge (all buckets belong to one + // logical window) — reuse the same merge path as Tumbling. + self.merge_precomputed_outputs(&values_map, true, agg_info.aggregation_type_for_value) } else { // Tumbling window: merge needed debug!("Tumbling window mode: Merging {} outputs", values_map.len()); @@ -576,13 +583,12 @@ impl SimpleEngine { }; let merge_duration = merge_start_time.elapsed(); + let did_merge = window_type == WindowType::Sliding + || do_merge + || agg_info.aggregation_type_for_value == AggregationType::DeltaSetAggregator; debug!( "[LATENCY] Precomputed output processing ({}): {:.2}ms, resulted in {} merged outputs", - if window_type == WindowType::Sliding { - "no merge" - } else { - "merge" - }, + if did_merge { "merge" } else { "no merge" }, merge_duration.as_secs_f64() * 1000.0, merged_values.len() ); @@ -817,8 +823,13 @@ impl SimpleEngine { /// Executes a pre-built DataFusion logical plan and returns results. /// - /// This is the shared execution kernel used by both `execute_plan` (for single-metric - /// queries) and the binary arithmetic dispatch path. + /// This was the entry point for the DataFusion-based binary arithmetic + /// dispatch path, cut over to a native implementation in #567. Unlike its + /// sibling `execute_plan` (still called by DataFusion-path tests), this + /// function has zero callers anywhere in the repo, including tests — it + /// is genuinely dead code, kept only in case the native cutover needs to + /// be reverted. + #[allow(dead_code)] pub async fn execute_logical_plan( &self, logical_plan: datafusion::logical_expr::LogicalPlan, @@ -1068,7 +1079,12 @@ impl SimpleEngine { let mut merged = HashMap::with_capacity(precomputed_outputs_map.len()); for (key, timestamped_buckets) in precomputed_outputs_map.iter() { - if !timestamped_buckets.is_empty() { + if timestamped_buckets.is_empty() { + warn!( + "Store returned key {:?} with no precompute buckets; skipping", + key + ); + } else { // Extract just the buckets (without timestamps) for merging let precomputes: Vec> = timestamped_buckets .iter() @@ -1080,7 +1096,7 @@ impl SimpleEngine { debug!(" Merging accumulators (should_merge=true)"); #[cfg(feature = "extra_debugging")] let merge_start = Instant::now(); - match self.merge_accumulators(&precomputes) { + match self.merge_accumulators(precomputes) { Ok(merged_accumulator) => { #[cfg(feature = "extra_debugging")] let merge_duration = merge_start.elapsed(); @@ -1123,21 +1139,21 @@ impl SimpleEngine { /// This follows the Python merge_accumulators approach fn merge_accumulators( &self, - accumulators: &[Box], + accumulators: Vec>, ) -> Result, AccumulatorError> { if accumulators.is_empty() { return Err(AccumulatorError::EmptySlice); } if accumulators.len() == 1 { - return Ok(accumulators[0].clone_boxed_core()); + return Ok(accumulators.into_iter().next().unwrap()); } // Try to use optimized batch merge for KLL accumulators if accumulators[0].get_accumulator_type() == AggregationType::DatasketchesKLL { use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; - match DatasketchesKLLAccumulator::merge_multiple(accumulators) { + match DatasketchesKLLAccumulator::merge_multiple(&accumulators) { Ok(merged) => return Ok(Box::new(merged)), Err(e) => { warn!( @@ -1153,7 +1169,7 @@ impl SimpleEngine { if accumulators[0].get_accumulator_type() == AggregationType::CountMinSketch { use crate::precompute_operators::count_min_sketch_accumulator::CountMinSketchAccumulator; - match CountMinSketchAccumulator::merge_multiple(accumulators) { + match CountMinSketchAccumulator::merge_multiple(&accumulators) { Ok(merged) => return Ok(Box::new(merged)), Err(e) => { warn!( diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index a9a2d546..8b37f843 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>( } } +/// 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( + 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(move |&rhs_val| { + let value = SimpleEngine::apply_range_binary_op(op, lhs_elem.value, rhs_val); + InstantVectorElement::new(lhs_elem.labels, value) + }) + }) + .collect(), + ) +} + +/// 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( + 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`. @@ -334,7 +395,7 @@ impl SimpleEngine { /// Recursively unwraps `Paren`, then structurally resolves a leaf PromQL /// arm (i.e. not `Binary` or `NumberLiteral`) to its `QueryConfig` and /// base `QueryExecutionContext`. Shared leaf-resolution step for both - /// `build_arm_logical_plan` (instant) and `build_arm_range_context` (range). + /// `evaluate_binary_arm` (instant) and `build_arm_range_context` (range). /// /// Returns `None` for `Binary` arms (caller handles recursion) and /// `NumberLiteral` arms (caller handles scalars). @@ -357,44 +418,72 @@ impl SimpleEngine { } } - /// Recursively builds a DataFusion logical plan for one arm of a binary - /// arithmetic expression. + /// Recursively evaluates one arm of a binary arithmetic expression via + /// the native pipeline (`execute_query_pipeline`). /// /// - Leaf arm (supported PromQL pattern): resolved via `resolve_arm_leaf_context`, - /// returning its `to_logical_plan()` together with the output label names. - /// - Binary arm: recursively build both sub-arms and combine with - /// `build_binary_vector_plan`. + /// executed through `execute_query_pipeline`. + /// - Binary arm: recursively evaluate both sub-arms and combine with + /// `combine_vector_vector`. Nested `Binary` arms are combined as + /// vector-vector only — a scalar inside a nested arm (e.g. `(a+5)*b`) + /// is not supported (tracked separately, same as before this cutover). /// - Scalar literal: returns `None` (handled by the caller separately). - fn build_arm_logical_plan( + fn evaluate_binary_arm( &self, arm_ast: &promql_parser::parser::Expr, time: f64, - ) -> Option<(datafusion::logical_expr::LogicalPlan, Vec)> { - use crate::engines::logical::plan_builder::build_binary_vector_plan; + ) -> Option<(Vec, Vec)> { use promql_parser::parser::Expr; match arm_ast { Expr::NumberLiteral(_) => None, // caller handles scalars - Expr::Paren(paren) => self.build_arm_logical_plan(&paren.expr, time), + Expr::Paren(paren) => self.evaluate_binary_arm(&paren.expr, time), Expr::Binary(binary) => { // Nested binary expression — recurse on both sides - let (lhs_plan, lhs_labels) = self.build_arm_logical_plan(&binary.lhs, time)?; - let (rhs_plan, _) = self.build_arm_logical_plan(&binary.rhs, time)?; - let combined = - build_binary_vector_plan(lhs_plan, rhs_plan, &binary.op, lhs_labels.clone()) - .ok()?; + let (lhs_results, lhs_labels) = self.evaluate_binary_arm(&binary.lhs, time)?; + let (rhs_results, rhs_labels) = self.evaluate_binary_arm(&binary.rhs, time)?; + let combined = combine_vector_vector( + 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)?; - let plan = ctx.to_logical_plan().ok()?; - Some((plan, label_names)) + // Unlike DataFusion's PrecomputedSummaryReadExec (which streamed + // whatever rows existed, including zero, so a currently-empty arm + // used to return 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!( + "Binary-expr arm for metric '{}' failed ({}) — \ + 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)) } } } - /// Handles a binary arithmetic PromQL expression by building a combined - /// DataFusion plan (vector–vector join or scalar projection) and executing it. + /// Handles a binary arithmetic PromQL expression via the native pipeline + /// (vector–vector join or scalar combine). /// /// Returns `None` if any arm is not acceleratable (caller falls back to Prometheus). fn handle_binary_expr_promql( @@ -402,7 +491,6 @@ impl SimpleEngine { ast: &promql_parser::parser::Expr, time: f64, ) -> Option<(KeyByLabelNames, QueryResult)> { - use crate::engines::logical::plan_builder::{build_binary_vector_plan, build_scalar_plan}; use promql_parser::parser::Expr; let query_time = Self::convert_query_time_to_data_time(time); @@ -417,40 +505,21 @@ impl SimpleEngine { let op = &binary.op; if let Some((scalar, vector_arm, scalar_on_left)) = detect_scalar_arm(lhs, rhs) { - let (vector_plan, label_names) = self.build_arm_logical_plan(vector_arm, time)?; - let combined = - build_scalar_plan(vector_plan, scalar, op, scalar_on_left, label_names.clone()) - .ok()?; - let results = tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.execute_logical_plan( - combined, - label_names.clone(), - "", - &Statistic::Sum, - )) - }) - .ok()?; + let (vector_results, label_names) = self.evaluate_binary_arm(vector_arm, time)?; + let combined = combine_scalar(vector_results, scalar, op, scalar_on_left); return Some(( KeyByLabelNames::new(label_names), - QueryResult::vector(results, query_time), + QueryResult::vector(combined, query_time), )); } // Vector–vector - let (lhs_plan, lhs_labels) = self.build_arm_logical_plan(lhs, time)?; - let (rhs_plan, _) = self.build_arm_logical_plan(rhs, time)?; - let combined = build_binary_vector_plan(lhs_plan, rhs_plan, op, lhs_labels.clone()).ok()?; - let results = tokio::task::block_in_place(|| { - tokio::runtime::Handle::current().block_on(self.execute_logical_plan( - combined, - lhs_labels.clone(), - "", - &Statistic::Sum, - )) - }) - .ok()?; + let (lhs_results, lhs_labels) = self.evaluate_binary_arm(lhs, time)?; + let (rhs_results, rhs_labels) = self.evaluate_binary_arm(rhs, time)?; + let combined = + combine_vector_vector(lhs_results, &lhs_labels, rhs_results, &rhs_labels, op)?; let output_labels = KeyByLabelNames::new(lhs_labels); - Some((output_labels, QueryResult::vector(results, query_time))) + Some((output_labels, QueryResult::vector(combined, query_time))) } /// Applies a PromQL binary arithmetic operator to two f64 values. @@ -475,7 +544,7 @@ impl SimpleEngine { /// arithmetic expression. /// /// Leaf resolution (Paren-unwrap + structural config lookup) is shared - /// with `build_arm_logical_plan` via `resolve_arm_leaf_context`. Note this + /// with `evaluate_binary_arm` via `resolve_arm_leaf_context`. Note this /// does not support nested `Binary` arms (e.g. `(a+b)*c` over a range) — /// tracked separately in #516. fn build_arm_range_context( @@ -582,9 +651,17 @@ impl SimpleEngine { return Some((KeyByLabelNames::new(labels), QueryResult::matrix(combined))); } - // Vector-vector: evaluate both arms, join by label key, apply op per matching timestamp + // Vector-vector: evaluate both arms, join by label key, apply op per matching timestamp. + // Reject mismatched label sets up front — same guard as the instant-query + // combine_vector_vector, and for the same reason: positional + // KeyByLabelValues equality below is only safe once the label *names* + // match (they're canonically sorted by KeyByLabelNames::new(), so two + // arms with the same label set always order their values the same way). let (lhs_ctx, lhs_labels) = self.build_arm_range_context(lhs, start, end, step)?; - let (rhs_ctx, _) = self.build_arm_range_context(rhs, start, end, step)?; + let (rhs_ctx, rhs_labels) = self.build_arm_range_context(rhs, start, end, step)?; + if lhs_labels != rhs_labels { + return None; + } let lhs_results = self.execute_range_query_pipeline(&lhs_ctx).ok()?; let rhs_results = self.execute_range_query_pipeline(&rhs_ctx).ok()?; @@ -1474,4 +1551,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(format!("{TOPK_QUERY} + 0"), QUERY_TIME) + .expect("binary-expr-wrapped topk should still resolve"); + + 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/datafusion/mod.rs b/asap-query-engine/src/tests/datafusion/mod.rs index 143cc0a4..fd946643 100644 --- a/asap-query-engine/src/tests/datafusion/mod.rs +++ b/asap-query-engine/src/tests/datafusion/mod.rs @@ -7,7 +7,6 @@ pub mod accumulator_serde_tests; pub mod dispatch_arithmetic_tests; pub mod plan_builder_binary_tests; pub mod plan_builder_regression_tests; -pub mod plan_execution_arithmetic_tests; pub mod plan_execution_dual_input_tests; pub mod plan_execution_temporal_tests; pub mod plan_execution_tests; diff --git a/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs b/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs index a2fca523..cdf67201 100644 --- a/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/datafusion/range_query_arithmetic_tests.rs @@ -36,18 +36,22 @@ mod tests { /// `engine_factories::create_engine_two_metrics` (single timestamp, instant /// queries only), this inserts one bucket per `(timestamp, value)` pair so /// range queries have more than one output sample to join across. + #[allow(clippy::too_many_arguments)] fn create_range_engine_two_metrics( metric_a: &str, + labels_a: Vec<&str>, data_a: TimeSeriesData, query_a: &str, metric_b: &str, + labels_b: Vec<&str>, data_b: TimeSeriesData, query_b: &str, ) -> SimpleEngine { - let labels = vec!["host".to_string()]; + let labels_a: Vec = labels_a.iter().map(|s| s.to_string()).collect(); + let labels_b: Vec = labels_b.iter().map(|s| s.to_string()).collect(); let mut aggregation_configs = HashMap::new(); - for (id, metric) in [(1u64, metric_a), (2u64, metric_b)] { + for (id, metric, labels) in [(1u64, metric_a, &labels_a), (2u64, metric_b, &labels_b)] { aggregation_configs.insert( id, AggregationConfig { @@ -91,8 +95,8 @@ mod tests { } let promql_schema = PromQLSchema::new() - .add_metric(metric_a.to_string(), KeyByLabelNames::new(labels.clone())) - .add_metric(metric_b.to_string(), KeyByLabelNames::new(labels)); + .add_metric(metric_a.to_string(), KeyByLabelNames::new(labels_a)) + .add_metric(metric_b.to_string(), KeyByLabelNames::new(labels_b)); let inference_config = InferenceConfig { schema: SchemaConfig::PromQL(promql_schema), @@ -143,9 +147,11 @@ mod tests { let data_requests = host_a_series([(1000, 200.0), (2000, 300.0)]); let engine = create_range_engine_two_metrics( "errors_total", + vec!["host"], data_errors, "sum(errors_total) by (host)", "requests_total", + vec!["host"], data_requests, "sum(requests_total) by (host)", ); @@ -174,9 +180,11 @@ mod tests { let data_b = host_a_series([(1000, 20.0), (2000, 25.0)]); let engine = create_range_engine_two_metrics( "metric_a", + vec!["host"], data_a, "sum(metric_a) by (host)", "metric_b", + vec!["host"], data_b, "sum(metric_b) by (host)", ); @@ -199,10 +207,12 @@ mod tests { let data_a = host_a_series([(1000, 5.0), (2000, 6.0)]); let engine = create_range_engine_two_metrics( "metric_a", + vec!["host"], data_a, "sum(metric_a) by (host)", // second metric not used but the helper requires it; empty data. "dummy", + vec!["host"], vec![], "sum(dummy) by (host)", ); @@ -225,9 +235,11 @@ mod tests { let data_a = host_a_series([(1000, 0.9), (2000, 0.75)]); let engine = create_range_engine_two_metrics( "metric_a", + vec!["host"], data_a, "sum(metric_a) by (host)", "dummy", + vec!["host"], vec![], "sum(dummy) by (host)", ); @@ -243,4 +255,35 @@ mod tests { assert!((by_ts[&1000] - 0.1).abs() < 1e-10); assert!((by_ts[&2000] - 0.25).abs() < 1e-10); } + + // Regression test: handle_binary_expr_range_promql's vector-vector join + // used to match purely on positional KeyByLabelValues equality (rhs + // labels discarded), unlike the instant-query combine_vector_vector, + // which rejects a join between arms grouped by different label sets. Two + // arms grouped by disjoint labels ((host) vs (region)) that happen to + // produce the same value could silently join into a wrong-but-plausible + // result across the whole range. + #[tokio::test(flavor = "multi_thread")] + async fn test_range_vector_vector_mismatched_label_sets_return_none() { + let data_a = host_a_series([(1000, 10.0), (2000, 15.0)]); + let data_b = host_a_series([(1000, 10.0), (2000, 15.0)]); + let engine = create_range_engine_two_metrics( + "metric_a", + vec!["host"], + data_a, + "sum(metric_a) by (host)", + "metric_b", + vec!["region"], + data_b, + "sum(metric_b) by (region)", + ); + + let query = "sum(metric_a) by (host) + sum(metric_b) by (region)"; + let result = engine.handle_range_query_promql(query.to_string(), 1.0, 2.0, 1.0); + assert!( + result.is_none(), + "BUG: arms grouped by different label sets must not join, even when their \ + values coincide, got {result:?}" + ); + } } diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 9fbe6fc1..6ac84281 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -3,6 +3,9 @@ pub mod clickhouse_forwarding_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests; +pub mod native_binary_arithmetic_plan_tests; +pub mod native_binary_instant_tests; +pub mod native_pipeline_merge_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod sql_pattern_matching_tests; diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs b/asap-query-engine/src/tests/native_binary_arithmetic_plan_tests.rs similarity index 99% rename from asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs rename to asap-query-engine/src/tests/native_binary_arithmetic_plan_tests.rs index 91d1356d..69a119ea 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_arithmetic_tests.rs +++ b/asap-query-engine/src/tests/native_binary_arithmetic_plan_tests.rs @@ -2,7 +2,8 @@ //! //! Verify that binary arithmetic queries (vector/vector and scalar/vector) //! produce numerically correct results when executed end-to-end through -//! `handle_binary_expr_promql` via DataFusion. +//! `handle_binary_expr_promql`, natively as of #567's Stage 3 cutover +//! (previously via DataFusion). #[cfg(test)] mod 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 00000000..4f0bd30c --- /dev/null +++ b/asap-query-engine/src/tests/native_binary_instant_tests.rs @@ -0,0 +1,469 @@ +//! PromQL binary-expr instant query tests, native execution (issue #567). +//! +//! `handle_query_promql`'s binary-arithmetic path (`handle_binary_expr_promql` +//! → `evaluate_binary_arm` → `combine_vector_vector`/`combine_scalar`) runs +//! natively as of #567's Stage 3 cutover — no more DataFusion involved. +//! These tests were originally written to compare the native path against +//! the (now-removed) DataFusion path before the cutover landed; they now +//! assert the native path's results directly. + +#[cfg(test)] +mod tests { + use crate::data_model::{AggregationType, KeyByLabelValues, WindowType}; + 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_multi_timestamp_with_window, + 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 binary_expr_vector_vector_all_ops() { + // `^` excluded: needs its own accumulator setup (see the dedicated + // power-operator 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 (_, qr) = engine + .handle_query_promql(query, QUERY_TIME) + .unwrap_or_else(|| panic!("query failed for op {op}")); + + let values = vector_values(qr); + assert_eq!(values.len(), 1, "op {op}"); + assert!( + (values[0].1 - expected).abs() < 1e-10, + "op {op}: expected {expected}, got {}", + values[0].1 + ); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_power_operator_computes_correctly() { + 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 (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("query failed"); + let values = vector_values(qr); + assert!((values[0].1 - 1024.0).abs() < 1e-6, "2^10 = 1024"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_scalar_both_orderings() { + 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 (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .unwrap_or_else(|| panic!("query failed for {query}")); + let values = vector_values(qr); + assert!((values[0].1 - 700.0).abs() < 1e-10, "{query}"); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_nested_binary() { + // (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 (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("query failed"); + let values = vector_values(qr); + assert!((values[0].1 - 90.0).abs() < 1e-10, "(10+20)*3 = 90"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_no_data_falls_back_to_none() { + // metric_a is configured (schema + pattern match) but has zero + // precomputed data. Accepted behavior change (#567, kept from + // DataFusion's now-removed empty-result behavior): falls back to + // Prometheus (returns None) for the whole expression — see + // evaluate_binary_arm's leaf branch, which warns loudly when this + // happens. + 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 result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + assert!( + result.is_none(), + "arm with no current precomputed data falls back to Prometheus" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_unsupported_arm_returns_none() { + // foo() is not a supported PromQL function -> arm lookup fails -> None + // (graceful fallback to Prometheus). + 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()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_vector_vector_dual_population() { + // event_frequency is dual-population (CountMinSketch values + + // DeltaSetAggregator keys) -- confirms the leaf's keys_query + // resolution works through the binary-expr path. 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 (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("query failed"); + assert!(!sorted(vector_values(qr)).is_empty()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_sliding_window_end_to_end_merges_correctly() { + // Ties Stage 1's sliding-bucket merge fix (#570) to the actual + // production entrypoint this issue changes: 2 buckets for the same + // key under one Sliding exact window must both be merged, not just + // the first, when reached through a real binary-expr query. + let data = vec![ + ( + 1_000_000, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ), + ( + 1_000_000, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(5.0)) as Box, + ), + ]; + let leaf_query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + leaf_query, + 1_000, // window_size_ms, matches the fixed 1000ms bucket width + WindowType::Sliding, + ); + + let query = format!("{leaf_query} + 0"); + let (_, qr) = engine + .handle_query_promql(query, QUERY_TIME) + .expect("query failed"); + let values = vector_values(qr); + assert_eq!(values.len(), 1); + assert!( + (values[0].1 - 15.0).abs() < 1e-10, + "expected both sliding-window buckets merged into 15.0, got {}", + values[0].1 + ); + } + + #[tokio::test] + async fn binary_expr_works_on_current_thread_runtime() { + // Default (single-threaded) tokio runtime, not `flavor = "multi_thread"` + // like every other test in this file. The old DataFusion path's + // tokio::task::block_in_place(...block_on(...)) wrapper panics on a + // current-thread runtime — this test only passes because that + // wrapper is actually gone (#567 Stage 3 cutover), not because of + // any value it computes. + 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)", + ); + + let result = + engine.handle_query_promql("sum(errors_total) by (host) * 2".to_string(), QUERY_TIME); + assert!(result.is_some()); + } + + // --- Regression tests: combine_vector_vector must reject a join between + // arms with different label sets rather than silently matching on + // positional KeyByLabelValues equality alone (raw Vec of values, + // no label names attached -- see key_by_label_values.rs). This mirrors + // DataFusion's build_binary_vector_plan, which fails to resolve a join + // column that only exists on one side. These tests originally compared + // against the DataFusion path to prove the divergence before the fix; + // now that combine_vector_vector checks label-set equality directly and + // DataFusion is no longer in the production path (#567 Stage 3), they + // assert the fixed behavior directly. + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_mismatched_label_sets_with_colliding_values_returns_none() { + // metric_a grouped by (host), metric_b grouped by (region) -- disjoint + // label sets -- but both happen to produce the value "us-east". A + // value-only join would spuriously match ["us-east"] == ["us-east"]; + // must return None instead. + 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 result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + + assert!( + result.is_none(), + "mismatched label sets must return None, not a spurious value-matched join: {result:?}" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_mismatched_label_sets_multi_label_full_collision_returns_none() { + // Same case 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 check 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 result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + + assert!( + result.is_none(), + "mismatched label sets must return None, not a spurious value-matched join: {result:?}" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_mismatched_label_sets_non_colliding_values_returns_none() { + // Same disjoint label sets (host vs region), but this time the values + // don't collide ("us-east" vs "eu-west") either -- must still return + // None because the label sets themselves don't match, not because no + // values happened to match. + 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 result = engine.handle_query_promql(query.to_string(), QUERY_TIME); + + assert!( + result.is_none(), + "mismatched label sets must return None even when no values happen to match: {result:?}" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn binary_expr_same_label_set_no_matching_values_returns_empty() { + // Control case: identical label sets (host on both sides) but + // disjoint values -- should resolve to an empty (not None) result. + // This isolates the label-set check from ordinary "no match" cases. + 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 (_, qr) = engine + .handle_query_promql(query.to_string(), QUERY_TIME) + .expect("query failed"); + + assert_eq!(vector_values(qr), Vec::new()); + } +} diff --git a/asap-query-engine/src/tests/native_pipeline_merge_tests.rs b/asap-query-engine/src/tests/native_pipeline_merge_tests.rs new file mode 100644 index 00000000..044b11a2 --- /dev/null +++ b/asap-query-engine/src/tests/native_pipeline_merge_tests.rs @@ -0,0 +1,174 @@ +//! Native pipeline merge tests (issue #567, Stage 1). +//! +//! `execute_and_merge_store_queries`'s Sliding-window branch +//! (`simple_engine/mod.rs`) must merge every precomputed bucket returned for +//! a key, not just the first one. The store's `query_precomputed_output_exact` +//! can legitimately return more than one bucket for the same key under one +//! exact window (see `per_key.rs::query_precomputed_output_exact`), and +//! DataFusion's `SummaryMergeMultipleExec` already merges all of them +//! correctly — native must match. + +use crate::data_model::{AggregationType, WindowType}; +use crate::engines::query_result::InstantVectorElement; +use crate::engines::simple_engine::SimpleEngine; +use crate::precompute_operators::sum_accumulator::SumAccumulator; +use crate::tests::test_utilities::engine_factories::create_engine_multi_timestamp_with_window; + +const QUERY_TIME: f64 = 1000.0; // -> data time 1_000_000ms, see convert_query_time_to_data_time +const DATA_TIME: u64 = 1_000_000; +const SLIDING_WINDOW_MS: u64 = 1_000; // matches create_engine_multi_timestamp_with_window's fixed bucket width + +/// Runs a query through the native pipeline (`execute_query_pipeline`), the +/// same path `execute_and_merge_store_queries` is reached from. +fn execute_native( + engine: &SimpleEngine, + query: &str, + query_time_sec: f64, +) -> Vec { + let context = engine + .build_query_execution_context_promql(query.to_string(), query_time_sec) + .expect("Failed to build context"); + engine + .execute_query_pipeline(&context, false, false) + .expect("execute_query_pipeline failed") +} + +#[tokio::test] +async fn sliding_single_bucket_returns_its_value() { + let data = vec![( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(42.0)) as Box, + )]; + let query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + SLIDING_WINDOW_MS, + WindowType::Sliding, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1); + assert!((results[0].value - 42.0).abs() < 1e-10); +} + +#[tokio::test] +async fn sliding_two_buckets_for_same_key_are_merged_not_dropped() { + // Two precomputed buckets land under the same key and the same exact + // window (both at DATA_TIME). Today's code takes the first and warns; + // it must merge both. + let data = vec![ + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ), + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(5.0)) as Box, + ), + ]; + let query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + SLIDING_WINDOW_MS, + WindowType::Sliding, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1, "expected one merged result for host-a"); + assert!( + (results[0].value - 15.0).abs() < 1e-10, + "expected both buckets merged into 15.0, got {}", + results[0].value + ); +} + +#[tokio::test] +async fn sliding_bucket_count_mismatch_still_returns_merged_result() { + // 3 buckets (not just 2) for one key: generalizes #2 beyond the + // exactly-one-extra case, and confirms a mismatch never errors/drops — + // it merges everything and only warns. + let data = vec![ + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ), + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(5.0)) as Box, + ), + ( + DATA_TIME, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(3.0)) as Box, + ), + ]; + let query = "sum_over_time(http_requests[1s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + SLIDING_WINDOW_MS, + WindowType::Sliding, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1); + assert!( + (results[0].value - 18.0).abs() < 1e-10, + "expected all 3 buckets merged into 18.0, got {}", + results[0].value + ); +} + +#[tokio::test] +async fn tumbling_multi_bucket_merge_unaffected_by_sliding_fix() { + // Regression guard: the Sliding-branch edit lives in the same `if` as + // the Tumbling branch below it — prove Tumbling's (already-correct) + // multi-timestamp merge is untouched, through the native pipeline. + let timestamps = [996_000u64, 997_000, 998_000, 999_000, 1_000_000]; + let data = timestamps + .iter() + .map(|&ts| { + ( + ts, + Some(vec!["host-a".to_string()]), + Box::new(SumAccumulator::with_sum(10.0)) as Box, + ) + }) + .collect(); + let query = "sum_over_time(http_requests[5s])"; + let engine = create_engine_multi_timestamp_with_window( + "http_requests", + AggregationType::Sum, + vec!["host"], + data, + query, + // window_size_ms < query range so do_merge=true. Equal (5s window, + // 5s range) hits a separate, pre-existing panic — see #569, not this stage. + 1_000, + WindowType::Tumbling, + ); + + let results = execute_native(&engine, query, QUERY_TIME); + assert_eq!(results.len(), 1); + assert!( + (results[0].value - 50.0).abs() < 1e-10, + "expected 5 timestamps merged into 50.0, got {}", + results[0].value + ); +} diff --git a/docs/design-252-arithmetic-operators.md b/docs/design-252-arithmetic-operators.md index 541f2fc2..fc5f22d5 100644 --- a/docs/design-252-arithmetic-operators.md +++ b/docs/design-252-arithmetic-operators.md @@ -1,5 +1,14 @@ # Design: PromQL Arithmetic Operator Acceleration (Issue #252) +> **Superseded (#567):** the DataFusion execution path described below +> (`execute_plan`/`execute_logical_plan`, the `Join + Projection` plan) was +> replaced by a purely native execution path in `handle_binary_expr_promql` / +> `evaluate_binary_arm` / `combine_vector_vector` / `combine_scalar` +> (`asap-query-engine/src/engines/simple_engine/promql.rs`). The DataFusion +> code is kept only for its own dedicated tests and is no longer reachable +> from production. This doc is retained as the historical record of the +> original design decision. + ## Problem ASAPQuery accelerates PromQL queries by pre-computing sketches over streaming data and serving answers from those sketches at query time, bypassing the underlying TSDB for supported query patterns.