Prune EXISTS subqueries during expression reduction

asenac · Jan 5, 2024 · cc1e6ed · cc1e6ed
1 parent fa52205
commit cc1e6ed
Show file tree

Hide file tree

Showing 5 changed files with 208 additions and 130 deletions.
diff --git a/src/query_graph/optimizer/rules/expression_reduction.rs b/src/query_graph/optimizer/rules/expression_reduction.rs
@@ -6,7 +6,7 @@ use crate::{
         properties::row_type,
         NodeId, QueryGraph, QueryNode,
     },
-    scalar_expr::reduction::reduce_expr_recursively,
+    scalar_expr::reduction::reduce_and_prune_exists_subplans_recursively,
 };
 
 pub struct ExpressionReductionRule;
@@ -20,23 +20,21 @@ impl SingleReplacementRule for ExpressionReductionRule {
         let new_node = match query_graph.node(node_id) {
             QueryNode::Project { outputs, input } => {
                 let row_type = row_type(query_graph, *input);
-                query_graph.project(
-                    *input,
-                    outputs
-                        .iter()
-                        .map(|e| reduce_expr_recursively(e, &query_graph, &row_type))
-                        .collect_vec(),
-                )
+                let input = *input;
+                let mut outputs = outputs.clone();
+                outputs.iter_mut().for_each(|e| {
+                    *e = reduce_and_prune_exists_subplans_recursively(e, query_graph, &row_type)
+                });
+                query_graph.project(input, outputs)
             }
             QueryNode::Filter { conditions, input } => {
                 let row_type = row_type(query_graph, *input);
-                query_graph.filter(
-                    *input,
-                    conditions
-                        .iter()
-                        .map(|e| reduce_expr_recursively(e, &query_graph, &row_type))
-                        .collect_vec(),
-                )
+                let input = *input;
+                let mut conditions = conditions.clone();
+                conditions.iter_mut().for_each(|e| {
+                    *e = reduce_and_prune_exists_subplans_recursively(e, query_graph, &row_type)
+                });
+                query_graph.filter(input, conditions)
             }
             QueryNode::Join {
                 join_type,
@@ -51,14 +49,18 @@ impl SingleReplacementRule for ExpressionReductionRule {
                     .chain(right_row_type.iter())
                     .cloned()
                     .collect_vec();
+                let left = *left;
+                let right = *right;
+                let join_type = join_type.clone();
+                let mut conditions = conditions.clone();
+                conditions.iter_mut().for_each(|e| {
+                    *e = reduce_and_prune_exists_subplans_recursively(e, query_graph, &row_type)
+                });
                 query_graph.add_node(QueryNode::Join {
-                    join_type: join_type.clone(),
-                    conditions: conditions
-                        .iter()
-                        .map(|e| reduce_expr_recursively(e, &query_graph, &row_type))
-                        .collect_vec(),
-                    left: *left,
-                    right: *right,
+                    join_type,
+                    conditions,
+                    left,
+                    right,
                 })
             }
             _ => node_id,

diff --git a/src/scalar_expr/reduction.rs b/src/scalar_expr/reduction.rs
@@ -2,9 +2,12 @@
 
 use itertools::Itertools;
 
-use crate::{data_type::DataType, query_graph::QueryGraph};
+use crate::{
+    data_type::DataType,
+    query_graph::{properties::num_columns, QueryGraph},
+};
 
-use super::{rewrite::rewrite_expr_pre_post, NaryOp, ScalarExpr, ScalarExprRef};
+use super::{rewrite::rewrite_expr_pre_post, NaryOp, ScalarExpr, ScalarExprRef, Subquery};
 
 /// Reduce the given expression recursively. Keeps trying until the expression cannot
 /// be reduced any further.
@@ -75,3 +78,43 @@ pub fn reduce_expr(
     }
     None
 }
+
+pub fn reduce_and_prune_exists_subplans_recursively(
+    expr: &ScalarExprRef,
+    query_graph: &mut QueryGraph,
+    row_type: &[DataType],
+) -> ScalarExprRef {
+    rewrite_expr_pre_post(
+        &mut |curr_expr: &ScalarExprRef| {
+            prune_exists_subplan(curr_expr, query_graph)
+                .or_else(|| reduce_expr(curr_expr, query_graph, row_type))
+        },
+        &expr,
+    )
+}
+
+pub fn prune_exists_subplan(
+    expr: &ScalarExprRef,
+    query_graph: &mut QueryGraph,
+) -> Option<ScalarExprRef> {
+    if let ScalarExpr::ExistsSubquery { subquery } = expr.as_ref() {
+        if num_columns(query_graph, subquery.root) > 0 {
+            // Skip the root node
+            let subquery_plan = query_graph.node(subquery.root).get_input(0);
+            let correlation = subquery.correlation.clone();
+            let project = query_graph.project(subquery_plan, vec![]);
+            let new_subquery_root = query_graph.add_subquery(project);
+            return Some(
+                ScalarExpr::ExistsSubquery {
+                    subquery: Subquery {
+                        root: new_subquery_root,
+                        correlation,
+                    },
+                }
+                .into(),
+            );
+        }
+    }
+
+    None
+}