mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-15 08:22:38 +00:00
refactor(promql): centralize PromQL function registration in promql::functions (#9500)
* refactor(promql): centralize PromQL function registration in promql::functions Move the PromQL scalar UDF and UDAF registrations out of the query plan decoder into a single `promql::functions::register_promql_functions` entry point, so plan decoding no longer has to enumerate every PromQL function implementation. The registered function set, aliases (e.g. `prom_holt_winters`) and registration order are unchanged. Add a registration test pinning the full set of `prom_*` scalar functions and the `prom_holt_winters` alias, plus a decode regression test covering the centralized registration path. Fixes #9493 Signed-off-by: fengyue-xve <1759519137@qq.com> * refactor(promql): register Ieee UDFs centrally and address review Rebase onto main picked up #9479, which registered the Ieee comparison and sqrt UDFs from the decoder directly. Register them from `register_promql_functions` instead so the single entry point stays complete, and add the seven `prom_eq`..`prom_sqrt` names to the registration test. Also per review: drop the doc paragraph that explains why the function was moved, drop the redundant `test_serializer_decode_centralized_promql_registration` test and the registration comment in the decoder, move `use std::sync::Arc` up with the imports, and simplify the registration test assertions. Signed-off-by: 风月 <1759519137@qq.com> Signed-off-by: fengyue-xve <1759519137@qq.com> --------- Signed-off-by: fengyue-xve <1759519137@qq.com> Signed-off-by: 风月 <1759519137@qq.com>
This commit is contained in:
fengyue-xve
@@ -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::<false>::scalar_udf()));
|
||||
let _ = session_state.register_udf(Arc::new(IDelta::<true>::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"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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::<false>::scalar_udf()));
|
||||
let _ = session_state.register_udf(Arc::new(IDelta::<true>::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};
|
||||
|
||||
Reference in New Issue
Block a user