feat: support PromQL scalar (#3693)

GreptimeTeam · Apr 19, 2024 · d077892 · d077892
1 parent cfed466
commit d077892
Show file tree

Hide file tree

Showing 8 changed files with 1,104 additions and 12 deletions.
diff --git a/src/common/substrait/src/extension_serializer.rs b/src/common/substrait/src/extension_serializer.rs
@@ -19,7 +19,7 @@ use datafusion::execution::registry::SerializerRegistry;
 use datafusion_common::DataFusionError;
 use datafusion_expr::UserDefinedLogicalNode;
 use promql::extension_plan::{
-    EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize,
+    EmptyMetric, InstantManipulate, RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize,
 };
 
 pub struct ExtensionSerializer;
@@ -50,6 +50,13 @@ impl SerializerRegistry for ExtensionSerializer {
                     .expect("Failed to downcast to RangeManipulate");
                 Ok(range_manipulate.serialize())
             }
+            name if name == ScalarCalculate::name() => {
+                let scalar_calculate = node
+                    .as_any()
+                    .downcast_ref::<ScalarCalculate>()
+                    .expect("Failed to downcast to ScalarCalculate");
+                Ok(scalar_calculate.serialize())
+            }
             name if name == SeriesDivide::name() => {
                 let series_divide = node
                     .as_any()
@@ -92,6 +99,10 @@ impl SerializerRegistry for ExtensionSerializer {
                 let series_divide = SeriesDivide::deserialize(bytes)?;
                 Ok(Arc::new(series_divide))
             }
+            name if name == ScalarCalculate::name() => {
+                let scalar_calculate = ScalarCalculate::deserialize(bytes)?;
+                Ok(Arc::new(scalar_calculate))
+            }
             name if name == EmptyMetric::name() => Err(DataFusionError::Substrait(
                 "EmptyMetric should not be deserialized".to_string(),
             )),

diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs
@@ -18,6 +18,7 @@ mod instant_manipulate;
 mod normalize;
 mod planner;
 mod range_manipulate;
+mod scalar_calculate;
 mod series_divide;
 #[cfg(test)]
 mod test_util;
@@ -30,6 +31,7 @@ pub use instant_manipulate::{InstantManipulate, InstantManipulateExec, InstantMa
 pub use normalize::{SeriesNormalize, SeriesNormalizeExec, SeriesNormalizeStream};
 pub use planner::PromExtensionPlanner;
 pub use range_manipulate::{RangeManipulate, RangeManipulateExec, RangeManipulateStream};
+pub use scalar_calculate::ScalarCalculate;
 pub use series_divide::{SeriesDivide, SeriesDivideExec, SeriesDivideStream};
 pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctOnStream};
 

diff --git a/src/promql/src/extension_plan/planner.rs b/src/promql/src/extension_plan/planner.rs
@@ -21,7 +21,7 @@ use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
 
-use super::{HistogramFold, UnionDistinctOn};
+use super::{HistogramFold, ScalarCalculate, UnionDistinctOn};
 use crate::extension_plan::{
     EmptyMetric, InstantManipulate, RangeManipulate, SeriesDivide, SeriesNormalize,
 };
@@ -48,6 +48,8 @@ impl ExtensionPlanner for PromExtensionPlanner {
             Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
         } else if let Some(node) = node.as_any().downcast_ref::<EmptyMetric>() {
             Ok(Some(node.to_execution_plan(session_state, planner)?))
+        } else if let Some(node) = node.as_any().downcast_ref::<ScalarCalculate>() {
+            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())?))
         } else if let Some(node) = node.as_any().downcast_ref::<HistogramFold>() {
             Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
         } else if let Some(node) = node.as_any().downcast_ref::<UnionDistinctOn>() {