feat(promql): bind the remaining range kernels to their PromQL names

`RangeFunction` already names 26 functions and every one of them answers: the
kernels for `first_over_time`, `mad_over_time`, the four `ts_of_*` functions,
`predict_linear` and `double_exponential_smoothing` are all in `range_eval`. What
was missing was the compiler's name table, so these queries were refused as names
with no kernel although the kernel existed.

Six of them are the functions Prometheus marks experimental. This crate implements
them unconditionally — `RangeFunction` names them without a gate and the kernels
compute them — so binding them is the same decision as binding `rate` rather than
a new policy. If a gate is ever wanted it belongs where a query is admitted, not
scattered across the compiler's name table.

The two parameterised ones needed a second look at where the selector sits, which
is not one place in PromQL: `quantile_over_time(0.9, x[5m])` writes its quantile
before the selector, while `predict_linear(x[5m], 300)` and
`double_exponential_smoothing(x[5m], 0.5, 0.1)` write theirs after it. A single
"parameters come first" rule read the selector as a constant for the second group
and rejected an ordinary query, so the signature now states the selector's
position as data instead of relying on a convention.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
(cherry picked from commit 8b15ab65c0)

src/query/src/promql/native_compile.rs

@@ -21,7 +21,9 @@
 //! ```text
 //! <selector>                                        e.g. x
 //! <range function>(<selector>[<range>])             e.g. rate(x[5m])
-//! <range function>(<constant>, <selector>[<range>]) e.g. quantile_over_time(0.9, x[5m])
+//! <range function>(<constant>[, <constant>], <selector>[<range>])
+//!                                                   e.g. quantile_over_time(0.9, x[5m]),
+//!                                                        double_exponential_smoothing(x[5m], 0.5, 0.5)
 //! <math function>(<expression>[, <constant>]*)      e.g. abs(rate(x[5m])), clamp_min(x, 0)
 //! <expression> <operator> <constant>                e.g. x * 100, x > bool 0
 //! <constant> <operator> <expression>                e.g. 1 / x, 0 < x
@@ -472,7 +474,7 @@ fn window_call(
 ) -> Result<NativeNode> {
     // 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;
+    let expected = signature.arity();
     if args.len() != expected {
         return Err(not_implemented(&format!(
             "a call to `{name}` with {} argument(s), where it takes {expected}: {}",
@@ -480,17 +482,24 @@ fn window_call(
             signature.arguments()
         )));
     }
-    let (parameters, selector) = args.split_at(signature.parameters());
-    let function = match signature {
-        RangeSignature::Selector(function) => function,
-        RangeSignature::Scalar(bind) => bind(constant_argument(name, &parameters[0], 0)?),
-    };
-    let argument = &selector[0];
+    // The selector sits where the signature says it does, and everything else is a constant. Reading
+    // it by position rather than by "the last argument" is what lets `predict_linear(x[5m], 300)` and
+    // `quantile_over_time(0.9, x[5m])` share one path.
+    let selector_index = signature.selector_index();
+    let argument = &args[selector_index];
     let PromExpr::MatrixSelector(MatrixSelector { vs, range }) = unparen(argument) else {
         return Err(not_implemented(&format!(
-            "a call to `{name}` whose argument is `{argument}`"
+            "a call to `{name}` whose range selector is `{argument}`, where it has to be a range \
+             selector like `x[5m]`"
         )));
     };
+    let constants: Vec<Box<PromExpr>> = args
+        .iter()
+        .enumerate()
+        .filter(|(position, _)| *position != selector_index)
+        .map(|(_, argument)| argument.clone())
+        .collect();
+    let function = signature.bind(name, &constants)?;
     window_selection(context, function, vs, range)
 }
 
@@ -727,34 +736,84 @@ 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.
+/// What a call to a range function has to be: where its range selector sits, how many constants it
+/// reads, and how they bind.
+///
+/// The selector's position is a variant rather than a convention because PromQL does not put it in
+/// one place: `quantile_over_time(0.9, x[5m])` writes its quantile before the selector, while
+/// `predict_linear(x[5m], 300)` and `double_exponential_smoothing(x[5m], 0.5, 0.1)` write theirs
+/// after it. A single "parameters come first" rule would read the selector as a constant in the
+/// second group and reject an ordinary query, so each shape states its own layout.
 #[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 constant before its range selector, and the function that constant
-    /// binds, as in `quantile_over_time(0.9, x[5m])`.
-    ///
-    /// The constant is kept out of the variant name because only the call site knows it.
-    Scalar(fn(f64) -> RangeFunction),
+    /// One constant, then the range selector, as in `quantile_over_time(0.9, x[5m])`.
+    ConstantBeforeSelector(fn(f64) -> RangeFunction),
+    /// The range selector, then one constant, as in `predict_linear(x[5m], 300)`.
+    ConstantAfterSelector(fn(f64) -> RangeFunction),
+    /// The range selector, then two constants, as in
+    /// `double_exponential_smoothing(x[5m], 0.5, 0.1)`.
+    TwoConstantsAfterSelector(fn(f64, f64) -> RangeFunction),
 }
 
 impl RangeSignature {
-    /// How many scalar parameters the function reads before its range selector.
-    fn parameters(self) -> usize {
+    /// The position of the range selector among the call's arguments.
+    fn selector_index(self) -> usize {
+        match self {
+            Self::Selector(_)
+            | Self::ConstantAfterSelector(_)
+            | Self::TwoConstantsAfterSelector(_) => 0,
+            Self::ConstantBeforeSelector(_) => 1,
+        }
+    }
+
+    /// How many constants the function reads.
+    fn constants(self) -> usize {
         match self {
             Self::Selector(_) => 0,
-            Self::Scalar(_) => 1,
+            Self::ConstantBeforeSelector(_) | Self::ConstantAfterSelector(_) => 1,
+            Self::TwoConstantsAfterSelector(_) => 2,
         }
     }
 
+    /// How many arguments the call writes, its range selector counted in.
+    fn arity(self) -> usize {
+        self.constants() + 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",
+            Self::ConstantBeforeSelector(_) => "a number and its range selector",
+            Self::ConstantAfterSelector(_) => "its range selector and a number",
+            Self::TwoConstantsAfterSelector(_) => "its range selector and two numbers",
         }
     }
+
+    /// Binds the function over the `constants` its call wrote, in the order they were written.
+    ///
+    /// The count is the signature's, so a caller that reached here has already checked it: reading
+    /// `constants[0]` and `constants[1]` cannot panic on a call this signature accepted.
+    fn bind(self, name: &str, constants: &[Box<PromExpr>]) -> Result<RangeFunction> {
+        // Each constant is folded by [`constant_argument`], which names the argument it found if the
+        // query wrote something only known at run time.
+        let mut folded = constants
+            .iter()
+            .enumerate()
+            .map(|(position, argument)| constant_argument(name, argument, position));
+        Ok(match self {
+            Self::Selector(function) => function,
+            Self::ConstantBeforeSelector(bind) | Self::ConstantAfterSelector(bind) => {
+                bind(folded.next().expect("one constant")?)
+            }
+            Self::TwoConstantsAfterSelector(bind) => bind(
+                folded.next().expect("first constant")?,
+                folded.next().expect("second constant")?,
+            ),
+        })
+    }
 }
 
 /// The signature the PromQL name calls, or `None` when this compiler has no range kernel for it.
@@ -765,7 +824,17 @@ impl RangeSignature {
 /// 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)),
+        "quantile_over_time" => Some(RangeSignature::ConstantBeforeSelector(
+            RangeFunction::QuantileOverTime,
+        )),
+        // These two write their constants after the selector, which is the opposite order to
+        // `quantile_over_time`.
+        "predict_linear" => Some(RangeSignature::ConstantAfterSelector(
+            RangeFunction::PredictLinear,
+        )),
+        "double_exponential_smoothing" => Some(RangeSignature::TwoConstantsAfterSelector(
+            |smoothing, trend| RangeFunction::DoubleExponentialSmoothing { smoothing, trend },
+        )),
         _ => range_function(name).map(RangeSignature::Selector),
     }
 }
@@ -1026,6 +1095,16 @@ fn range_function(name: &str) -> Option<RangeFunction> {
         "resets" => Some(RangeFunction::Resets),
         "changes" => Some(RangeFunction::Changes),
         "deriv" => Some(RangeFunction::Deriv),
+        // Prometheus marks these six as experimental, and this crate implements every one of them
+        // unconditionally: `RangeFunction` names them without a gate and the kernels answer them.
+        // Binding them here is therefore the same decision as binding `rate`, not a new policy —
+        // if a gate is ever wanted it belongs where a query is admitted, not scattered here.
+        "first_over_time" => Some(RangeFunction::FirstOverTime),
+        "mad_over_time" => Some(RangeFunction::MadOverTime),
+        "ts_of_first_over_time" => Some(RangeFunction::TsOfFirstOverTime),
+        "ts_of_last_over_time" => Some(RangeFunction::TsOfLastOverTime),
+        "ts_of_min_over_time" => Some(RangeFunction::TsOfMinOverTime),
+        "ts_of_max_over_time" => Some(RangeFunction::TsOfMaxOverTime),
         _ => None,
     }
 }

src/query/src/promql/native_compile/tests.rs

@@ -262,6 +262,15 @@ fn the_supported_range_functions_map_to_their_kernels() {
         ("resets", RangeFunction::Resets),
         ("changes", RangeFunction::Changes),
         ("deriv", RangeFunction::Deriv),
+        // The six Prometheus marks experimental. `RangeFunction` names all of them and the kernels
+        // answer all of them, so the compiler binds them like any other name; a test per name is
+        // what keeps a rename or a drop from passing as a typo.
+        ("first_over_time", RangeFunction::FirstOverTime),
+        ("mad_over_time", RangeFunction::MadOverTime),
+        ("ts_of_first_over_time", RangeFunction::TsOfFirstOverTime),
+        ("ts_of_last_over_time", RangeFunction::TsOfLastOverTime),
+        ("ts_of_min_over_time", RangeFunction::TsOfMinOverTime),
+        ("ts_of_max_over_time", RangeFunction::TsOfMaxOverTime),
     ];
 
     for (name, kernel) in cases {
@@ -878,13 +887,10 @@ fn an_unsupported_shape_is_rejected_and_names_what_it_found() {
         // function arguments, so `rate(x)` never gets past parsing.
         (
             "rate(x[5m:])",
-            "a call to `rate` whose argument is `x[5m:]`",
-        ),
-        // A function with no kernel, whether or not it takes a constant argument.
-        (
-            "predict_linear(x[5m], 300)",
-            "a call to `predict_linear`, which has no native kernel",
+            "a call to `rate` whose range selector is `x[5m:]`",
         ),
+        // A function with no kernel. `predict_linear` used to stand here as the parameterised case
+        // and is now bound, so what is left is a name this compiler has no kernel for at all.
         (
             "absent(x)",
             "a call to `absent`, which has no native kernel",
@@ -1869,3 +1875,168 @@ fn the_batch_acceptance_queries_compile() {
         );
     }
 }
+
+/// `predict_linear` reads its horizon before its range selector, so the number the query wrote is the
+/// horizon the kernel extrapolates to. The value is asserted apart from the window, because a
+/// compiler that bound the horizon on the wrong side of the call would still produce a window.
+#[test]
+fn a_predict_linear_binds_its_horizon_and_its_window() {
+    let tree = compile(
+        &eval_stmt("predict_linear(x[5m], 300)"),
+        raw_schema(),
+        SeriesOrdering::KeyThenTime,
+    )
+    .unwrap_or_else(|error| panic!("`predict_linear(x[5m], 300)` was refused: {error}"));
+
+    let (_, op) = root_window(&tree);
+    assert_eq!(op.function, RangeFunction::PredictLinear(300.0));
+    assert_eq!(op.window.range_ms, 300_000, "the window was not bound");
+    assert!(op.input_grid.is_none(), "raw input has no grid to declare");
+}
+
+/// `double_exponential_smoothing` reads two constants, and both reach the kernel in the order they
+/// were written: swapping them would smooth with the trend factor and trend with the smoothing one,
+/// which is a different series rather than an error.
+///
+/// The pair is asserted with two different numbers, because a compiler that bound the second
+/// constant twice would agree with a test that passed the same number twice.
+#[test]
+fn a_double_exponential_smoothing_binds_both_factors_in_order() {
+    let cases = [
+        ("double_exponential_smoothing(x[5m], 0.5, 0.1)", 0.5, 0.1),
+        ("double_exponential_smoothing(x[5m], 0.1, 0.5)", 0.1, 0.5),
+        // The constants are folded like every other constant this compiler binds.
+        (
+            "double_exponential_smoothing(x[5m], 1 / 2, 0.1 + 0.2)",
+            0.5,
+            0.30000000000000004,
+        ),
+    ];
+
+    for (query, smoothing, trend) in cases {
+        let tree = compile(&eval_stmt(query), raw_schema(), SeriesOrdering::KeyThenTime)
+            .unwrap_or_else(|error| panic!("`{query}` was refused: {error}"));
+        let (_, op) = root_window(&tree);
+        assert_eq!(
+            op.function,
+            RangeFunction::DoubleExponentialSmoothing { smoothing, trend },
+            "`{query}` bound its factors wrongly"
+        );
+        assert_eq!(op.window.range_ms, 300_000, "`{query}` lost its window");
+    }
+}
+
+/// The kernel needs both smoothing factors, so a call that wrote one is refused and the rejection
+/// says what the function reads. The parser enforces the same count, so the only way to hand the
+/// compiler such a call is to edit a tree it accepted: the same shape as the arity test above,
+/// applied to a function whose constants follow its selector.
+#[test]
+fn a_smoothing_with_one_factor_is_refused_and_names_what_it_reads() {
+    let stmt = with_call_arguments("double_exponential_smoothing(x[5m], 0.5, 0.1)", |args| {
+        args.pop();
+    });
+    let error = compile(&stmt, raw_schema(), SeriesOrdering::KeyThenTime)
+        .expect_err("a smoothing with one factor compiled, and it should not have");
+    assert!(
+        matches!(error, DataFusionError::NotImplemented(_)),
+        "the refusal is not NotImplemented: {error}"
+    );
+    let message = error.to_string();
+    assert!(
+        message.contains("with 2 argument(s), where it takes 3"),
+        "the refusal did not name the count the function takes: {message}"
+    );
+    assert!(
+        message.contains("its range selector and two numbers"),
+        "the refusal did not say what the function reads: {message}"
+    );
+}
+
+/// A parameter computed at query time is refused for a two-constant function too, and the rejection
+/// names which argument it found: a smoothing factor that only exists at run time has no field in the
+/// tree, the same as a quantile.
+#[test]
+fn a_computed_smoothing_factor_is_refused_and_names_its_position() {
+    let cases = [
+        (
+            "double_exponential_smoothing(x[5m], time(), 0.1)",
+            "first argument",
+        ),
+        (
+            "double_exponential_smoothing(x[5m], 0.5, time())",
+            "second argument",
+        ),
+        ("predict_linear(x[5m], time())", "first argument"),
+    ];
+
+    for (query, position) in cases {
+        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(position),
+            "`{query}` did not name which argument it found: {message}"
+        );
+        assert!(
+            message.contains("constant expression"),
+            "`{query}` did not say what the argument has to be: {message}"
+        );
+    }
+}
+
+/// The parameterised functions keep the anchor and the offset a function without parameters binds,
+/// because the parameters decide what folds the window rather than which window is folded.
+#[test]
+fn a_two_constant_function_binds_an_anchor_and_an_offset_too() {
+    let tree = compile(
+        &eval_stmt("double_exponential_smoothing(x[5m] @ 1234 offset 1m, 0.5, 0.1)"),
+        raw_schema(),
+        SeriesOrdering::KeyThenTime,
+    )
+    .unwrap_or_else(|error| panic!("a smoothed window with an anchor was refused: {error}"));
+
+    let (_, op) = root_window(&tree);
+    assert_eq!(
+        op.function,
+        RangeFunction::DoubleExponentialSmoothing {
+            smoothing: 0.5,
+            trend: 0.1,
+        }
+    );
+    assert_eq!(op.window.range_ms, 300_000);
+    assert_eq!(op.window.anchor_ms, Some(1_234_000), "the anchor was lost");
+    assert_eq!(op.window.offset_ms, 60_000, "the offset was lost");
+}
+
+/// The two parameterised functions nest under the rest of the compiler, so a smoothed window can be
+/// compared or aggregated like any other: an aggregate over a smoothing is a tree, not a shape the
+/// compiler knows by name.
+#[test]
+fn a_parameterised_window_nests_under_a_comparison_and_an_aggregate() {
+    let tree = compile(
+        &eval_stmt("sum by(job)(predict_linear(x[5m], 300) > 0)"),
+        raw_schema(),
+        SeriesOrdering::KeyThenTime,
+    )
+    .unwrap_or_else(|error| {
+        panic!("`sum by(job)(predict_linear(x[5m], 300) > 0)` refused: {error}")
+    });
+
+    let (input, aggregate) = aggregate_parts(tree.root());
+    assert_eq!(aggregate.kind, AggregateKind::Sum);
+    let (left, right, matching) = binary_parts(input);
+    assert_eq!(matching.op, BinaryOp::Gt);
+    assert!(!matching.return_bool, "`>` without `bool` filters");
+    assert_eq!(constant(right), Some(0.0));
+    let NativeExpr::Window { op, .. } = &left.expr else {
+        panic!(
+            "the comparison's left operand is not a window: {:?}",
+            left.expr
+        );
+    };
+    assert_eq!(op.function, RangeFunction::PredictLinear(300.0));
+}