promql/extension_plan/
planner.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
15use std::sync::Arc;
16
17use async_trait::async_trait;
18use datafusion::catalog::Session;
19use datafusion::error::Result as DfResult;
20use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
21use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNode};
22use datafusion::physical_plan::ExecutionPlan;
23use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
24
25use crate::extension_plan::{
26    Absent, EmptyMetric, HistogramFold, InstantManipulate, RangeManipulate, ScalarCalculate,
27    SeriesDivide, SeriesNormalize, UnionDistinctOn,
28};
29
30pub struct PromExtensionPlanner;
31
32#[async_trait]
33impl ExtensionPlanner for PromExtensionPlanner {
34    async fn plan_extension(
35        &self,
36        planner: &dyn PhysicalPlanner,
37        node: &dyn UserDefinedLogicalNode,
38        _logical_inputs: &[&LogicalPlan],
39        physical_inputs: &[Arc<dyn ExecutionPlan>],
40        session: &dyn Session,
41        planning_ctx: &PhysicalPlanningContext,
42    ) -> DfResult<Option<Arc<dyn ExecutionPlan>>> {
43        if let Some(node) = node.as_any().downcast_ref::<SeriesNormalize>() {
44            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
45        } else if let Some(node) = node.as_any().downcast_ref::<InstantManipulate>() {
46            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
47        } else if let Some(node) = node.as_any().downcast_ref::<RangeManipulate>() {
48            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
49        } else if let Some(node) = node.as_any().downcast_ref::<SeriesDivide>() {
50            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
51        } else if let Some(node) = node.as_any().downcast_ref::<EmptyMetric>() {
52            Ok(Some(node.to_execution_plan(
53                session,
54                planner,
55                planning_ctx,
56            )?))
57        } else if let Some(node) = node.as_any().downcast_ref::<ScalarCalculate>() {
58            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())?))
59        } else if let Some(node) = node.as_any().downcast_ref::<HistogramFold>() {
60            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
61        } else if let Some(node) = node.as_any().downcast_ref::<UnionDistinctOn>() {
62            Ok(Some(node.to_execution_plan(
63                physical_inputs[0].clone(),
64                physical_inputs[1].clone(),
65            )))
66        } else if let Some(node) = node.as_any().downcast_ref::<Absent>() {
67            Ok(Some(node.to_execution_plan(physical_inputs[0].clone())))
68        } else {
69            Ok(None)
70        }
71    }
72}
promql/extension_plan/planner.rs

promql/extension_plan/
planner.rs
