Skip to content

Commit 7c2eb98

Browse files
milindsrivastava1997claude
milindsrivastava1997
and
claude
authored
fix(query-engine): merge all sliding-window buckets per key instead of taking first (#570)
* fix(query-engine): merge all sliding-window buckets per key instead of taking first execute_and_merge_store_queries kept only the first precomputed bucket per key for Sliding-window queries, discarding the rest when the store returned more than expected. DataFusion's SummaryMergeMultipleExec already merges all of them correctly, so binary-expr queries (still on DataFusion) don't hit this. #567 will move binary-expr onto this native path, so this native/DataFusion behavior gap needed closing first. Part of #567 Stage 1. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * fix(query-engine): dedupe sliding-window merge path, fix double-clone and stale latency label - Sliding branch now delegates to merge_precomputed_outputs (do_merge=true) instead of hand-rolling its own extract/merge/insert loop, removing the duplication with the Tumbling branch's merge path. - merge_accumulators now takes ownership of the accumulator Vec so its single-element shortcut can move the value out instead of re-cloning it on top of the clone already done to build the Vec. - The [LATENCY] log's merge/no-merge label was hardcoded off window_type and said "no merge" even when merge_accumulators was in fact called; it now reflects whether merging actually occurs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * fix(query-engine): warn instead of silently dropping keys with no precompute buckets merge_precomputed_outputs silently skipped any key whose timestamped_buckets list was empty, with no log signal — unlike the "found N, expected 1" mismatch case a few lines up in the Sliding caller, which does warn. Since this function is shared by Sliding, Tumbling, and the keys-merge path, the warn now covers all three instead of being Sliding-only. Also files #575 to compute EXPECTED_BUCKETS_PER_KEY instead of hardcoding it to 1, since #554 will make >1 legitimate whenever a sliding-window query's range exceeds the window size. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 808e09d commit 7c2eb98

3 files changed

Lines changed: 213 additions & 27 deletions

File tree

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

Lines changed: 38 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -549,22 +549,29 @@ impl SimpleEngine {
549549
};
550550

551551
let merged_values = if plan.values_query.is_exact_query {
552-
// Sliding window: no merge needed, extract buckets from timestamped data
553-
debug!("Sliding window mode: Skipping merge (expecting 1 precompute per key)");
554-
values_map
555-
.into_iter()
556-
.map(|(key, timestamped_buckets)| {
557-
if timestamped_buckets.len() != 1 {
558-
warn!(
559-
"Sliding window expected 1 precompute per key, found {}. Using first.",
560-
timestamped_buckets.len()
561-
);
562-
}
563-
// Extract bucket from timestamped tuple
564-
let (_, bucket) = timestamped_buckets.into_iter().next().unwrap();
565-
(key, bucket.as_ref().clone_boxed_core())
566-
})
567-
.collect()
552+
// Sliding window: expected exactly 1 precompute per key today
553+
// (ponytail: hardcoded, #554 will make >1 legitimate — don't
554+
// block on it). The store can legitimately return more than
555+
// expected for one exact window; merge whatever came back
556+
// instead of arbitrarily keeping the first and dropping the
557+
// rest (see #567).
558+
const EXPECTED_BUCKETS_PER_KEY: usize = 1;
559+
debug!("Sliding window mode: merging {} keys", values_map.len());
560+
for timestamped_buckets in values_map.values() {
561+
if timestamped_buckets.is_empty() {
562+
continue;
563+
}
564+
if timestamped_buckets.len() != EXPECTED_BUCKETS_PER_KEY {
565+
warn!(
566+
"Sliding window expected {} precompute(s) per key, found {}. Merging all.",
567+
EXPECTED_BUCKETS_PER_KEY,
568+
timestamped_buckets.len()
569+
);
570+
}
571+
}
572+
// Sliding windows always merge (all buckets belong to one
573+
// logical window) — reuse the same merge path as Tumbling.
574+
self.merge_precomputed_outputs(&values_map, true, agg_info.aggregation_type_for_value)
568575
} else {
569576
// Tumbling window: merge needed
570577
debug!("Tumbling window mode: Merging {} outputs", values_map.len());
@@ -576,13 +583,12 @@ impl SimpleEngine {
576583
};
577584

578585
let merge_duration = merge_start_time.elapsed();
586+
let did_merge = window_type == WindowType::Sliding
587+
|| do_merge
588+
|| agg_info.aggregation_type_for_value == AggregationType::DeltaSetAggregator;
579589
debug!(
580590
"[LATENCY] Precomputed output processing ({}): {:.2}ms, resulted in {} merged outputs",
581-
if window_type == WindowType::Sliding {
582-
"no merge"
583-
} else {
584-
"merge"
585-
},
591+
if did_merge { "merge" } else { "no merge" },
586592
merge_duration.as_secs_f64() * 1000.0,
587593
merged_values.len()
588594
);
@@ -1068,7 +1074,12 @@ impl SimpleEngine {
10681074
let mut merged = HashMap::with_capacity(precomputed_outputs_map.len());
10691075

10701076
for (key, timestamped_buckets) in precomputed_outputs_map.iter() {
1071-
if !timestamped_buckets.is_empty() {
1077+
if timestamped_buckets.is_empty() {
1078+
warn!(
1079+
"Store returned key {:?} with no precompute buckets; skipping",
1080+
key
1081+
);
1082+
} else {
10721083
// Extract just the buckets (without timestamps) for merging
10731084
let precomputes: Vec<Box<dyn AggregateCore>> = timestamped_buckets
10741085
.iter()
@@ -1080,7 +1091,7 @@ impl SimpleEngine {
10801091
debug!(" Merging accumulators (should_merge=true)");
10811092
#[cfg(feature = "extra_debugging")]
10821093
let merge_start = Instant::now();
1083-
match self.merge_accumulators(&precomputes) {
1094+
match self.merge_accumulators(precomputes) {
10841095
Ok(merged_accumulator) => {
10851096
#[cfg(feature = "extra_debugging")]
10861097
let merge_duration = merge_start.elapsed();
@@ -1123,21 +1134,21 @@ impl SimpleEngine {
11231134
/// This follows the Python merge_accumulators approach
11241135
fn merge_accumulators(
11251136
&self,
1126-
accumulators: &[Box<dyn crate::data_model::AggregateCore>],
1137+
accumulators: Vec<Box<dyn crate::data_model::AggregateCore>>,
11271138
) -> Result<Box<dyn crate::data_model::AggregateCore>, AccumulatorError> {
11281139
if accumulators.is_empty() {
11291140
return Err(AccumulatorError::EmptySlice);
11301141
}
11311142

11321143
if accumulators.len() == 1 {
1133-
return Ok(accumulators[0].clone_boxed_core());
1144+
return Ok(accumulators.into_iter().next().unwrap());
11341145
}
11351146

11361147
// Try to use optimized batch merge for KLL accumulators
11371148
if accumulators[0].get_accumulator_type() == AggregationType::DatasketchesKLL {
11381149
use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator;
11391150

1140-
match DatasketchesKLLAccumulator::merge_multiple(accumulators) {
1151+
match DatasketchesKLLAccumulator::merge_multiple(&accumulators) {
11411152
Ok(merged) => return Ok(Box::new(merged)),
11421153
Err(e) => {
11431154
warn!(
@@ -1153,7 +1164,7 @@ impl SimpleEngine {
11531164
if accumulators[0].get_accumulator_type() == AggregationType::CountMinSketch {
11541165
use crate::precompute_operators::count_min_sketch_accumulator::CountMinSketchAccumulator;
11551166

1156-
match CountMinSketchAccumulator::merge_multiple(accumulators) {
1167+
match CountMinSketchAccumulator::merge_multiple(&accumulators) {
11571168
Ok(merged) => return Ok(Box::new(merged)),
11581169
Err(e) => {
11591170
warn!(

‎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_pipeline_merge_tests;
67
pub mod prometheus_forwarding_tests;
78
pub mod query_equivalence_tests;
89
pub mod sql_pattern_matching_tests;

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

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
//! Native pipeline merge tests (issue #567, Stage 1).
2+
//!
3+
//! `execute_and_merge_store_queries`'s Sliding-window branch
4+
//! (`simple_engine/mod.rs`) must merge every precomputed bucket returned for
5+
//! a key, not just the first one. The store's `query_precomputed_output_exact`
6+
//! can legitimately return more than one bucket for the same key under one
7+
//! exact window (see `per_key.rs::query_precomputed_output_exact`), and
8+
//! DataFusion's `SummaryMergeMultipleExec` already merges all of them
9+
//! correctly — native must match.
10+
11+
use crate::data_model::{AggregationType, WindowType};
12+
use crate::engines::query_result::InstantVectorElement;
13+
use crate::engines::simple_engine::SimpleEngine;
14+
use crate::precompute_operators::sum_accumulator::SumAccumulator;
15+
use crate::tests::test_utilities::engine_factories::create_engine_multi_timestamp_with_window;
16+
17+
const QUERY_TIME: f64 = 1000.0; // -> data time 1_000_000ms, see convert_query_time_to_data_time
18+
const DATA_TIME: u64 = 1_000_000;
19+
const SLIDING_WINDOW_MS: u64 = 1_000; // matches create_engine_multi_timestamp_with_window's fixed bucket width
20+
21+
/// Runs a query through the native pipeline (`execute_query_pipeline`), the
22+
/// same path `execute_and_merge_store_queries` is reached from.
23+
fn execute_native(
24+
engine: &SimpleEngine,
25+
query: &str,
26+
query_time_sec: f64,
27+
) -> Vec<InstantVectorElement> {
28+
let context = engine
29+
.build_query_execution_context_promql(query.to_string(), query_time_sec)
30+
.expect("Failed to build context");
31+
engine
32+
.execute_query_pipeline(&context, false, false)
33+
.expect("execute_query_pipeline failed")
34+
}
35+
36+
#[tokio::test]
37+
async fn sliding_single_bucket_returns_its_value() {
38+
let data = vec![(
39+
DATA_TIME,
40+
Some(vec!["host-a".to_string()]),
41+
Box::new(SumAccumulator::with_sum(42.0)) as Box<dyn crate::AggregateCore>,
42+
)];
43+
let query = "sum_over_time(http_requests[1s])";
44+
let engine = create_engine_multi_timestamp_with_window(
45+
"http_requests",
46+
AggregationType::Sum,
47+
vec!["host"],
48+
data,
49+
query,
50+
SLIDING_WINDOW_MS,
51+
WindowType::Sliding,
52+
);
53+
54+
let results = execute_native(&engine, query, QUERY_TIME);
55+
assert_eq!(results.len(), 1);
56+
assert!((results[0].value - 42.0).abs() < 1e-10);
57+
}
58+
59+
#[tokio::test]
60+
async fn sliding_two_buckets_for_same_key_are_merged_not_dropped() {
61+
// Two precomputed buckets land under the same key and the same exact
62+
// window (both at DATA_TIME). Today's code takes the first and warns;
63+
// it must merge both.
64+
let data = vec![
65+
(
66+
DATA_TIME,
67+
Some(vec!["host-a".to_string()]),
68+
Box::new(SumAccumulator::with_sum(10.0)) as Box<dyn crate::AggregateCore>,
69+
),
70+
(
71+
DATA_TIME,
72+
Some(vec!["host-a".to_string()]),
73+
Box::new(SumAccumulator::with_sum(5.0)) as Box<dyn crate::AggregateCore>,
74+
),
75+
];
76+
let query = "sum_over_time(http_requests[1s])";
77+
let engine = create_engine_multi_timestamp_with_window(
78+
"http_requests",
79+
AggregationType::Sum,
80+
vec!["host"],
81+
data,
82+
query,
83+
SLIDING_WINDOW_MS,
84+
WindowType::Sliding,
85+
);
86+
87+
let results = execute_native(&engine, query, QUERY_TIME);
88+
assert_eq!(results.len(), 1, "expected one merged result for host-a");
89+
assert!(
90+
(results[0].value - 15.0).abs() < 1e-10,
91+
"expected both buckets merged into 15.0, got {}",
92+
results[0].value
93+
);
94+
}
95+
96+
#[tokio::test]
97+
async fn sliding_bucket_count_mismatch_still_returns_merged_result() {
98+
// 3 buckets (not just 2) for one key: generalizes #2 beyond the
99+
// exactly-one-extra case, and confirms a mismatch never errors/drops —
100+
// it merges everything and only warns.
101+
let data = vec![
102+
(
103+
DATA_TIME,
104+
Some(vec!["host-a".to_string()]),
105+
Box::new(SumAccumulator::with_sum(10.0)) as Box<dyn crate::AggregateCore>,
106+
),
107+
(
108+
DATA_TIME,
109+
Some(vec!["host-a".to_string()]),
110+
Box::new(SumAccumulator::with_sum(5.0)) as Box<dyn crate::AggregateCore>,
111+
),
112+
(
113+
DATA_TIME,
114+
Some(vec!["host-a".to_string()]),
115+
Box::new(SumAccumulator::with_sum(3.0)) as Box<dyn crate::AggregateCore>,
116+
),
117+
];
118+
let query = "sum_over_time(http_requests[1s])";
119+
let engine = create_engine_multi_timestamp_with_window(
120+
"http_requests",
121+
AggregationType::Sum,
122+
vec!["host"],
123+
data,
124+
query,
125+
SLIDING_WINDOW_MS,
126+
WindowType::Sliding,
127+
);
128+
129+
let results = execute_native(&engine, query, QUERY_TIME);
130+
assert_eq!(results.len(), 1);
131+
assert!(
132+
(results[0].value - 18.0).abs() < 1e-10,
133+
"expected all 3 buckets merged into 18.0, got {}",
134+
results[0].value
135+
);
136+
}
137+
138+
#[tokio::test]
139+
async fn tumbling_multi_bucket_merge_unaffected_by_sliding_fix() {
140+
// Regression guard: the Sliding-branch edit lives in the same `if` as
141+
// the Tumbling branch below it — prove Tumbling's (already-correct)
142+
// multi-timestamp merge is untouched, through the native pipeline.
143+
let timestamps = [996_000u64, 997_000, 998_000, 999_000, 1_000_000];
144+
let data = timestamps
145+
.iter()
146+
.map(|&ts| {
147+
(
148+
ts,
149+
Some(vec!["host-a".to_string()]),
150+
Box::new(SumAccumulator::with_sum(10.0)) as Box<dyn crate::AggregateCore>,
151+
)
152+
})
153+
.collect();
154+
let query = "sum_over_time(http_requests[5s])";
155+
let engine = create_engine_multi_timestamp_with_window(
156+
"http_requests",
157+
AggregationType::Sum,
158+
vec!["host"],
159+
data,
160+
query,
161+
// window_size_ms < query range so do_merge=true. Equal (5s window,
162+
// 5s range) hits a separate, pre-existing panic — see #569, not this stage.
163+
1_000,
164+
WindowType::Tumbling,
165+
);
166+
167+
let results = execute_native(&engine, query, QUERY_TIME);
168+
assert_eq!(results.len(), 1);
169+
assert!(
170+
(results[0].value - 50.0).abs() < 1e-10,
171+
"expected 5 timestamps merged into 50.0, got {}",
172+
results[0].value
173+
);
174+
}

0 commit comments

Comments
 (0)