query/
promql.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod error;
16pub mod label_values;
17pub mod planner;
18
19use datafusion_common::tree_node::{TreeNode as _, TreeNodeRecursion};
20use datafusion_expr::{Extension, LogicalPlan};
21use promql::extension_plan::{
22    Absent, EmptyMetric, HistogramFold, InstantManipulate, RangeManipulate, ScalarCalculate,
23    SeriesDivide, SeriesNormalize, UnionDistinctOn,
24};
25
26/// Returns true if the plan contains PromQL-specific extension plan nodes.
27pub fn plan_contains_promql_extension(plan: &LogicalPlan) -> bool {
28    let mut found = false;
29    let _ = plan.apply(|node| {
30        if is_promql_extension_plan(node) {
31            found = true;
32            Ok(TreeNodeRecursion::Stop)
33        } else {
34            Ok(TreeNodeRecursion::Continue)
35        }
36    });
37    found
38}
39
40fn is_promql_extension_plan(plan: &LogicalPlan) -> bool {
41    let LogicalPlan::Extension(Extension { node }) = plan else {
42        return false;
43    };
44
45    node.as_any().is::<Absent>()
46        || node.as_any().is::<EmptyMetric>()
47        || node.as_any().is::<HistogramFold>()
48        || node.as_any().is::<InstantManipulate>()
49        || node.as_any().is::<RangeManipulate>()
50        || node.as_any().is::<ScalarCalculate>()
51        || node.as_any().is::<SeriesDivide>()
52        || node.as_any().is::<SeriesNormalize>()
53        || node.as_any().is::<UnionDistinctOn>()
54}
55
56#[cfg(test)]
57mod tests {
58    use std::sync::Arc;
59
60    use datafusion_common::DFSchema;
61    use datafusion_expr::{EmptyRelation, Extension, LogicalPlanBuilder, col};
62
63    use super::*;
64
65    #[test]
66    fn plan_contains_promql_extension_returns_true_for_promql_extension() {
67        let plan = empty_metric_plan();
68
69        assert!(plan_contains_promql_extension(&plan));
70    }
71
72    #[test]
73    fn plan_contains_promql_extension_returns_true_for_nested_promql_extension() {
74        let plan = LogicalPlanBuilder::from(empty_metric_plan())
75            .project(vec![col("ts")])
76            .unwrap()
77            .build()
78            .unwrap();
79
80        assert!(plan_contains_promql_extension(&plan));
81    }
82
83    #[test]
84    fn plan_contains_promql_extension_returns_false_for_non_promql_plan() {
85        let plan = LogicalPlan::EmptyRelation(EmptyRelation {
86            produce_one_row: false,
87            schema: Arc::new(DFSchema::empty()),
88        });
89
90        assert!(!plan_contains_promql_extension(&plan));
91    }
92
93    fn empty_metric_plan() -> LogicalPlan {
94        let empty_metric = EmptyMetric::new(
95            0,
96            10_000,
97            5_000,
98            "ts".to_string(),
99            "greptime_value".to_string(),
100            None,
101        )
102        .unwrap();
103
104        LogicalPlan::Extension(Extension {
105            node: Arc::new(empty_metric),
106        })
107    }
108}
query/promql.rs

query/
promql.rs
