Skip to content

fix(query-engine): range queries expand keys_query and merge same-timestamp buckets #582

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
milindsrivastava1997 merged 3 commits into from
Aug 24, 2026
+520 −98
Merged

fix(query-engine): range queries expand keys_query and merge same-timestamp buckets #582

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4c59e9b
test(query-engine): add red tests for range-query keys_query/is_exact…
milindsrivastava1997 Aug 23, 2026
fc6080e
fix(query-engine): range queries expand keys_query and merge same-tim…
milindsrivastava1997 Aug 23, 2026
7b9f084
fix(query-engine): hard-error on unresolvable range keys, dedup keys …
milindsrivastava1997 Aug 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 145 additions & 98 deletions asap-query-engine/src/engines/simple_engine/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -594,37 +594,50 @@ impl SimpleEngine {
);

// Query and merge keys if needed
let merged_keys = if let Some(keys_params) = &plan.keys_query {
let keys_store_query_start_time = Instant::now();
let keys_map = self.execute_store_query(keys_params).map_err(|e| {
warn!("Error querying store for keys: {}", e);
e
})?;
debug!(
"[LATENCY] Keys store query (metric: {}, agg: {}): {}ms",
&keys_params.metric,
keys_params.aggregation_id,
keys_store_query_start_time.elapsed().as_millis()
);
debug!("Keys query returned {} unique keys", keys_map.len());
let merged_keys = self.fetch_and_merge_keys(
&plan.keys_query,
agg_info.aggregation_type_for_key,
do_merge,
)?;

let keys_merge_start_time = Instant::now();
let merged = self.merge_precomputed_outputs(
&keys_map,
do_merge,
agg_info.aggregation_type_for_key,
);
debug!(
"[LATENCY] Keys merge operation: {:.2}ms, resulted in {} merged outputs",
keys_merge_start_time.elapsed().as_secs_f64() * 1000.0,
merged.len()
);
Some(merged)
} else {
None
Ok((merged_values, merged_keys))
}

/// Fetches and merges the keys side of a dual-population query plan, if
/// present. Shared by `execute_and_merge_store_queries` (instant) and
/// `execute_range_query_pipeline` (range) — both build a `StoreQueryPlan`
/// that may carry a separate `keys_query` and need it merged the same way.
fn fetch_and_merge_keys(
&self,
keys_query: &Option<StoreQueryParams>,
key_aggregation_type: AggregationType,
do_merge: bool,
) -> Result<Option<MergedOutputsMap>, String> {
let Some(keys_params) = keys_query else {
return Ok(None);
};

Ok((merged_values, merged_keys))
let keys_store_query_start_time = Instant::now();
let keys_map = self.execute_store_query(keys_params).map_err(|e| {
warn!("Error querying store for keys: {}", e);
e
})?;
debug!(
"[LATENCY] Keys store query (metric: {}, agg: {}): {}ms",
&keys_params.metric,
keys_params.aggregation_id,
keys_store_query_start_time.elapsed().as_millis()
);
debug!("Keys query returned {} unique keys", keys_map.len());

let keys_merge_start_time = Instant::now();
let merged = self.merge_precomputed_outputs(&keys_map, do_merge, key_aggregation_type);
debug!(
"[LATENCY] Keys merge operation: {:.2}ms, resulted in {} merged outputs",
keys_merge_start_time.elapsed().as_secs_f64() * 1000.0,
merged.len()
);
Ok(Some(merged))
}

/// Collects all results based on whether keys are separate or not
Expand Down Expand Up @@ -1113,12 +1126,27 @@ impl SimpleEngine {
}
}
} else {
assert_eq!(
precomputes.len(),
1,
"Spatial queries should have exactly 1 precompute per key"
);
merged.insert(key.clone(), precomputes[0].clone_boxed_core());
// Spatial queries (do_merge=false) normally see exactly 1
// precompute per key. A range query's widened fetch can
// surface more than expected even when do_merge was
// computed false for the base instant range, and future
// Sliding-window support may make >1 legitimate here too
// — warn and merge instead of asserting/panicking.
if precomputes.len() != 1 {
warn!(
"Spatial query expected 1 precompute per key {:?}, found {}. Merging anyway.",
key,
precomputes.len()
);
}
match self.merge_accumulators(precomputes) {
Ok(merged_accumulator) => {
merged.insert(key.clone(), merged_accumulator);
}
Err(e) => {
warn!("Failed to merge accumulators for key {:?}: {}", key, e);
}
}
}
}
}
Expand Down Expand Up @@ -1456,6 +1484,16 @@ impl SimpleEngine {
all_data.values().map(|v| v.len()).sum::<usize>()
);

// Dual-population metrics (separate key/value aggregations) need the
// keys side merged once up front, then expanded per value group below
// — mirrors execute_and_merge_store_queries/collect_results_separate_keys
// (#580: this pipeline previously never read keys_query at all).
let merged_keys = self.fetch_and_merge_keys(
&context.base.store_plan.keys_query,
context.base.agg_info.aggregation_type_for_key,
context.base.do_merge,
)?;

let mut results: HashMap<KeyByLabelValues, RangeVectorElement> = HashMap::new();

// Determine accumulator type for merger selection
Expand All @@ -1467,6 +1505,8 @@ impl SimpleEngine {
let end_ms = context.range_params.end;
let buckets_per_step = context.buckets_per_step;
let lookback_bucket_count = context.lookback_bucket_count;
let tumbling_window_ms = context.tumbling_window_ms;
let lookback_ms = (lookback_bucket_count as u64) * tumbling_window_ms;

let window_mode = if buckets_per_step <= lookback_bucket_count {
"sliding (slide <= size)"
Expand All @@ -1479,42 +1519,58 @@ impl SimpleEngine {
start_ms,
end_ms,
step_ms,
context.tumbling_window_ms,
tumbling_window_ms,
buckets_per_step,
lookback_bucket_count,
window_mode
);

// Process each key independently
for (key_opt, timestamped_buckets) in &all_data {
let key = match key_opt {
Some(k) => k.clone(),
None => continue, // Skip None keys for now
};

// Build lookup: bucket_start_timestamp -> bucket for O(1) access
let bucket_map: HashMap<u64, &dyn AggregateCore> = timestamped_buckets
// Resolve, for every value group, which output label-keys it serves:
// its own key for single-population metrics, or every key the merged
// keys aggregation expands it to for dual-population metrics. Mirrors
// collect_results_separate_keys exactly, including its error
// semantics — an unresolvable key set fails the whole range query
// (so callers fall back to Prometheus) instead of silently returning
// a partial result. See #582 review.
let groups: Vec<(
&Vec<crate::stores::TimestampedBucket>,
Vec<KeyByLabelValues>,
)> = match &merged_keys {
Some(keys_map) => keys_map
.iter()
.map(|((start, _), bucket)| (*start, bucket.as_ref()))
.collect();

debug!(
"Key {:?}: built bucket_map with {} entries, timestamps: {:?}",
key,
bucket_map.len(),
bucket_map.keys().collect::<Vec<_>>()
);

// Create result element for this key
let mut element = RangeVectorElement::new(key.clone());
.map(|(group_key, keys_precompute)| {
let timestamped_buckets = all_data
.get(group_key)
.ok_or_else(|| format!("No value for key: {:?}", group_key))?;
let expansion_keys = keys_precompute
.get_keys()
.ok_or_else(|| "Keys required for separate aggregation".to_string())?;
Ok((timestamped_buckets, expansion_keys))
})
.collect::<Result<Vec<_>, String>>()?,
None => all_data
.iter()
.filter_map(|(group_key, buckets)| {
group_key.as_ref().map(|k| (buckets, vec![k.clone()]))
})
.collect(),
};

// Calculate window parameters
let tumbling_window_ms = context.tumbling_window_ms;
let lookback_ms = (lookback_bucket_count as u64) * tumbling_window_ms;
// Process each value group independently
for (timestamped_buckets, expansion_keys) in groups {
// Build lookup: bucket_start_timestamp -> all buckets sharing
// that start. A Sliding aggregation can legitimately return more
// than one bucket per start timestamp (#567/#570) — every one of
// them must be merged, not just the last one collected here.
let mut bucket_map: HashMap<u64, Vec<&dyn AggregateCore>> = HashMap::new();
for ((start, _), bucket) in timestamped_buckets {
bucket_map.entry(*start).or_default().push(bucket.as_ref());
}

debug!(
"Key {:?}: range [{}, {}], step={}, lookback_ms={}, tumbling_window_ms={}",
key, start_ms, end_ms, step_ms, lookback_ms, tumbling_window_ms
"Group with {} start-timestamps, expands to keys: {:?}",
bucket_map.len(),
expansion_keys
);

// Iterate by OUTPUT timestamp, not by bucket index
Expand All @@ -1529,10 +1585,10 @@ impl SimpleEngine {

let mut t = window_start;
while t < current_time {
if let Some(bucket) = bucket_map.get(&t) {
window_buckets.push((*bucket).clone_boxed_core());
if let Some(buckets) = bucket_map.get(&t) {
window_buckets.extend(buckets.iter().map(|b| b.clone_boxed_core()));
}
// If bucket missing at timestamp t, just skip it (partial data is okay)
// If no bucket at timestamp t, just skip it (partial data is okay)
t += tumbling_window_ms;
}

Expand All @@ -1543,56 +1599,47 @@ impl SimpleEngine {

match merger.get_merged() {
Ok(merged) => {
// Query statistic and emit sample at current_time
match self.query_precompute_for_statistic(
merged.as_ref(),
&context.base.metadata.statistic_to_compute,
&Some(key.clone()),
&context.base.metadata.query_kwargs,
) {
Ok(value) => {
debug!(
"Key {:?}: emitting sample (t={}, value={})",
key, current_time, value
);
element.add_sample(current_time, value);
}
Err(e) => {
debug!(
"Failed to query statistic at t={} for key {:?}: {}",
current_time, key, e
);
// Query statistic and emit a sample at current_time
// for every expanded key sharing this value group.
for key in &expansion_keys {
match self.query_precompute_for_statistic(
merged.as_ref(),
&context.base.metadata.statistic_to_compute,
&Some(key.clone()),
&context.base.metadata.query_kwargs,
) {
Ok(value) => {
results
.entry(key.clone())
.or_insert_with(|| RangeVectorElement::new(key.clone()))
.add_sample(current_time, value);
}
Err(e) => {
debug!(
"Failed to query statistic at t={} for key {:?}: {}",
current_time, key, e
);
}
}
}
}
Err(e) => {
debug!(
"Failed to get merged result at t={} for key {:?}: {}",
current_time, key, e
"Failed to get merged result at t={} for keys {:?}: {}",
current_time, expansion_keys, e
);
}
}
} else {
// No data at all for this window - skip sample
debug!(
"Key {:?}: skipping sample at {} - no data in window [{}, {})",
key, current_time, window_start, current_time
"Keys {:?}: skipping sample at {} - no data in window [{}, {})",
expansion_keys, current_time, window_start, current_time
);
}

current_time += step_ms;
}

debug!(
"Key {:?}: finished with {} samples",
key,
element.samples.len()
);

// Only include keys with samples
if !element.samples.is_empty() {
results.insert(key, element);
}
}

// Convert to Vec
Expand Down
1 change: 1 addition & 0 deletions asap-query-engine/src/tests/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ 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 native_range_query_tests;
pub mod prometheus_forwarding_tests;
pub mod query_equivalence_tests;
pub mod sql_pattern_matching_tests;
Expand Down
Loading
Loading