From f02dc0e27420590351615eac1db6aa1ffaa8b79e Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Thu, 9 Nov 2023 20:44:50 +0800 Subject: [PATCH] feat: run dist planner after simplify expression rule (#2723) Signed-off-by: Ruihang Xia --- src/query/src/dist_plan/analyzer.rs | 8 ++++++++ src/query/src/query_engine/state.rs | 6 +++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 06aee5844d45..a7f4fa5b4c26 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -21,6 +21,8 @@ use datafusion_common::tree_node::{RewriteRecursion, Transformed, TreeNode, Tree use datafusion_expr::expr::{Exists, InSubquery}; use datafusion_expr::{col, Expr, LogicalPlan, LogicalPlanBuilder, Subquery}; use datafusion_optimizer::analyzer::AnalyzerRule; +use datafusion_optimizer::simplify_expressions::SimplifyExpressions; +use datafusion_optimizer::{OptimizerContext, OptimizerRule}; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use table::metadata::TableType; use table::table::adapter::DfTableProviderAdapter; @@ -42,6 +44,12 @@ impl AnalyzerRule for DistPlannerAnalyzer { plan: LogicalPlan, _config: &ConfigOptions, ) -> datafusion_common::Result { + // preprocess the input plan + let optimizer_context = OptimizerContext::new(); + let plan = SimplifyExpressions::new() + .try_optimize(&plan, &optimizer_context)? + .unwrap_or(plan); + let plan = plan.transform(&Self::inspect_plan_with_subquery)?; let mut rewriter = PlanRewriter::default(); let result = plan.rewrite(&mut rewriter)?; diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index b5dc35b3e9b4..c80ecadc9fb8 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -80,13 +80,13 @@ impl QueryEngineState { let session_config = SessionConfig::new().with_create_default_catalog_and_schema(false); // Apply the type conversion rule first. let mut analyzer = Analyzer::new(); - if with_dist_planner { - analyzer.rules.insert(0, Arc::new(DistPlannerAnalyzer)); - } analyzer.rules.insert(0, Arc::new(TypeConversionRule)); analyzer.rules.insert(0, Arc::new(StringNormalizationRule)); Self::remove_analyzer_rule(&mut analyzer.rules, CountWildcardRule {}.name()); analyzer.rules.insert(0, Arc::new(CountWildcardRule {})); + if with_dist_planner { + analyzer.rules.push(Arc::new(DistPlannerAnalyzer)); + } let mut optimizer = Optimizer::new(); optimizer.rules.push(Arc::new(OrderHintRule));