diff --git a/asap-common/dependencies/rs/datafusion_summary_library/src/sketch_operators.rs b/asap-common/dependencies/rs/datafusion_summary_library/src/sketch_operators.rs index 64fe3946..7d51333d 100644 --- a/asap-common/dependencies/rs/datafusion_summary_library/src/sketch_operators.rs +++ b/asap-common/dependencies/rs/datafusion_summary_library/src/sketch_operators.rs @@ -180,11 +180,21 @@ pub enum InferOperation { /// Extract maximum value from MinMax/HydraMinMax accumulator ExtractMax, - /// Extract increase value from Increase/HydraIncrease accumulator - ExtractIncrease, - - /// Extract rate (increase / time_range) from Increase accumulator - ExtractRate, + /// Extract increase value from Increase/HydraIncrease accumulator. + /// + /// The optional payload carries the range-vector boundaries + /// `(range_start_ms, range_end_ms)` of the instant query being evaluated. + /// When present, the accumulator reproduces Prometheus `extrapolatedRate` + /// (counter-reset correction + extrapolation); when `None` (callers without a + /// range vector, e.g. SQL/Elastic), it falls back to the plain reset-corrected + /// increase. Embedded here so it survives the logical→physical boundary, the + /// same way `Quantile`/`TopK` carry their parameters. + ExtractIncrease(Option<(i64, i64)>), + + /// Extract rate (increase / time_range) from Increase accumulator. Carries the + /// same optional `(range_start_ms, range_end_ms)` payload as + /// [`InferOperation::ExtractIncrease`]. + ExtractRate(Option<(i64, i64)>), // ======================================================================== // Sketch operations @@ -262,8 +272,8 @@ impl fmt::Display for InferOperation { InferOperation::ExtractCount => write!(f, "EXTRACT_COUNT"), InferOperation::ExtractMin => write!(f, "EXTRACT_MIN"), InferOperation::ExtractMax => write!(f, "EXTRACT_MAX"), - InferOperation::ExtractIncrease => write!(f, "EXTRACT_INCREASE"), - InferOperation::ExtractRate => write!(f, "EXTRACT_RATE"), + InferOperation::ExtractIncrease(_) => write!(f, "EXTRACT_INCREASE"), + InferOperation::ExtractRate(_) => write!(f, "EXTRACT_RATE"), // Sketch operations InferOperation::CountDistinct => write!(f, "COUNT_DISTINCT"), InferOperation::Quantile(p) => write!(f, "QUANTILE({:.4})", *p as f64 / 10000.0), @@ -1112,8 +1122,8 @@ impl SummaryInfer { InferOperation::ExtractCount => DataType::Float64, InferOperation::ExtractMin => DataType::Float64, InferOperation::ExtractMax => DataType::Float64, - InferOperation::ExtractIncrease => DataType::Float64, - InferOperation::ExtractRate => DataType::Float64, + InferOperation::ExtractIncrease(_) => DataType::Float64, + InferOperation::ExtractRate(_) => DataType::Float64, // Sketch operations InferOperation::CountDistinct => DataType::UInt64, InferOperation::Quantile(_) | InferOperation::Median => DataType::Float64, diff --git a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs index 04110285..c10877ae 100644 --- a/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs +++ b/asap-common/dependencies/rs/promql_utilities/src/query_logics/enums.rs @@ -3,6 +3,16 @@ use std::fmt; use std::str::FromStr; use tracing::debug; +/// `query_kwargs` key carrying the range-vector start boundary (epoch ms) for a +/// `rate()`/`increase()` evaluation. Used to drive Prometheus-style extrapolation +/// in `IncreaseAccumulator::query`. When absent, the accumulator falls back to a +/// non-extrapolated, reset-corrected result. +pub const RANGE_START_MS_KWARG: &str = "range_start_ms"; + +/// `query_kwargs` key carrying the range-vector end boundary (epoch ms) for a +/// `rate()`/`increase()` evaluation. See [`RANGE_START_MS_KWARG`]. +pub const RANGE_END_MS_KWARG: &str = "range_end_ms"; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum QueryTreatmentType { Exact, diff --git a/asap-query-engine/src/engines/logical/plan_builder.rs b/asap-query-engine/src/engines/logical/plan_builder.rs index 813f0241..7f20a82d 100644 --- a/asap-query-engine/src/engines/logical/plan_builder.rs +++ b/asap-query-engine/src/engines/logical/plan_builder.rs @@ -15,7 +15,9 @@ use datafusion_summary_library::{ InferOperation, PrecomputedSummaryRead, SketchType, SummaryInfer, SummaryMergeMultiple, }; use promql_parser::parser::token::{self, T_ADD, T_DIV, T_MOD, T_MUL, T_POW, T_SUB}; -use promql_utilities::query_logics::enums::{AggregationType, Statistic}; +use promql_utilities::query_logics::enums::{ + AggregationType, Statistic, RANGE_END_MS_KWARG, RANGE_START_MS_KWARG, +}; use std::sync::Arc; use crate::engines::simple_engine::{QueryExecutionContext, StoreQueryParams}; @@ -189,8 +191,8 @@ impl QueryExecutionContext { Statistic::Min => Ok(InferOperation::ExtractMin), Statistic::Max => Ok(InferOperation::ExtractMax), Statistic::Count => Ok(InferOperation::ExtractCount), - Statistic::Increase => Ok(InferOperation::ExtractIncrease), - Statistic::Rate => Ok(InferOperation::ExtractRate), + Statistic::Increase => Ok(InferOperation::ExtractIncrease(self.range_bounds_ms())), + Statistic::Rate => Ok(InferOperation::ExtractRate(self.range_bounds_ms())), Statistic::Quantile => { // Extract quantile parameter from query_kwargs let q = self @@ -215,6 +217,30 @@ impl QueryExecutionContext { } } + /// Parse the range-vector boundaries `(range_start_ms, range_end_ms)` from the + /// query kwargs. + /// + /// The PromQL layer populates these for rate()/increase() instant queries via + /// `build_query_kwargs_promql`; embedding them in the `InferOperation` carries + /// them across the logical→physical boundary so `IncreaseAccumulator` can + /// extrapolate. Returns `None` when absent (non-range callers), leaving the + /// accumulator on its reset-corrected fallback. + fn range_bounds_ms(&self) -> Option<(i64, i64)> { + let start = self + .metadata + .query_kwargs + .get(RANGE_START_MS_KWARG)? + .parse::() + .ok()?; + let end = self + .metadata + .query_kwargs + .get(RANGE_END_MS_KWARG)? + .parse::() + .ok()?; + Some((start, end)) + } + /// Map aggregation type to SketchType (SummaryType) for the value accumulator fn map_aggregation_type_to_summary_type(&self) -> Result { Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_value) diff --git a/asap-query-engine/src/engines/physical/accumulator_serde.rs b/asap-query-engine/src/engines/physical/accumulator_serde.rs index 6b7e480e..242a62b8 100644 --- a/asap-query-engine/src/engines/physical/accumulator_serde.rs +++ b/asap-query-engine/src/engines/physical/accumulator_serde.rs @@ -12,8 +12,8 @@ use datafusion_summary_library::SketchType; use crate::data_model::{MultipleSubpopulationAggregate, SingleSubpopulationAggregate}; use crate::precompute_operators::{ CountMinSketchAccumulator, DatasketchesKLLAccumulator, DeltaSetAggregatorAccumulator, - HllAccumulator, HydraKllSketchAccumulator, MultipleIncreaseAccumulator, MultipleSumAccumulator, - SetAggregatorAccumulator, SumAccumulator, + HllAccumulator, HydraKllSketchAccumulator, IncreaseAccumulator, MultipleIncreaseAccumulator, + MultipleSumAccumulator, SetAggregatorAccumulator, SumAccumulator, }; use crate::AggregateCore; @@ -40,9 +40,12 @@ pub fn deserialize_accumulator( })?; Ok(Box::new(acc)) } - SketchType::Increase => Err(DataFusionError::NotImplemented( - "Increase Arroyo deserialization not implemented".to_string(), - )), + SketchType::Increase => { + let acc = IncreaseAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { + DataFusionError::Internal(format!("Failed to deserialize Increase: {}", e)) + })?; + Ok(Box::new(acc)) + } SketchType::MinMax => Err(DataFusionError::NotImplemented( "MinMax Arroyo deserialization not implemented".to_string(), )), @@ -151,6 +154,9 @@ pub fn serialize_accumulator_arroyo(acc: &dyn AggregateCore) -> Vec { if let Some(set_acc) = acc.as_any().downcast_ref::() { return set_acc.serialize_to_bytes_arroyo(); } + if let Some(inc_acc) = acc.as_any().downcast_ref::() { + return inc_acc.serialize_to_bytes_arroyo(); + } if let Some(inc_acc) = acc.as_any().downcast_ref::() { return inc_acc.serialize_to_bytes_arroyo(); } @@ -178,9 +184,12 @@ pub fn deserialize_single_subpopulation( })?; Ok(Box::new(acc)) } - SketchType::Increase => Err(DataFusionError::NotImplemented( - "Increase Arroyo deserialization not implemented".to_string(), - )), + SketchType::Increase => { + let acc = IncreaseAccumulator::deserialize_from_bytes_arroyo(bytes).map_err(|e| { + DataFusionError::Internal(format!("Failed to deserialize Increase: {}", e)) + })?; + Ok(Box::new(acc)) + } SketchType::MinMax => Err(DataFusionError::NotImplemented( "MinMax Arroyo deserialization not implemented".to_string(), )), diff --git a/asap-query-engine/src/engines/physical/summary_infer_exec.rs b/asap-query-engine/src/engines/physical/summary_infer_exec.rs index 52885528..1825685b 100644 --- a/asap-query-engine/src/engines/physical/summary_infer_exec.rs +++ b/asap-query-engine/src/engines/physical/summary_infer_exec.rs @@ -24,7 +24,7 @@ use datafusion::physical_plan::{ }; use datafusion_summary_library::{InferOperation, SketchType, SummaryInfer}; use futures::stream; -use promql_utilities::query_logics::enums::Statistic; +use promql_utilities::query_logics::enums::{Statistic, RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; use std::any::Any; use std::collections::HashMap; use std::fmt; @@ -150,8 +150,8 @@ impl SummaryInferExec { InferOperation::ExtractCount => Statistic::Count, InferOperation::ExtractMin => Statistic::Min, InferOperation::ExtractMax => Statistic::Max, - InferOperation::ExtractIncrease => Statistic::Increase, - InferOperation::ExtractRate => Statistic::Rate, + InferOperation::ExtractIncrease(_) => Statistic::Increase, + InferOperation::ExtractRate(_) => Statistic::Rate, InferOperation::CountDistinct => Statistic::Cardinality, InferOperation::Quantile(_) | InferOperation::Median => Statistic::Quantile, InferOperation::TopK(_) => Statistic::Topk, @@ -164,6 +164,12 @@ impl SummaryInferExec { /// /// Some operations embed parameters (e.g. Quantile embeds the quantile value, /// TopK embeds k) that accumulators need via the kwargs HashMap. + /// + /// For `ExtractRate`/`ExtractIncrease` the embedded `(range_start_ms, + /// range_end_ms)` boundaries are re-materialized into + /// [`RANGE_START_MS_KWARG`]/[`RANGE_END_MS_KWARG`] so `IncreaseAccumulator` + /// applies Prometheus extrapolation instead of the reset-corrected fallback. + /// Without a payload (`None`), no kwargs are produced and the fallback stands. fn infer_op_to_kwargs(op: &InferOperation) -> Option> { match op { InferOperation::Quantile(q_u16) => { @@ -182,6 +188,13 @@ impl SummaryInferExec { kwargs.insert("k".to_string(), k.to_string()); Some(kwargs) } + InferOperation::ExtractIncrease(Some((range_start_ms, range_end_ms))) + | InferOperation::ExtractRate(Some((range_start_ms, range_end_ms))) => { + let mut kwargs = HashMap::new(); + kwargs.insert(RANGE_START_MS_KWARG.to_string(), range_start_ms.to_string()); + kwargs.insert(RANGE_END_MS_KWARG.to_string(), range_end_ms.to_string()); + Some(kwargs) + } _ => None, } } diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index 501f7218..ae4d8b2e 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -24,7 +24,8 @@ use promql_utilities::ast_matching::{PromQLPattern, PromQLPatternBuilder}; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::get_is_collapsable; use promql_utilities::query_logics::enums::{ - AggregationOperator, AggregationType, PromQLFunction, Statistic, + AggregationOperator, AggregationType, PromQLFunction, Statistic, RANGE_END_MS_KWARG, + RANGE_START_MS_KWARG, }; use serde_json::Value; @@ -1501,12 +1502,19 @@ impl SimpleEngine { key, start_ms, end_ms, step_ms, lookback_ms, tumbling_window_ms ); + // Per-step query kwargs: clone the base map once, then refresh the + // range-vector boundaries each step (rate/increase extrapolation needs + // [window_start, current_time) per emitted point). + let mut step_kwargs = context.base.metadata.query_kwargs.clone(); + // Iterate by OUTPUT timestamp, not by bucket index let mut current_time = start_ms; while current_time <= end_ms { // Window covers [current_time - lookback_ms, current_time) // This means we look at buckets that START within this range let window_start = current_time.saturating_sub(lookback_ms); + step_kwargs.insert(RANGE_START_MS_KWARG.to_string(), window_start.to_string()); + step_kwargs.insert(RANGE_END_MS_KWARG.to_string(), current_time.to_string()); // Collect all AVAILABLE buckets in this window (skip missing ones) let mut window_buckets: Vec> = Vec::new(); @@ -1532,7 +1540,7 @@ impl SimpleEngine { merged.as_ref(), &context.base.metadata.statistic_to_compute, &Some(key.clone()), - &context.base.metadata.query_kwargs, + &step_kwargs, ) { Ok(value) => { debug!( diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index a9a2d546..1ead33f1 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -14,7 +14,7 @@ use asap_types::query_requirements::build_query_requirements_promql; use asap_types::PromQLSchema; use promql_utilities::ast_matching::PromQLMatchResult; use promql_utilities::data_model::KeyByLabelNames; -use promql_utilities::query_logics::enums::Statistic; +use promql_utilities::query_logics::enums::{Statistic, RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; use promql_utilities::query_logics::parsing::get_metric_and_spatial_filter; use std::collections::HashMap; use std::time::Instant; @@ -129,6 +129,7 @@ impl SimpleEngine { &self, statistic: &Statistic, match_result: &PromQLMatchResult, + timestamps: &QueryTimestamps, ) -> Result, String> { let mut query_kwargs = HashMap::new(); @@ -145,6 +146,19 @@ impl SimpleEngine { debug!("Extracted k value: {:?}", k); query_kwargs.insert("k".to_string(), k); } + Statistic::Rate | Statistic::Increase => { + // Pass the range-vector boundaries so IncreaseAccumulator can apply + // Prometheus extrapolation. For an instant `rate(x[5m])` evaluated + // at T, `timestamps` spans [T - 5m, T]. + query_kwargs.insert( + RANGE_START_MS_KWARG.to_string(), + timestamps.start_timestamp.to_string(), + ); + query_kwargs.insert( + RANGE_END_MS_KWARG.to_string(), + timestamps.end_timestamp.to_string(), + ); + } _ => {} } @@ -286,7 +300,7 @@ impl SimpleEngine { } let query_kwargs = self - .build_query_kwargs_promql(&statistic_to_compute, match_result) + .build_query_kwargs_promql(&statistic_to_compute, match_result, ×tamps) .map_err(|e| { warn!("{}", e); e diff --git a/asap-query-engine/src/precompute_operators/increase_accumulator.rs b/asap-query-engine/src/precompute_operators/increase_accumulator.rs index 6c8c0293..ffbcdc97 100644 --- a/asap-query-engine/src/precompute_operators/increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/increase_accumulator.rs @@ -6,19 +6,62 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; -use promql_utilities::query_logics::enums::Statistic; +use promql_utilities::query_logics::enums::{Statistic, RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; + +/// Sample count assumed when deserializing a blob that predates the +/// `sample_count` field. Two (the minimum for a defined increase) keeps +/// extrapolation well-defined instead of dividing by `sample_count - 1 == 0`. +const DEFAULT_LEGACY_SAMPLE_COUNT: u64 = 2; + +/// MessagePack ("Arroyo") wire form of an [`IncreaseAccumulator`]: the window +/// endpoints plus the reset correction and sample count, with the measurements +/// flattened to `f64`. +/// +/// The field order/types intentionally mirror the per-key `MeasurementData` +/// struct inside `MultipleIncreaseAccumulator`, so a single-population Increase +/// blob is byte-identical to one MultipleIncrease entry (rmp-serde encodes +/// structs as positional arrays). `#[serde(default)]` on the trailing fields +/// lets legacy 4-field blobs (from before reset support) still decode. +#[derive(Serialize, Deserialize)] +struct ArroyoIncrease { + starting_measurement: f64, + starting_timestamp: i64, + last_seen_measurement: f64, + last_seen_timestamp: i64, + #[serde(default)] + counter_reset_correction: f64, + #[serde(default)] + sample_count: u64, +} -/// Accumulator for tracking increases in counter metrics -/// Stores the starting and last seen measurements with timestamps +/// Accumulator for tracking increases in counter metrics. +/// +/// Stores the starting and last-seen measurements with timestamps, plus the +/// running counter-reset correction and the number of samples observed. The +/// extra state is what lets `query` reproduce Prometheus `rate()`/`increase()` +/// semantics (counter-reset correction + extrapolation) even though only the +/// window endpoints are retained — see [`IncreaseAccumulator::update`] and +/// [`IncreaseAccumulator::extrapolated`]. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct IncreaseAccumulator { pub starting_measurement: Measurement, pub starting_timestamp: i64, pub last_seen_measurement: Measurement, pub last_seen_timestamp: i64, + /// Sum of the values observed just before each counter reset within this + /// window (Prometheus `counterCorrection`). Added back to `last - start` so + /// the increase stays monotonic across restarts. + pub counter_reset_correction: f64, + /// Number of samples folded into this accumulator. Needed to estimate the + /// average inter-sample interval during extrapolation. + pub sample_count: u64, } impl IncreaseAccumulator { + /// Construct from the first observed sample (`starting == last_seen`), with a + /// zero reset correction and a sample count of 1. The per-sample updaters in + /// `accumulator_factory` create via this constructor and then call + /// [`update`](Self::update) for every subsequent sample. pub fn new( starting_measurement: Measurement, starting_timestamp: i64, @@ -30,12 +73,159 @@ impl IncreaseAccumulator { starting_timestamp, last_seen_measurement, last_seen_timestamp, + counter_reset_correction: 0.0, + sample_count: 1, + } + } + + /// Construct with every field explicit. Used only by deserializers, which + /// must restore the persisted correction and sample count rather than the + /// first-sample defaults of [`new`](Self::new). + pub fn new_full( + starting_measurement: Measurement, + starting_timestamp: i64, + last_seen_measurement: Measurement, + last_seen_timestamp: i64, + counter_reset_correction: f64, + sample_count: u64, + ) -> Self { + Self { + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + counter_reset_correction, + sample_count, } } pub fn update(&mut self, measurement: Measurement, timestamp: i64) { + // Counter-reset detection: a sample lower than its predecessor means the + // counter restarted, so add the pre-reset value back (Prometheus + // `counterCorrection`). This must happen here, per sample — only the + // window endpoints survive, so a reset between them is unrecoverable at + // query time. + if measurement.value < self.last_seen_measurement.value { + self.counter_reset_correction += self.last_seen_measurement.value; + } self.last_seen_measurement = measurement; self.last_seen_timestamp = timestamp; + self.sample_count += 1; + } + + /// Reset-corrected raw increase over the observed samples: `last - start` + /// plus the accumulated counter-reset correction. + pub fn corrected_increase(&self) -> f64 { + self.last_seen_measurement.value - self.starting_measurement.value + + self.counter_reset_correction + } + + /// Prometheus `extrapolatedRate`: extrapolate the reset-corrected increase to + /// the range-vector boundaries `[range_start_ms, range_end_ms]`. When + /// `is_rate` is true the result is divided by the range duration (per-second + /// rate); otherwise it is the extrapolated increase. + /// + /// This mirrors `extrapolatedRate` in Prometheus `promql/functions.go` + /// (`isCounter == true`, the rate/increase case) step for step, including the + /// order in which `durationToStart` is clamped (threshold first, then the + /// counter zero-point clamp) and the `factor = 1.0` fallback when the sampled + /// interval is zero. ASAP does not track per-sample created timestamps, so the + /// start-timestamp / anchored / smoothed branches are not reproduced. + /// + /// Returns `None` when fewer than two samples were observed (Prometheus emits + /// no point in that case, absent a created-timestamp reset) or when the query + /// range is degenerate. + pub fn extrapolated( + &self, + is_rate: bool, + range_start_ms: i64, + range_end_ms: i64, + ) -> Option { + // Without two samples (and lacking created-timestamp tracking), Prometheus + // returns no point. + if self.sample_count < 2 { + return None; + } + let num_samples_minus_one = (self.sample_count - 1) as f64; + + let first_value = self.starting_measurement.value; + let result_value = self.corrected_increase(); + + // Duration between first/last samples and the boundary of the range. + let mut duration_to_start = (self.starting_timestamp - range_start_ms) as f64 / 1000.0; + let mut duration_to_end = (range_end_ms - self.last_seen_timestamp) as f64 / 1000.0; + + let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; + let average_duration_between_samples = sampled_interval / num_samples_minus_one; + let extrapolation_threshold = average_duration_between_samples * 1.1; + + // If the first sample is close enough to the lower boundary, extrapolate to + // it; otherwise only extrapolate half an average gap. + if duration_to_start >= extrapolation_threshold { + duration_to_start = average_duration_between_samples / 2.0; + } + // Counters cannot be negative: if the series has a positive slope, never + // extrapolate the start earlier than where the counter would have been 0. + // (Prometheus applies this AFTER the threshold clamp above.) + let mut duration_to_zero = duration_to_start; + if result_value > 0.0 && first_value >= 0.0 { + duration_to_zero = sampled_interval * (first_value / result_value); + } + if duration_to_zero < duration_to_start { + duration_to_start = duration_to_zero; + } + + if duration_to_end >= extrapolation_threshold { + duration_to_end = average_duration_between_samples / 2.0; + } + + let mut factor = 1.0; + if sampled_interval != 0.0 { + factor = (sampled_interval + duration_to_start + duration_to_end) / sampled_interval; + } + if is_rate { + let range_seconds = (range_end_ms - range_start_ms) as f64 / 1000.0; + if range_seconds <= 0.0 { + return None; + } + factor /= range_seconds; + } + Some(result_value * factor) + } + + /// Order-independent pairwise merge of two windows that together form a + /// contiguous (tumbling) range. Keeps the earliest start and latest + /// last-seen, sums both corrections and sample counts, and adds a *boundary* + /// counter-reset correction when the later window's first value is below the + /// earlier window's last value (a reset across the window seam). Folding this + /// left-to-right over time-ordered windows is correct because the running + /// result always carries the latest segment's `last_seen`. + /// + /// Assumes non-overlapping windows (the only kind the engine produces — see + /// the tumbling-window merge loop in `simple_engine`), so the window with the + /// earlier start also has the earlier `last_seen`. + pub fn merge_pair(a: &IncreaseAccumulator, b: &IncreaseAccumulator) -> IncreaseAccumulator { + let (earlier, later) = if a.starting_timestamp <= b.starting_timestamp { + (a, b) + } else { + (b, a) + }; + + let boundary_correction = + if later.starting_measurement.value < earlier.last_seen_measurement.value { + earlier.last_seen_measurement.value + } else { + 0.0 + }; + + IncreaseAccumulator::new_full( + earlier.starting_measurement.clone(), + earlier.starting_timestamp, + later.last_seen_measurement.clone(), + later.last_seen_timestamp, + earlier.counter_reset_correction + later.counter_reset_correction + boundary_correction, + earlier.sample_count + later.sample_count, + ) } pub fn deserialize_from_json(data: &Value) -> Result> { @@ -49,12 +239,19 @@ impl IncreaseAccumulator { let last_seen_timestamp = data["last_seen_timestamp"] .as_i64() .ok_or("Missing or invalid 'last_seen_timestamp' field")?; + // Tolerate legacy blobs written before these fields existed. + let counter_reset_correction = data["counter_reset_correction"].as_f64().unwrap_or(0.0); + let sample_count = data["sample_count"] + .as_u64() + .unwrap_or(DEFAULT_LEGACY_SAMPLE_COUNT); - Ok(Self::new( + Ok(Self::new_full( starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, + counter_reset_correction, + sample_count, )) } @@ -131,12 +328,91 @@ impl IncreaseAccumulator { buffer[offset + 6], buffer[offset + 7], ]); + offset += 8; - Ok(Self::new( + // Counter-reset correction (f64) and sample count (u64) were appended + // after the original four fields. Tolerate legacy buffers that lack them. + let counter_reset_correction = if buffer.len() >= offset + 8 { + let v = f64::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; + v + } else { + 0.0 + }; + let sample_count = if buffer.len() >= offset + 8 { + 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], + ]) + } else { + DEFAULT_LEGACY_SAMPLE_COUNT + }; + + Ok(Self::new_full( starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, + counter_reset_correction, + sample_count, + )) + } + + /// Serialize to the Arroyo-compatible MessagePack form (see [`ArroyoIncrease`]). + /// + /// This is the format the DataFusion read path emits for single-population + /// Increase sketches (via `serialize_accumulator_arroyo`), so downstream + /// operators can reconstruct the accumulator with + /// [`deserialize_from_bytes_arroyo`](Self::deserialize_from_bytes_arroyo). + pub fn serialize_to_bytes_arroyo(&self) -> Vec { + let wire = ArroyoIncrease { + starting_measurement: self.starting_measurement.value, + starting_timestamp: self.starting_timestamp, + last_seen_measurement: self.last_seen_measurement.value, + last_seen_timestamp: self.last_seen_timestamp, + counter_reset_correction: self.counter_reset_correction, + sample_count: self.sample_count, + }; + rmp_serde::to_vec(&wire).expect("Failed to serialize IncreaseAccumulator to MessagePack") + } + + /// Deserialize from the Arroyo-compatible MessagePack form. Mirrors + /// `MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo`, including the + /// legacy handling where an absent `sample_count` (decoded as 0) falls back to + /// [`DEFAULT_LEGACY_SAMPLE_COUNT`] so extrapolation stays well-defined. + pub fn deserialize_from_bytes_arroyo( + buffer: &[u8], + ) -> Result> { + let wire: ArroyoIncrease = rmp_serde::from_slice(buffer).map_err(|e| { + format!("Failed to deserialize IncreaseAccumulator from MessagePack: {e}") + })?; + let sample_count = if wire.sample_count == 0 { + DEFAULT_LEGACY_SAMPLE_COUNT + } else { + wire.sample_count + }; + Ok(Self::new_full( + Measurement::new(wire.starting_measurement), + wire.starting_timestamp, + Measurement::new(wire.last_seen_measurement), + wire.last_seen_timestamp, + wire.counter_reset_correction, + sample_count, )) } } @@ -148,6 +424,8 @@ 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, + "counter_reset_correction": self.counter_reset_correction, + "sample_count": self.sample_count, }) } @@ -171,6 +449,11 @@ impl SerializableToSink for IncreaseAccumulator { // Last seen timestamp buffer.extend_from_slice(&self.last_seen_timestamp.to_le_bytes()); + // Counter-reset correction (f64) and sample count (u64), appended last so + // legacy readers/decoders that stop after the four original fields still work. + buffer.extend_from_slice(&self.counter_reset_correction.to_le_bytes()); + buffer.extend_from_slice(&self.sample_count.to_le_bytes()); + buffer } } @@ -183,20 +466,14 @@ impl MergeableAccumulator for IncreaseAccumulator { return Err("No accumulators to merge".into()); } - let mut result = accumulators[0].clone(); + // Sort by start time so boundary counter-reset corrections are evaluated + // between temporally adjacent windows, then fold left-to-right. + let mut sorted = accumulators; + sorted.sort_by_key(|a| a.starting_timestamp); - for acc in &accumulators[1..] { - // Use the earlier starting point - if acc.starting_timestamp < result.starting_timestamp { - result.starting_measurement = acc.starting_measurement.clone(); - result.starting_timestamp = acc.starting_timestamp; - } - - // Use the later last seen point - if acc.last_seen_timestamp > result.last_seen_timestamp { - result.last_seen_measurement = acc.last_seen_measurement.clone(); - result.last_seen_timestamp = acc.last_seen_timestamp; - } + let mut result = sorted[0].clone(); + for acc in &sorted[1..] { + result = IncreaseAccumulator::merge_pair(&result, acc); } Ok(result) @@ -235,8 +512,8 @@ impl AggregateCore for IncreaseAccumulator { .downcast_ref::() .ok_or("Failed to downcast to IncreaseAccumulator")?; - // Use the existing merge_accumulators method - let merged = Self::merge_accumulators(vec![self.clone(), other_increase.clone()])?; + // Boundary-aware pairwise merge (order-independent). + let merged = Self::merge_pair(self, other_increase); Ok(Box::new(merged)) } @@ -253,10 +530,11 @@ impl AggregateCore for IncreaseAccumulator { &self, statistic: promql_utilities::query_logics::enums::Statistic, _key: &Option, - _query_kwargs: &std::collections::HashMap, + query_kwargs: &std::collections::HashMap, ) -> Result> { use crate::data_model::SingleSubpopulationAggregate; - self.query(statistic, None) + // Forward the kwargs so range-vector boundaries reach `query`. + self.query(statistic, Some(query_kwargs)) } } @@ -266,25 +544,43 @@ impl SingleSubpopulationAggregate for IncreaseAccumulator { statistic: Statistic, query_kwargs: Option<&HashMap>, ) -> Result> { - // IncreaseAccumulator doesn't use query_kwargs, assert it's None - if query_kwargs.is_some() { - return Err("IncreaseAccumulator does not support query parameters".into()); + let is_rate = match statistic { + Statistic::Increase => false, + Statistic::Rate => true, + _ => { + return Err( + format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into(), + ) + } + }; + + // When the PromQL layer supplies the range-vector boundaries, reproduce + // Prometheus `extrapolatedRate` (counter-reset correction + extrapolation). + let range_bounds = query_kwargs.and_then(|kw| { + let start = kw.get(RANGE_START_MS_KWARG)?.parse::().ok()?; + let end = kw.get(RANGE_END_MS_KWARG)?.parse::().ok()?; + Some((start, end)) + }); + + if let Some((range_start_ms, range_end_ms)) = range_bounds { + return self + .extrapolated(is_rate, range_start_ms, range_end_ms) + .ok_or_else(|| -> Box { + "Insufficient samples for rate/increase extrapolation".into() + }); } - match statistic { - Statistic::Increase => { - Ok(self.last_seen_measurement.value - self.starting_measurement.value) - } - Statistic::Rate => { - // Convert to per second; timestamps are in milliseconds - let time_diff = (self.last_seen_timestamp - self.starting_timestamp) as f64; - if time_diff <= 0.0 { - return Err("Invalid time difference for rate calculation".into()); - } - let value_diff = self.last_seen_measurement.value - self.starting_measurement.value; - Ok(value_diff / time_diff * 1000.0) + // Fallback (e.g. SQL/Elastic callers with no range vector): reset-corrected + // increase, with rate divided by the sampled interval as before. + let increase = self.corrected_increase(); + if is_rate { + let time_diff = (self.last_seen_timestamp - self.starting_timestamp) as f64; + if time_diff <= 0.0 { + return Err("Invalid time difference for rate calculation".into()); } - _ => Err(format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into()), + Ok(increase / time_diff * 1000.0) + } else { + Ok(increase) } } @@ -479,4 +775,252 @@ mod tests { assert_eq!(acc.type_name(), "IncreaseAccumulator"); } + + use promql_utilities::query_logics::enums::{RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; + + /// Feed a sample stream into an accumulator the way the ingest updater does: + /// `new` on the first sample, then `update` for each subsequent one. + fn feed(samples: &[(i64, f64)]) -> IncreaseAccumulator { + let (ts0, v0) = samples[0]; + let mut acc = + IncreaseAccumulator::new(Measurement::new(v0), ts0, Measurement::new(v0), ts0); + for &(ts, v) in &samples[1..] { + acc.update(Measurement::new(v), ts); + } + acc + } + + #[test] + fn test_reset_within_window_tracks_correction() { + // 10 -> 100 -> 5 (reset, +100) -> 30. Increase = 30 - 10 + 100 = 120. + let acc = feed(&[(1000, 10.0), (2000, 100.0), (3000, 5.0), (4000, 30.0)]); + assert_eq!(acc.counter_reset_correction, 100.0); + assert_eq!(acc.sample_count, 4); + assert_eq!(acc.corrected_increase(), 120.0); + // Fallback query (no range kwargs) returns the reset-corrected increase. + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, None).unwrap(), + 120.0 + ); + } + + #[test] + fn test_multiple_resets_within_window() { + // 10 ->50 ->5(+50) ->40 ->2(+40) ->20. Increase = 20 - 10 + 50 + 40 = 100. + let acc = feed(&[ + (1000, 10.0), + (2000, 50.0), + (3000, 5.0), + (4000, 40.0), + (5000, 2.0), + (6000, 20.0), + ]); + assert_eq!(acc.counter_reset_correction, 90.0); + assert_eq!(acc.corrected_increase(), 100.0); + } + + #[test] + fn test_fallback_rate_uses_corrected_increase() { + // Reset-corrected increase 120 over 3000 ms => 40/s. + let acc = feed(&[(1000, 10.0), (2000, 100.0), (3000, 5.0), (4000, 30.0)]); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, None).unwrap(), + 40.0 + ); + } + + #[test] + fn test_merge_boundary_reset_correction() { + // Window A: 10 -> 20. Window B: 5 -> 30 (B starts below A's last => reset). + // Increase = 30 - 10 + boundary(20) = 40. + let a = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000); + let b = IncreaseAccumulator::new(Measurement::new(5.0), 3000, Measurement::new(30.0), 4000); + + let merged = IncreaseAccumulator::merge_pair(&a, &b); + assert_eq!(merged.starting_measurement.value, 10.0); + assert_eq!(merged.last_seen_measurement.value, 30.0); + assert_eq!(merged.counter_reset_correction, 20.0); + assert_eq!(merged.sample_count, 2); + assert_eq!(merged.corrected_increase(), 40.0); + + // Order-independent: merging B,A yields the same result. + let merged_rev = IncreaseAccumulator::merge_pair(&b, &a); + assert_eq!(merged_rev.corrected_increase(), 40.0); + assert_eq!(merged_rev.starting_measurement.value, 10.0); + assert_eq!(merged_rev.last_seen_measurement.value, 30.0); + + // No reset at the boundary (B starts >= A's last) => no boundary correction. + let c = + IncreaseAccumulator::new(Measurement::new(25.0), 3000, Measurement::new(40.0), 4000); + let merged_no_reset = IncreaseAccumulator::merge_pair(&a, &c); + assert_eq!(merged_no_reset.counter_reset_correction, 0.0); + assert_eq!(merged_no_reset.corrected_increase(), 30.0); + } + + #[test] + fn test_merge_accumulators_sorts_and_sums() { + // Three contiguous windows fed out of order; corrections + counts sum. + let w1 = feed(&[(1000, 0.0), (2000, 40.0)]); // +40, count 2 + let w2 = feed(&[(3000, 50.0), (4000, 10.0), (5000, 60.0)]); // reset +50, count 3 + let w3 = feed(&[(6000, 70.0), (7000, 90.0)]); // +20, count 2 + let merged = + >::merge_accumulators( + vec![w3.clone(), w1.clone(), w2.clone()], + ) + .unwrap(); + // boundaries: w1.last(40) -> w2.start(50): no reset. w2.last(60) -> w3.start(70): no reset. + // total correction = 0 + 50 + 0 (intra w2) + boundaries(0) = 50. + assert_eq!(merged.starting_measurement.value, 0.0); + assert_eq!(merged.last_seen_measurement.value, 90.0); + assert_eq!(merged.counter_reset_correction, 50.0); + assert_eq!(merged.sample_count, 7); + // increase = 90 - 0 + 50 = 140. + assert_eq!(merged.corrected_increase(), 140.0); + } + + #[test] + fn test_extrapolation_parity() { + // 6 samples spanning [5s, 55s] (sampledInterval 50s, avgGap 10s) inside the + // range [0, 60s]. Hand-computed Prometheus extrapolatedRate: + // durationToStart = 5s (clamped by durationToZero = 50*(10/60)=8.33 -> 5) + // durationToEnd = 5s, threshold = 11s + // extrapolateTo = 50 + 5 + 5 = 60 ; factor = 60/50 = 1.2 + // increase = 60 * 1.2 = 72 ; rate = 72 / 60 = 1.2 + let acc = IncreaseAccumulator::new_full( + Measurement::new(10.0), + 5_000, + Measurement::new(70.0), + 55_000, + 0.0, + 6, + ); + let inc = acc.extrapolated(false, 0, 60_000).unwrap(); + let rate = acc.extrapolated(true, 0, 60_000).unwrap(); + assert!((inc - 72.0).abs() < 1e-9, "increase = {inc}"); + assert!((rate - 1.2).abs() < 1e-9, "rate = {rate}"); + + // Reached via query() with range kwargs. + let mut kwargs = HashMap::new(); + kwargs.insert(RANGE_START_MS_KWARG.to_string(), "0".to_string()); + kwargs.insert(RANGE_END_MS_KWARG.to_string(), "60000".to_string()); + let q = + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) + .unwrap(); + assert!((q - 72.0).abs() < 1e-9, "query increase = {q}"); + } + + #[test] + fn test_extrapolation_threshold_then_zero_clamp_order() { + // First sample far from the lower boundary (durationToStart 50s >> threshold + // 11s) so Prometheus clamps it to avgGap/2 = 5s FIRST, then the zero-point + // clamp (durationToZero = 50*(10/62.5) = 8s) does NOT apply (8 >= 5). + // factor = (50 + 5 + 0)/50 = 1.1 ; increase = 62.5 * 1.1 = 68.75. + // (Applying the zero clamp before the threshold clamp would wrongly give + // 72.5 — this test pins the Prometheus ordering.) + let acc = IncreaseAccumulator::new_full( + Measurement::new(10.0), + 60_000, + Measurement::new(72.5), + 110_000, + 0.0, + 6, + ); + let inc = acc.extrapolated(false, 10_000, 110_000).unwrap(); + assert!((inc - 68.75).abs() < 1e-9, "increase = {inc}"); + } + + #[test] + fn test_extrapolation_zero_clamp_engages() { + // Counter starts near zero (2) relative to its increase (100), far from the + // boundary. After the threshold clamp (durationToStart -> 5s), the zero + // clamp applies: durationToZero = 50*(2/100) = 1s < 5s. + // factor = (50 + 1 + 0)/50 = 1.02 ; increase = 100 * 1.02 = 102. + let acc = IncreaseAccumulator::new_full( + Measurement::new(2.0), + 60_000, + Measurement::new(102.0), + 110_000, + 0.0, + 6, + ); + let inc = acc.extrapolated(false, 10_000, 110_000).unwrap(); + assert!((inc - 102.0).abs() < 1e-9, "increase = {inc}"); + } + + #[test] + fn test_extrapolation_requires_two_samples() { + let acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(10.0), 1000); + assert_eq!(acc.sample_count, 1); + assert!(acc.extrapolated(false, 0, 60_000).is_none()); + } + + #[test] + fn test_serde_roundtrip_new_fields() { + let acc = feed(&[(1000, 10.0), (2000, 100.0), (3000, 5.0), (4000, 30.0)]); + // JSON + let json = acc.serialize_to_json(); + let back = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!(back.counter_reset_correction, 100.0); + assert_eq!(back.sample_count, 4); + // Bytes + let bytes = acc.serialize_to_bytes(); + let back_b = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(back_b.counter_reset_correction, 100.0); + assert_eq!(back_b.sample_count, 4); + } + + #[test] + fn test_legacy_bytes_decode_with_defaults() { + // A buffer that ends after the original four fields (correction + count + // absent) must decode with defaults instead of erroring. + let acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); + let mut bytes = acc.serialize_to_bytes(); + bytes.truncate(bytes.len() - 16); // drop counter_reset_correction (8) + sample_count (8) + let back = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(back.counter_reset_correction, 0.0); + assert_eq!(back.sample_count, DEFAULT_LEGACY_SAMPLE_COUNT); + assert_eq!(back.last_seen_measurement.value, 25.0); + } + + #[test] + fn test_arroyo_roundtrip_preserves_all_fields() { + // A window with a reset so every field is non-default. + let acc = feed(&[(1000, 10.0), (2000, 100.0), (3000, 5.0), (4000, 30.0)]); + let bytes = acc.serialize_to_bytes_arroyo(); + let back = IncreaseAccumulator::deserialize_from_bytes_arroyo(&bytes).unwrap(); + assert_eq!(back.starting_measurement.value, 10.0); + assert_eq!(back.starting_timestamp, 1000); + assert_eq!(back.last_seen_measurement.value, 30.0); + assert_eq!(back.last_seen_timestamp, 4000); + assert_eq!(back.counter_reset_correction, 100.0); + assert_eq!(back.sample_count, 4); + assert_eq!(back.corrected_increase(), 120.0); + } + + #[test] + fn test_arroyo_legacy_4field_blob_defaults_sample_count() { + // A legacy 4-field MessagePack blob (from before reset support) must decode + // with sample_count defaulted so extrapolation stays well-defined. + #[derive(serde::Serialize)] + struct LegacyFour { + starting_measurement: f64, + starting_timestamp: i64, + last_seen_measurement: f64, + last_seen_timestamp: i64, + } + let bytes = rmp_serde::to_vec(&LegacyFour { + starting_measurement: 10.0, + starting_timestamp: 1000, + last_seen_measurement: 40.0, + last_seen_timestamp: 2000, + }) + .unwrap(); + let back = IncreaseAccumulator::deserialize_from_bytes_arroyo(&bytes).unwrap(); + assert_eq!(back.counter_reset_correction, 0.0); + assert_eq!(back.sample_count, DEFAULT_LEGACY_SAMPLE_COUNT); + assert_eq!(back.corrected_increase(), 30.0); + } } 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 6ee70402..fccc0b5f 100644 --- a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs @@ -23,6 +23,13 @@ struct MeasurementData { starting_timestamp: i64, last_seen_measurement: f64, last_seen_timestamp: i64, + // rmp-serde encodes structs as positional arrays, so `#[serde(default)]` is + // what lets legacy 4-field blobs (from the Arroyo UDF before reset support) + // still decode — the trailing fields are filled with their defaults. + #[serde(default)] + counter_reset_correction: f64, + #[serde(default)] + sample_count: u64, } impl MultipleIncreaseAccumulator { @@ -94,7 +101,8 @@ impl MultipleIncreaseAccumulator { // Calculate consumed bytes for IncreaseAccumulator // Structure: starting_measurement_len(4) + starting_measurement + starting_timestamp(8) + - // last_seen_measurement_len(4) + last_seen_measurement + last_seen_timestamp(8) + // last_seen_measurement_len(4) + last_seen_measurement + last_seen_timestamp(8) + + // counter_reset_correction(8) + sample_count(8) let starting_measurement_len = u32::from_le_bytes([ buffer[offset], buffer[offset + 1], @@ -108,7 +116,7 @@ impl MultipleIncreaseAccumulator { buffer[offset + 4 + starting_measurement_len + 8 + 3], ]) as usize; let consumed_bytes = - 4 + starting_measurement_len + 8 + 4 + last_seen_measurement_len + 8; + 4 + starting_measurement_len + 8 + 4 + last_seen_measurement_len + 8 + 8 + 8; offset += consumed_bytes; accumulator.increases.insert(key, increase_data); @@ -139,12 +147,22 @@ 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; + // `sample_count == 0` means the field was absent (legacy UDF output); + // fall back to 2 so extrapolation stays well-defined. A genuine + // single-sample window reports count 1 and is handled as "no rate". + let sample_count = if values.sample_count == 0 { + 2 + } else { + values.sample_count + }; - let increase_accumulator = IncreaseAccumulator::new( + let increase_accumulator = IncreaseAccumulator::new_full( starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, + values.counter_reset_correction, + sample_count, ); accumulator.increases.insert(key_obj, increase_accumulator); @@ -169,6 +187,8 @@ 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, + counter_reset_correction: increase_acc.counter_reset_correction, + sample_count: increase_acc.sample_count, }, ); } @@ -257,19 +277,13 @@ impl AggregateCore for MultipleIncreaseAccumulator { .downcast_ref::() .ok_or("Failed to downcast to MultipleIncreaseAccumulator")?; - // Clone self once, then merge other's data in-place + // Clone self once, then merge other's data in-place. Per key, delegate to + // the boundary-aware IncreaseAccumulator merge so counter-reset correction + // and sample counts are combined consistently. let mut merged = self.clone(); for (key, data) in &other_multiple_increase.increases { if let Some(existing_data) = merged.increases.get_mut(key) { - // Merge in-place: take earliest start, latest end - if data.starting_timestamp < existing_data.starting_timestamp { - existing_data.starting_measurement = data.starting_measurement.clone(); - existing_data.starting_timestamp = data.starting_timestamp; - } - if data.last_seen_timestamp > existing_data.last_seen_timestamp { - existing_data.last_seen_measurement = data.last_seen_measurement.clone(); - existing_data.last_seen_timestamp = data.last_seen_timestamp; - } + *existing_data = IncreaseAccumulator::merge_pair(existing_data, data); } else { merged.increases.insert(key.clone(), data.clone()); } @@ -305,14 +319,15 @@ impl MultipleSubpopulationAggregate for MultipleIncreaseAccumulator { &self, statistic: Statistic, key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, + query_kwargs: Option<&HashMap>, ) -> Result> { let data = self .increases .get(key) .ok_or_else(|| format!("Key {key} not found in MultipleIncreaseAccumulator"))?; - data.query(statistic, None) + // Forward kwargs so the inner accumulator receives the range boundaries. + data.query(statistic, query_kwargs) } fn clone_boxed(&self) -> Box { @@ -333,16 +348,9 @@ impl MergeableAccumulator for MultipleIncreaseAccum for accumulator in accumulators { for (key, data) in accumulator.increases { if let Some(existing_data) = result.increases.get_mut(&key) { - // Merge in-place without cloning existing_data - // Take the earliest start time and latest end time - if data.starting_timestamp < existing_data.starting_timestamp { - existing_data.starting_measurement = data.starting_measurement; - existing_data.starting_timestamp = data.starting_timestamp; - } - if data.last_seen_timestamp > existing_data.last_seen_timestamp { - existing_data.last_seen_measurement = data.last_seen_measurement; - existing_data.last_seen_timestamp = data.last_seen_timestamp; - } + // Boundary-aware merge: earliest start, latest last-seen, summed + // corrections + sample counts, plus a seam reset correction. + *existing_data = IncreaseAccumulator::merge_pair(existing_data, &data); } else { result.increases.insert(key, data); } @@ -519,6 +527,105 @@ mod tests { assert_eq!(keys.len(), 1); } + #[test] + fn test_arroyo_roundtrip_preserves_reset_fields() { + let mut acc = MultipleIncreaseAccumulator::new(); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + // A window with a reset: correction 100, four samples. + let inc = IncreaseAccumulator::new_full( + Measurement::new(10.0), + 1000, + Measurement::new(30.0), + 4000, + 100.0, + 4, + ); + acc.update(key.clone(), inc); + + let bytes = acc.serialize_to_bytes_arroyo(); + let back = MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo(&bytes).unwrap(); + let got = back.increases.get(&key).unwrap(); + assert_eq!(got.counter_reset_correction, 100.0); + assert_eq!(got.sample_count, 4); + assert_eq!(got.corrected_increase(), 120.0); + } + + #[test] + fn test_arroyo_legacy_blob_defaults() { + // A legacy 4-field MeasurementData blob (no correction / count) must decode + // with safe defaults rather than failing. + #[derive(serde::Serialize)] + struct LegacyMeasurementData { + starting_measurement: f64, + starting_timestamp: i64, + last_seen_measurement: f64, + last_seen_timestamp: i64, + } + let mut legacy: HashMap = HashMap::new(); + legacy.insert( + "web".to_string(), + LegacyMeasurementData { + starting_measurement: 10.0, + starting_timestamp: 1000, + last_seen_measurement: 25.0, + last_seen_timestamp: 2000, + }, + ); + let bytes = rmp_serde::to_vec(&legacy).unwrap(); + let back = MultipleIncreaseAccumulator::deserialize_from_bytes_arroyo(&bytes).unwrap(); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let got = back.increases.get(&key).unwrap(); + assert_eq!(got.counter_reset_correction, 0.0); + assert_eq!(got.sample_count, 2); // legacy default + assert_eq!(got.corrected_increase(), 15.0); + } + + #[test] + fn test_query_with_range_kwargs_extrapolates() { + use promql_utilities::query_logics::enums::{RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; + let mut acc = MultipleIncreaseAccumulator::new(); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + // Same numbers as IncreaseAccumulator::test_extrapolation_parity => increase 72. + acc.update( + key.clone(), + IncreaseAccumulator::new_full( + Measurement::new(10.0), + 5_000, + Measurement::new(70.0), + 55_000, + 0.0, + 6, + ), + ); + let mut kwargs = HashMap::new(); + kwargs.insert(RANGE_START_MS_KWARG.to_string(), "0".to_string()); + kwargs.insert(RANGE_END_MS_KWARG.to_string(), "60000".to_string()); + let v = acc.query(Statistic::Increase, &key, Some(&kwargs)).unwrap(); + assert!((v - 72.0).abs() < 1e-9, "increase = {v}"); + } + + #[test] + fn test_merge_keyed_boundary_reset() { + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut a = MultipleIncreaseAccumulator::new(); + a.update( + key.clone(), + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000), + ); + let mut b = MultipleIncreaseAccumulator::new(); + b.update( + key.clone(), + IncreaseAccumulator::new(Measurement::new(5.0), 3000, Measurement::new(30.0), 4000), + ); + + let merged = + MultipleIncreaseAccumulator::merge_accumulators(vec![a.clone(), b.clone()]).unwrap(); + let got = merged.increases.get(&key).unwrap(); + // boundary reset (5 < 20) => +20, increase = 30 - 10 + 20 = 40. + assert_eq!(got.corrected_increase(), 40.0); + assert_eq!(got.sample_count, 2); + } + // #[test] // fn test_multiple_increase_accumulator_arroyo_deserialization() { // // Create test data in Arroyo MessagePack format diff --git a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs b/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs index c90e8e56..cf556fb7 100644 --- a/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs +++ b/asap-query-engine/src/tests/datafusion/accumulator_serde_tests.rs @@ -211,26 +211,48 @@ mod tests { assert_eq!(keys.len(), 1); } - // ======================================================================== - // Not-implemented types - // ======================================================================== - #[test] - fn test_not_implemented_increase() { - let bytes = vec![1, 2, 3, 4]; - let result = deserialize_accumulator(&bytes, &SketchType::Increase); - assert!(result.is_err()); - let err_msg = match result { - Err(e) => format!("{}", e), - Ok(_) => panic!("Expected error"), - }; + fn test_round_trip_increase() { + use promql_utilities::query_logics::enums::{RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; + + // A 2-sample counter window (built the way ingest does), round-tripped + // through the Arroyo MessagePack form. + let mut acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(10.0), 1000); + acc.update(Measurement::new(70.0), 3000); + let bytes = serialize_accumulator_arroyo(&acc); + + let restored = deserialize_accumulator(&bytes, &SketchType::Increase).unwrap(); + assert_eq!(restored.get_accumulator_type(), AggregationType::Increase); + + // Fallback (no range vector): reset-corrected increase = 70 - 10 = 60. + let single = deserialize_single_subpopulation(&bytes, &SketchType::Increase).unwrap(); + let increase = single.query(Statistic::Increase, None).unwrap(); assert!( - err_msg.contains("Not") || err_msg.contains("not"), - "Expected NotImplemented error, got: {}", - err_msg + (increase - 60.0).abs() < 1e-10, + "expected 60, got {increase}" + ); + + // The sample_count field survives the round-trip, so extrapolation is + // well-defined and matches the source accumulator exactly. + let mut kwargs = HashMap::new(); + kwargs.insert(RANGE_START_MS_KWARG.to_string(), "0".to_string()); + kwargs.insert(RANGE_END_MS_KWARG.to_string(), "4000".to_string()); + let extrapolated = deserialize_single_subpopulation(&bytes, &SketchType::Increase) + .unwrap() + .query(Statistic::Increase, Some(&kwargs)) + .unwrap(); + let expected = acc.extrapolated(false, 0, 4000).unwrap(); + assert!( + (extrapolated - expected).abs() < 1e-9, + "round-tripped extrapolation {extrapolated} != source {expected}" ); } + // ======================================================================== + // Not-implemented types + // ======================================================================== + #[test] fn test_not_implemented_minmax() { let bytes = vec![1, 2, 3, 4]; diff --git a/asap-query-engine/src/tests/datafusion/mod.rs b/asap-query-engine/src/tests/datafusion/mod.rs index 143cc0a4..f14dd5c4 100644 --- a/asap-query-engine/src/tests/datafusion/mod.rs +++ b/asap-query-engine/src/tests/datafusion/mod.rs @@ -12,4 +12,6 @@ pub mod plan_execution_dual_input_tests; pub mod plan_execution_temporal_tests; pub mod plan_execution_tests; pub mod range_query_arithmetic_tests; +pub mod rate_increase_extrapolation_tests; +pub mod rate_increase_synthetic_tests; pub mod structural_matching_tests; diff --git a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs b/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs index 8764a806..7d9c6ab2 100644 --- a/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_builder_regression_tests.rs @@ -143,6 +143,7 @@ mod tests { #[test] fn test_rate_to_extract_rate() { use datafusion_summary_library::InferOperation; + // No range kwargs -> no embedded bounds -> reset-corrected fallback at query time. let ctx = create_context( Statistic::Rate, AggregationType::Increase, @@ -151,7 +152,29 @@ mod tests { ); assert!(matches!( ctx.map_statistic_to_infer_operation().unwrap(), - InferOperation::ExtractRate + InferOperation::ExtractRate(None) + )); + } + + #[test] + fn test_rate_with_range_kwargs_embeds_bounds() { + use datafusion_summary_library::InferOperation; + use promql_utilities::query_logics::enums::{RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; + // The binary-arithmetic instant path supplies range-vector boundaries; they + // must be embedded in the InferOperation so extrapolation survives to the + // physical SummaryInferExec. + let mut kwargs = HashMap::new(); + kwargs.insert(RANGE_START_MS_KWARG.to_string(), "1000".to_string()); + kwargs.insert(RANGE_END_MS_KWARG.to_string(), "301000".to_string()); + let ctx = create_context( + Statistic::Rate, + AggregationType::Increase, + vec!["host"], + kwargs, + ); + assert!(matches!( + ctx.map_statistic_to_infer_operation().unwrap(), + InferOperation::ExtractRate(Some((1000, 301000))) )); } diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs index 649b2b24..a7cb2671 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_dual_input_tests.rs @@ -29,15 +29,20 @@ mod tests { let key = KeyByLabelValues { labels: key_labels.iter().map(|s| s.to_string()).collect(), }; - increases.insert( - key, - IncreaseAccumulator::new( - Measurement::new(start_val), - 0, - Measurement::new(end_val), - 10, - ), + // Build the accumulator the way ingest does: first sample via `new`, + // then a second via `update`. This yields sample_count == 2, so the + // accumulator is valid for Prometheus extrapolation (rate/increase + // needs >= 2 points) — the path now exercised end-to-end via the + // DataFusion plan. A single-sample `new(start, .., end, ..)` would be a + // physically impossible state that real ingest never produces. + let mut acc = IncreaseAccumulator::new( + Measurement::new(start_val), + 0, + Measurement::new(start_val), + 0, ); + acc.update(Measurement::new(end_val), 10); + increases.insert(key, acc); } MultipleIncreaseAccumulator::new_with_increases(increases) } @@ -105,6 +110,88 @@ mod tests { ); } + /// Regression guard for the counter-reset/extrapolation fix on the DataFusion + /// physical path (the plumbing shared by the binary-arithmetic instant path): + /// the range-vector boundaries must survive the logical→physical boundary so + /// `IncreaseAccumulator` applies Prometheus `extrapolatedRate` instead of the + /// reset-corrected fallback. Previously they were dropped and the fallback was + /// silently used. Asserts the exact extrapolated value, computed from the same + /// boundaries the engine derives for the query. + #[tokio::test] + async fn test_self_keyed_multiple_increase_applies_extrapolation() { + use crate::data_model::SingleSubpopulationAggregate; + use promql_utilities::query_logics::enums::{ + Statistic, RANGE_END_MS_KWARG, RANGE_START_MS_KWARG, + }; + + const QUERY_TIME: f64 = 1000.0; + + // A single-sub-key MultipleIncrease with a 2-sample counter window inside + // the query range [990s, 1000s], built the way ingest does (new + update). + let build_inner = || { + let mut a = IncreaseAccumulator::new( + Measurement::new(10.0), + 992_000, + Measurement::new(10.0), + 992_000, + ); + a.update(Measurement::new(70.0), 997_000); + a + }; + let key = KeyByLabelValues { + labels: vec!["host-a".to_string(), "endpoint-1".to_string()], + }; + let mut increases = HashMap::new(); + increases.insert(key, build_inner()); + let acc = MultipleIncreaseAccumulator::new_with_increases(increases); + + let engine = create_engine_single_pop_with_aggregated( + "http_requests_total", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(acc))], + "increase(http_requests_total[10s])", + ); + + // Read the actual range-vector boundaries the engine derives for the query. + let ctx = engine + .build_query_execution_context_promql( + "increase(http_requests_total[10s])".to_string(), + QUERY_TIME, + ) + .expect("build context"); + let start_ms: i64 = ctx.metadata.query_kwargs[RANGE_START_MS_KWARG] + .parse() + .unwrap(); + let end_ms: i64 = ctx.metadata.query_kwargs[RANGE_END_MS_KWARG] + .parse() + .unwrap(); + + // Expected = Prometheus extrapolated increase over those boundaries. It must + // differ from the naive reset-corrected fallback, else the check is vacuous. + let expected = build_inner() + .extrapolated(false, start_ms, end_ms) + .expect("extrapolation should produce a value"); + let fallback = + SingleSubpopulationAggregate::query(&build_inner(), Statistic::Increase, None).unwrap(); + assert!( + (expected - fallback).abs() > 1e-9, + "test setup should distinguish extrapolation ({expected}) from fallback ({fallback})" + ); + + let results = + execute_new_plan(&engine, "increase(http_requests_total[10s])", QUERY_TIME).await; + assert_eq!(results.len(), 1, "expected 1 sub-key result"); + assert!( + (results[0].value - expected).abs() < 1e-6, + "physical path must apply extrapolation: got {}, expected {} (fallback would be {})", + results[0].value, + expected, + fallback + ); + } + // ======================================================================== // HydraKLL + DeltaSetAggregator (quantile dual-input) // ======================================================================== diff --git a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs b/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs index 8f942905..1c0b7290 100644 --- a/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs +++ b/asap-query-engine/src/tests/datafusion/plan_execution_tests.rs @@ -536,12 +536,20 @@ mod tests { // ======================================================================== #[tokio::test] - async fn test_execute_plan_not_implemented_increase() { - // IncreaseAccumulator has no arroyo serde, so execute_plan should fail. - // Uses a bare OnlyTemporal query (not `sum(increase(...))`): Increase is - // never collapsable with any spatial aggregation (see get_is_collapsable), - // so a spatial wrapper would no longer match any pattern at all (#508). - let inc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(100.0), 10); + async fn test_execute_plan_increase_single_pop_extrapolates() { + // Single-population IncreaseAccumulator now round-trips through the + // DataFusion path (arroyo/MessagePack serde), so execute_plan succeeds and + // applies Prometheus extrapolation. A 2-sample window [10@992s → 70@997s] + // over the query range [990s, 1000s] extrapolates to 106 (fallback: 60). + // (Uses a bare OnlyTemporal query: post-#508 Increase is never collapsable + // with a spatial aggregation, so `sum(increase(...))` matches no pattern.) + let mut inc = IncreaseAccumulator::new( + Measurement::new(10.0), + 992_000, + Measurement::new(10.0), + 992_000, + ); + inc.update(Measurement::new(70.0), 997_000); let engine = create_engine_single_pop( "http_requests_total", @@ -559,9 +567,13 @@ mod tests { .expect("Should build context"); let result = engine.execute_plan(&context).await; + assert!(result.is_ok(), "execute_plan failed: {:?}", result.err()); + let results = result.unwrap(); + assert_eq!(results.len(), 1, "expected 1 result, got {}", results.len()); assert!( - result.is_err(), - "execute_plan should fail for IncreaseAccumulator (no arroyo serde)" + (results[0].value - 106.0).abs() < 1e-6, + "expected extrapolated 106, got {}", + results[0].value ); } diff --git a/asap-query-engine/src/tests/datafusion/rate_increase_extrapolation_tests.rs b/asap-query-engine/src/tests/datafusion/rate_increase_extrapolation_tests.rs new file mode 100644 index 00000000..92f87877 --- /dev/null +++ b/asap-query-engine/src/tests/datafusion/rate_increase_extrapolation_tests.rs @@ -0,0 +1,406 @@ +//! Cross-path correctness tests for `rate()` / `increase()` counter-reset +//! correction + Prometheus extrapolation. +//! +//! `rate()`/`increase()` are served by three distinct execution paths, and all +//! three must reproduce Prometheus `extrapolatedRate` (counter-reset correction +//! followed by extrapolation to the range-vector boundaries) rather than the +//! plain reset-corrected fallback: +//! +//! 1. **Instant, single expression** — `handle_query_promql("increase(m[w])")` +//! → `execute_query_pipeline` (in-memory accumulator query). +//! 2. **Range query** — `handle_range_query_promql(...)` +//! → `execute_range_query_pipeline` (per-step range boundaries). +//! 3. **Instant, binary arithmetic** — `handle_query_promql("increase(m[w]) * k")` +//! → DataFusion `SummaryInferExec` (bounds embedded in `InferOperation`). +//! +//! The oracle for "correct" is [`IncreaseAccumulator::extrapolated`], which is +//! pinned to hand-computed Prometheus values in `increase_accumulator.rs` +//! (`test_extrapolation_parity` etc.). Each test asserts the path reproduces that +//! value AND that it differs from the non-extrapolated fallback, so a regression +//! back to the old behavior fails loudly. + +#[cfg(test)] +mod tests { + use crate::data_model::{ + AggregationType, KeyByLabelValues, Measurement, SingleSubpopulationAggregate, WindowType, + }; + use crate::engines::query_result::{InstantVectorElement, QueryResult}; + use crate::precompute_operators::{IncreaseAccumulator, MultipleIncreaseAccumulator}; + use crate::tests::test_utilities::engine_factories::*; + use crate::AggregateCore; + use promql_utilities::query_logics::enums::{ + Statistic, RANGE_END_MS_KWARG, RANGE_START_MS_KWARG, + }; + use std::collections::HashMap; + + const QUERY_TIME: f64 = 1000.0; // seconds → 1_000_000 ms + const WINDOW: &str = "[10s]"; + + /// A monotonic 2-sample counter window [10@992s → 70@997s] (no reset), built + /// the way ingest does (`new` then `update` ⇒ sample_count == 2). + fn counter_no_reset() -> IncreaseAccumulator { + let mut a = IncreaseAccumulator::new( + Measurement::new(10.0), + 992_000, + Measurement::new(10.0), + 992_000, + ); + a.update(Measurement::new(70.0), 997_000); + a + } + + /// A 4-sample counter window with one reset: 10 → 50 → 5 (reset, +50) → 30. + /// Reset-corrected increase = 30 - 10 + 50 = 70. + fn counter_with_reset() -> IncreaseAccumulator { + let mut a = IncreaseAccumulator::new( + Measurement::new(10.0), + 992_000, + Measurement::new(10.0), + 992_000, + ); + a.update(Measurement::new(50.0), 994_000); + a.update(Measurement::new(5.0), 995_000); // reset: +50 correction + a.update(Measurement::new(30.0), 997_000); + a + } + + fn multi_pop(inner: IncreaseAccumulator) -> MultipleIncreaseAccumulator { + let key = KeyByLabelValues { + labels: vec!["host-a".to_string(), "endpoint-1".to_string()], + }; + let mut m = HashMap::new(); + m.insert(key, inner); + MultipleIncreaseAccumulator::new_with_increases(m) + } + + fn values(qr: QueryResult) -> Vec { + match qr { + QueryResult::Vector(iv) => iv.values, + other => panic!("expected instant vector, got {other:?}"), + } + } + + /// Range boundaries the engine derives for `(m[10s])` at `QUERY_TIME`. + fn engine_bounds( + engine: &crate::engines::simple_engine::SimpleEngine, + query: &str, + ) -> (i64, i64) { + let ctx = engine + .build_query_execution_context_promql(query.to_string(), QUERY_TIME) + .expect("build context"); + ( + ctx.metadata.query_kwargs[RANGE_START_MS_KWARG] + .parse() + .unwrap(), + ctx.metadata.query_kwargs[RANGE_END_MS_KWARG] + .parse() + .unwrap(), + ) + } + + // ======================================================================== + // Anchor: the extrapolation oracle produces the expected absolute values. + // ======================================================================== + + /// Ground-truth anchor independent of the query engine: over [990s, 1000s] + /// the no-reset counter extrapolates to increase 106 / rate 10.6, well away + /// from the reset-corrected fallback of 60. + #[test] + fn oracle_anchor_values() { + let inc = counter_no_reset() + .extrapolated(false, 990_000, 1_000_000) + .unwrap(); + let rate = counter_no_reset() + .extrapolated(true, 990_000, 1_000_000) + .unwrap(); + assert!((inc - 106.0).abs() < 1e-9, "increase = {inc}"); + assert!((rate - 10.6).abs() < 1e-9, "rate = {rate}"); + let fallback = + SingleSubpopulationAggregate::query(&counter_no_reset(), Statistic::Increase, None) + .unwrap(); + assert_eq!(fallback, 60.0, "reset-corrected fallback"); + } + + // ======================================================================== + // Path 1: instant single-expression (pipeline) + // ======================================================================== + + #[tokio::test(flavor = "multi_thread")] + async fn path1_instant_increase_extrapolates() { + let engine = create_engine_single_pop_with_aggregated( + "m", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(multi_pop(counter_no_reset())))], + &format!("increase(m{WINDOW})"), + ); + let (s, e) = engine_bounds(&engine, &format!("increase(m{WINDOW})")); + let expected = counter_no_reset().extrapolated(false, s, e).unwrap(); + + let out = values( + engine + .handle_query_promql(format!("increase(m{WINDOW})"), QUERY_TIME) + .expect("path #1 should serve increase") + .1, + ); + assert_eq!(out.len(), 1); + assert!((out[0].value - expected).abs() < 1e-6); + assert!((out[0].value - 106.0).abs() < 1e-6, "got {}", out[0].value); + } + + // ======================================================================== + // Path 2: range query (pipeline, per-step boundaries) + // ======================================================================== + + /// Range query over two steps. The pipeline skips `None` keys, so this uses a + /// single-population Increase bucket with an explicit key. Each output point + /// extrapolates over its own window: T=1000s → [990s,1000s] → 106; T=999s → + /// [989s,999s] → 94. + #[tokio::test(flavor = "multi_thread")] + async fn path2_range_increase_extrapolates() { + // Bucket start timestamp is 992_000 (factory stores [ts-1000, ts]). + let engine = create_engine_multi_timestamp_with_window( + "m", + AggregationType::Increase, + vec!["host"], + vec![( + 993_000u64, + Some(vec!["host-a".to_string()]), + Box::new(counter_no_reset()) as Box, + )], + &format!("increase(m{WINDOW})"), + 1000, + WindowType::Tumbling, + ); + + let expected_t1000 = counter_no_reset() + .extrapolated(false, 990_000, 1_000_000) + .unwrap(); + let expected_t999 = counter_no_reset() + .extrapolated(false, 989_000, 999_000) + .unwrap(); + + let (_, qr) = engine + .handle_range_query_promql(format!("increase(m{WINDOW})"), 999.0, 1000.0, 1.0) + .expect("path #2 should serve range increase"); + let series = match qr { + QueryResult::Matrix(rv) => rv.values, + other => panic!("expected matrix, got {other:?}"), + }; + assert_eq!(series.len(), 1, "one series"); + let samples: Vec = series[0].samples.iter().map(|s| s.value).collect(); + assert_eq!(samples.len(), 2, "two output steps: {samples:?}"); + assert!( + (samples[0] - expected_t999).abs() < 1e-6, + "T=999s: {samples:?}" + ); + assert!( + (samples[1] - expected_t1000).abs() < 1e-6, + "T=1000s: {samples:?}" + ); + // Anchor + proof it is extrapolation, not the flat fallback of 60. + assert!((samples[1] - 106.0).abs() < 1e-6); + assert!((samples[0] - 94.0).abs() < 1e-6); + } + + // ======================================================================== + // Path 3: instant binary arithmetic (DataFusion / SummaryInferExec) — the + // path this fix repairs. Single-pop Increase has no arroyo deserializer, so + // the executable case is the multi-population accumulator. + // ======================================================================== + + #[tokio::test(flavor = "multi_thread")] + async fn path3_binary_increase_extrapolates() { + let engine = create_engine_single_pop_with_aggregated( + "m", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(multi_pop(counter_no_reset())))], + &format!("increase(m{WINDOW})"), + ); + let (s, e) = engine_bounds(&engine, &format!("increase(m{WINDOW})")); + let expected = counter_no_reset().extrapolated(false, s, e).unwrap(); + + // `* 1` forces the binary-arithmetic instant path (DataFusion) while + // leaving the extrapolated value unchanged. + let out = values( + engine + .handle_query_promql(format!("increase(m{WINDOW}) * 1"), QUERY_TIME) + .expect("path #3 should serve binary increase") + .1, + ); + assert_eq!(out.len(), 1); + assert!((out[0].value - expected).abs() < 1e-6); + assert!((out[0].value - 106.0).abs() < 1e-6, "got {}", out[0].value); + } + + #[tokio::test(flavor = "multi_thread")] + async fn path3_binary_rate_extrapolates_with_arithmetic() { + let engine = create_engine_single_pop_with_aggregated( + "m", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(multi_pop(counter_no_reset())))], + &format!("rate(m{WINDOW})"), + ); + let (s, e) = engine_bounds(&engine, &format!("rate(m{WINDOW})")); + let expected_rate = counter_no_reset().extrapolated(true, s, e).unwrap(); + + let out = values( + engine + .handle_query_promql(format!("rate(m{WINDOW}) * 2"), QUERY_TIME) + .expect("path #3 should serve binary rate") + .1, + ); + assert_eq!(out.len(), 1); + assert!( + (out[0].value - expected_rate * 2.0).abs() < 1e-6, + "got {}, expected {}", + out[0].value, + expected_rate * 2.0 + ); + assert!((out[0].value - 21.2).abs() < 1e-6, "got {}", out[0].value); + } + + // ======================================================================== + // Cross-path consistency + counter-reset correctness + // ======================================================================== + + /// The instant pipeline (#1) and the binary DataFusion path (#3) must produce + /// the identical extrapolated value for the same sub-expression — the property + /// that was violated before this fix (#3 silently used the flat fallback). + #[tokio::test(flavor = "multi_thread")] + async fn paths_1_and_3_agree() { + let make = || { + create_engine_single_pop_with_aggregated( + "m", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(multi_pop(counter_no_reset())))], + &format!("increase(m{WINDOW})"), + ) + }; + let engine = make(); + let p1 = values( + engine + .handle_query_promql(format!("increase(m{WINDOW})"), QUERY_TIME) + .unwrap() + .1, + ); + let p3 = values( + engine + .handle_query_promql(format!("increase(m{WINDOW}) * 1"), QUERY_TIME) + .unwrap() + .1, + ); + assert_eq!(p1.len(), 1); + assert_eq!(p3.len(), 1); + assert!( + (p1[0].value - p3[0].value).abs() < 1e-9, + "instant #1 ({}) and binary #3 ({}) must agree", + p1[0].value, + p3[0].value + ); + } + + /// Counter-reset correction must be applied before extrapolation on the binary + /// path: with a mid-window reset the corrected increase is 70 (not 20), and the + /// path must reproduce the extrapolation of that corrected value. + #[tokio::test(flavor = "multi_thread")] + async fn path3_binary_increase_applies_reset_correction() { + let engine = create_engine_single_pop_with_aggregated( + "m", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(multi_pop(counter_with_reset())))], + &format!("increase(m{WINDOW})"), + ); + let (s, e) = engine_bounds(&engine, &format!("increase(m{WINDOW})")); + + // Sanity: the reset correction is in the corrected increase (70, not 20), + // and the extrapolation of a reset-corrected series differs from a naive + // last-first (20) as well as from the reset-corrected-but-flat value (70). + assert_eq!(counter_with_reset().corrected_increase(), 70.0); + let expected = counter_with_reset().extrapolated(false, s, e).unwrap(); + assert!( + expected > 70.0, + "extrapolated {expected} should exceed flat 70" + ); + + let out = values( + engine + .handle_query_promql(format!("increase(m{WINDOW}) * 1"), QUERY_TIME) + .expect("path #3 should serve binary increase with reset") + .1, + ); + assert_eq!(out.len(), 1); + assert!( + (out[0].value - expected).abs() < 1e-6, + "got {}, expected {}", + out[0].value, + expected + ); + } + + // ======================================================================== + // Single-population Increase in a binary expression (DataFusion path) + // ======================================================================== + + /// Single-population `IncreaseAccumulator` is deserializable on the DataFusion + /// path (arroyo/MessagePack serde), so a binary expression with a single-pop + /// increase arm is accelerated and extrapolates — identically to the same + /// single-pop increase served by the instant pipeline (#1). This is the + /// scenario the arroyo-serde addition unlocked; before it, the binary form + /// returned `None` and fell back to Prometheus. + #[tokio::test(flavor = "multi_thread")] + async fn singlepop_increase_binary_extrapolates_like_instant() { + let engine = create_engine_single_pop( + "m", + AggregationType::Increase, + vec!["host"], + vec![( + Some(vec!["host-a".to_string()]), + Box::new(counter_no_reset()), + )], + &format!("increase(m{WINDOW})"), + ); + + let (s, e) = engine_bounds(&engine, &format!("increase(m{WINDOW})")); + let expected = counter_no_reset().extrapolated(false, s, e).unwrap(); + + // Instant pipeline (#1) extrapolates correctly for single-pop Increase. + let instant = values( + engine + .handle_query_promql(format!("increase(m{WINDOW})"), QUERY_TIME) + .expect("instant single-pop increase should work") + .1, + ); + assert_eq!(instant.len(), 1); + assert!((instant[0].value - expected).abs() < 1e-6); + + // Binary arithmetic (#3, DataFusion) is now accelerated and returns the + // same extrapolated value — no Prometheus fallback. + let binary = values( + engine + .handle_query_promql(format!("increase(m{WINDOW}) * 1"), QUERY_TIME) + .expect("single-pop increase in a binary expr should now accelerate") + .1, + ); + assert_eq!(binary.len(), 1); + assert!((binary[0].value - expected).abs() < 1e-6); + assert!( + (binary[0].value - instant[0].value).abs() < 1e-9, + "binary must match instant" + ); + assert!( + (binary[0].value - 106.0).abs() < 1e-6, + "got {}", + binary[0].value + ); + } +} diff --git a/asap-query-engine/src/tests/datafusion/rate_increase_synthetic_tests.rs b/asap-query-engine/src/tests/datafusion/rate_increase_synthetic_tests.rs new file mode 100644 index 00000000..92915fe3 --- /dev/null +++ b/asap-query-engine/src/tests/datafusion/rate_increase_synthetic_tests.rs @@ -0,0 +1,625 @@ +//! Synthetic differential correctness tests for `rate()` / `increase()`. +//! +//! The PR requirement is that ASAP reproduces Prometheus counter semantics: +//! per-sample counter-reset correction, then `extrapolatedRate` to the +//! range-vector boundaries. ASAP only retains the window *endpoints* plus the +//! running reset correction and sample count — but Prometheus's own +//! `extrapolatedRate` likewise derives its interval estimate from just +//! `(first_ts, last_ts, N)` and its value from `(first, last, correction)`, so +//! ASAP should match Prometheus *exactly* for any sample stream. +//! +//! These tests verify that against an INDEPENDENT reference reimplementation of +//! Prometheus's algorithm that consumes the raw `(timestamp, value)` samples +//! directly (not ASAP's reduced accumulator state). The reference is first +//! anchored to hand-computed Prometheus values, then used as the oracle across +//! hundreds of generated series (varied reset patterns, spacing, magnitude, +//! length, and range alignment). A discrepancy on any case is a real bug in the +//! accumulator's update / merge / extrapolation logic. + +#[cfg(test)] +mod tests { + use crate::data_model::{AggregationType, KeyByLabelValues, Measurement}; + use crate::engines::query_result::QueryResult; + use crate::precompute_operators::{IncreaseAccumulator, MultipleIncreaseAccumulator}; + use crate::tests::test_utilities::engine_factories::*; + use promql_utilities::query_logics::enums::{RANGE_END_MS_KWARG, RANGE_START_MS_KWARG}; + use std::collections::HashMap; + + // ======================================================================== + // Independent reference (operates on raw samples, not on the accumulator) + // ======================================================================== + + /// Prometheus `counterCorrection`: sum of the pre-reset value each time the + /// series drops. The reset-corrected increase is `last - first + correction`. + fn reference_corrected_increase(samples: &[(i64, f64)]) -> f64 { + let mut correction = 0.0; + let mut last = samples[0].1; + for &(_, v) in &samples[1..] { + if v < last { + correction += last; + } + last = v; + } + samples[samples.len() - 1].1 - samples[0].1 + correction + } + + /// Independent reimplementation of Prometheus `extrapolatedRate` + /// (promql/functions.go, `isCounter == true`) over the raw sample stream. + /// Deliberately written from the Prometheus algorithm, not from ASAP's code. + fn reference_extrapolated( + samples: &[(i64, f64)], + range_start_ms: i64, + range_end_ms: i64, + is_rate: bool, + ) -> Option { + if samples.len() < 2 { + return None; + } + let first_value = samples[0].1; + let result_value = reference_corrected_increase(samples); + + let first_ts = samples[0].0; + let last_ts = samples[samples.len() - 1].0; + + let mut duration_to_start = (first_ts - range_start_ms) as f64 / 1000.0; + let mut duration_to_end = (range_end_ms - last_ts) as f64 / 1000.0; + + let sampled_interval = (last_ts - first_ts) as f64 / 1000.0; + let average_duration_between_samples = sampled_interval / (samples.len() - 1) as f64; + let extrapolation_threshold = average_duration_between_samples * 1.1; + + if duration_to_start >= extrapolation_threshold { + duration_to_start = average_duration_between_samples / 2.0; + } + if result_value > 0.0 && first_value >= 0.0 { + let duration_to_zero = sampled_interval * (first_value / result_value); + if duration_to_zero < duration_to_start { + duration_to_start = duration_to_zero; + } + } + if duration_to_end >= extrapolation_threshold { + duration_to_end = average_duration_between_samples / 2.0; + } + + let mut factor = 1.0; + if sampled_interval != 0.0 { + factor = (sampled_interval + duration_to_start + duration_to_end) / sampled_interval; + } + if is_rate { + let range_seconds = (range_end_ms - range_start_ms) as f64 / 1000.0; + if range_seconds <= 0.0 { + return None; + } + factor /= range_seconds; + } + Some(result_value * factor) + } + + /// Feed a sample stream through the accumulator exactly as ingest does: + /// `new` on the first sample, then `update` for each subsequent one. + fn feed(samples: &[(i64, f64)]) -> IncreaseAccumulator { + let (ts0, v0) = samples[0]; + let mut acc = + IncreaseAccumulator::new(Measurement::new(v0), ts0, Measurement::new(v0), ts0); + for &(ts, v) in &samples[1..] { + acc.update(Measurement::new(v), ts); + } + acc + } + + fn close(a: f64, b: f64) -> bool { + (a - b).abs() <= 1e-9 * (1.0 + a.abs().max(b.abs())) + } + + // ======================================================================== + // Anchor the reference to externally-known Prometheus values, so the oracle + // itself is trustworthy before we test ASAP against it. + // ======================================================================== + + #[test] + fn reference_matches_hand_computed_prometheus() { + // From increase_accumulator::test_extrapolation_parity: 6 samples spanning + // [5s,55s] inside [0,60s] => increase 72, rate 1.2. + let s: Vec<(i64, f64)> = (0..6) + .map(|i| (5_000 + i * 10_000, 10.0 + i as f64 * 12.0)) + .collect(); + assert!(close( + reference_extrapolated(&s, 0, 60_000, false).unwrap(), + 72.0 + )); + assert!(close( + reference_extrapolated(&s, 0, 60_000, true).unwrap(), + 1.2 + )); + + // The [10s] window used throughout the cross-path tests: 10@992s, 70@997s + // inside [990s,1000s] => increase 106, rate 10.6. + let s2 = [(992_000i64, 10.0), (997_000i64, 70.0)]; + assert!(close( + reference_extrapolated(&s2, 990_000, 1_000_000, false).unwrap(), + 106.0 + )); + assert!(close( + reference_extrapolated(&s2, 990_000, 1_000_000, true).unwrap(), + 10.6 + )); + + // counterCorrection anchor: 10 -> 100 -> 5(reset) -> 30 => 30-10+100 = 120. + let s3 = [(0i64, 10.0), (1i64, 100.0), (2i64, 5.0), (3i64, 30.0)]; + assert!(close(reference_corrected_increase(&s3), 120.0)); + } + + // ======================================================================== + // Deterministic pseudo-random generator (no Math::random; reproducible) + // ======================================================================== + + struct Lcg(u64); + impl Lcg { + fn next_u64(&mut self) -> u64 { + // Numerical-Recipes LCG constants. + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + self.0 + } + fn unit(&mut self) -> f64 { + (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64 + } + fn range(&mut self, lo: usize, hi: usize) -> usize { + lo + (self.next_u64() as usize) % (hi - lo + 1) + } + } + + /// Generate a synthetic counter series: `n` samples starting at `start_ts`, + /// each gap drawn from `[min_gap, max_gap]` ms (uneven spacing allowed), + /// monotonically increasing with occasional resets to a small value. + fn synthetic_series( + rng: &mut Lcg, + start_ts: i64, + n: usize, + min_gap: usize, + max_gap: usize, + reset_prob: f64, + ) -> Vec<(i64, f64)> { + let mut samples = Vec::with_capacity(n); + let mut value = rng.unit() * 50.0; // >= 0 + let mut ts = start_ts; + for _ in 0..n { + samples.push((ts, value)); + if rng.unit() < reset_prob { + value = rng.unit() * 10.0; // counter restart (drops below current) + } else { + value += rng.unit() * 25.0; // positive increment + } + ts += rng.range(min_gap, max_gap) as i64; + } + samples + } + + /// A set of range configurations relative to the series span, to exercise the + /// tight case, the half-gap extrapolation, and the threshold clamp. + fn range_configs(samples: &[(i64, f64)]) -> Vec<(i64, i64)> { + let first = samples[0].0; + let last = samples[samples.len() - 1].0; + let span = (last - first).max(1); + let gap = span / (samples.len() as i64 - 1).max(1); + vec![ + (first, last), // tight: no outward extrapolation + (first - gap / 2, last + gap / 2), // within threshold both sides + (first - gap, last + gap), // ~ at threshold + (first - 5 * span, last + 5 * span), // far: threshold clamp both sides + ] + } + + // ======================================================================== + // Differential property test: single window + // ======================================================================== + + #[test] + fn synthetic_single_window_matches_prometheus() { + let mut rng = Lcg(0x9E3779B97F4A7C15); + let mut cases = 0usize; + for iter in 0..1500 { + let n = rng.range(2, 24); + let uneven = iter % 3 == 0; + let (min_gap, max_gap) = if uneven { + (1_000, 30_000) + } else { + (15_000, 15_000) + }; + let reset_prob = [0.0, 0.1, 0.25, 0.5][iter % 4]; + let samples = synthetic_series(&mut rng, 1_000_000, n, min_gap, max_gap, reset_prob); + if samples[samples.len() - 1].0 == samples[0].0 { + continue; // skip degenerate zero-span (guarded elsewhere) + } + + let acc = feed(&samples); + + // Reset correction must match the independent counterCorrection. + let ref_inc = reference_corrected_increase(&samples); + assert!( + close(acc.corrected_increase(), ref_inc), + "corrected_increase mismatch: acc={} ref={} samples={:?}", + acc.corrected_increase(), + ref_inc, + samples + ); + // sample_count must equal the number of samples fed. + assert_eq!(acc.sample_count, samples.len() as u64); + + // Extrapolation must match Prometheus for every range config & stat. + for (rs, re) in range_configs(&samples) { + for is_rate in [false, true] { + let got = acc.extrapolated(is_rate, rs, re); + let want = reference_extrapolated(&samples, rs, re, is_rate); + match (got, want) { + (Some(g), Some(w)) => assert!( + close(g, w), + "extrapolated mismatch (is_rate={is_rate}) got={g} want={w} \ + range=[{rs},{re}] samples={samples:?}" + ), + (None, None) => {} + (g, w) => panic!( + "presence mismatch got={g:?} want={w:?} range=[{rs},{re}] samples={samples:?}" + ), + } + cases += 1; + } + } + } + assert!(cases > 5_000, "expected broad coverage, only {cases} cases"); + } + + // ======================================================================== + // Named edge cases (guaranteed coverage of the tricky branches) + // ======================================================================== + + #[test] + fn synthetic_edge_cases_match_prometheus() { + let g = 15_000i64; // even 15s spacing + let cases: Vec> = vec![ + // Monotonic, no reset. + (0..5).map(|i| (i * g, 10.0 + 7.0 * i as f64)).collect(), + // Reset on the 2nd sample. + vec![(0, 100.0), (g, 5.0), (2 * g, 25.0), (3 * g, 40.0)], + // Reset on the last sample. + vec![(0, 10.0), (g, 40.0), (2 * g, 90.0), (3 * g, 3.0)], + // Multiple consecutive resets. + vec![ + (0, 50.0), + (g, 5.0), + (2 * g, 3.0), + (3 * g, 1.0), + (4 * g, 60.0), + ], + // Counter starting exactly at 0 (zero-point clamp with first_value 0). + (0..6).map(|i| (i * g, 12.0 * i as f64)).collect(), + // Near-zero start relative to large increase (zero clamp engages). + vec![(0, 2.0), (g, 40.0), (2 * g, 80.0), (3 * g, 102.0)], + // Flat counter (no increase, no reset) => increase 0. + (0..4).map(|i| (i * g, 42.0)).collect(), + // Exactly two samples (minimum for a defined rate). + vec![(0, 10.0), (g, 70.0)], + // Uneven spacing with a reset. + vec![(0, 10.0), (2_000, 15.0), (2_500, 4.0), (30_000, 25.0)], + ]; + + for samples in &cases { + let acc = feed(samples); + assert!( + close( + acc.corrected_increase(), + reference_corrected_increase(samples) + ), + "corrected_increase edge mismatch for {samples:?}" + ); + for (rs, re) in range_configs(samples) { + for is_rate in [false, true] { + let got = acc.extrapolated(is_rate, rs, re); + let want = reference_extrapolated(samples, rs, re, is_rate); + assert_eq!( + got.is_some(), + want.is_some(), + "presence mismatch for {samples:?}" + ); + if let (Some(g), Some(w)) = (got, want) { + assert!( + close(g, w), + "edge extrapolated mismatch (is_rate={is_rate}) got={g} want={w} \ + range=[{rs},{re}] samples={samples:?}" + ); + } + } + } + } + } + + // ======================================================================== + // Merge path: windowed ingest (tumbling buckets) then merge must reproduce + // the same result as a single window over the full series — including + // counter resets that straddle a window boundary. + // ======================================================================== + + fn merge_fold(mut windows: Vec) -> IncreaseAccumulator { + let mut acc = windows.remove(0); + for w in windows { + acc = IncreaseAccumulator::merge_pair(&acc, &w); + } + acc + } + + #[test] + fn synthetic_merge_matches_full_window() { + let mut rng = Lcg(0xD1B54A32D192ED03); + for iter in 0..800 { + let n = rng.range(4, 20); + let samples = synthetic_series( + &mut rng, + 1_000_000, + n, + 10_000, + 10_000, + [0.0, 0.2, 0.4][iter % 3], + ); + if samples[samples.len() - 1].0 == samples[0].0 { + continue; + } + + // Split into contiguous, non-overlapping windows at random cut points. + let n_windows = rng.range(1, (n - 1).max(1)); + let mut cuts: Vec = (1..n).collect(); + // pick `n_windows-1` interior cut points deterministically + let mut chosen = Vec::new(); + for _ in 0..n_windows.saturating_sub(1) { + if cuts.is_empty() { + break; + } + let idx = rng.range(0, cuts.len() - 1); + chosen.push(cuts.remove(idx)); + } + chosen.sort_unstable(); + chosen.dedup(); + + let mut windows = Vec::new(); + let mut start = 0usize; + for &cut in &chosen { + windows.push(feed(&samples[start..cut])); + start = cut; + } + windows.push(feed(&samples[start..])); + + let merged = merge_fold(windows); + + // The merged accumulator must equal a single window over all samples. + assert!( + close( + merged.corrected_increase(), + reference_corrected_increase(&samples) + ), + "merged corrected_increase mismatch: got={} ref={} cuts={:?} samples={:?}", + merged.corrected_increase(), + reference_corrected_increase(&samples), + chosen, + samples + ); + assert_eq!(merged.sample_count, samples.len() as u64); + assert_eq!(merged.starting_timestamp, samples[0].0); + assert_eq!(merged.last_seen_timestamp, samples[samples.len() - 1].0); + + for (rs, re) in range_configs(&samples) { + let got = merged.extrapolated(false, rs, re); + let want = reference_extrapolated(&samples, rs, re, false); + if let (Some(g), Some(w)) = (got, want) { + assert!( + close(g, w), + "merged extrapolated mismatch got={g} want={w} cuts={chosen:?} samples={samples:?}" + ); + } + } + } + } + + // ======================================================================== + // End-to-end: synthetic series through the real query paths. + // ======================================================================== + + fn instant_values(qr: QueryResult) -> Vec { + match qr { + QueryResult::Vector(iv) => iv.values.into_iter().map(|e| e.value).collect(), + other => panic!("expected vector, got {other:?}"), + } + } + + /// Drive several synthetic series through the instant pipeline (#1) and the + /// binary-arithmetic DataFusion path (#3) as single-population increase, and + /// confirm both equal the independent Prometheus reference over the exact + /// range boundaries the engine derives. + #[tokio::test(flavor = "multi_thread")] + async fn synthetic_end_to_end_paths_match_reference() { + const QT: f64 = 1000.0; // eval at t=1_000_000 ms + let g = 20_000i64; + // A few representative series that fit inside a [100s] window ending at QT. + let series: Vec> = vec![ + // monotonic + (0..5) + .map(|i| (910_000 + i * g, 10.0 + 30.0 * i as f64)) + .collect(), + // with a mid-window reset + vec![ + (910_000, 10.0), + (930_000, 90.0), + (950_000, 5.0), + (970_000, 45.0), + (990_000, 80.0), + ], + // near-zero start (zero clamp) + vec![(915_000, 1.0), (945_000, 50.0), (975_000, 130.0)], + ]; + + for samples in &series { + let acc = feed(samples); + let engine = create_engine_single_pop( + "m", + AggregationType::Increase, + vec!["host"], + vec![(Some(vec!["host-a".to_string()]), Box::new(acc))], + "increase(m[100s])", + ); + + // Range boundaries the engine derives for this query. + let ctx = engine + .build_query_execution_context_promql("increase(m[100s])".to_string(), QT) + .unwrap(); + let rs: i64 = ctx.metadata.query_kwargs[RANGE_START_MS_KWARG] + .parse() + .unwrap(); + let re: i64 = ctx.metadata.query_kwargs[RANGE_END_MS_KWARG] + .parse() + .unwrap(); + let want_inc = reference_extrapolated(samples, rs, re, false).unwrap(); + + // Path #1: instant pipeline. + let p1 = instant_values( + engine + .handle_query_promql("increase(m[100s])".to_string(), QT) + .expect("instant increase") + .1, + ); + assert_eq!(p1.len(), 1); + assert!( + close(p1[0], want_inc), + "path#1 got={} want={} samples={:?}", + p1[0], + want_inc, + samples + ); + + // Path #3: binary-arithmetic DataFusion path (× 1 keeps the value). + let p3 = instant_values( + engine + .handle_query_promql("increase(m[100s]) * 1".to_string(), QT) + .expect("binary increase") + .1, + ); + assert_eq!(p3.len(), 1); + assert!( + close(p3[0], want_inc), + "path#3 got={} want={} samples={:?}", + p3[0], + want_inc, + samples + ); + + // rate() via the binary path scaled by 2. Needs an engine whose config + // structurally matches `rate(...)` (config lookup is by function form), + // served by the same underlying Increase accumulator. + let rate_engine = create_engine_single_pop( + "m", + AggregationType::Increase, + vec!["host"], + vec![(Some(vec!["host-a".to_string()]), Box::new(feed(samples)))], + "rate(m[100s])", + ); + let want_rate = reference_extrapolated(samples, rs, re, true).unwrap(); + let p3r = instant_values( + rate_engine + .handle_query_promql("rate(m[100s]) * 2".to_string(), QT) + .expect("binary rate") + .1, + ); + assert_eq!(p3r.len(), 1); + assert!( + close(p3r[0], want_rate * 2.0), + "path#3 rate got={} want={} samples={:?}", + p3r[0], + want_rate * 2.0, + samples + ); + } + } + + /// Multi-population synthetic coverage: several sub-keys, each an independent + /// synthetic counter, extrapolated through the DataFusion path and checked + /// per-key against the reference. + #[tokio::test(flavor = "multi_thread")] + async fn synthetic_multipop_paths_match_reference() { + const QT: f64 = 1000.0; + let mut rng = Lcg(0x2545F4914F6CDD1D); + + // Build sub-keys with distinct synthetic series inside [900s, 1000s]. + // Key 0 is an explicit reset series so the multi-pop path is always + // exercised over counter-reset correction, not just monotonic counters. + let mut per_key_samples: Vec<(String, Vec<(i64, f64)>)> = vec![( + "host-0".to_string(), + vec![ + (905_000, 10.0), + (930_000, 90.0), + (955_000, 5.0), // reset + (980_000, 45.0), + ], + )]; + let mut increases = HashMap::new(); + for k in 1..5 { + let n = rng.range(3, 8); + let samples = synthetic_series(&mut rng, 905_000, n, 8_000, 8_000, 0.4); + // keep inside the 100s window + let samples: Vec<(i64, f64)> = samples + .into_iter() + .filter(|(t, _)| *t <= 1_000_000) + .collect(); + if samples.len() < 2 { + continue; + } + per_key_samples.push((format!("host-{k}"), samples)); + } + for (host, samples) in &per_key_samples { + let ep = host.replace("host-", "ep-"); + let key = KeyByLabelValues { + labels: vec![host.clone(), ep], + }; + increases.insert(key, feed(samples)); + } + let acc = MultipleIncreaseAccumulator::new_with_increases(increases); + + let engine = create_engine_single_pop_with_aggregated( + "m", + AggregationType::MultipleIncrease, + vec![], + vec!["host", "endpoint"], + vec![(None, Box::new(acc))], + "increase(m[100s])", + ); + let ctx = engine + .build_query_execution_context_promql("increase(m[100s])".to_string(), QT) + .unwrap(); + let rs: i64 = ctx.metadata.query_kwargs[RANGE_START_MS_KWARG] + .parse() + .unwrap(); + let re: i64 = ctx.metadata.query_kwargs[RANGE_END_MS_KWARG] + .parse() + .unwrap(); + + let result = engine + .handle_query_promql("increase(m[100s]) * 1".to_string(), QT) + .expect("multipop binary increase"); + let elems = match result.1 { + QueryResult::Vector(iv) => iv.values, + other => panic!("expected vector, got {other:?}"), + }; + assert_eq!(elems.len(), per_key_samples.len(), "one result per sub-key"); + + for (host, samples) in &per_key_samples { + let want = reference_extrapolated(samples, rs, re, false).unwrap(); + let got = elems + .iter() + .find(|e| e.labels.labels.first().map(|s| s.as_str()) == Some(host.as_str())) + .unwrap_or_else(|| panic!("missing result for {host}")) + .value; + assert!( + close(got, want), + "multipop {host}: got={got} want={want} samples={samples:?}" + ); + } + } +} diff --git a/asap-summary-ingest/templates/hashed_key_udfs/multipleincrease_.rs b/asap-summary-ingest/templates/hashed_key_udfs/multipleincrease_.rs index 9f1225c2..a7153d44 100644 --- a/asap-summary-ingest/templates/hashed_key_udfs/multipleincrease_.rs +++ b/asap-summary-ingest/templates/hashed_key_udfs/multipleincrease_.rs @@ -15,36 +15,55 @@ struct MeasurementData { starting_timestamp: i64, last_seen_measurement: f64, last_seen_timestamp: i64, + // Prometheus `counterCorrection`: sum of pre-reset values within the window. + // Field order MUST match IncreaseAccumulator's MeasurementData on the query + // side (rmp-serde encodes structs positionally). + counter_reset_correction: f64, + // Number of samples folded into this window (for extrapolation). + sample_count: u64, } #[udf] fn multipleincrease_(keys: Vec, values: Vec, timestamps: Vec) -> Option> { - // Create a new hashmap to store measurement data with timestamps - let mut per_key_storage: HashMap = HashMap::new(); - - // Iterate through the keys, values, and timestamps + // Group all (timestamp, value) samples per key first, because counter-reset + // detection requires processing them in timestamp order and the input vectors + // are not guaranteed to be ordered. + let mut per_key_samples: HashMap> = HashMap::new(); for (i, &key) in keys.iter().enumerate() { if i < values.len() && i < timestamps.len() { - let value = values[i]; - let timestamp = timestamps[i]; + per_key_samples + .entry(key) + .or_default() + .push((timestamps[i], values[i])); + } + } - let entry = per_key_storage.entry(key).or_insert(MeasurementData { - starting_measurement: value, - starting_timestamp: timestamp, - last_seen_measurement: value, - last_seen_timestamp: timestamp, - }); + let mut per_key_storage: HashMap = HashMap::new(); + for (key, mut samples) in per_key_samples { + samples.sort_by_key(|&(ts, _)| ts); - // Update last seen measurement and timestamp - entry.last_seen_measurement = value; - entry.last_seen_timestamp = timestamp; + let (first_ts, first_val) = samples[0]; + let mut entry = MeasurementData { + starting_measurement: first_val, + starting_timestamp: first_ts, + last_seen_measurement: first_val, + last_seen_timestamp: first_ts, + counter_reset_correction: 0.0, + sample_count: 1, + }; - // If this timestamp is earlier than our current starting timestamp, update starting values - //if timestamp < entry.starting_timestamp { - // entry.starting_measurement = value; - // entry.starting_timestamp = timestamp; - //} + for &(ts, value) in &samples[1..] { + // A drop below the previous value is a counter reset; add the + // pre-reset value back so the increase stays monotonic. + if value < entry.last_seen_measurement { + entry.counter_reset_correction += entry.last_seen_measurement; + } + entry.last_seen_measurement = value; + entry.last_seen_timestamp = ts; + entry.sample_count += 1; } + + per_key_storage.insert(key, entry); } let mut buf = Vec::new(); diff --git a/asap-summary-ingest/templates/udfs/multipleincrease_.rs b/asap-summary-ingest/templates/udfs/multipleincrease_.rs index 3ba2b429..c9c3dea4 100644 --- a/asap-summary-ingest/templates/udfs/multipleincrease_.rs +++ b/asap-summary-ingest/templates/udfs/multipleincrease_.rs @@ -15,36 +15,55 @@ struct MeasurementData { starting_timestamp: i64, last_seen_measurement: f64, last_seen_timestamp: i64, + // Prometheus `counterCorrection`: sum of pre-reset values within the window. + // Field order MUST match IncreaseAccumulator's MeasurementData on the query + // side (rmp-serde encodes structs positionally). + counter_reset_correction: f64, + // Number of samples folded into this window (for extrapolation). + sample_count: u64, } #[udf] fn multipleincrease_(keys: Vec<&str>, values: Vec, timestamps: Vec) -> Option> { - // Create a new hashmap to store measurement data with timestamps - let mut per_key_storage: HashMap = HashMap::new(); - - // Iterate through the keys, values, and timestamps + // Group all (timestamp, value) samples per key first, because counter-reset + // detection requires processing them in timestamp order and the input vectors + // are not guaranteed to be ordered. + let mut per_key_samples: HashMap> = HashMap::new(); for (i, &key) in keys.iter().enumerate() { if i < values.len() && i < timestamps.len() { - let value = values[i]; - let timestamp = timestamps[i]; + per_key_samples + .entry(key.to_string()) + .or_default() + .push((timestamps[i], values[i])); + } + } - let entry = per_key_storage.entry(key.to_string()).or_insert(MeasurementData { - starting_measurement: value, - starting_timestamp: timestamp, - last_seen_measurement: value, - last_seen_timestamp: timestamp, - }); + let mut per_key_storage: HashMap = HashMap::new(); + for (key, mut samples) in per_key_samples { + samples.sort_by_key(|&(ts, _)| ts); - // Update last seen measurement and timestamp - entry.last_seen_measurement = value; - entry.last_seen_timestamp = timestamp; + let (first_ts, first_val) = samples[0]; + let mut entry = MeasurementData { + starting_measurement: first_val, + starting_timestamp: first_ts, + last_seen_measurement: first_val, + last_seen_timestamp: first_ts, + counter_reset_correction: 0.0, + sample_count: 1, + }; - // If this timestamp is earlier than our current starting timestamp, update starting values - //if timestamp < entry.starting_timestamp { - // entry.starting_measurement = value; - // entry.starting_timestamp = timestamp; - //} + for &(ts, value) in &samples[1..] { + // A drop below the previous value is a counter reset; add the + // pre-reset value back so the increase stays monotonic. + if value < entry.last_seen_measurement { + entry.counter_reset_correction += entry.last_seen_measurement; + } + entry.last_seen_measurement = value; + entry.last_seen_timestamp = ts; + entry.sample_count += 1; } + + per_key_storage.insert(key, entry); } let mut buf = Vec::new();