Push down projection to prewhere scan to prune output columns.

Author: RinChanNOWWW

src/query/legacy-planners/src/plan_node_extras.rs

@@ -68,8 +68,10 @@ impl Debug for Projection {
 
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq)]
 pub struct PrewhereInfo {
-    /// column indices of the table used for prewhere
-    pub need_columns: Projection,
+    /// columns to be ouput be prewhere scan
+    pub output_columns: Projection,
+    /// columns used for prewhere
+    pub prewhere_columns: Projection,
     /// remain_columns = scan.columns - need_columns
     pub remain_columns: Projection,
     /// filter for prewhere

src/query/service/src/sql/executor/physical_plan_builder.rs

@@ -121,8 +121,16 @@ impl PhysicalPlanBuilder {
                     if column.path_indices.is_some() {
                         has_inner_column = true;
                     }
-                    let name = column.name.clone();
-                    name_mapping.insert(name, index.to_string());
+                    if let Some(prewhere) = &scan.prewhere {
+                        // if there is a prewhere optimization,
+                        // we can prune `PhysicalScan`'s ouput schema.
+                        if prewhere.output_columns.contains(index) {
+                            name_mapping.insert(column.name.clone(), index.to_string());
+                        }
+                    } else {
+                        let name = column.name.clone();
+                        name_mapping.insert(name, index.to_string());
+                    }
                 }
 
                 let table_entry = metadata.table(scan.table_index);
@@ -424,13 +432,19 @@ impl PhysicalPlanBuilder {
 
                 let remain_columns = scan
                     .columns
-                    .difference(&prewhere.columns)
+                    .difference(&prewhere.prewhere_columns)
                     .copied()
                     .collect::<HashSet<usize>>();
-                let need_columns = Self::build_projection(
+                let output_columns = Self::build_projection(
+                    &metadata,
+                    table_schema,
+                    &prewhere.output_columns,
+                    has_inner_column,
+                );
+                let prewhere_columns = Self::build_projection(
                     &metadata,
                     table_schema,
-                    &prewhere.columns,
+                    &prewhere.prewhere_columns,
                     has_inner_column,
                 );
                 let remain_columns = Self::build_projection(
@@ -441,7 +455,8 @@ impl PhysicalPlanBuilder {
                 );
 
                 Ok::<PrewhereInfo, ErrorCode>(PrewhereInfo {
-                    need_columns,
+                    output_columns,
+                    prewhere_columns,
                     remain_columns,
                     filter,
                 })

src/query/service/src/sql/planner/metadata.rs

@@ -230,6 +230,14 @@ impl Metadata {
         }
         table_index
     }
+
+    pub fn find_smallest_column_within(&self, indices: &[usize]) -> usize {
+        let entries = indices
+            .iter()
+            .map(|i| self.column(*i).clone())
+            .collect::<Vec<_>>();
+        find_smallest_column(entries.as_slice())
+    }
 }
 
 pub fn optimize_remove_count_args(name: &str, distinct: bool, args: &[&Expr]) -> bool {

src/query/service/src/sql/planner/optimizer/heuristic/mod.rs

@@ -14,16 +14,17 @@
 
 mod decorrelate;
 mod implement;
+mod prewhere_optimization;
 mod prune_columns;
 mod rule_list;
 mod subquery_rewriter;
-mod where_optimizer;
 
 use std::sync::Arc;
 
 use common_exception::Result;
 use once_cell::sync::Lazy;
 
+use self::prewhere_optimization::PrewhereOptimizer;
 use super::rule::RuleID;
 use super::ColumnSet;
 use crate::sessions::TableContext;
@@ -93,13 +94,13 @@ impl HeuristicOptimizer {
     }
 
     fn post_optimize(&mut self, s_expr: SExpr) -> Result<SExpr> {
+        let prewhere_optimizer = PrewhereOptimizer::new(self.metadata.clone());
+        let s_expr = prewhere_optimizer.prewhere_optimize(s_expr)?;
+
         let pruner = prune_columns::ColumnPruner::new(self.metadata.clone());
         let require_columns: ColumnSet =
             self.bind_context.columns.iter().map(|c| c.index).collect();
-        let s_expr = pruner.prune_columns(&s_expr, require_columns)?;
-
-        let where_opt = where_optimizer::WhereOptimizer::new(self.metadata.clone());
-        where_opt.where_optimize(s_expr)
+        pruner.prune_columns(&s_expr, require_columns)
     }
 
     pub fn optimize(&mut self, s_expr: SExpr) -> Result<SExpr> {

src/query/service/src/sql/planner/optimizer/heuristic/where_optimizer.rs renamed to src/query/service/src/sql/planner/optimizer/heuristic/prewhere_optimization.rs

@@ -24,12 +24,12 @@ use crate::sql::plans::RelOp;
 use crate::sql::plans::Scalar;
 use crate::sql::MetadataRef;
 
-pub struct WhereOptimizer {
+pub struct PrewhereOptimizer {
     metadata: MetadataRef,
     pattern: SExpr,
 }
 
-impl WhereOptimizer {
+impl PrewhereOptimizer {
     pub fn new(metadata: MetadataRef) -> Self {
         Self {
             metadata,
@@ -97,7 +97,7 @@ impl WhereOptimizer {
         )
     }
 
-    pub fn where_optimize(&self, s_expr: SExpr) -> Result<SExpr> {
+    pub fn prewhere_optimize(&self, s_expr: SExpr) -> Result<SExpr> {
         let rel_op = s_expr.plan();
         if s_expr.match_pattern(&self.pattern) {
             let mut filter: Filter = s_expr.plan().clone().try_into()?;
@@ -131,7 +131,8 @@ impl WhereOptimizer {
                 None
             } else {
                 Some(Prewhere {
-                    columns: prewhere_columns,
+                    output_columns: get.columns.clone(),
+                    prewhere_columns,
                     predicates: prewhere_pred,
                 })
             };
@@ -149,7 +150,7 @@ impl WhereOptimizer {
             let children = s_expr
                 .children()
                 .iter()
-                .map(|expr| self.where_optimize(expr.clone()))
+                .map(|expr| self.prewhere_optimize(expr.clone()))
                 .collect::<Result<Vec<_>>>()?;
             Ok(SExpr::create(rel_op.clone(), children, None, None))
         }

src/query/service/src/sql/planner/optimizer/heuristic/prune_columns.rs

@@ -76,7 +76,29 @@ impl ColumnPruner {
     fn keep_required_columns(&self, expr: &SExpr, mut required: ColumnSet) -> Result<SExpr> {
         match expr.plan() {
             RelOperator::LogicalGet(p) => {
+                let mut prewhere = p.prewhere.clone();
                 let mut used: ColumnSet = required.intersection(&p.columns).cloned().collect();
+                if let Some(ref mut pw) = prewhere {
+                    debug_assert!(
+                        pw.prewhere_columns.is_subset(&p.columns),
+                        "prewhere columns should be a subset of scan columns"
+                    );
+                    // `used` is the columns which prewhere scan needs to output for its upper operator.
+                    if used.is_empty() {
+                        let prewhere_columns =
+                            pw.prewhere_columns.iter().map(|&i| i).collect::<Vec<_>>();
+                        let smallest_index = self
+                            .metadata
+                            .read()
+                            .find_smallest_column_within(prewhere_columns.as_slice());
+                        used.insert(smallest_index);
+                    }
+                    pw.output_columns = used.clone();
+                    // `prune_columns` is after `prewhere_optimize`,
+                    // so we need to add prewhere columns to scan columns.
+                    used = used.union(&pw.prewhere_columns).cloned().collect();
+                }
+
                 if used.is_empty() {
                     let columns = self.metadata.read().columns_by_table_index(p.table_index);
                     let smallest_index = find_smallest_column(&columns);
@@ -90,7 +112,7 @@ impl ColumnPruner {
                     limit: p.limit,
                     order_by: p.order_by.clone(),
                     statistics: p.statistics,
-                    prewhere: p.prewhere.clone(),
+                    prewhere,
                 })))
             }
             RelOperator::LogicalInnerJoin(p) => {

src/query/service/src/sql/planner/plans/logical_get.rs

@@ -29,7 +29,11 @@ use crate::sql::plans::SortItem;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Prewhere {
-    pub columns: ColumnSet,
+    // columns needed to be output after prewhere scan
+    pub output_columns: ColumnSet,
+    // columns needed to conduct prewhere filter
+    pub prewhere_columns: ColumnSet,
+    // prewhere filter predicates
     pub predicates: Vec<Scalar>,
 }
 

src/query/service/src/storages/storage_table_read_plan.rs

@@ -15,11 +15,13 @@
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
+use common_datavalues::DataField;
 use common_exception::Result;
 use common_legacy_planners::Extras;
 use common_legacy_planners::Projection;
 use common_legacy_planners::ReadDataSourcePlan;
 use common_legacy_planners::SourceInfo;
+use common_meta_app::schema::TableInfo;
 
 use crate::sessions::QueryContext;
 use crate::storages::Table;
@@ -57,29 +59,14 @@ impl ToReadDataSourcePlan for dyn Table {
         let description = statistics.get_description(table_info);
 
         let scan_fields = match (self.benefit_column_prune(), &push_downs) {
-            (true, Some(push_downs)) => match &push_downs.projection {
-                Some(projection) => match projection {
-                    Projection::Columns(ref indices) => {
-                        if indices.len() < table_info.schema().fields().len() {
-                            let fields = indices
-                                .iter()
-                                .map(|i| table_info.schema().field(*i).clone());
-
-                            Some((indices.iter().cloned().zip(fields)).collect::<BTreeMap<_, _>>())
-                        } else {
-                            None
-                        }
-                    }
-                    Projection::InnerColumns(ref path_indices) => {
-                        let column_ids: Vec<usize> = path_indices.keys().cloned().collect();
-                        let new_schema = table_info.schema().inner_project(path_indices);
-                        Some(
-                            (column_ids.iter().cloned().zip(new_schema.fields().clone()))
-                                .collect::<BTreeMap<_, _>>(),
-                        )
-                    }
+            (true, Some(push_downs)) => match &push_downs.prewhere {
+                Some(prewhere) => {
+                    extract_scan_fields_from_projection(table_info, &prewhere.output_columns)
+                }
+                _ => match &push_downs.projection {
+                    Some(projection) => extract_scan_fields_from_projection(table_info, projection),
+                    _ => None,
                 },
-                _ => None,
             },
             _ => None,
         };
@@ -98,3 +85,30 @@ impl ToReadDataSourcePlan for dyn Table {
         })
     }
 }
+
+fn extract_scan_fields_from_projection(
+    table_info: &TableInfo,
+    projection: &Projection,
+) -> Option<BTreeMap<usize, DataField>> {
+    match projection {
+        Projection::Columns(ref indices) => {
+            if indices.len() < table_info.schema().fields().len() {
+                let fields = indices
+                    .iter()
+                    .map(|i| table_info.schema().field(*i).clone());
+
+                Some((indices.iter().cloned().zip(fields)).collect::<BTreeMap<_, _>>())
+            } else {
+                None
+            }
+        }
+        Projection::InnerColumns(ref path_indices) => {
+            let column_ids: Vec<usize> = path_indices.keys().cloned().collect();
+            let new_schema = table_info.schema().inner_project(path_indices);
+            Some(
+                (column_ids.iter().cloned().zip(new_schema.fields().clone()))
+                    .collect::<BTreeMap<_, _>>(),
+            )
+        }
+    }
+}

src/query/storages/fuse/src/operations/read.rs

@@ -90,7 +90,7 @@ impl FuseTable {
 
         let (output_reader, prewhere_reader, prewhere_filter, remain_reader) =
             if let Some(prewhere) = self.prewhere_of_push_downs(&plan.push_downs) {
-                let prewhere_schema = prewhere.need_columns.project_schema(&table_schema);
+                let prewhere_schema = prewhere.prewhere_columns.project_schema(&table_schema);
                 let prewhere_schema = Arc::new(prewhere_schema);
                 let expr_field = prewhere.filter.to_data_field(&prewhere_schema)?;
                 let expr_schema = DataSchemaRefExt::create(vec![expr_field]);
@@ -103,9 +103,10 @@ impl FuseTable {
                     vec![prewhere.filter.clone()],
                     false,
                 )?;
-
+                let output_reader =
+                    self.create_block_reader(&ctx, prewhere.output_columns.clone())?;
                 let prewhere_reader =
-                    self.create_block_reader(&ctx, prewhere.need_columns.clone())?;
+                    self.create_block_reader(&ctx, prewhere.prewhere_columns.clone())?;
                 let remain_reader = if prewhere.remain_columns.is_empty() {
                     None
                 } else {