fix(query-engine): merge all sliding-window buckets per key instead of taking first

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()
         );
@@ -1068,7 +1074,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<Box<dyn AggregateCore>> = timestamped_buckets
                     .iter()
@@ -1080,7 +1091,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 +1134,21 @@ impl SimpleEngine {
     /// This follows the Python merge_accumulators approach
     fn merge_accumulators(
         &self,
-        accumulators: &[Box<dyn crate::data_model::AggregateCore>],
+        accumulators: Vec<Box<dyn crate::data_model::AggregateCore>>,
     ) -> Result<Box<dyn crate::data_model::AggregateCore>, 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 +1164,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!(

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_pipeline_merge_tests;
 pub mod prometheus_forwarding_tests;
 pub mod query_equivalence_tests;
 pub mod sql_pattern_matching_tests;

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<InstantVectorElement> {
+    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<dyn crate::AggregateCore>,
+    )];
+    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<dyn crate::AggregateCore>,
+        ),
+        (
+            DATA_TIME,
+            Some(vec!["host-a".to_string()]),
+            Box::new(SumAccumulator::with_sum(5.0)) as Box<dyn crate::AggregateCore>,
+        ),
+    ];
+    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<dyn crate::AggregateCore>,
+        ),
+        (
+            DATA_TIME,
+            Some(vec!["host-a".to_string()]),
+            Box::new(SumAccumulator::with_sum(5.0)) as Box<dyn crate::AggregateCore>,
+        ),
+        (
+            DATA_TIME,
+            Some(vec!["host-a".to_string()]),
+            Box::new(SumAccumulator::with_sum(3.0)) as Box<dyn crate::AggregateCore>,
+        ),
+    ];
+    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<dyn crate::AggregateCore>,
+            )
+        })
+        .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
+    );
+}