rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 06:46:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(promql): evaluate range windows over subquery results

Browse Source 
Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
Dennis Zhuang
2026-10-17 17:30:40 +08:00
parent 7dbcac094a
commit 536b807889
2 changed files with 259 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/promql/src/range_eval.rs
+107 -19
View File
@@ -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<SeriesSchema>,
output: GridSeriesBuilder,
/// The grid of evaluated input series, such as the inner result of a subquery; `None` for
/// raw input.
input_grid: Option<GridSpec>,
series: Option<SeriesState>,
fields: Vec<FieldSamples>,
row: Vec<Option<f64>>,
}
struct SeriesState {
primary_key: Buffer,
/// `None` for evaluated input, which arrives as one complete series.
primary_key: Option<Buffer>,
metric: MetricName,
last_ts: i64,
/// Next grid index to evaluate.
@@ -146,11 +150,6 @@ impl RangeEvaluator {
kernel: WindowKernel,
step_invariant: bool,
) -> Result<Self> {
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<Option<Series>> {
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<Option<Series>> {
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<Option<Series>> {
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 &timestamp in chunk.time.as_buffer().iter() {
let mut samples: Vec<_> = fields.iter().map(|field| field.iter()).collect();
for &timestamp 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<i64> {
fold_offset_of(grid.spec(), window)
}
fn fold_offset_of(grid: GridSpec, window: RangeWindow) -> Result<i64> {
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<i64> {
))
})
}
/// 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<Option<GridSpec>> {
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)))
}
src/promql/src/range_eval/tests.rs
+152 -30
View File
@@ -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<Row> {
.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<Row>, Vec<Range<usize>>) {
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<Range<u
}
/// Applies the probe to the valid samples of every reference window.
fn reference(rows: &[Row], grid: &EvalGrid, window: RangeWindow, fields: usize) -> 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<RangeFunction> {
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<EvalGrid>,
fields: usize,
host: u64,
) -> (Option<Series>, Vec<Row>) {
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<i64> = ((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));
}