feat: run dist planner after simplify expression rule (#2723)

Signed-off-by: Ruihang Xia <[email protected]>

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<LogicalPlan> {
+        // 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)?;

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));