From d6f004aebb9abc76a4937d7d5cc43833a4987418 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sun, 18 Oct 2026 00:24:15 +0800 Subject: [PATCH] fix(promql): reject subquery grids longer than u64 and pin anchored predict_linear Signed-off-by: Dennis Zhuang --- src/promql/src/range_eval.rs | 5 +-- src/promql/src/range_eval/tests.rs | 51 ++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/promql/src/range_eval.rs b/src/promql/src/range_eval.rs index 6fbb9a442d2..8f3df3315b6 100644 --- a/src/promql/src/range_eval.rs +++ b/src/promql/src/range_eval.rs @@ -561,11 +561,12 @@ pub fn subquery_grid( if first > last { return Ok(None); } - let len = (last - first) / step + 1; + let len = u64::try_from((last - first) / step + 1) + .map_err(|_| DataFusionError::Plan("subquery grid has too many points".into()))?; let first = i64::try_from(first).map_err(|_| { DataFusionError::Plan(format!("subquery grid start {first} ms is out of range")) })?; - GridSpec::try_new(first, step_ms, len as u64) + GridSpec::try_new(first, step_ms, len) .map(Some) .map_err(|error| DataFusionError::External(Box::new(error))) } diff --git a/src/promql/src/range_eval/tests.rs b/src/promql/src/range_eval/tests.rs index 6ba09593ded..a0cc750f3e3 100644 --- a/src/promql/src/range_eval/tests.rs +++ b/src/promql/src/range_eval/tests.rs @@ -1062,3 +1062,54 @@ fn samples_after_the_last_window_are_validated_but_not_retained() { Err(DataFusionError::Execution(_)) )); } + +/// An anchored `predict_linear` regresses the one anchored window around each grid point, as +/// Prometheus does with `enh.Ts`; over `y = t` that predicts `T` whatever the offset. +#[test] +fn anchored_predict_linear_regresses_around_each_point() { + let rows: Vec = (0..=60) + .step_by(5) + .map(|ts| (ts, vec![Some(ts as f64)])) + .collect(); + for offset_ms in [0, 10] { + let window = RangeWindow { + range_ms: 30, + offset_ms, + anchor_ms: Some(50), + }; + let mut evaluator = RangeEvaluator::try_new( + &raw_schema(1), + grid(100, 10, 3), + window, + RangeFunction::PredictLinear(0.0), + ) + .unwrap(); + let series = evaluator.push(&chunk(&rows, 1, 0, true)).unwrap().unwrap(); + let values: Vec<_> = (0..3) + .map(|index| match series.fields.as_slice()[0].sample(index) { + SampleRef::Float(value) => value, + _ => panic!("expected a float"), + }) + .collect(); + for (value, expected) in values.into_iter().zip([100.0, 110.0, 120.0]) { + assert!( + (value - expected).abs() < 1e-9, + "offset {offset_ms}: {value}" + ); + } + } +} + +#[test] +fn rejects_a_subquery_grid_longer_than_u64() { + let outer = GridSpec::try_new(i64::MIN, i64::MAX, 3).unwrap(); + let window = RangeWindow { + range_ms: 7, + offset_ms: -6, + anchor_ms: None, + }; + assert!(matches!( + subquery_grid(outer, window, 1), + Err(DataFusionError::Plan(_)) + )); +}