From 536b807889ed76455cc87cac779b41fdf2e11a45 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sat, 17 Oct 2026 17:30:40 +0800 Subject: [PATCH] feat(promql): evaluate range windows over subquery results Signed-off-by: Dennis Zhuang --- src/promql/src/range_eval.rs | 126 +++++++++++++++++--- src/promql/src/range_eval/tests.rs | 182 ++++++++++++++++++++++++----- 2 files changed, 259 insertions(+), 49 deletions(-) diff --git a/src/promql/src/range_eval.rs b/src/promql/src/range_eval.rs index 1dbbfccf129..06a73e83903 100644 --- a/src/promql/src/range_eval.rs +++ b/src/promql/src/range_eval.rs @@ -21,8 +21,8 @@ use datafusion::arrow::buffer::Buffer; use datafusion::error::{DataFusionError, Result}; use datatypes::series::MetricName; use datatypes::series::raw::RawSeriesChunk; -use datatypes::series::samples::{SampleRef, Series, SeriesSchema}; -use datatypes::series::time::EvalGrid; +use datatypes::series::samples::{SampleRef, Series, SeriesFields, SeriesSchema}; +use datatypes::series::time::{EvalGrid, GridSpec, TimeAxis}; use crate::grid_series::GridSeriesBuilder; @@ -72,13 +72,17 @@ pub struct RangeEvaluator { kernel: WindowKernel, output_schema: Arc, output: GridSeriesBuilder, + /// The grid of evaluated input series, such as the inner result of a subquery; `None` for + /// raw input. + input_grid: Option, series: Option, fields: Vec, row: Vec>, } struct SeriesState { - primary_key: Buffer, + /// `None` for evaluated input, which arrives as one complete series. + primary_key: Option, metric: MetricName, last_ts: i64, /// Next grid index to evaluate. @@ -146,11 +150,6 @@ impl RangeEvaluator { kernel: WindowKernel, step_invariant: bool, ) -> Result { - if input.eval_grid.is_some() { - return Err(DataFusionError::Plan( - "range evaluation input must be a raw series schema".into(), - )); - } if window.range_ms <= 0 { return Err(DataFusionError::Plan(format!( "range must be positive, got {} ms", @@ -173,6 +172,7 @@ impl RangeEvaluator { replay, kernel, output: GridSeriesBuilder::try_new(output_schema.clone())?, + input_grid: input.eval_grid.as_ref().map(|grid| grid.spec()), output_schema, series: None, fields: input @@ -196,8 +196,13 @@ impl RangeEvaluator { /// /// The evaluator must not be reused after an error. pub fn push(&mut self, chunk: &RawSeriesChunk) -> Result> { + if self.input_grid.is_some() { + return Err(DataFusionError::Execution( + "raw series chunk pushed to an evaluator of evaluated series".into(), + )); + } self.check_chunk(chunk)?; - self.ingest(chunk)?; + self.ingest(chunk.time.as_buffer(), &chunk.fields, true)?; let last_ts = chunk.time.as_buffer()[chunk.time.len() - 1]; if !chunk.is_last { // Later chunks only hold samples after `last_ts`. @@ -206,8 +211,44 @@ impl RangeEvaluator { ))?; return Ok(None); } + self.finish_series() + } + + /// Evaluates one complete evaluated series, such as the inner result of a subquery: the + /// value at each grid point is a sample at that time. Stale markers are kept as values, + /// since the input is computed. + pub fn push_series(&mut self, series: &Series) -> Result> { + let TimeAxis::Grid(range) = &series.time else { + return Err(DataFusionError::Execution( + "evaluated range input must be a grid series".into(), + )); + }; + if self.input_grid != Some(range.grid().spec()) + || series.fields.len() != self.fields.len() + || series.fields.iter().any(|field| field.len() != range.len()) + { + return Err(DataFusionError::Execution( + "evaluated series does not match the input schema".into(), + )); + } + if self.series.is_some() { + return Err(DataFusionError::Execution( + "evaluated series pushed inside a raw series".into(), + )); + } + self.series = Some(SeriesState { + primary_key: None, + metric: series.metric.clone(), + last_ts: 0, + cursor: 0, + }); + self.ingest(&range.timestamps(), &series.fields, false)?; + self.finish_series() + } + + fn finish_series(&mut self) -> Result> { self.evaluate(None)?; - let state = self.series.take().expect("checked by check_chunk"); + let state = self.series.take().expect("set by the caller"); for field in &mut self.fields { field.timestamps.clear(); field.values.clear(); @@ -227,6 +268,7 @@ impl RangeEvaluator { } fn check_chunk(&mut self, chunk: &RawSeriesChunk) -> Result<()> { + let primary_key = Some(chunk.primary_key.as_slice()); let time = chunk.time.as_buffer(); if time.is_empty() || chunk.fields.len() != self.fields.len() @@ -247,14 +289,14 @@ impl RangeEvaluator { ) })?; self.series = Some(SeriesState { - primary_key: chunk.primary_key.clone(), + primary_key: Some(chunk.primary_key.clone()), metric, last_ts: time[time.len() - 1], cursor: 0, }); } Some(state) => { - if state.primary_key.as_slice() != chunk.primary_key.as_slice() + if state.primary_key.as_ref().map(|key| key.as_slice()) != primary_key || chunk.metric.as_ref() != Some(&state.metric) { return Err(DataFusionError::Execution( @@ -273,18 +315,25 @@ impl RangeEvaluator { Ok(()) } - fn ingest(&mut self, chunk: &RawSeriesChunk) -> Result<()> { - let state = self.series.as_ref().expect("checked by check_chunk"); + /// Retains the samples of `fields` at `timestamps`; with `drop_stale`, a row holding a stale + /// marker in any field is dropped. + fn ingest( + &mut self, + timestamps: &[i64], + fields: &SeriesFields, + drop_stale: bool, + ) -> Result<()> { + let state = self.series.as_ref().expect("set by the caller"); if state.cursor >= self.points_read() { // Every window that reads samples has been evaluated. return Ok(()); } - let mut samples: Vec<_> = chunk.fields.iter().map(|field| field.iter()).collect(); - for ×tamp in chunk.time.as_buffer().iter() { + let mut samples: Vec<_> = fields.iter().map(|field| field.iter()).collect(); + for ×tamp in timestamps { self.row.clear(); let mut stale = false; for field in &mut samples { - match field.next().expect("field length checked by check_chunk") { + match field.next().expect("field length checked by the caller") { SampleRef::Missing => self.row.push(None), SampleRef::Float(value) => { stale |= is_prometheus_stale_nan(value); @@ -297,7 +346,7 @@ impl RangeEvaluator { } } } - if stale { + if stale && drop_stale { continue; } // Like RangeManipulate, any kept sample whose shifted time leaves i64 fails the series. @@ -442,12 +491,16 @@ fn evaluate_row( /// The offset samples are shifted by: the selector offset, or for `@` the offset that places /// the anchored window at the first grid point. fn fold_offset(grid: &EvalGrid, window: RangeWindow) -> Result { + fold_offset_of(grid.spec(), window) +} + +fn fold_offset_of(grid: GridSpec, window: RangeWindow) -> Result { let Some(anchor) = window.anchor_ms else { return Ok(window.offset_ms); }; anchor .checked_sub(window.offset_ms) - .and_then(|end| grid.spec().start_ms.checked_sub(end)) + .and_then(|end| grid.start_ms.checked_sub(end)) .ok_or_else(|| { DataFusionError::Plan(format!( "@ {anchor} ms with offset {} ms is out of range", @@ -455,3 +508,38 @@ fn fold_offset(grid: &EvalGrid, window: RangeWindow) -> Result { )) }) } + +/// The inner grid of a subquery read by `window` from the outer grid `outer`: the multiples of +/// `step_ms` after `start - F - range` up to the last window end, where `F` is the fold offset +/// of `window`; with `@`, the only window end is `start - F`. `None` when no multiple falls in +/// the windows. +pub fn subquery_grid( + outer: GridSpec, + window: RangeWindow, + step_ms: i64, +) -> Result> { + if step_ms <= 0 { + return Err(DataFusionError::Plan(format!( + "subquery step must be positive, got {step_ms} ms" + ))); + } + let fold_offset = i128::from(fold_offset_of(outer, window)?); + let step = i128::from(step_ms); + let start = i128::from(outer.start_ms); + let last_eval = match window.anchor_ms { + Some(_) => start, + None => start + i128::from(outer.step_ms) * (i128::from(outer.len) - 1), + }; + let first = ((start - fold_offset - i128::from(window.range_ms)).div_euclid(step) + 1) * step; + let last = last_eval - fold_offset; + if first > last { + return Ok(None); + } + let len = (last - first) / step + 1; + 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) + .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 2fde06b00a3..8002447fd53 100644 --- a/src/promql/src/range_eval/tests.rs +++ b/src/promql/src/range_eval/tests.rs @@ -30,19 +30,20 @@ use datatypes::series::histogram::{ BucketSideView, CounterResetHint, HistogramBatchBuilder, HistogramView, }; use datatypes::series::raw::{RawSeriesChunk, RawTimestamps}; -use datatypes::series::samples::{FieldSpec, Samples, SeriesFields, SeriesSchema}; +use datatypes::series::samples::{FieldSpec, Samples, Series, SeriesFields, SeriesSchema}; use datatypes::series::time::{EvalGrid, GridSpec}; use datatypes::series::{Label, MetricName}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; -use super::{FloatWindow, RangeEvaluator, RangeFunction, RangeWindow, WindowKernel}; +use super::{FloatWindow, RangeEvaluator, RangeFunction, RangeWindow, WindowKernel, subquery_grid}; use crate::functions::{ AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, StddevOverTime, StdvarOverTime, SumOverTime, }; -use crate::grid_series::reference::{Dense, assert_series_eq, special_value}; +use crate::grid_series::GridSeriesBuilder; +use crate::grid_series::reference::{Dense, assert_series_eq, random_dense, special_value}; use crate::range_array::RangeArray; /// One raw row: source timestamp and a value per field, `None` when missing. @@ -195,20 +196,23 @@ fn random_rows(rng: &mut StdRng, fields: usize) -> Vec { .collect() } -/// Rows without stale markers and, for each grid point, the range of those rows in its window, -/// found by scanning all rows independently of chunking and cursors. +/// Rows on the evaluation timeline, without stale rows when `drop_stale`, and for each grid +/// point the range of those rows in its window, found by scanning all rows independently of +/// chunking and cursors. fn reference_windows( rows: &[Row], grid: &EvalGrid, window: RangeWindow, + drop_stale: bool, ) -> (Vec, Vec>) { let kept: Vec<_> = rows .iter() .filter(|(_, values)| { - !values - .iter() - .flatten() - .any(|value| value.to_bits() == PROMETHEUS_STALE_NAN_BITS) + !drop_stale + || !values + .iter() + .flatten() + .any(|value| value.to_bits() == PROMETHEUS_STALE_NAN_BITS) }) .map(|(timestamp, values)| (timestamp + fold_offset(grid, window), values.clone())) .collect(); @@ -250,8 +254,14 @@ fn windows_of(rows: &[Row], grid: &EvalGrid, window: RangeWindow) -> Vec Dense { - let (kept, windows) = reference_windows(rows, grid, window); +fn reference( + rows: &[Row], + grid: &EvalGrid, + window: RangeWindow, + fields: usize, + drop_stale: bool, +) -> Dense { + let (kept, windows) = reference_windows(rows, grid, window, drop_stale); grid.timestamps() .iter() .zip(windows) @@ -291,8 +301,9 @@ fn udf_reference( grid: &EvalGrid, window: RangeWindow, fields: usize, + drop_stale: bool, ) -> Dense { - let (kept, _) = reference_windows(rows, grid, window); + let (kept, _) = reference_windows(rows, grid, window, drop_stale); let mut dense = vec![vec![None; fields]; grid.timestamps().len()]; for field in 0..fields { let rows: Vec<_> = if matches!( @@ -455,7 +466,7 @@ fn evaluator_matches_brute_force_reference() { let mut evaluator = probe_evaluator(fields, grid.clone(), window); for host in 0..3 { let rows = random_rows(&mut rng, fields); - let expected = reference(&rows, &grid, window, fields); + let expected = reference(&rows, &grid, window, fields, true); let mut result = None; for piece in random_chunks(&mut rng, &rows, fields, host) { assert!( @@ -470,9 +481,9 @@ fn evaluator_matches_brute_force_reference() { } } -#[test] -fn functions_match_range_udfs() { - for function in [ +/// Every function, with parameters inside, at and outside their valid ranges. +fn all_functions() -> Vec { + vec![ RangeFunction::Rate, RangeFunction::Increase, RangeFunction::Delta, @@ -519,7 +530,12 @@ fn functions_match_range_udfs() { smoothing: f64::NAN, trend: 0.1, }, - ] { + ] +} + +#[test] +fn functions_match_range_udfs() { + for function in all_functions() { for seed in 0..300 { let mut rng = StdRng::seed_from_u64(seed); let fields = rng.random_range(1..=3); @@ -529,7 +545,7 @@ fn functions_match_range_udfs() { .unwrap(); for host in 0..3 { let rows = random_rows(&mut rng, fields); - let expected = udf_reference(function, &rows, &grid, window, fields); + let expected = udf_reference(function, &rows, &grid, window, fields, true); let mut result = None; for piece in random_chunks(&mut rng, &rows, fields, host) { result = evaluator.push(&piece).unwrap(); @@ -593,18 +609,6 @@ fn rejects_invalid_parameters() { Err(DataFusionError::Plan(_)) )); - let mut evaluated = raw_schema(1); - evaluated.eval_grid = Some(grid(0, 10, 4)); - let window = RangeWindow { - range_ms: 10, - offset_ms: 0, - anchor_ms: None, - }; - assert!(matches!( - RangeEvaluator::try_new(&evaluated, grid(0, 10, 4), window, RangeFunction::Rate), - Err(DataFusionError::Plan(_)) - )); - let far_offset = RangeWindow { range_ms: 10, offset_ms: i64::MAX - 20, @@ -756,3 +760,121 @@ fn retains_only_samples_later_windows_need() { assert!(evaluator.fields[0].timestamps.len() <= 60); } } + +/// A random evaluated series on `grid`; values may carry the stale bit pattern, which computed +/// input keeps as an ordinary NaN. +fn random_evaluated( + rng: &mut StdRng, + grid: &Arc, + fields: usize, + host: u64, +) -> (Option, Vec) { + let mut dense = random_dense(rng, grid.timestamps().len(), fields); + for value in dense.iter_mut().flatten().flatten() { + if rng.random_bool(0.05) { + *value = stale(); + } + } + let schema = Arc::new(SeriesSchema { + fields: raw_schema(fields).fields, + eval_grid: Some(grid.clone()), + }); + let mut builder = GridSeriesBuilder::try_new(schema).unwrap(); + for (index, row) in dense.iter().enumerate() { + builder.push(index, row).unwrap(); + } + let rows = grid + .timestamps() + .iter() + .zip(&dense) + .filter(|(_, row)| row.iter().any(Option::is_some)) + .map(|(timestamp, row)| (*timestamp, row.clone())) + .collect(); + (builder.finish(metric(host)).unwrap(), rows) +} + +#[test] +fn subquery_series_match_references() { + for seed in 0..300 { + let mut rng = StdRng::seed_from_u64(seed); + let fields = rng.random_range(1..=3); + let (outer, window) = random_window(&mut rng); + let step = rng.random_range(1..4_000); + let Some(inner) = subquery_grid(outer.spec(), window, step).unwrap() else { + continue; + }; + let inner = Arc::new(EvalGrid::try_new(inner).unwrap()); + let schema = SeriesSchema { + fields: raw_schema(fields).fields, + eval_grid: Some(inner.clone()), + }; + let (Some(series), rows) = random_evaluated(&mut rng, &inner, fields, seed) else { + continue; + }; + let mut probe = + RangeEvaluator::with_kernel(&schema, outer.clone(), window, Box::new(probe), true) + .unwrap(); + let expected = reference(&rows, &outer, window, fields, false); + assert_series_eq( + probe.push_series(&series).unwrap().as_ref(), + &expected, + &metric(seed), + ); + for function in all_functions() { + let mut evaluator = + RangeEvaluator::try_new(&schema, outer.clone(), window, function).unwrap(); + let expected = udf_reference(function, &rows, &outer, window, fields, false); + assert_series_eq( + evaluator.push_series(&series).unwrap().as_ref(), + &expected, + &metric(seed), + ); + } + } +} + +#[test] +fn subquery_grid_covers_every_step_multiple_in_the_windows() { + for seed in 0..2_000 { + let mut rng = StdRng::seed_from_u64(seed); + let (outer, window) = random_window(&mut rng); + let spec = outer.spec(); + let step = rng.random_range(1..4_000); + let fold = fold_offset(&outer, window); + let last_eval = match window.anchor_ms { + Some(_) => spec.start_ms, + None => *outer.timestamps().last().unwrap(), + }; + let lower = spec.start_ms - fold - window.range_ms; + let expected: Vec = ((lower.div_euclid(step) - 1) * step..=last_eval - fold) + .step_by(step as usize) + .filter(|point| *point > lower && point.rem_euclid(step) == 0) + .collect(); + let actual = subquery_grid(spec, window, step) + .unwrap() + .map(|grid| EvalGrid::try_new(grid).unwrap().timestamps().to_vec()) + .unwrap_or_default(); + assert_eq!(actual, expected, "seed {seed}"); + } +} + +#[test] +fn rejects_input_of_the_other_kind() { + let window = RangeWindow { + range_ms: 10, + offset_ms: 0, + anchor_ms: None, + }; + let evaluated = SeriesSchema { + fields: raw_schema(1).fields, + eval_grid: Some(grid(0, 10, 4)), + }; + let mut evaluator = + RangeEvaluator::try_new(&evaluated, grid(0, 10, 4), window, RangeFunction::Rate).unwrap(); + execution_error(evaluator.push(&chunk(&[(1, vec![Some(1.0)])], 1, 0, true))); + + let mut builder = GridSeriesBuilder::try_new(Arc::new(evaluated)).unwrap(); + builder.push(0, &[Some(1.0)]).unwrap(); + let series = builder.finish(metric(0)).unwrap().unwrap(); + execution_error(probe_evaluator(1, grid(0, 10, 4), window).push_series(&series)); +}