diff --git a/src/query/src/promql/native_compile.rs b/src/query/src/promql/native_compile.rs index c58409e5932..9389a3fad17 100644 --- a/src/query/src/promql/native_compile.rs +++ b/src/query/src/promql/native_compile.rs @@ -21,7 +21,9 @@ //! ```text //! e.g. x //! ([]) e.g. rate(x[5m]) -//! (, []) e.g. quantile_over_time(0.9, x[5m]) +//! ([, ], []) +//! e.g. quantile_over_time(0.9, x[5m]), +//! double_exponential_smoothing(x[5m], 0.5, 0.5) //! ([, ]*) e.g. abs(rate(x[5m])), clamp_min(x, 0) //! e.g. x * 100, x > bool 0 //! e.g. 1 / x, 0 < x @@ -472,7 +474,7 @@ fn window_call( ) -> Result { // 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, ¶meters[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> = 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]) -> Result { + // 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 { 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 { "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, } } diff --git a/src/query/src/promql/native_compile/tests.rs b/src/query/src/promql/native_compile/tests.rs index 26d9d2ec27d..ce6e617380a 100644 --- a/src/query/src/promql/native_compile/tests.rs +++ b/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)); +}