feat: support common subexpression elimination

src/query/service/src/physical_plans/physical_aggregate_final.rs

@@ -208,23 +208,20 @@ impl PhysicalPlanBuilder {
         &mut self,
         s_expr: &SExpr,
         agg: &Aggregate,
-        mut required: ColumnSet,
+        required: ColumnSet,
         stat_info: PlanStatsInfo,
     ) -> Result<PhysicalPlan> {
         // 1. Prune unused Columns.
         let mut used = vec![];
         for item in &agg.aggregate_functions {
             if required.contains(&item.index) {
-                required.extend(item.scalar.used_columns());
                 used.push(item.clone());
             }
         }
 
-        agg.group_items.iter().for_each(|i| {
-            // If the group item comes from a complex expression, we only include the final
-            // column index here. The used columns will be included in its EvalScalar child.
-            required.insert(i.index);
-        });
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
 
         // single key without aggregation
         if agg.group_items.is_empty() && used.is_empty() {
@@ -245,7 +242,7 @@ impl PhysicalPlanBuilder {
         };
 
         // 2. Build physical plan.
-        let input = self.build(s_expr.child(0)?, required).await?;
+        let input = self.build(s_expr.child(0)?, child_required).await?;
         let input_schema = input.output_schema()?;
         let group_items = agg.group_items.iter().map(|v| v.index).collect::<Vec<_>>();
 

src/query/service/src/physical_plans/physical_async_func.rs

@@ -132,23 +132,26 @@ impl PhysicalPlanBuilder {
         &mut self,
         s_expr: &SExpr,
         async_func_plan: &databend_common_sql::plans::AsyncFunction,
-        mut required: ColumnSet,
+        required: ColumnSet,
         stat_info: PlanStatsInfo,
     ) -> Result<PhysicalPlan> {
         // 1. Prune unused Columns.
         let mut used = vec![];
         for item in async_func_plan.items.iter() {
             if required.contains(&item.index) {
-                required.extend(item.scalar.used_columns());
                 used.push(item.clone());
             }
         }
 
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
+
         // 2. Build physical plan.
         if used.is_empty() {
-            return self.build(s_expr.child(0)?, required).await;
+            return self.build(s_expr.child(0)?, child_required).await;
         }
-        let input = self.build(s_expr.child(0)?, required).await?;
+        let input = self.build(s_expr.child(0)?, child_required).await?;
         let input_schema = input.output_schema()?;
 
         let async_func_descs = used

src/query/service/src/physical_plans/physical_cte_consumer.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::any::Any;
+use std::collections::HashMap;
 
 use databend_common_exception::Result;
 use databend_common_expression::DataField;
@@ -93,11 +94,23 @@ impl PhysicalPlanBuilder {
         cte_consumer: &databend_common_sql::plans::MaterializedCTERef,
         stat_info: PlanStatsInfo,
     ) -> Result<PhysicalPlan> {
+        let def_to_ref = cte_consumer
+            .column_mapping
+            .iter()
+            .map(|(k, v)| (*v, *k))
+            .collect::<HashMap<_, _>>();
+        let cte_output_columns: Vec<_> = self
+            .cte_required_columns
+            .get(&cte_consumer.cte_name)
+            .unwrap()
+            .iter()
+            .map(|c| def_to_ref.get(c).unwrap())
+            .collect();
         let mut fields = Vec::new();
         let metadata = self.metadata.read();
 
-        for index in &cte_consumer.output_columns {
-            let column = metadata.column(*index);
+        for index in cte_output_columns.iter() {
+            let column = metadata.column(**index);
             let data_type = column.data_type();
             fields.push(DataField::new(&index.to_string(), data_type));
         }

src/query/service/src/physical_plans/physical_eval_scalar.rs

@@ -192,7 +192,7 @@ impl PhysicalPlanBuilder {
         &mut self,
         s_expr: &SExpr,
         eval_scalar: &databend_common_sql::plans::EvalScalar,
-        mut required: ColumnSet,
+        required: ColumnSet,
         stat_info: PlanStatsInfo,
     ) -> Result<PhysicalPlan> {
         // 1. Prune unused Columns.
@@ -204,19 +204,20 @@ impl PhysicalPlanBuilder {
                 continue;
             }
             used.push(s.clone());
-            s.scalar.used_columns().iter().for_each(|c| {
-                required.insert(*c);
-            })
         }
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
+
         // 2. Build physical plan.
         if used.is_empty() {
-            self.build(s_expr.child(0)?, required).await
+            self.build(s_expr.child(0)?, child_required).await
         } else {
             let child = s_expr.child(0)?;
             let input = if let Some(new_child) = self.try_eliminate_flatten_columns(&used, child)? {
-                self.build(&new_child, required).await?
+                self.build(&new_child, child_required.clone()).await?
             } else {
-                self.build(child, required).await?
+                self.build(child, child_required).await?
             };
 
             let column_projections: HashSet<usize> = column_projections

src/query/service/src/physical_plans/physical_exchange.rs

@@ -97,17 +97,15 @@ impl PhysicalPlanBuilder {
         &mut self,
         s_expr: &SExpr,
         exchange: &databend_common_sql::plans::Exchange,
-        mut required: ColumnSet,
+        required: ColumnSet,
     ) -> Result<PhysicalPlan> {
         // 1. Prune unused Columns.
-        if let databend_common_sql::plans::Exchange::Hash(exprs) = exchange {
-            for expr in exprs {
-                required.extend(expr.used_columns());
-            }
-        }
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
 
         // 2. Build physical plan.
-        let input = self.build(s_expr.child(0)?, required).await?;
+        let input = self.build(s_expr.child(0)?, child_required).await?;
         let input_schema = input.output_schema()?;
         let mut keys = vec![];
         let mut allow_adjust_parallelism = true;

src/query/service/src/physical_plans/physical_expression_scan.rs

@@ -117,7 +117,10 @@ impl PhysicalPlanBuilder {
         scan: &databend_common_sql::plans::ExpressionScan,
         required: ColumnSet,
     ) -> Result<PhysicalPlan> {
-        let input = self.build(s_expr.child(0)?, required).await?;
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
+        let input = self.build(s_expr.child(0)?, child_required).await?;
         let input_schema = input.output_schema()?;
 
         let values = scan

src/query/service/src/physical_plans/physical_filter.rs

@@ -137,12 +137,12 @@ impl PhysicalPlanBuilder {
         stat_info: PlanStatsInfo,
     ) -> Result<PhysicalPlan> {
         // 1. Prune unused Columns.
-        let used = filter.predicates.iter().fold(required.clone(), |acc, v| {
-            acc.union(&v.used_columns()).cloned().collect()
-        });
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
 
         // 2. Build physical plan.
-        let input = self.build(s_expr.child(0)?, used).await?;
+        let input = self.build(s_expr.child(0)?, child_required).await?;
         required = required
             .union(self.metadata.read().get_retained_column())
             .cloned()

src/query/service/src/physical_plans/physical_join.rs

@@ -150,28 +150,10 @@ impl PhysicalPlanBuilder {
                 others_required.insert(*column);
             }
         }
-
-        // Include columns referenced in left conditions and right conditions.
-        let left_required: ColumnSet = join
-            .equi_conditions
-            .iter()
-            .fold(required.clone(), |acc, v| {
-                acc.union(&v.left.used_columns()).cloned().collect()
-            })
-            .union(&others_required)
-            .cloned()
-            .collect();
-        let right_required: ColumnSet = join
-            .equi_conditions
-            .iter()
-            .fold(required.clone(), |acc, v| {
-                acc.union(&v.right.used_columns()).cloned().collect()
-            })
-            .union(&others_required)
-            .cloned()
-            .collect();
-        let left_required = left_required.union(&others_required).cloned().collect();
-        let right_required = right_required.union(&others_required).cloned().collect();
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let left_required = child_required.remove(0);
+        let right_required = child_required.remove(0);
 
         // 2. Build physical plan.
         // Choose physical join type by join conditions

src/query/service/src/physical_plans/physical_limit.rs

@@ -149,7 +149,10 @@ impl PhysicalPlanBuilder {
         }
 
         // 2. Build physical plan.
-        let input_plan = self.build(s_expr.child(0)?, required).await?;
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
+        let input_plan = self.build(s_expr.child(0)?, child_required).await?;
         if limit.before_exchange || limit.lazy_columns.is_empty() || !support_lazy_materialize {
             return Ok(PhysicalPlan::new(Limit {
                 input: input_plan,

src/query/service/src/physical_plans/physical_materialized_cte.rs

@@ -16,9 +16,10 @@ use std::any::Any;
 
 use databend_common_exception::Result;
 use databend_common_expression::DataSchemaRef;
-use databend_common_sql::optimizer::ir::RelExpr;
+use databend_common_pipeline_transforms::TransformPipelineHelper;
+use databend_common_sql::evaluator::BlockOperator;
+use databend_common_sql::evaluator::CompoundBlockOperator;
 use databend_common_sql::optimizer::ir::SExpr;
-use databend_common_sql::ColumnBinding;
 
 use crate::physical_plans::explain::PlanStatsInfo;
 use crate::physical_plans::format::MaterializedCTEFormatter;
@@ -38,7 +39,7 @@ pub struct MaterializedCTE {
     pub stat_info: Option<PlanStatsInfo>,
     pub input: PhysicalPlan,
     pub cte_name: String,
-    pub cte_output_columns: Option<Vec<ColumnBinding>>,
+    pub cte_output_columns: Option<Vec<usize>>,
     pub ref_count: usize,
     pub channel_size: Option<usize>,
     pub meta: PhysicalPlanMeta,
@@ -95,13 +96,20 @@ impl IPhysicalPlan for MaterializedCTE {
 
         let input_schema = self.input.output_schema()?;
         if let Some(output_columns) = &self.cte_output_columns {
-            PipelineBuilder::build_result_projection(
-                &builder.func_ctx,
-                input_schema,
-                output_columns,
-                &mut builder.main_pipeline,
-                false,
-            )?;
+            let mut projections = Vec::with_capacity(output_columns.len());
+            for index in output_columns {
+                projections.push(input_schema.index_of(index.to_string().as_str())?);
+            }
+            let num_input_columns = input_schema.num_fields();
+            builder.main_pipeline.add_transformer(|| {
+                CompoundBlockOperator::new(
+                    vec![BlockOperator::Project {
+                        projection: projections.clone(),
+                    }],
+                    builder.func_ctx.clone(),
+                    num_input_columns,
+                )
+            });
         }
 
         builder.main_pipeline.try_resize(1)?;
@@ -123,20 +131,19 @@ impl PhysicalPlanBuilder {
         materialized_cte: &databend_common_sql::plans::MaterializedCTE,
         stat_info: PlanStatsInfo,
     ) -> Result<PhysicalPlan> {
-        let required = match &materialized_cte.cte_output_columns {
-            Some(o) => o.iter().map(|c| c.index).collect(),
-            None => RelExpr::with_s_expr(s_expr.child(0)?)
-                .derive_relational_prop()?
-                .output_columns
-                .clone(),
-        };
-        let input = self.build(s_expr.child(0)?, required).await?;
+        let required = self
+            .cte_required_columns
+            .get(&materialized_cte.cte_name)
+            .unwrap()
+            .clone();
+        let cte_output_columns = Some(required.iter().copied().collect());
+        let input = self.build_physical_plan(s_expr.child(0)?, required).await?;
         Ok(PhysicalPlan::new(MaterializedCTE {
             plan_id: 0,
             stat_info: Some(stat_info),
             input,
             cte_name: materialized_cte.cte_name.clone(),
-            cte_output_columns: materialized_cte.cte_output_columns.clone(),
+            cte_output_columns,
             ref_count: materialized_cte.ref_count,
             channel_size: materialized_cte.channel_size,
             meta: PhysicalPlanMeta::new("MaterializedCTE"),

src/query/service/src/physical_plans/physical_mutation.rs

@@ -297,7 +297,10 @@ impl PhysicalPlanBuilder {
         let udf_col_num = required_udf_ids.len();
         required.extend(required_udf_ids);
 
-        let mut plan = self.build(s_expr.child(0)?, required).await?;
+        let mut child_required = self.derive_child_required_columns(s_expr, &required)?;
+        debug_assert_eq!(child_required.len(), s_expr.arity());
+        let child_required = child_required.remove(0);
+        let mut plan = self.build(s_expr.child(0)?, child_required).await?;
         if *no_effect {
             return Ok(plan);
         }