From 3ad5f963285e51166fd8789920e4b3bf2cf60423 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 20 Oct 2026 19:26:29 +0800 Subject: [PATCH] feat(promql): compile a parameterised range function (#9743) * feat(promql): compile a parameterised range function `quantile_over_time(q, x[r])` now compiles: the quantile is bound as a number literal and the range selector is bound by the same path a function without parameters takes. The two shapes share one binder, so the range, the offset, the anchor and the grid cannot drift between them. The quantile is not range-checked. Prometheus binds any value and the kernel answers `-Infinity` below zero and `+Infinity` above one, so refusing it here would be this compiler inventing a rule the engine does not have. A parameter in parentheses binds the same quantile: the planner already looks through the parentheses when it reads a literal, so a query written `quantile_over_time((0.9), x[5m])` keeps working on this path too. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): cover a parenthesised selector of a parameterised function The unparen rule applies to the selector a parameterised function binds, not only to the one a plain range function binds: a query that writes its range selector in parentheses asks for the same tree. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(promql): name the real run-time parameter, and reuse `unparen` The doc said `1 + 1` is computed at run time. It is not: `try_build_literal_expr` folds `Paren`, `Unary` and literal-only `Binary` in the existing planner, so `quantile_over_time(1 - 0.1, x[5m])` and `quantile_over_time(-(0.9), x[5m])` plan today and reach this compiler unfolded. Refusing them is the agreed behaviour for this stage, not a run-time-only case, and the doc now says so; `time()` is the real run-time example. A test pins both folded forms as refused, so implementing folding later is a deliberate change rather than a silent one. `scalar_parameter` open-coded the parenthesis loop `unparen` already implements, so the rule for transparent wrappers now lives in one place. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/promql/native_compile.rs | 122 +++++++++- src/query/src/promql/native_compile/tests.rs | 227 ++++++++++++++++++- 2 files changed, 336 insertions(+), 13 deletions(-) diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index 831f10eab6d..88e5023aa42 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -16,11 +16,12 @@ //! //! This is the front half of the native evaluation model: it reads the parsed statement and writes //! down what it means as bound nodes. It runs nothing, and nothing calls it yet, so the subset it -//! accepts is deliberately two shapes: +//! accepts is deliberately three shapes: //! //! ```text -//! e.g. x -//! ([]) e.g. rate(x[5m]) +//! e.g. x +//! ([]) e.g. rate(x[5m]) +//! (, []) e.g. quantile_over_time(0.9, x[5m]) //! ``` //! //! which become `Input` (raw samples) → `Instant` (the last sample within the statement's lookback) @@ -36,6 +37,7 @@ //! another. use std::sync::Arc; +use std::time::Duration; use datafusion::error::{DataFusionError, Result}; use datatypes::series::SeriesOrdering; @@ -46,14 +48,17 @@ use promql::native::plan::{ }; use promql::range_eval::{RangeFunction, RangeWindow}; use promql_parser::parser::{ - AtModifier, Call, EvalStmt, Expr as PromExpr, MatrixSelector, Offset, VectorSelector, + AtModifier, Call, EvalStmt, Expr as PromExpr, MatrixSelector, NumberLiteral, Offset, + VectorSelector, }; use crate::promql::label_values::signed_millis_since_epoch; /// The shapes this compiler accepts, quoted in every rejection. -const SUPPORTED: &str = "a bare selector or `([])`, for example \ - `x` or `rate(x[5m])`"; +const SUPPORTED: &str = "a bare selector, a range function over a selector — \ + `([])` — or a parameterised one, \ + `(, [])`, for example `x`, \ + `rate(x[5m])` or `quantile_over_time(0.9, x[5m])`"; /// Compiles `stmt` into a tree that reads `input` and either selects from it or folds it with the /// statement's one window. @@ -90,19 +95,31 @@ pub fn compile( let PromExpr::Call(Call { func, args }) = unparen(&stmt.expr) else { return Err(not_implemented(&format!("the expression `{}`", stmt.expr))); }; - let Some(function) = range_function(func.name) else { + let Some(signature) = range_signature(func.name) else { return Err(not_implemented(&format!( "a call to `{}`, which has no native range kernel", func.name ))); }; - let [argument] = args.args.as_slice() else { + // The count is the signature's, so the rejection names what this function takes rather than + // what one argument happens to be. + let expected = signature.parameters() + 1; + if args.args.len() != expected { return Err(not_implemented(&format!( - "a call to `{}` with {} argument(s)", + "a call to `{}` with {} argument(s), where it takes {expected}: {}", func.name, - args.args.len() + args.args.len(), + signature.arguments() ))); + } + // The parameters come first and the range selector is last, for both signatures: a function + // without parameters has no `parameters` to read. + let (parameters, selector) = args.args.split_at(signature.parameters()); + let function = match signature { + RangeSignature::Selector(function) => function, + RangeSignature::Scalar(bind) => bind(scalar_parameter(func.name, ¶meters[0])?), }; + let argument = &selector[0]; let PromExpr::MatrixSelector(MatrixSelector { vs, range }) = unparen(argument) else { return Err(not_implemented(&format!( "a call to `{}` whose argument is `{}`", @@ -119,6 +136,22 @@ pub fn compile( series may not interleave", )); } + window_selection(stmt, function, vs, range, input, ordering) +} + +/// Binds `function` over the range selector `vs[range]` into a window over the raw input. +/// +/// Both shapes reach this: the range, the offset, the anchor, the grid and the output are the same +/// whichever function reads them, so a parameterised function only decides *what* folds the window, +/// never *which* window is folded. +fn window_selection( + stmt: &EvalStmt, + function: RangeFunction, + vs: &VectorSelector, + range: &Duration, + input: Arc, + ordering: SeriesOrdering, +) -> Result { // An `@` anchor is step-invariant: every step folds the same samples, which the window's // `anchor_ms` tells `RangeEvaluator` to replay over the grid. let range_ms = range_millis(range.as_millis())?; @@ -313,6 +346,75 @@ fn grid_error(error: datatypes::error::Error) -> DataFusionError { DataFusionError::Plan(format!("the statement's grid is invalid: {error}")) } +/// What a call to a range function has to be: the parameters it reads, and how to bind them. +#[derive(Clone, Copy)] +enum RangeSignature { + /// A function that reads only its range selector, as in `rate(x[5m])`. + Selector(RangeFunction), + /// A function that reads one number literal before its range selector, and the function that + /// number binds, as in `quantile_over_time(0.9, x[5m])`. + /// + /// The number is kept out of the variant name because only the call site knows it. + Scalar(fn(f64) -> RangeFunction), +} + +impl RangeSignature { + /// How many scalar parameters the function reads before its range selector. + fn parameters(self) -> usize { + match self { + Self::Selector(_) => 0, + Self::Scalar(_) => 1, + } + } + + /// What the call's arguments are, for the arity rejection to name them. + fn arguments(self) -> &'static str { + match self { + Self::Selector(_) => "its range selector", + Self::Scalar(_) => "a number and its range selector", + } + } +} + +/// The signature the PromQL name calls, or `None` when this compiler has no kernel for it. +/// +/// The map is explicit on purpose, as in [`range_function`]: a name that is not listed is rejected, +/// so supporting one more function is a decision made in one place rather than a consequence of a +/// name looking familiar. A parameterised function is named here rather than in the map, because +/// its name alone cannot produce the variant the tree binds. +fn range_signature(name: &str) -> Option { + match name { + "quantile_over_time" => Some(RangeSignature::Scalar(RangeFunction::QuantileOverTime)), + _ => range_function(name).map(RangeSignature::Selector), + } +} + +/// The number a parameterised range function was given, as the literal it was written as. +/// +/// A number literal is the only parameter this compiler binds. A parameter whose value is only +/// known at run time, `time()` being the real case, would make the kernel's constant a query-time +/// value, which the tree has no field for. +/// +/// A parameter written as an expression the existing planner folds to a literal is refused here +/// and is a gap this stage leaves open on purpose: `quantile_over_time(1 - 0.1, x[5m])` and +/// `quantile_over_time(-(0.9), x[5m])` plan today through `try_build_literal_expr`, which folds +/// `Paren`, `Unary` and literal-only `Binary`, and reach this function still unfolded. They are +/// refused rather than answered from the wrong number, and folding them here would be the work +/// to close the gap. +/// +/// A quantile outside `[0, 1]` is not refused here: Prometheus binds it and the kernel reports +/// `-Infinity` below zero and `+Infinity` above one, so refusing it would be this compiler +/// inventing a rule the query engine does not have. +fn scalar_parameter(name: &str, parameter: &PromExpr) -> Result { + match unparen(parameter) { + PromExpr::NumberLiteral(NumberLiteral { val }) => Ok(*val), + found => Err(not_implemented(&format!( + "a call to `{name}` whose first argument is `{found}`, where it has to be a number \ + literal" + ))), + } +} + /// The range function a PromQL name denotes, or `None` when this compiler has no kernel for it. /// /// The map is explicit on purpose: a name that is not listed here is rejected, so supporting one diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 1918ce3ec2e..5e375858111 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/src/query/src/promql/native_compile/tests.rs @@ -22,7 +22,7 @@ use promql::native::plan::{ BoundInstant, BoundWindow, NativeExpr, NativeNode, NativeTree, NativeType, }; use promql::range_eval::RangeFunction; -use promql_parser::parser::{EvalStmt, Expr as PromExpr}; +use promql_parser::parser::{EvalStmt, Expr as PromExpr, NumberLiteral}; use super::compile; @@ -69,6 +69,19 @@ fn with_selector_range(range: Duration) -> EvalStmt { stmt } +/// `query`'s statement with its call's arguments edited by `edit`. +/// +/// The parser checks arity, so a call with the wrong number of arguments cannot be parsed: the only +/// way to hand one to the compiler is to edit a tree the parser accepted. +fn with_call_arguments(query: &str, edit: impl FnOnce(&mut Vec>)) -> EvalStmt { + let mut stmt = eval_stmt(query); + let PromExpr::Call(call) = &mut stmt.expr else { + unreachable!("`{query}` parses to a call"); + }; + edit(&mut call.args.args); + stmt +} + /// The window at the root of `tree`, for the range-function shape. fn root_window(tree: &NativeTree) -> (&NativeNode, &BoundWindow) { let NativeExpr::Window { input, op } = &tree.root().expr else { @@ -169,6 +182,214 @@ fn the_supported_range_functions_map_to_their_kernels() { } } +/// `quantile_over_time(0.9, x[5m])` binds the quantile the query wrote, over the same window a +/// function without parameters would fold: a parameter decides what folds the window, not which +/// window is folded. +#[test] +fn a_quantile_over_time_binds_its_quantile_and_its_window() { + let tree = compile( + &eval_stmt("quantile_over_time(0.9, x[5m])"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + + let (_, op) = root_window(&tree); + assert_eq!(op.function, RangeFunction::QuantileOverTime(0.9)); + assert_eq!(op.window.range_ms, 300_000); + assert_eq!(op.window.offset_ms, 0); + assert_eq!(op.window.anchor_ms, None); + assert!(op.input_grid.is_none(), "raw input has no grid to declare"); + + // The statement's grid: 0, 5s, ..., 100s. + let spec = op.output_grid.spec(); + assert_eq!( + (spec.start_ms, spec.step_ms, spec.len), + (0, 5_000, 21), + "the window does not land on the statement's grid" + ); +} + +/// A quantile outside `[0, 1]` compiles, because it is not an error in Prometheus: the kernel +/// returns `-Infinity` below zero, `+Infinity` above one, and a real quantile at either bound. A +/// compiler that refused it would invent a rule the query engine does not have, so the number is +/// bound exactly as written. +#[test] +fn a_quantile_outside_the_unit_interval_is_bound_as_given() { + let cases = [("-1.0", -1.0), ("2.0", 2.0), ("0.0", 0.0), ("1.0", 1.0)]; + + for (written, quantile) in cases { + let query = format!("quantile_over_time({written}, x[5m])"); + let tree = compile( + &eval_stmt(&query), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap_or_else(|error| panic!("`{query}` did not compile: {error}")); + let (_, op) = root_window(&tree); + assert_eq!( + op.function, + RangeFunction::QuantileOverTime(quantile), + "`{query}` did not bind the quantile it wrote" + ); + assert_eq!(op.window.range_ms, 300_000, "`{query}` lost its window"); + } +} + +/// A parameter in parentheses binds the same quantile, because it is the same query: the planner +/// already looks through the parentheses when it reads a literal, so the native compiler does too. +#[test] +fn a_parenthesised_quantile_binds_as_written() { + let tree = compile( + &eval_stmt("quantile_over_time((0.9), x[5m])"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .expect("a parenthesised literal is the same quantile"); + + let (_, op) = root_window(&tree); + assert_eq!(op.function, RangeFunction::QuantileOverTime(0.9)); +} + +/// The selector of a parameterised function is looked through its parentheses too, so the range +/// function and the parameter are both read from a statement that was merely typed with grouping. +#[test] +fn a_parenthesised_selector_binds_as_written() { + let bare = compile( + &eval_stmt("quantile_over_time(0.9, x[5m])"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .expect("the bare selector binds"); + + for query in [ + "quantile_over_time(0.9, (x[5m]))", + "quantile_over_time((0.9), (x[5m]))", + "(quantile_over_time(0.9, (x[5m])))", + ] { + let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) + .unwrap_or_else(|error| panic!("`{query}` was refused: {error}")); + assert_eq!( + tree.to_string(), + bare.to_string(), + "`{query}` did not compile to the tree `quantile_over_time(0.9, x[5m])` does" + ); + } +} + +/// A parameter the planner folds to a literal is still refused, and that gap is deliberate: the +/// kernel's constant is a field of the tree, so binding `1 - 0.1` means folding it here rather than +/// carrying an expression. A refusal is the agreed behaviour for this stage; an answer computed +/// from the wrong number would not be. +#[test] +fn a_parameter_folded_from_an_expression_is_refused() { + for query in [ + "quantile_over_time(1 - 0.1, x[5m])", + "quantile_over_time(-(0.9), x[5m])", + ] { + let error = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) + .expect_err(&format!("`{query}` compiled, and it should not have")); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "`{query}` did not report NotImplemented: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("number literal"), + "`{query}` did not say what the argument has to be: {message}" + ); + } +} + +/// A parameter computed at query time is refused, and the error names it: the tree binds a kernel's +/// constant, so a quantile that would only exist at run time has no field to go in. +#[test] +fn a_computed_quantile_is_refused_and_names_what_it_found() { + // `time()` is a scalar, so the parser accepts it as the quantile; the compiler binds literals. + let query = "quantile_over_time(time(), x[5m])"; + let error = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime) + .expect_err(&format!("`{query}` compiled, and it should not have")); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "`{query}` did not report NotImplemented: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("whose first argument is `time()`"), + "`{query}` did not name the argument it found: {message}" + ); + assert!( + message.contains("number literal"), + "`{query}` did not say what the argument has to be: {message}" + ); +} + +/// A parameterised window binds its anchor and its offset exactly as an unparameterised one does. +/// The two are asserted apart because the runtime folds `anchor - offset`, so a binding that +/// dropped either number would still compile a window, and only these assertions would notice. +#[test] +fn a_quantile_over_time_binds_an_anchor_and_an_offset_apart() { + let tree = compile( + &eval_stmt("quantile_over_time(0.9, x[5m] @ 1234 offset 1m)"), + raw_schema(), + SeriesOrdering::KeyThenTime, + ) + .unwrap(); + + let (_, op) = root_window(&tree); + assert_eq!(op.function, RangeFunction::QuantileOverTime(0.9)); + assert_eq!(op.window.range_ms, 300_000, "the window was not bound"); + assert_eq!( + op.window.anchor_ms, + Some(1_234_000), + "the anchor was not bound" + ); + assert_eq!(op.window.offset_ms, 60_000, "the offset was not bound"); +} + +/// A call with the wrong number of arguments names the count its function takes, rather than +/// reporting whatever the compiler does with one argument. The parser checks arity too, so a +/// mismatched call can only reach the compiler as an edited tree. +#[test] +fn a_call_with_the_wrong_number_of_arguments_names_the_count_it_takes() { + let cases = [ + // `quantile_over_time` takes its quantile and its range selector: two arguments. + ( + with_call_arguments("quantile_over_time(0.9, x[5m])", |args| { + args.remove(1); + }), + "quantile_over_time", + "where it takes 2", + ), + // `rate` takes one argument, so nothing can be passed in front of its range selector. + ( + with_call_arguments("rate(x[5m])", |args| { + args.push(Box::new(PromExpr::NumberLiteral(NumberLiteral::new(1.0)))); + }), + "rate", + "where it takes 1", + ), + ]; + + for (stmt, name, expected) in cases { + let error = compile(&stmt, raw_schema(), SeriesOrdering::KeyThenTime) + .expect_err(&format!("`{name}` with the wrong argument count compiled")); + assert!( + matches!(error, DataFusionError::NotImplemented(_)), + "`{name}` did not report NotImplemented: {error}" + ); + let message = error.to_string(); + assert!( + message.contains(&format!("a call to `{name}`")), + "`{name}` did not name itself: {message}" + ); + assert!( + message.contains(expected), + "`{name}` did not name the count it takes: {message}" + ); + } +} + /// `x` becomes an instant selection over the raw input, on the statement's grid, with the /// statement's lookback and the input's fields. #[test] @@ -503,8 +724,8 @@ fn an_unsupported_shape_is_rejected_and_names_what_it_found() { ), // A function with no kernel, whether or not it takes a constant argument. ( - "quantile_over_time(0.9, x[5m])", - "a call to `quantile_over_time`", + "predict_linear(x[5m], 300)", + "a call to `predict_linear`, which has no native range kernel", ), ( "absent(x)",