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

@@ -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
-//! <selector>                              e.g. x
-//! <range function>(<selector>[<range>])   e.g. rate(x[5m])
+//! <selector>                                       e.g. x
+//! <range function>(<selector>[<range>])            e.g. rate(x[5m])
+//! <range function>(<number>, <selector>[<range>])  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 `<range function>(<selector>[<range>])`, for example \
-                         `x` or `rate(x[5m])`";
+const SUPPORTED: &str = "a bare selector, a range function over a selector — \
+                         `<range function>(<selector>[<range>])` — or a parameterised one, \
+                         `<range function>(<number>, <selector>[<range>])`, 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, &parameters[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<SeriesSchema>,
+    ordering: SeriesOrdering,
+) -> Result<NativeTree> {
     // 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<RangeSignature> {
+    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<f64> {
+    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

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<Box<PromExpr>>)) -> 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)",