fix(query-engine): counter-reset correction + Prometheus extrapolation for rate()/increase() (#476)

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,

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,

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::<i64>()
+            .ok()?;
+        let end = self
+            .metadata
+            .query_kwargs
+            .get(RANGE_END_MS_KWARG)?
+            .parse::<i64>()
+            .ok()?;
+        Some((start, end))
+    }
+
     /// Map aggregation type to SketchType (SummaryType) for the value accumulator
     fn map_aggregation_type_to_summary_type(&self) -> Result<SketchType, DataFusionError> {
         Self::agg_type_to_sketch_type(self.agg_info.aggregation_type_for_value)

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<u8> {
     if let Some(set_acc) = acc.as_any().downcast_ref::<SetAggregatorAccumulator>() {
         return set_acc.serialize_to_bytes_arroyo();
     }
+    if let Some(inc_acc) = acc.as_any().downcast_ref::<IncreaseAccumulator>() {
+        return inc_acc.serialize_to_bytes_arroyo();
+    }
     if let Some(inc_acc) = acc.as_any().downcast_ref::<MultipleIncreaseAccumulator>() {
         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(),
         )),

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<HashMap<String, String>> {
         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,
         }
     }

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<Box<dyn AggregateCore>> = 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!(

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<HashMap<String, String>, 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, &timestamps)
             .map_err(|e| {
                 warn!("{}", e);
                 e