diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index 6703505c277..6067deabac1 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -30,6 +30,8 @@ mod round; mod test_util; mod vector_matching; +use std::sync::Arc; + pub use aggr_over_time::{ AbsentOverTime, AvgOverTime, CountOverTime, LastOverTime, MaxOverTime, MinOverTime, PresentOverTime, StddevOverTime, StdvarOverTime, SumOverTime, @@ -39,6 +41,7 @@ use datafusion::arrow::array::{ ArrayRef, DictionaryArray, Float64Array, TimestampMillisecondArray, }; use datafusion::error::DataFusionError; +use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::physical_plan::ColumnarValue; use datatypes::arrow::array::Array; use datatypes::arrow::datatypes::{DataType, Int64Type}; @@ -232,10 +235,95 @@ pub(crate) fn linear_regression_slices( (Some(slope), Some(intercept)) } +/// Registers every PromQL scalar function (UDF) and aggregate function (UDAF) +/// implemented in this module into `session_state`. +pub fn register_promql_functions(session_state: &mut SessionState) { + let _ = session_state.register_udaf(quantile_udaf()); + + for udf in IeeeComparison::all() { + let _ = session_state.register_udf(Arc::new(udf)); + } + let _ = session_state.register_udf(Arc::new(IeeeSqrt::scalar_udf())); + + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); + // Legacy Prometheus name of the double exponential smoothing function. + let double_exponential_smoothing_udf = + DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); + let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); + + for udf in [ + NativeHistogramAbsentOverTime::scalar_udf(), + NativeHistogramAdd::scalar_udf(), + NativeHistogramAvg::scalar_udf(), + NativeHistogramAvgOverTime::scalar_udf(), + NativeHistogramChanges::scalar_udf(), + NativeHistogramCount::scalar_udf(), + NativeHistogramCountOverTime::scalar_udf(), + NativeHistogramDelta::scalar_udf(), + NativeHistogramDivScalar::scalar_udf(), + NativeHistogramDrop::bool_false_udf(String::new(), None), + NativeHistogramDrop::bool_true_udf(String::new(), None), + NativeHistogramDrop::float_null_udf(String::new(), None), + NativeHistogramEq::scalar_udf(), + NativeHistogramFraction::scalar_udf(), + NativeHistogramIDelta::scalar_udf(), + NativeHistogramIRate::scalar_udf(), + NativeHistogramIncrease::scalar_udf(), + NativeHistogramLastOverTime::scalar_udf(), + MixedRange::float_udf(None), + MixedRange::histogram_udf(None), + NativeHistogramMulScalar::scalar_udf(), + NativeHistogramNeg::scalar_udf(), + NativeHistogramNotEq::scalar_udf(), + NativeHistogramPresentOverTime::scalar_udf(), + NativeHistogramQuantile::scalar_udf(), + NativeHistogramRate::scalar_udf(), + NativeHistogramResets::scalar_udf(), + NativeHistogramScalarMul::scalar_udf(), + NativeHistogramStddev::scalar_udf(), + NativeHistogramStdvar::scalar_udf(), + NativeHistogramSub::scalar_udf(), + NativeHistogramSum::scalar_udf(), + NativeHistogramSumOverTime::scalar_udf(), + NativeHistogramToString::scalar_udf(), + PromqlFloatToString::scalar_udf(), + ] { + let _ = session_state.register_udf(Arc::new(udf)); + } + for udaf in [ + NativeHistogramAggAvg::aggregate_udf(), + NativeHistogramAggSum::aggregate_udf(), + ] { + let _ = session_state.register_udaf(Arc::new(udaf)); + } +} + #[cfg(test)] mod test { + use std::collections::BTreeSet; use std::sync::Arc; + use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::ColumnarValue; use datatypes::arrow::array::Int64Array; use datatypes::arrow::datatypes::Int64Type; @@ -381,4 +469,100 @@ mod test { assert_eq!(extracted.get_offset_length(0), Some((0, 2))); assert_eq!(extracted.get_offset_length(1), Some((1, 2))); } + + #[test] + fn register_promql_functions_registers_all_promql_functions() { + let mut session_state = SessionStateBuilder::new().with_default_features().build(); + register_promql_functions(&mut session_state); + + // Every PromQL scalar function is registered under its `prom_` name; + // DataFusion built-ins never use the `prom_` prefix. + let registered: BTreeSet<&str> = session_state + .scalar_functions() + .keys() + .filter(|name| name.starts_with("prom_")) + .map(String::as_str) + .collect(); + let expected: BTreeSet<&str> = [ + "prom_absent_over_time", + "prom_avg_over_time", + "prom_changes", + "prom_count_over_time", + "prom_delta", + "prom_deriv", + "prom_double_exponential_smoothing", + "prom_eq", + "prom_float_to_string", + "prom_ge", + "prom_gt", + "prom_holt_winters", + "prom_idelta", + "prom_increase", + "prom_irate", + "prom_last_over_time", + "prom_le", + "prom_lt", + "prom_max_over_time", + "prom_min_over_time", + "prom_mixed_range_float", + "prom_mixed_range_histogram", + "prom_native_histogram_absent_over_time", + "prom_native_histogram_add", + "prom_native_histogram_avg", + "prom_native_histogram_avg_over_time", + "prom_native_histogram_changes", + "prom_native_histogram_count", + "prom_native_histogram_count_over_time", + "prom_native_histogram_delta", + "prom_native_histogram_div_scalar", + "prom_native_histogram_drop_bool", + "prom_native_histogram_drop_float", + "prom_native_histogram_eq", + "prom_native_histogram_fraction", + "prom_native_histogram_idelta", + "prom_native_histogram_increase", + "prom_native_histogram_irate", + "prom_native_histogram_keep_bool", + "prom_native_histogram_last_over_time", + "prom_native_histogram_mul_scalar", + "prom_native_histogram_neg", + "prom_native_histogram_not_eq", + "prom_native_histogram_present_over_time", + "prom_native_histogram_quantile", + "prom_native_histogram_rate", + "prom_native_histogram_resets", + "prom_native_histogram_scalar_mul", + "prom_native_histogram_stddev", + "prom_native_histogram_stdvar", + "prom_native_histogram_sub", + "prom_native_histogram_sum", + "prom_native_histogram_sum_over_time", + "prom_native_histogram_to_string", + "prom_ne", + "prom_predict_linear", + "prom_present_over_time", + "prom_quantile_over_time", + "prom_rate", + "prom_resets", + "prom_round", + "prom_sqrt", + "prom_stddev_over_time", + "prom_stdvar_over_time", + "prom_sum_over_time", + ] + .into_iter() + .collect(); + assert_eq!(registered, expected); + + for name in [ + QUANTILE_NAME, + NativeHistogramAggAvg::name(), + NativeHistogramAggSum::name(), + ] { + assert!( + session_state.aggregate_functions().contains_key(name), + "{name} should be registered" + ); + } + } } diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index ced633382d3..76f2bab125d 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -30,21 +30,7 @@ use datafusion::execution::{FunctionRegistry, SessionStateBuilder}; use datafusion::logical_expr::{Extension, LogicalPlan}; use datafusion_expr::{Expr, UserDefinedLogicalNode}; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; -use promql::functions::{ - AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, IeeeComparison, IeeeSqrt, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, - NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, - NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, - NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, - NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, - NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, - NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, - NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, - NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, - NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, - Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, -}; +use promql::functions::register_promql_functions; use prost::Message; use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; @@ -333,83 +319,7 @@ fn register_greptime_functions( .context(RegisterUdfSnafu { name })?; } - let _ = session_state.register_udaf(quantile_udaf()); - for udf in IeeeComparison::all() { - let _ = session_state.register_udf(Arc::new(udf)); - } - let _ = session_state.register_udf(Arc::new(IeeeSqrt::scalar_udf())); - - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); - let double_exponential_smoothing_udf = - DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); - let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); - - for udf in [ - NativeHistogramAbsentOverTime::scalar_udf(), - NativeHistogramAdd::scalar_udf(), - NativeHistogramAvg::scalar_udf(), - NativeHistogramAvgOverTime::scalar_udf(), - NativeHistogramChanges::scalar_udf(), - NativeHistogramCount::scalar_udf(), - NativeHistogramCountOverTime::scalar_udf(), - NativeHistogramDelta::scalar_udf(), - NativeHistogramDivScalar::scalar_udf(), - NativeHistogramDrop::bool_false_udf(String::new(), None), - NativeHistogramDrop::bool_true_udf(String::new(), None), - NativeHistogramDrop::float_null_udf(String::new(), None), - NativeHistogramEq::scalar_udf(), - NativeHistogramFraction::scalar_udf(), - NativeHistogramIDelta::scalar_udf(), - NativeHistogramIRate::scalar_udf(), - NativeHistogramIncrease::scalar_udf(), - NativeHistogramLastOverTime::scalar_udf(), - MixedRange::float_udf(None), - MixedRange::histogram_udf(None), - NativeHistogramMulScalar::scalar_udf(), - NativeHistogramNeg::scalar_udf(), - NativeHistogramNotEq::scalar_udf(), - NativeHistogramPresentOverTime::scalar_udf(), - NativeHistogramQuantile::scalar_udf(), - NativeHistogramRate::scalar_udf(), - NativeHistogramResets::scalar_udf(), - NativeHistogramScalarMul::scalar_udf(), - NativeHistogramStddev::scalar_udf(), - NativeHistogramStdvar::scalar_udf(), - NativeHistogramSub::scalar_udf(), - NativeHistogramSum::scalar_udf(), - NativeHistogramSumOverTime::scalar_udf(), - NativeHistogramToString::scalar_udf(), - PromqlFloatToString::scalar_udf(), - ] { - let _ = session_state.register_udf(Arc::new(udf)); - } - - for udaf in [ - NativeHistogramAggAvg::aggregate_udf(), - NativeHistogramAggSum::aggregate_udf(), - ] { - let _ = session_state.register_udaf(Arc::new(udaf)); - } + register_promql_functions(session_state); Ok(()) } @@ -471,6 +381,10 @@ mod tests { use datatypes::data_type::DataType; use futures::stream::BoxStream; use promql::extension_plan::{RangeManipulate, SeriesNormalize, UnionDistinctOn}; + use promql::functions::{ + MixedRange, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramCount, + NativeHistogramDrop, PromqlFloatToString, + }; use session::context::QueryContext; use table::TableRef; use table::metadata::{TableId, TableInfoRef};