flow/
df_optimizer.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
15//! Datafusion optimizer for flow plan
16
17#![warn(unused)]
18
19use std::collections::HashSet;
20use std::sync::Arc;
21
22use common_error::ext::BoxedError;
23use common_telemetry::debug;
24use datafusion::config::ConfigOptions;
25use datafusion::error::DataFusionError;
26use datafusion::optimizer::analyzer::type_coercion::TypeCoercion;
27use datafusion::optimizer::common_subexpr_eliminate::CommonSubexprEliminate;
28use datafusion::optimizer::optimize_projections::OptimizeProjections;
29use datafusion::optimizer::simplify_expressions::SimplifyExpressions;
30use datafusion::optimizer::{Analyzer, AnalyzerRule, Optimizer, OptimizerContext};
31use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor};
32use query::optimizer::count_wildcard::CountWildcardToTimeIndexRule;
33use session::context::QueryContextRef;
34use snafu::ResultExt;
35
36/// note here we are using the `substrait_proto_df` crate from the `substrait` module and
37/// rename it to `substrait_proto`
38use crate::error::{DatafusionSnafu, Error, ExternalSnafu};
39
40// TODO(discord9): use `Analyzer` to manage rules if more `AnalyzerRule` is needed
41pub async fn apply_df_optimizer(
42    plan: datafusion_expr::LogicalPlan,
43    query_ctx: &QueryContextRef,
44) -> Result<datafusion_expr::LogicalPlan, Error> {
45    let cfg = query_ctx.create_config_options();
46    let analyzer = Analyzer::with_rules(vec![
47        Arc::new(CountWildcardToTimeIndexRule),
48        Arc::new(CheckGroupByRule::new()),
49        Arc::new(TypeCoercion::new()),
50    ]);
51    let plan = analyzer
52        .execute_and_check(plan, &cfg, |p, r| {
53            debug!("After apply rule {}, get plan: \n{:?}", r.name(), p);
54        })
55        .context(DatafusionSnafu {
56            context: "Fail to apply analyzer",
57        })?;
58
59    let mut ctx = OptimizerContext::new();
60    let scheduled_time = query::options::parse_scheduled_time_datetime(&query_ctx.extensions())
61        .map_err(BoxedError::new)
62        .context(ExternalSnafu)?;
63    if let Some(dt) = scheduled_time {
64        ctx = ctx.with_query_execution_start_time(dt);
65    }
66    let optimizer = Optimizer::with_rules(vec![
67        Arc::new(OptimizeProjections::new()),
68        Arc::new(CommonSubexprEliminate::new()),
69        Arc::new(SimplifyExpressions::new()),
70    ]);
71    let plan = optimizer
72        .optimize(plan, &ctx, |_, _| {})
73        .context(DatafusionSnafu {
74            context: "Fail to apply optimizer",
75        })?;
76
77    Ok(plan)
78}
79
80/// This rule check all group by exprs, and make sure they are also in select clause in a aggr query
81#[derive(Debug)]
82struct CheckGroupByRule {}
83
84impl CheckGroupByRule {
85    pub fn new() -> Self {
86        Self {}
87    }
88}
89
90impl AnalyzerRule for CheckGroupByRule {
91    fn analyze(
92        &self,
93        plan: datafusion_expr::LogicalPlan,
94        _config: &ConfigOptions,
95    ) -> datafusion_common::Result<datafusion_expr::LogicalPlan> {
96        let transformed = plan
97            .transform_up_with_subqueries(check_group_by_analyzer)?
98            .data;
99        Ok(transformed)
100    }
101
102    fn name(&self) -> &str {
103        "check_groupby"
104    }
105}
106
107/// make sure everything in group by's expr is in select
108fn check_group_by_analyzer(
109    plan: datafusion_expr::LogicalPlan,
110) -> Result<Transformed<datafusion_expr::LogicalPlan>, DataFusionError> {
111    if let datafusion_expr::LogicalPlan::Projection(proj) = &plan
112        && let datafusion_expr::LogicalPlan::Aggregate(aggr) = proj.input.as_ref()
113    {
114        let mut found_column_used = FindColumn::new();
115        proj.expr
116            .iter()
117            .map(|i| i.visit(&mut found_column_used))
118            .count();
119        for expr in aggr.group_expr.iter() {
120            if !found_column_used
121                .names_for_alias
122                .contains(&expr.name_for_alias()?)
123            {
124                return Err(DataFusionError::Plan(format!(
125                    "Expect {} expr in group by also exist in select list, but select list only contain {:?}",
126                    expr.name_for_alias()?,
127                    found_column_used.names_for_alias
128                )));
129            }
130        }
131    }
132
133    Ok(Transformed::no(plan))
134}
135
136/// Find all column names in a plan
137#[derive(Debug, Default)]
138struct FindColumn {
139    names_for_alias: HashSet<String>,
140}
141
142impl FindColumn {
143    fn new() -> Self {
144        Default::default()
145    }
146}
147
148impl TreeNodeVisitor<'_> for FindColumn {
149    type Node = datafusion_expr::Expr;
150    fn f_down(
151        &mut self,
152        node: &datafusion_expr::Expr,
153    ) -> Result<TreeNodeRecursion, DataFusionError> {
154        if let datafusion_expr::Expr::Column(_) = node {
155            self.names_for_alias.insert(node.name_for_alias()?);
156        }
157        Ok(TreeNodeRecursion::Continue)
158    }
159}
flow/df_optimizer.rs

flow/
df_optimizer.rs
