From a82945a63dc76349b2d8d577e902f69b94fa8199 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 1 Sep 2026 18:08:10 -0400 Subject: [PATCH 1/3] feat(query-engine): extrapolate rate and increase at boundaries --- asap-query-engine/src/data_model/traits.rs | 33 ++ .../src/engines/simple_engine/mod.rs | 28 +- .../src/engines/simple_engine/promql.rs | 1 + .../increase_accumulator.rs | 311 +++++++++++++++++- .../multiple_increase_accumulator.rs | 94 +++++- .../src/tests/native_range_query_tests.rs | 71 ++++ 6 files changed, 528 insertions(+), 10 deletions(-) diff --git a/asap-query-engine/src/data_model/traits.rs b/asap-query-engine/src/data_model/traits.rs index 064619f..b388578 100644 --- a/asap-query-engine/src/data_model/traits.rs +++ b/asap-query-engine/src/data_model/traits.rs @@ -6,6 +6,25 @@ use promql_utilities::query_logics::enums::{AggregationType, Statistic}; pub use asap_types::traits::SerializableToSink; +/// Exact time boundaries of the range vector being evaluated. +/// +/// Timestamps are milliseconds since the Unix epoch, matching the query +/// engine's data timestamps. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct QueryBounds { + pub start_timestamp: i64, + pub end_timestamp: i64, +} + +impl QueryBounds { + pub const fn new(start_timestamp: i64, end_timestamp: i64) -> Self { + Self { + start_timestamp, + end_timestamp, + } + } +} + /// Core trait for all aggregates containing shared functionality /// This trait provides common operations like serialization, cloning, and type identification pub trait AggregateCore: SerializableToSink + Send + Sync { @@ -44,6 +63,20 @@ pub trait AggregateCore: SerializableToSink + Send + Sync { key: &Option, query_kwargs: &HashMap, ) -> Result>; + + /// Dispatch a statistic query with exact range-vector boundaries. + /// + /// Accumulators that need Prometheus range semantics override this + /// method. Other accumulators retain their existing query behavior. + fn query_statistic_with_bounds( + &self, + statistic: Statistic, + key: &Option, + query_kwargs: &HashMap, + _bounds: &QueryBounds, + ) -> Result> { + self.query_statistic(statistic, key, query_kwargs) + } } /// Trait for accumulators that support a single subpopulation diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 43d2dda..e583957 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -3,7 +3,7 @@ mod promql; mod sql; use crate::data_model::{ - AggregationIdInfo, InferenceConfig, KeyByLabelValues, QueryConfig, QueryLanguage, + AggregationIdInfo, InferenceConfig, KeyByLabelValues, QueryBounds, QueryConfig, QueryLanguage, StreamingConfig, }; use crate::engines::query_result::{InstantVectorElement, QueryResult}; @@ -109,6 +109,8 @@ pub struct RangeQueryExecutionContext { /// list, rather than a start/end/step triple that only ever meant /// something for range. pub output_timestamps: Vec, + /// Exact range-vector duration used for extrapolation, in milliseconds. + pub query_range_ms: u64, /// Number of buckets per step (step / tumbling_window) pub buckets_per_step: usize, /// Number of buckets in lookback window @@ -641,6 +643,7 @@ impl SimpleEngine { ..base_context }, output_timestamps: vec![query_time], + query_range_ms: lookback_ms, // Placeholder: no real "step" for a single instant point. Only // feeds a debug-log string today -- not type-enforced, recheck // before using it for anything functional. @@ -975,6 +978,7 @@ impl SimpleEngine { fallback_key: &Option, statistic: &Statistic, query_kwargs: &HashMap, + query_bounds: Option<&QueryBounds>, ) -> Vec<(Option, f64)> { let Some(value_precompute) = value_precompute else { warn!( @@ -1011,6 +1015,7 @@ impl SimpleEngine { statistic, &key, query_kwargs, + query_bounds, ) { Ok(value) => Some((key, value)), Err(e) => { @@ -1512,6 +1517,7 @@ impl SimpleEngine { group_key, statistic, query_kwargs, + None, ) { unformatted_results.insert(key, value); } @@ -1544,6 +1550,7 @@ impl SimpleEngine { group_key, statistic, query_kwargs, + None, ) { unformatted_results.insert(key, value); } @@ -1558,8 +1565,14 @@ impl SimpleEngine { statistic: &Statistic, key: &Option, query_kwargs: &HashMap, + query_bounds: Option<&QueryBounds>, ) -> Result> { - precompute.query_statistic(*statistic, key, query_kwargs) + match query_bounds { + Some(bounds) => { + precompute.query_statistic_with_bounds(*statistic, key, query_kwargs, bounds) + } + None => precompute.query_statistic(*statistic, key, query_kwargs), + } } // ============================================================ @@ -2124,6 +2137,16 @@ impl SimpleEngine { // (#581). One loop shape for topk and non-topk alike, rather than // maintaining two. for ¤t_time in &context.output_timestamps { + let current_time_i64 = i64::try_from(current_time) + .map_err(|_| "Output timestamp exceeds signed timestamp range".to_string())?; + let query_range_ms = i64::try_from(context.query_range_ms) + .map_err(|_| "Query range exceeds signed timestamp range".to_string())?; + let query_bounds = QueryBounds::new( + current_time_i64 + .checked_sub(query_range_ms) + .ok_or("Query range underflows timestamp range".to_string())?, + current_time_i64, + ); // This timestamp's (key, value) pairs across every group, // collected before insertion into `results` so a topk query can // rank/truncate them as one step-local set (#581 stage E.3 -- @@ -2314,6 +2337,7 @@ impl SimpleEngine { &fallback_key, &context.base.metadata.statistic_to_compute, &context.base.metadata.query_kwargs, + Some(&query_bounds), ) { // A fully unlabeled result (fallback_key was None and // the value accumulator has no self-keys) has no diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 897d352..f124329 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -661,6 +661,7 @@ impl SimpleEngine { // per-step loop's `current_time` sequence exactly: start_ms, // start_ms+step_ms, ..., the last value <= end_ms. output_timestamps: (start_ms..=end_ms).step_by(step_ms as usize).collect(), + query_range_ms: lookback_ms, buckets_per_step, lookback_bucket_count, tumbling_window_ms, diff --git a/asap-query-engine/src/precompute_operators/increase_accumulator.rs b/asap-query-engine/src/precompute_operators/increase_accumulator.rs index b32dd4a..de10869 100644 --- a/asap-query-engine/src/precompute_operators/increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/increase_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ - AggregateCore, AggregationType, Measurement, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, SingleSubpopulationAggregateFactory, + AggregateCore, AggregationType, Measurement, MergeableAccumulator, QueryBounds, + SerializableToSink, SingleSubpopulationAggregate, SingleSubpopulationAggregateFactory, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -8,7 +8,7 @@ use std::collections::HashMap; use promql_utilities::query_logics::enums::Statistic; -pub(crate) const INCREASE_BINARY_FORMAT_MAGIC: [u8; 4] = *b"INC6"; +pub(crate) const INCREASE_BINARY_FORMAT_MAGIC: [u8; 4] = *b"INC7"; pub(crate) const RESET_RECORD_BYTES: usize = std::mem::size_of::() + std::mem::size_of::(); @@ -27,6 +27,7 @@ pub struct IncreaseAccumulator { pub starting_timestamp: i64, pub last_seen_measurement: Measurement, pub last_seen_timestamp: i64, + pub sample_count: u64, pub counter_reset_adjustment: f64, pub counter_reset_events: Vec, } @@ -38,11 +39,32 @@ impl IncreaseAccumulator { last_seen_measurement: Measurement, last_seen_timestamp: i64, ) -> Self { + Self::new_with_sample_count( + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + 1, + ) + } + + pub fn new_with_sample_count( + starting_measurement: Measurement, + starting_timestamp: i64, + last_seen_measurement: Measurement, + last_seen_timestamp: i64, + sample_count: u64, + ) -> Self { + assert!( + sample_count > 0, + "IncreaseAccumulator sample count must be positive" + ); Self { starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, + sample_count, counter_reset_adjustment: 0.0, counter_reset_events: Vec::new(), } @@ -59,6 +81,10 @@ impl IncreaseAccumulator { } self.last_seen_measurement = measurement; self.last_seen_timestamp = timestamp; + self.sample_count = self + .sample_count + .checked_add(1) + .expect("IncreaseAccumulator sample count overflow"); } fn increase(&self) -> f64 { @@ -66,6 +92,70 @@ impl IncreaseAccumulator { + self.counter_reset_adjustment } + pub fn query_with_bounds( + &self, + statistic: Statistic, + bounds: &QueryBounds, + ) -> Result> { + if bounds.start_timestamp >= bounds.end_timestamp { + return Err("Query range must have a positive duration".into()); + } + if self.sample_count < 2 { + return Err("At least two samples are required".into()); + } + + let sampled_interval = self + .last_seen_timestamp + .checked_sub(self.starting_timestamp) + .ok_or("Observed sample timestamps overflowed")? as f64; + if sampled_interval <= 0.0 { + return Err("Observed samples must span a positive duration".into()); + } + + let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; + let extrapolation_threshold = average_sample_interval * 1.1; + let mut duration_to_start = + self.starting_timestamp + .checked_sub(bounds.start_timestamp) + .ok_or("Start boundary duration overflowed")? as f64; + let mut duration_to_end = bounds + .end_timestamp + .checked_sub(self.last_seen_timestamp) + .ok_or("End boundary duration overflowed")? as f64; + + if duration_to_start >= extrapolation_threshold { + duration_to_start = average_sample_interval / 2.0; + } + if duration_to_end >= extrapolation_threshold { + duration_to_end = average_sample_interval / 2.0; + } + + let increase = self.increase(); + if increase > 0.0 && self.starting_measurement.value >= 0.0 { + let duration_to_zero = sampled_interval * (self.starting_measurement.value / increase); + if duration_to_zero < duration_to_start { + duration_to_start = duration_to_zero; + } + } + + let factor = (sampled_interval + duration_to_start + duration_to_end) / sampled_interval; + let query_interval = bounds + .end_timestamp + .checked_sub(bounds.start_timestamp) + .ok_or("Query range duration overflowed")? as f64; + let result = match statistic { + Statistic::Increase => increase * factor, + Statistic::Rate => increase * factor / query_interval * 1000.0, + _ => { + return Err( + format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into(), + ) + } + }; + + Ok(result) + } + fn add_reset_event(&mut self, event: CounterResetEvent) { if self .counter_reset_events @@ -95,6 +185,12 @@ impl IncreaseAccumulator { let last_seen_timestamp = data["last_seen_timestamp"] .as_i64() .ok_or("Missing or invalid 'last_seen_timestamp' field")?; + let sample_count = data["sample_count"] + .as_u64() + .ok_or("Missing or invalid 'sample_count' field")?; + if sample_count == 0 { + return Err("Sample count must be positive".into()); + } let counter_reset_adjustment = data["counter_reset_adjustment"] .as_f64() .ok_or("Missing or invalid 'counter_reset_adjustment' field")?; @@ -106,6 +202,7 @@ impl IncreaseAccumulator { last_seen_measurement, last_seen_timestamp, ); + accumulator.sample_count = sample_count; accumulator.counter_reset_adjustment = counter_reset_adjustment; accumulator.counter_reset_events = counter_reset_events; Ok(accumulator) @@ -195,6 +292,24 @@ impl IncreaseAccumulator { ]); offset += 8; + if buffer.len() < offset + 8 { + return Err("Buffer too short for sample count".into()); + } + let sample_count = u64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + if sample_count == 0 { + return Err("Sample count must be positive".into()); + } + if buffer.len() < offset + 8 { return Err("Buffer too short for counter reset adjustment".into()); } @@ -262,6 +377,7 @@ impl IncreaseAccumulator { last_seen_measurement, last_seen_timestamp, ); + accumulator.sample_count = sample_count; accumulator.counter_reset_adjustment = counter_reset_adjustment; accumulator.counter_reset_events = counter_reset_events; Ok((accumulator, offset)) @@ -275,6 +391,7 @@ impl SerializableToSink for IncreaseAccumulator { "starting_timestamp": self.starting_timestamp, "last_seen_measurement": self.last_seen_measurement.serialize_to_json(), "last_seen_timestamp": self.last_seen_timestamp, + "sample_count": self.sample_count, "counter_reset_adjustment": self.counter_reset_adjustment, "counter_reset_events": self.counter_reset_events, }) @@ -300,6 +417,7 @@ impl SerializableToSink for IncreaseAccumulator { // Last seen timestamp and total reset adjustment buffer.extend_from_slice(&self.last_seen_timestamp.to_le_bytes()); + buffer.extend_from_slice(&self.sample_count.to_le_bytes()); buffer.extend_from_slice(&self.counter_reset_adjustment.to_le_bytes()); buffer.extend_from_slice(&(self.counter_reset_events.len() as u32).to_le_bytes()); for event in &self.counter_reset_events { @@ -324,6 +442,11 @@ impl MergeableAccumulator for IncreaseAccumulator { let mut result = accumulators.remove(0); for acc in accumulators { + result.sample_count = result + .sample_count + .checked_add(acc.sample_count) + .ok_or("Sample count overflow while merging IncreaseAccumulator")?; + // Adjacent accumulators represent consecutive portions of the same // counter. A decrease at their boundary is also a reset. if acc.starting_timestamp >= result.last_seen_timestamp @@ -405,6 +528,16 @@ impl AggregateCore for IncreaseAccumulator { use crate::data_model::SingleSubpopulationAggregate; self.query(statistic, None) } + + fn query_statistic_with_bounds( + &self, + statistic: Statistic, + _key: &Option, + _query_kwargs: &HashMap, + bounds: &QueryBounds, + ) -> Result> { + self.query_with_bounds(statistic, bounds) + } } impl SingleSubpopulationAggregate for IncreaseAccumulator { @@ -481,6 +614,165 @@ impl SingleSubpopulationAggregateFactory for IncreaseAccumulatorFactory { mod tests { use super::*; + #[test] + fn no_boundary_gap_keeps_increase_and_scales_rate_to_requested_range() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(100.0), + 1_000, + Measurement::new(100.0), + 1_000, + ); + acc.update(Measurement::new(110.0), 2_000); + acc.update(Measurement::new(120.0), 3_000); + + let bounds = crate::data_model::QueryBounds::new(1_000, 3_000); + + assert_eq!( + acc.query_with_bounds(Statistic::Increase, &bounds).unwrap(), + 20.0 + ); + assert_eq!( + acc.query_with_bounds(Statistic::Rate, &bounds).unwrap(), + 10.0 + ); + } + + #[test] + fn large_boundary_gaps_are_limited_to_half_average_sample_interval() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(100.0), + 1_000, + Measurement::new(100.0), + 1_000, + ); + acc.update(Measurement::new(110.0), 2_000); + acc.update(Measurement::new(120.0), 3_000); + + let bounds = crate::data_model::QueryBounds::new(-1_000, 5_000); + + assert_eq!( + acc.query_with_bounds(Statistic::Increase, &bounds).unwrap(), + 30.0 + ); + assert_eq!( + acc.query_with_bounds(Statistic::Rate, &bounds).unwrap(), + 5.0 + ); + } + + #[test] + fn exact_threshold_uses_half_an_average_sample_interval() { + let acc = IncreaseAccumulator::new_with_sample_count( + Measurement::new(100.0), + 1_000, + Measurement::new(110.0), + 2_000, + 2, + ); + let bounds = crate::data_model::QueryBounds::new(-100, 2_000); + + // The left gap is exactly 1.1 times the average interval, so the + // inclusive threshold must choose half an interval (500ms). + assert_eq!( + acc.query_with_bounds(Statistic::Increase, &bounds).unwrap(), + 15.0 + ); + } + + #[test] + fn irregular_sample_intervals_use_the_average_interval() { + let acc = IncreaseAccumulator::new_with_sample_count( + Measurement::new(100.0), + 0, + Measurement::new(125.0), + 1_700, + 3, + ); + let bounds = crate::data_model::QueryBounds::new(-700, 2_400); + + let expected_increase = 25.0 * 3_100.0 / 1_700.0; + let expected_rate = expected_increase / 3_100.0 * 1_000.0; + assert!( + (acc.query_with_bounds(Statistic::Increase, &bounds).unwrap() - expected_increase) + .abs() + < 1e-12 + ); + assert!( + (acc.query_with_bounds(Statistic::Rate, &bounds).unwrap() - expected_rate).abs() + < 1e-12 + ); + } + + #[test] + fn fewer_than_two_samples_and_degenerate_ranges_are_rejected() { + let single_sample = IncreaseAccumulator::new( + Measurement::new(100.0), + 1_000, + Measurement::new(100.0), + 1_000, + ); + assert!(single_sample + .query_with_bounds( + Statistic::Increase, + &crate::data_model::QueryBounds::new(0, 2_000) + ) + .is_err()); + + let two_samples = IncreaseAccumulator::new_with_sample_count( + Measurement::new(100.0), + 1_000, + Measurement::new(110.0), + 2_000, + 2, + ); + assert!(two_samples + .query_with_bounds( + Statistic::Rate, + &crate::data_model::QueryBounds::new(2_000, 2_000) + ) + .is_err()); + } + + #[test] + fn counter_resets_are_corrected_before_boundary_extrapolation() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(100.0), + 1_000, + Measurement::new(100.0), + 1_000, + ); + acc.update(Measurement::new(150.0), 2_000); + acc.update(Measurement::new(10.0), 3_000); + acc.update(Measurement::new(60.0), 4_000); + + let bounds = crate::data_model::QueryBounds::new(0, 5_000); + + let increase = acc.query_with_bounds(Statistic::Increase, &bounds).unwrap(); + let rate = acc.query_with_bounds(Statistic::Rate, &bounds).unwrap(); + assert!((increase - (110.0 * 5.0 / 3.0)).abs() < f64::EPSILON); + assert!((rate - (110.0 / 3.0)).abs() < f64::EPSILON); + } + + #[test] + fn counter_extrapolation_is_clamped_to_duration_to_zero() { + let acc = IncreaseAccumulator::new_with_sample_count( + Measurement::new(1.0), + 1_000, + Measurement::new(101.0), + 3_000, + 2, + ); + let bounds = crate::data_model::QueryBounds::new(-500, 3_000); + + // Without the counter-to-zero clamp, the 1.5s left gap would be + // extrapolated as if the counter had already been increasing before + // the first observed sample. Prometheus limits it to 20ms here. + assert_eq!( + acc.query_with_bounds(Statistic::Increase, &bounds).unwrap(), + 101.0 + ); + } + #[test] fn test_increase_accumulator_creation() { let starting_measurement = Measurement::new(10.0); @@ -599,6 +891,7 @@ mod tests { assert_eq!(merged.starting_timestamp, 500); assert_eq!(merged.last_seen_measurement.value, 30.0); assert_eq!(merged.last_seen_timestamp, 3000); + assert_eq!(merged.sample_count, 3); } #[test] @@ -644,11 +937,14 @@ mod tests { #[test] fn test_increase_accumulator_serialization() { - let acc = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); + let mut acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(10.0), 1000); + acc.update(Measurement::new(15.0), 1500); + acc.update(Measurement::new(25.0), 2000); // Test JSON serialization let json = acc.serialize_to_json(); + assert_eq!(json["sample_count"], 3); assert!(json.get("opaque_reset_adjustment").is_none()); assert!(json.get("opaque_reset_ranges").is_none()); let deserialized = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); @@ -662,10 +958,11 @@ mod tests { deserialized.last_seen_measurement.value ); assert_eq!(acc.last_seen_timestamp, deserialized.last_seen_timestamp); + assert_eq!(acc.sample_count, deserialized.sample_count); // Test byte serialization let bytes = acc.serialize_to_bytes(); - assert_eq!(&bytes[..4], b"INC6"); + assert_eq!(&bytes[..4], b"INC7"); let deserialized_bytes = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); assert_eq!( acc.starting_measurement.value, @@ -683,11 +980,13 @@ mod tests { acc.last_seen_timestamp, deserialized_bytes.last_seen_timestamp ); + assert_eq!(acc.sample_count, deserialized_bytes.sample_count); } #[test] fn test_deserialize_from_bytes_rejects_previous_formats() { assert!(IncreaseAccumulator::deserialize_from_bytes(b"INC2").is_err()); + assert!(IncreaseAccumulator::deserialize_from_bytes(b"INC6").is_err()); assert!(IncreaseAccumulator::deserialize_from_bytes(b"INC5").is_err()); } diff --git a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs index a513b25..c4ff1ca 100644 --- a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs @@ -1,6 +1,6 @@ use crate::data_model::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, SingleSubpopulationAggregate, + MultipleSubpopulationAggregate, QueryBounds, SerializableToSink, SingleSubpopulationAggregate, }; use crate::precompute_operators::{CounterResetEvent, IncreaseAccumulator}; use serde::{Deserialize, Serialize}; @@ -24,6 +24,7 @@ struct MeasurementData { starting_timestamp: i64, last_seen_measurement: f64, last_seen_timestamp: i64, + sample_count: u64, counter_reset_adjustment: f64, counter_reset_events: Vec, } @@ -122,6 +123,9 @@ impl MultipleIncreaseAccumulator { let starting_timestamp = values.starting_timestamp; let last_seen_measurement = Measurement::new(values.last_seen_measurement); let last_seen_timestamp = values.last_seen_timestamp; + if values.sample_count == 0 { + return Err("Sample count must be positive".into()); + } let mut increase_accumulator = IncreaseAccumulator::new( starting_measurement, @@ -129,6 +133,7 @@ impl MultipleIncreaseAccumulator { last_seen_measurement, last_seen_timestamp, ); + increase_accumulator.sample_count = values.sample_count; increase_accumulator.counter_reset_adjustment = values.counter_reset_adjustment; increase_accumulator.counter_reset_events = values.counter_reset_events; @@ -173,6 +178,7 @@ impl MultipleIncreaseAccumulator { starting_timestamp: increase_acc.starting_timestamp, last_seen_measurement: increase_acc.last_seen_measurement.value, last_seen_timestamp: increase_acc.last_seen_timestamp, + sample_count: increase_acc.sample_count, counter_reset_adjustment: increase_acc.counter_reset_adjustment, counter_reset_events: increase_acc.counter_reset_events.clone(), }, @@ -292,6 +298,23 @@ impl AggregateCore for MultipleIncreaseAccumulator { .ok_or("Key required for MultipleIncreaseAccumulator")?; self.query(statistic, key_val, Some(query_kwargs)) } + + fn query_statistic_with_bounds( + &self, + statistic: Statistic, + key: &Option, + _query_kwargs: &std::collections::HashMap, + bounds: &QueryBounds, + ) -> Result> { + let key_val = key + .as_ref() + .ok_or("Key required for MultipleIncreaseAccumulator")?; + let data = self + .increases + .get(key_val) + .ok_or_else(|| format!("Key {key_val} not found in MultipleIncreaseAccumulator"))?; + data.query_with_bounds(statistic, bounds) + } } impl MultipleSubpopulationAggregate for MultipleIncreaseAccumulator { @@ -454,6 +477,7 @@ mod tests { let merged_key1 = merged.increases.get(&key1).unwrap(); assert_eq!(merged_key1.starting_measurement.value, 10.0); // Earlier start assert_eq!(merged_key1.last_seen_measurement.value, 30.0); // Later end + assert_eq!(merged_key1.sample_count, 2); } #[test] @@ -462,7 +486,16 @@ mod tests { let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); + acc.update( + key.clone(), + IncreaseAccumulator::new_with_sample_count( + Measurement::new(10.0), + 1_000, + Measurement::new(25.0), + 2_000, + 2, + ), + ); // Test JSON serialization let json_value = acc.serialize_to_json(); @@ -472,6 +505,7 @@ mod tests { let deserialized_acc = deserialized.increases.get(&key).unwrap(); assert_eq!(deserialized_acc.starting_measurement.value, 10.0); assert_eq!(deserialized_acc.last_seen_measurement.value, 25.0); + assert_eq!(deserialized_acc.sample_count, 2); // Test binary serialization let bytes = acc.serialize_to_bytes(); @@ -482,6 +516,61 @@ mod tests { let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 25.0); + assert_eq!(deserialized_acc_bytes.sample_count, 2); + + let arroyo_round_trip = MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo( + &acc.serialize_to_bytes_arroyo(), + ) + .unwrap(); + assert_eq!( + arroyo_round_trip.increases.get(&key).unwrap().sample_count, + 2 + ); + } + + #[test] + fn keyed_queries_use_each_key_sample_count() { + let key_with_enough_samples = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let key_with_one_sample = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + let mut acc = MultipleIncreaseAccumulator::new(); + acc.update( + key_with_enough_samples.clone(), + IncreaseAccumulator::new_with_sample_count( + Measurement::new(100.0), + 1_000, + Measurement::new(110.0), + 2_000, + 2, + ), + ); + acc.update( + key_with_one_sample.clone(), + IncreaseAccumulator::new( + Measurement::new(100.0), + 1_000, + Measurement::new(100.0), + 1_000, + ), + ); + + let bounds = crate::data_model::QueryBounds::new(1_000, 2_000); + let query_kwargs = HashMap::new(); + assert!(acc + .query_statistic_with_bounds( + Statistic::Rate, + &Some(key_with_enough_samples), + &query_kwargs, + &bounds, + ) + .is_ok()); + assert!(acc + .query_statistic_with_bounds( + Statistic::Rate, + &Some(key_with_one_sample), + &query_kwargs, + &bounds, + ) + .is_err()); } #[test] @@ -525,6 +614,7 @@ mod tests { starting_timestamp: 0, last_seen_measurement: 60.0, last_seen_timestamp: 3_000, + sample_count: 4, counter_reset_adjustment: 150.0, counter_reset_events: vec![CounterResetEvent { timestamp: 2_000, diff --git a/asap-query-engine/src/tests/native_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs index 238cd08..258d003 100644 --- a/asap-query-engine/src/tests/native_range_query_tests.rs +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -28,6 +28,7 @@ mod tests { use crate::engines::query_result::{QueryResult, RangeVectorElement}; use crate::engines::simple_engine::SimpleEngine; use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::IncreaseAccumulator; use crate::precompute_operators::{ CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, DeltaSetAggregatorAccumulator, SetAggregatorAccumulator, @@ -78,6 +79,76 @@ mod tests { .is_some_and(|e| e.samples.iter().any(|s| s.timestamp == ts)) } + #[test] + fn range_rate_uses_exact_boundaries_for_each_output_step() { + let host = Some(vec!["host-a".to_string()]); + let data = vec![ + ( + 1_000_000, + host.clone(), + Box::new(IncreaseAccumulator::new_with_sample_count( + crate::data_model::Measurement::new(100.0), + 999_500, + crate::data_model::Measurement::new(110.0), + 1_000_000, + 2, + )) as Box, + ), + ( + 1_001_000, + host.clone(), + Box::new(IncreaseAccumulator::new_with_sample_count( + crate::data_model::Measurement::new(110.0), + 1_000_000, + crate::data_model::Measurement::new(120.0), + 1_001_000, + 2, + )) as Box, + ), + ( + 1_002_000, + host, + Box::new(IncreaseAccumulator::new_with_sample_count( + crate::data_model::Measurement::new(120.0), + 1_001_000, + crate::data_model::Measurement::new(130.0), + 1_002_000, + 2, + )) as Box, + ), + ]; + let engine = create_engine_multi_timestamp_with_window( + "http_requests_total", + AggregationType::Increase, + vec!["host"], + data, + "rate(http_requests_total[2s])", + 1_000, + 1_000, + WindowType::Tumbling, + ); + + let (_, result) = engine + .handle_range_query_promql( + "rate(http_requests_total[2s])".to_string(), + 1_000.0, + 1_002.0, + 1.0, + ) + .expect("range rate query failed"); + let elements = matrix_values(result); + let samples = &elements + .iter() + .find(|element| element.labels.labels.contains(&"host-a".to_string())) + .expect("host-a result missing") + .samples; + + assert_eq!(samples.len(), 3); + assert!((samples[0].value - 7.5).abs() < f64::EPSILON); + assert!((samples[1].value - (40.0 / 3.0)).abs() < 1e-12); + assert!((samples[2].value - 10.0).abs() < f64::EPSILON); + } + /// Checks every `(label_values, ts, expected_present, reason)` case /// against `elements` and reports ALL mismatches in one panic, instead of /// stopping at the first failing `assert!` -- each of these tests makes From c3eedc97e276df6be7e2096b061efb64dea86173 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 1 Sep 2026 20:09:32 -0400 Subject: [PATCH 2/3] fix(query-engine): reject off-grid sliding counter bounds --- .../src/engines/simple_engine/mod.rs | 14 +++++ .../increase_accumulator.rs | 4 ++ .../src/tests/native_range_query_tests.rs | 57 +++++++++++++++++++ 3 files changed, 75 insertions(+) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index e583957..2de0ccf 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -2137,6 +2137,20 @@ impl SimpleEngine { // (#581). One loop shape for topk and non-topk alike, rather than // maintaining two. for ¤t_time in &context.output_timestamps { + if context.window_type == WindowType::Sliding + && context.tumbling_window_ms > 0 + && matches!( + context.base.metadata.statistic_to_compute, + Statistic::Increase | Statistic::Rate + ) + && !current_time.is_multiple_of(context.tumbling_window_ms) + { + return Err(format!( + "Exact Prometheus counter bounds are unavailable for off-grid Sliding \ + timestamp {} (grid interval {}ms)", + current_time, context.tumbling_window_ms + )); + } let current_time_i64 = i64::try_from(current_time) .map_err(|_| "Output timestamp exceeds signed timestamp range".to_string())?; let query_range_ms = i64::try_from(context.query_range_ms) diff --git a/asap-query-engine/src/precompute_operators/increase_accumulator.rs b/asap-query-engine/src/precompute_operators/increase_accumulator.rs index de10869..585932c 100644 --- a/asap-query-engine/src/precompute_operators/increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/increase_accumulator.rs @@ -442,6 +442,10 @@ impl MergeableAccumulator for IncreaseAccumulator { let mut result = accumulators.remove(0); for acc in accumulators { + // Query-time merges receive disjoint pane/window observations. + // Reset events are deduplicated for defensive overlap handling, + // but arbitrary duplicate samples are not identifiable from this + // summary shape, so sample counts remain additive by contract. result.sample_count = result .sample_count .checked_add(acc.sample_count) diff --git a/asap-query-engine/src/tests/native_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs index 258d003..2ea8d55 100644 --- a/asap-query-engine/src/tests/native_range_query_tests.rs +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -149,6 +149,63 @@ mod tests { assert!((samples[2].value - 10.0).abs() < f64::EPSILON); } + #[test] + fn sliding_counter_range_rejects_off_grid_output_timestamps() { + let host = Some(vec!["host-a".to_string()]); + let data = vec![ + ( + 1_000, + host.clone(), + Box::new(IncreaseAccumulator::new_with_sample_count( + crate::data_model::Measurement::new(100.0), + 0, + crate::data_model::Measurement::new(110.0), + 1_000, + 2, + )) as Box, + ), + ( + 2_000, + host.clone(), + Box::new(IncreaseAccumulator::new_with_sample_count( + crate::data_model::Measurement::new(110.0), + 1_000, + crate::data_model::Measurement::new(120.0), + 2_000, + 2, + )) as Box, + ), + ( + 3_000, + host, + Box::new(IncreaseAccumulator::new_with_sample_count( + crate::data_model::Measurement::new(120.0), + 2_000, + crate::data_model::Measurement::new(130.0), + 3_000, + 2, + )) as Box, + ), + ]; + let engine = create_engine_multi_timestamp_with_window( + "http_requests_total", + AggregationType::Increase, + vec!["host"], + data, + "rate(http_requests_total[2s])", + 2_000, + 1_000, + WindowType::Sliding, + ); + + // The stored Sliding windows are aligned to the 1s grid, so native + // execution cannot represent the exact [t-2s, t] range at t=1.5s. + // Returning None lets the caller use an exact Prometheus fallback. + assert!(engine + .handle_range_query_promql("rate(http_requests_total[2s])".to_string(), 1.5, 2.5, 1.0,) + .is_none()); + } + /// Checks every `(label_values, ts, expected_present, reason)` case /// against `elements` and reports ALL mismatches in one panic, instead of /// stopping at the first failing `assert!` -- each of these tests makes From a1c9654e318f3284f7ac95a4bb63de4b3eab9237 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 1 Sep 2026 20:50:38 -0400 Subject: [PATCH 3/3] fix(query-engine): validate sliding counter bounds early --- .../src/engines/simple_engine/mod.rs | 34 +++++++++------ .../src/tests/native_range_query_tests.rs | 43 +++++++++++++------ 2 files changed, 50 insertions(+), 27 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 2de0ccf..c142397 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1883,6 +1883,26 @@ impl SimpleEngine { use crate::engines::query_result::RangeVectorElement; use crate::engines::window_merger::create_window_merger; + if context.window_type == WindowType::Sliding + && context.tumbling_window_ms > 0 + && matches!( + context.base.metadata.statistic_to_compute, + Statistic::Increase | Statistic::Rate + ) + { + if let Some(&off_grid_timestamp) = context + .output_timestamps + .iter() + .find(|&×tamp| !timestamp.is_multiple_of(context.tumbling_window_ms)) + { + return Err(format!( + "Exact Prometheus counter bounds are unavailable for off-grid Sliding \ + timestamp {} (grid interval {}ms)", + off_grid_timestamp, context.tumbling_window_ms + )); + } + } + let lookback_ms = (context.lookback_bucket_count as u64) * context.tumbling_window_ms; // Step 1: Fetch all data needed for the entire range. Sliding @@ -2137,20 +2157,6 @@ impl SimpleEngine { // (#581). One loop shape for topk and non-topk alike, rather than // maintaining two. for ¤t_time in &context.output_timestamps { - if context.window_type == WindowType::Sliding - && context.tumbling_window_ms > 0 - && matches!( - context.base.metadata.statistic_to_compute, - Statistic::Increase | Statistic::Rate - ) - && !current_time.is_multiple_of(context.tumbling_window_ms) - { - return Err(format!( - "Exact Prometheus counter bounds are unavailable for off-grid Sliding \ - timestamp {} (grid interval {}ms)", - current_time, context.tumbling_window_ms - )); - } let current_time_i64 = i64::try_from(current_time) .map_err(|_| "Output timestamp exceeds signed timestamp range".to_string())?; let query_range_ms = i64::try_from(context.query_range_ms) diff --git a/asap-query-engine/src/tests/native_range_query_tests.rs b/asap-query-engine/src/tests/native_range_query_tests.rs index 2ea8d55..8f96856 100644 --- a/asap-query-engine/src/tests/native_range_query_tests.rs +++ b/asap-query-engine/src/tests/native_range_query_tests.rs @@ -154,35 +154,35 @@ mod tests { let host = Some(vec!["host-a".to_string()]); let data = vec![ ( - 1_000, + 10_000, host.clone(), Box::new(IncreaseAccumulator::new_with_sample_count( crate::data_model::Measurement::new(100.0), - 0, + 9_000, crate::data_model::Measurement::new(110.0), - 1_000, + 10_000, 2, )) as Box, ), ( - 2_000, + 11_000, host.clone(), Box::new(IncreaseAccumulator::new_with_sample_count( crate::data_model::Measurement::new(110.0), - 1_000, + 10_000, crate::data_model::Measurement::new(120.0), - 2_000, + 11_000, 2, )) as Box, ), ( - 3_000, + 12_000, host, Box::new(IncreaseAccumulator::new_with_sample_count( crate::data_model::Measurement::new(120.0), - 2_000, + 11_000, crate::data_model::Measurement::new(130.0), - 3_000, + 12_000, 2, )) as Box, ), @@ -199,11 +199,28 @@ mod tests { ); // The stored Sliding windows are aligned to the 1s grid, so native - // execution cannot represent the exact [t-2s, t] range at t=1.5s. + // execution cannot represent the exact [t-2s, t] range at t=10.5s. // Returning None lets the caller use an exact Prometheus fallback. - assert!(engine - .handle_range_query_promql("rate(http_requests_total[2s])".to_string(), 1.5, 2.5, 1.0,) - .is_none()); + assert!( + engine + .handle_range_query_promql( + "rate(http_requests_total[2s])".to_string(), + 10.5, + 12.5, + 1.0, + ) + .is_none() + ); + + let (_, on_grid_result) = engine + .handle_range_query_promql("rate(http_requests_total[2s])".to_string(), 11.0, 12.0, 1.0) + .expect("on-grid Sliding counter query should use native execution"); + let on_grid_samples = matrix_values(on_grid_result) + .into_iter() + .find(|element| element.labels.labels.contains(&"host-a".to_string())) + .expect("host-a result missing") + .samples; + assert_eq!(on_grid_samples.len(), 2); } /// Checks every `(label_values, ts, expected_present, reason)` case