Merge pull request #7646 from RinChanNOWWW/prune-prewhere

refactor(query): push all filters to prewhere and prune columns for it.

Author: mergify[bot]

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,
@@ -48,59 +48,50 @@ impl WhereOptimizer {
         }
     }
 
-    fn collect_columns(expr: &Scalar, columns: &mut ColumnSet) {
+    fn collect_columns_impl(expr: &Scalar, columns: &mut ColumnSet) {
         match expr {
             Scalar::BoundColumnRef(column) => {
                 columns.insert(column.column.index);
             }
             Scalar::AndExpr(and) => {
-                Self::collect_columns(and.left.as_ref(), columns);
-                Self::collect_columns(and.right.as_ref(), columns);
+                Self::collect_columns_impl(and.left.as_ref(), columns);
+                Self::collect_columns_impl(and.right.as_ref(), columns);
             }
             Scalar::OrExpr(or) => {
-                Self::collect_columns(or.left.as_ref(), columns);
-                Self::collect_columns(or.right.as_ref(), columns);
+                Self::collect_columns_impl(or.left.as_ref(), columns);
+                Self::collect_columns_impl(or.right.as_ref(), columns);
             }
             Scalar::ComparisonExpr(cmp) => {
-                Self::collect_columns(cmp.left.as_ref(), columns);
-                Self::collect_columns(cmp.right.as_ref(), columns);
+                Self::collect_columns_impl(cmp.left.as_ref(), columns);
+                Self::collect_columns_impl(cmp.right.as_ref(), columns);
             }
             Scalar::FunctionCall(func) => {
                 for arg in func.arguments.iter() {
-                    Self::collect_columns(arg, columns);
+                    Self::collect_columns_impl(arg, columns);
                 }
             }
             Scalar::CastExpr(cast) => {
-                Self::collect_columns(cast.argument.as_ref(), columns);
+                Self::collect_columns_impl(cast.argument.as_ref(), columns);
             }
             // 1. ConstantExpr is not collected.
-            // 2. SubqueryExpr is not collected.
-            // 3. AggregateFunction will not appear in where clause.
+            // 2. SubqueryExpr and AggregateFunction will not appear in Filter-LogicalGet
             _ => {}
         }
     }
 
     // analyze if the expression can be moved to prewhere
-    fn analyze(expr: &Scalar, columns_to_scan: usize) -> (bool, ColumnSet) {
+    fn collect_columns(expr: &Scalar) -> ColumnSet {
         let mut columns = ColumnSet::new();
-
         // columns in subqueries are not considered
-        Self::collect_columns(expr, &mut columns);
+        Self::collect_columns_impl(expr, &mut columns);
 
-        // viable conditions:
-        // 1. Condition depend on some column. Constant expressions are not moved.
-        // 2. Do not move conditions involving all queried columns.
-        // 3. Only current table columns are considered. (This condition is always true in current Pattern (Filter -> LogicalGet)).
-        (
-            !columns.is_empty() && columns.len() < columns_to_scan,
-            columns,
-        )
+        columns
     }
 
-    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()?;
+            let filter: Filter = s_expr.plan().clone().try_into()?;
             let mut get: LogicalGet = s_expr.child(0)?.plan().clone().try_into()?;
             let metadata = self.metadata.read().clone();
 
@@ -112,44 +103,30 @@ impl WhereOptimizer {
 
             let mut prewhere_columns = ColumnSet::new();
             let mut prewhere_pred = Vec::new();
-            let mut remain_pred = Vec::new();
-
-            let columns_to_scan = get.columns.len();
 
             // filter.predicates are already splited by AND
             for pred in filter.predicates.iter() {
-                let (viable, columns) = Self::analyze(pred, columns_to_scan);
-                if viable {
-                    prewhere_pred.push(pred.clone());
-                    prewhere_columns.extend(&columns);
-                } else {
-                    remain_pred.push(pred.clone());
-                }
+                let columns = Self::collect_columns(pred);
+                prewhere_pred.push(pred.clone());
+                prewhere_columns.extend(&columns);
             }
 
             get.prewhere = if prewhere_pred.is_empty() {
                 None
             } else {
                 Some(Prewhere {
-                    columns: prewhere_columns,
+                    output_columns: get.columns.clone(),
+                    prewhere_columns,
                     predicates: prewhere_pred,
                 })
             };
 
-            if !remain_pred.is_empty() {
-                filter.predicates = remain_pred;
-                Ok(SExpr::create_unary(
-                    filter.into(),
-                    SExpr::create_leaf(get.into()),
-                ))
-            } else {
-                Ok(SExpr::create_leaf(get.into()))
-            }
+            Ok(SExpr::create_leaf(get.into()))
         } else {
             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().copied().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<_, _>>(),
+            )
+        }
+    }
+}