Fix correlated subquery optimizer rules (#1405)

* Fix correlated subquery optimizer rules

Part of RewritePullFilterThroughAggregation and
DependentSingleJoinToInnerJoin rules adds a group by
column to the aggregation from one of the sides of the predicate.
When selecting a column we need to make sure that the column is at
the same depth as the aggregation or deeper. If the group by column
is part of the outer query and the aggregation part of the inner
query, then the aggregation has no way of accessing the column. A
higher value for depth means deeper in the querer and a lower value
for depth means more shallow in the query. The current code was
always setting the column from the left side of the predicate as the
group by column. This was sometimes causing the group by column to
be more shallow than the aggregation itself which caused the query
to error out.

Fixes #1404

* Add comments

* Fix comments and variable naming

* Fix comments

* Remove dead code from DependentSingleJoinToInnerJoin

* Revert "Remove dead code from DependentSingleJoinToInnerJoin"

This reverts commit d1255a3.

* Extract common code from nested predicate rules

* Respond to PR comments

Co-authored-by: Wan Shen Lim <[email protected]>
Co-authored-by: Matt Butrovich <[email protected]>

Author: 3 people

script/testing/junit/traces/nested-query.test

@@ -174,6 +174,11 @@ SELECT item_name FROM shipment WHERE qty <= (SELECT SUM(amount) FROM part WHERE
 ----
 14 values hashing to eeafa4e65b8e56d2198e1a4703307d57
 
+query T nosort
+SELECT item_name FROM shipment WHERE qty <= (SELECT SUM(amount) FROM part WHERE shipment.pno = part.pno) order by item_name;
+----
+14 values hashing to eeafa4e65b8e56d2198e1a4703307d57
+
 query T nosort
 SELECT item_name FROM shipment WHERE qty < (SELECT SUM(amount) FROM part INNER JOIN shipment s on part.pno = s.pno) order by item_name;
 ----

src/include/optimizer/rules/unnesting_rules.h

@@ -1,6 +1,9 @@
 #pragma once
 
 #include <memory>
+#include <string>
+#include <tuple>
+#include <unordered_set>
 #include <vector>
 
 #include "optimizer/rule.h"
@@ -120,4 +123,19 @@ class DependentSingleJoinToInnerJoin : public Rule {
                  OptimizationContext *context) const override;
 };
 
+/**
+ * Given predicates associated with some aggregate, this function will extract the predicates that are correlated with
+ * an outer query, as well as the columns from those predicates that aren't correlated.
+ * @param predicates vector of predicates associated with an aggregate
+ * @param child_group_aliases_set the table alias set of the predicate's child node
+ * @param[out] correlated_predicates predicates which are correlated to an outer query
+ * @param[out] normal_predicates predicates which are not correlated to an outer query
+ * @param[out] new_group_cols columns from correlated predicates which are not correlated to an outer query. These
+ * will be used as group by columns in the nested aggregate
+ */
+void ExtractCorrelatedPredicatesWithAggregate(
+    const std::vector<AnnotatedExpression> &predicates, const std::unordered_set<std::string> &child_group_aliases_set,
+    std::vector<AnnotatedExpression> *correlated_predicates, std::vector<AnnotatedExpression> *normal_predicates,
+    std::vector<common::ManagedPointer<parser::AbstractExpression>> *new_groupby_cols);
+
 };  // namespace noisepage::optimizer

src/optimizer/rules/rewrite_rules.cpp

@@ -16,6 +16,7 @@
 #include "optimizer/optimizer_defs.h"
 #include "optimizer/physical_operators.h"
 #include "optimizer/properties.h"
+#include "optimizer/rules/unnesting_rules.h"
 #include "optimizer/util.h"
 #include "parser/expression_util.h"
 
@@ -493,21 +494,8 @@ void RewritePullFilterThroughAggregation::Transform(common::ManagedPointer<Abstr
   std::vector<AnnotatedExpression> correlated_predicates;
   std::vector<AnnotatedExpression> normal_predicates;
   std::vector<common::ManagedPointer<parser::AbstractExpression>> new_groupby_cols;
-  for (auto &predicate : predicates) {
-    if (OptimizerUtil::IsSubset(child_group_aliases_set, predicate.GetTableAliasSet())) {
-      normal_predicates.emplace_back(predicate);
-    } else {
-      // Correlated predicate, already in the form of
-      // (outer_relation.a = (expr))
-      correlated_predicates.emplace_back(predicate);
-      auto root_expr = predicate.GetExpr();
-      if (root_expr->GetChild(0)->GetDepth() < root_expr->GetDepth()) {
-        new_groupby_cols.emplace_back(root_expr->GetChild(1).Get());
-      } else {
-        new_groupby_cols.emplace_back(root_expr->GetChild(0).Get());
-      }
-    }
-  }
+  ExtractCorrelatedPredicatesWithAggregate(predicates, child_group_aliases_set, &correlated_predicates,
+                                           &normal_predicates, &new_groupby_cols);
 
   if (correlated_predicates.empty()) {
     // No need to pull

src/optimizer/rules/unnesting_rules.cpp

@@ -150,32 +150,11 @@ void DependentSingleJoinToInnerJoin::Transform(common::ManagedPointer<AbstractOp
   const auto &agg_group_aliases_set = memo.GetGroupByID(agg_group_id)->GetTableAliases();
   auto &filter_predicates = filter_expr->Contents()->GetContentsAs<LogicalFilter>()->GetPredicates();
 
-  std::vector<AnnotatedExpression> ancestor_predicates;
-  std::vector<AnnotatedExpression> descendant_predicates;
+  std::vector<AnnotatedExpression> correlated_predicates;
+  std::vector<AnnotatedExpression> normal_predicates;
   std::vector<common::ManagedPointer<parser::AbstractExpression>> new_groupby_cols;
-
-  // loop over all predicates check each of them if they refer table not contained in agg
-  // from RewritePullFilterThroughAggregation
-  for (auto &predicate : filter_predicates) {
-    if (OptimizerUtil::IsSubset(agg_group_aliases_set, predicate.GetTableAliasSet())) {
-      descendant_predicates.emplace_back(predicate);
-    } else {
-      // Correlated predicate, already in the form of
-      // (outer_relation.a = (expr))
-      ancestor_predicates.emplace_back(predicate);
-      auto root_expr = predicate.GetExpr();
-      // If the sub-query depth level of the first child is less than the current expression
-      // the first child is outer_relation.a and the second child is a (expr)
-      // The second child expression shall be evaluated as a part of the new aggregation before the new filter
-      if (root_expr->GetChild(0)->GetDepth() < root_expr->GetDepth()) {
-        new_groupby_cols.emplace_back(root_expr->GetChild(1).Get());
-      } else {
-        // Otherwise, the first child is a (expr) and the second child is outer_relation.a
-        // The first child expression shall be evaluated as a part of the new aggregation before the new filter
-        new_groupby_cols.emplace_back(root_expr->GetChild(0).Get());
-      }
-    }
-  }
+  ExtractCorrelatedPredicatesWithAggregate(filter_predicates, agg_group_aliases_set, &correlated_predicates,
+                                           &normal_predicates, &new_groupby_cols);
 
   // Create a new agg node
   auto aggregation = agg_expr->Contents()->GetContentsAs<LogicalAggregateAndGroupBy>();
@@ -193,10 +172,10 @@ void DependentSingleJoinToInnerJoin::Transform(common::ManagedPointer<AbstractOp
   // Create a new inner join node from single join
   std::vector<std::unique_ptr<AbstractOptimizerNode>> inner_node;
   inner_node.emplace_back(input->GetChildren()[0]->Copy());
-  if (!descendant_predicates.empty()) {
+  if (!normal_predicates.empty()) {
     std::vector<std::unique_ptr<AbstractOptimizerNode>> child_node;
     child_node.emplace_back(std::move(new_aggr));
-    auto filter = std::make_unique<OperatorNode>(LogicalFilter::Make(std::move(descendant_predicates))
+    auto filter = std::make_unique<OperatorNode>(LogicalFilter::Make(std::move(normal_predicates))
                                                      .RegisterWithTxnContext(context->GetOptimizerContext()->GetTxn()),
                                                  std::move(child_node), context->GetOptimizerContext()->GetTxn());
     inner_node.emplace_back(std::move(filter));
@@ -210,10 +189,10 @@ void DependentSingleJoinToInnerJoin::Transform(common::ManagedPointer<AbstractOp
   std::unique_ptr<OperatorNode> output;
   // Create new filter nodes
   // Construct a top filter if any
-  if (!ancestor_predicates.empty()) {
+  if (!correlated_predicates.empty()) {
     std::vector<std::unique_ptr<AbstractOptimizerNode>> root_node;
     root_node.emplace_back(std::move(new_inner));
-    output = std::make_unique<OperatorNode>(LogicalFilter::Make(std::move(ancestor_predicates))
+    output = std::make_unique<OperatorNode>(LogicalFilter::Make(std::move(correlated_predicates))
                                                 .RegisterWithTxnContext(context->GetOptimizerContext()->GetTxn()),
                                             std::move(root_node), context->GetOptimizerContext()->GetTxn());
 
@@ -223,4 +202,31 @@ void DependentSingleJoinToInnerJoin::Transform(common::ManagedPointer<AbstractOp
   transformed->emplace_back(std::move(output));
 }
 
+void ExtractCorrelatedPredicatesWithAggregate(
+    const std::vector<AnnotatedExpression> &predicates, const std::unordered_set<std::string> &child_group_aliases_set,
+    std::vector<AnnotatedExpression> *correlated_predicates, std::vector<AnnotatedExpression> *normal_predicates,
+    std::vector<common::ManagedPointer<parser::AbstractExpression>> *new_groupby_cols) {
+  for (auto &predicate : predicates) {
+    if (OptimizerUtil::IsSubset(child_group_aliases_set, predicate.GetTableAliasSet())) {
+      normal_predicates->emplace_back(predicate);
+    } else {
+      // Correlated predicate, predicate in nested query references column in outer query
+      correlated_predicates->emplace_back(predicate);
+      auto root_expr = predicate.GetExpr();
+      // See https://github.com/cmu-db/noisepage/issues/1404
+      // The higher the depth of an expression the deeper/more nested it is.
+      // If the sub-query depth level of the left side of the predicate is greater than the current expression then the
+      // left side doesn't references the outer query (i.e. the right side references the outer query).
+      // The left side shall be evaluated as a part of the new aggregation before the new filter
+      if (root_expr->GetChild(0)->GetDepth() > root_expr->GetDepth()) {
+        new_groupby_cols->emplace_back(root_expr->GetChild(0).Get());
+      } else {
+        // Otherwise, the right side of the predicate doesn't references the outer query (i.e. the left side references
+        // the outer query). The right side shall be evaluated as a part of the new aggregation before the new filter
+        new_groupby_cols->emplace_back(root_expr->GetChild(1).Get());
+      }
+    }
+  }
+}
+
 }  // namespace noisepage::optimizer