rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-27 02:10:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support arbitrary constant expression in PromQL function (#6315)

Browse Source 
* refactor holt_winters, predict_linear, quantile, round

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix clippy

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* some sqlness result

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* support some functions

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* make all sqlness cases pass

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix other sqlness cases

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* some refactor

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix clippy

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2025-06-16 23:12:27 +08:00
committed by GitHub
parent 2a3445c72c
commit be4e0d589e
10 changed files with 579 additions and 302 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
src/query/src/promql/planner.rs
View File

@@ -1435,27 +1435,22 @@ impl PromPlanner {
for arg in args {
match *arg.clone() {
PromExpr::Aggregate(_)
| PromExpr::Unary(_)
| PromExpr::Binary(_)
| PromExpr::Paren(_)
| PromExpr::Subquery(_)
PromExpr::Subquery(_)
| PromExpr::VectorSelector(_)
| PromExpr::MatrixSelector(_)
| PromExpr::Extension(_)
| PromExpr::Aggregate(_)
| PromExpr::Paren(_)
| PromExpr::Call(_) => {
if result.input.replace(*arg.clone()).is_some() {
MultipleVectorSnafu { expr: *arg.clone() }.fail()?;
}
}
PromExpr::NumberLiteral(NumberLiteral { val, .. }) => {
let scalar_value = ScalarValue::Float64(Some(val));
result.literals.push(DfExpr::Literal(scalar_value));
}
PromExpr::StringLiteral(StringLiteral { val, .. }) => {
let scalar_value = ScalarValue::Utf8(Some(val));
result.literals.push(DfExpr::Literal(scalar_value));
_ => {
let expr =
Self::get_param_as_literal_expr(&Some(Box::new(*arg.clone())), None, None)?;
result.literals.push(expr);
}
}
}
@@ -1507,7 +1502,13 @@ impl PromPlanner {
"stddev_over_time" => ScalarFunc::Udf(Arc::new(StddevOverTime::scalar_udf())),
"stdvar_over_time" => ScalarFunc::Udf(Arc::new(StdvarOverTime::scalar_udf())),
"quantile_over_time" => ScalarFunc::Udf(Arc::new(QuantileOverTime::scalar_udf())),
"predict_linear" => ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf())),
"predict_linear" => {
other_input_exprs[0] = DfExpr::Cast(Cast {
expr: Box::new(other_input_exprs[0].clone()),
data_type: ArrowDataType::Int64,
});
ScalarFunc::Udf(Arc::new(PredictLinear::scalar_udf()))
}
"holt_winters" => ScalarFunc::Udf(Arc::new(HoltWinters::scalar_udf())),
"time" => {
exprs.push(build_special_time_expr(

10
src/query/src/query_engine/default_serializer.rs
View File

@@ -29,9 +29,9 @@ use datafusion::logical_expr::LogicalPlan;
use datafusion_expr::UserDefinedLogicalNode;
use greptime_proto::substrait_extension::MergeScan as PbMergeScan;
use promql::functions::{
quantile_udaf, AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, IDelta,
Increase, LastOverTime, MaxOverTime, MinOverTime, PresentOverTime, Rate, Resets, Round,
StddevOverTime, StdvarOverTime, SumOverTime,
quantile_udaf, AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, HoltWinters,
IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, PredictLinear, PresentOverTime,
QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime,
};
use prost::Message;
use session::context::QueryContextRef;
@@ -161,7 +161,9 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder {
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()));
// TODO(ruihang): add quantile_over_time, predict_linear, holt_winters, round
let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf()));
let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf()));
let _ = session_state.register_udf(Arc::new(HoltWinters::scalar_udf()));
let logical_plan = DFLogicalSubstraitConvertor
.decode(message, session_state)