Merge pull request #1147 from chenboy/nested_query

Add unnesting for correlated subquery in where clause with aggregation

Author: linmagit

src/binder/bind_node_visitor.cpp

@@ -75,10 +75,8 @@ void BindNodeVisitor::Visit(parser::SelectStatement *node) {
     }
     select_element->DeriveSubqueryFlag();
 
-    // Recursively deduce expression value type
-    expression::ExpressionUtil::EvaluateExpression({ExprMap()},
-                                                   select_element.get());
-    // Recursively deduce expression name
+    // Traverse the expression to deduce expression value type and name
+    select_element->DeduceExpressionType();
     select_element->DeduceExpressionName();
     new_select_list.push_back(std::move(select_element));
   }

src/common/internal_types.cpp

@@ -978,6 +978,9 @@ std::string ExpressionTypeToString(ExpressionType type, bool short_str) {
     case ExpressionType::CAST: {
       return ("CAST");
     }
+    case ExpressionType::OPERATOR_IS_NOT_NULL: {
+      return ("IS_NOT_NULL");
+    }
     default: {
       throw ConversionException(StringUtil::Format(
           "No string conversion for ExpressionType value '%d'",

src/expression/aggregate_expression.cpp

@@ -17,7 +17,6 @@ namespace expression {
 
 const std::string AggregateExpression::GetInfo(int num_indent) const {
   std::ostringstream os;
-
   os << StringUtil::Indent(num_indent) << "Expression ::\n"
      << StringUtil::Indent(num_indent + 1) << "expression type = Aggregate,\n"
      << StringUtil::Indent(num_indent + 1) << "aggregate type = " << expr_name_

src/include/common/internal_types.h

@@ -1336,6 +1336,7 @@ enum class RuleType : uint32_t {
   MARK_JOIN_INNER_JOIN_TO_INNER_JOIN,
   MARK_JOIN_FILTER_TO_INNER_JOIN,
   PULL_FILTER_THROUGH_MARK_JOIN,
+  PULL_FILTER_THROUGH_AGGREGATION,
 
   // Place holder to generate number of rules compile time
   NUM_RULES

src/include/expression/aggregate_expression.h

@@ -79,16 +79,16 @@ class AggregateExpression : public AbstractExpression {
   }
 
   // Attribute binding
-  void PerformBinding(const std::vector<const planner::BindingContext *> &
-                          binding_contexts) override {
+  void PerformBinding(const std::vector<const planner::BindingContext *>
+                          &binding_contexts) override {
     const auto &context = binding_contexts[0];
     ai_ = context->Find(value_idx_);
     PL_ASSERT(ai_ != nullptr);
-    LOG_DEBUG("AggregateOutput Column ID %u.%u binds to AI %p (%s)", 0,
+    LOG_TRACE("AggregateOutput Column ID %u.%u binds to AI %p (%s)", 0,
               value_idx_, ai_, ai_->name.c_str());
   }
 
-  const planner::AttributeInfo* GetAttributeRef() const { return ai_; }
+  const planner::AttributeInfo *GetAttributeRef() const { return ai_; }
 
   inline void SetValueIdx(int value_idx) { value_idx_ = value_idx; }
 
@@ -130,7 +130,7 @@ class AggregateExpression : public AbstractExpression {
 
  private:
   int value_idx_ = -1;
-  const planner::AttributeInfo* ai_;
+  const planner::AttributeInfo *ai_;
 };
 
 }  // namespace expression

src/include/expression/expression_util.h

@@ -401,35 +401,43 @@ class ExpressionUtil {
   }
 
   /**
-   * TODO(boweic): this function may not be efficient, in the future we may want
-   * to add expressions to groups so that we do not need to walk through the
-   * expression tree when judging '==' each time
+   * @brief TODO(boweic): this function may not be efficient, in the future we
+   * may want to add expressions to groups so that we do not need to walk
+   * through the expression tree when judging '==' each time
    *
    * Convert all expression in the current expression tree that is in
    * child_expr_map to tuple value expression with corresponding column offset
    * of the input child tuple. This is used for handling projection
    * on situations like aggregate function (e.g. SELECT sum(a)+max(b) FROM ...
    * GROUP BY ...) when input columns contain sum(a) and sum(b). We need to
    * treat them as tuple value expression in the projection plan. This function
-   *should always be called before calling EvaluateExpression
+   * should always be called before calling EvaluateExpression
    *
    * Please notice that this function should only apply to copied expression
-   *since it would modify the current expression. We do not want to modify the
-   *original expression since it may be referenced in other places
+   * since it would modify the current expression. We do not want to modify the
+   * original expression since it may be referenced in other places
+   *
+   * @param expr The expression to modify
+   * @param child_expr_maps map from child column ids to expression
    */
   static void ConvertToTvExpr(AbstractExpression *expr,
-                              ExprMap &child_expr_map) {
+                              std::vector<ExprMap> child_expr_maps) {
+    if (expr == nullptr) {
+      return;
+    };
     for (size_t i = 0; i < expr->GetChildrenSize(); i++) {
       auto child_expr = expr->GetModifiableChild(i);
-      if (child_expr->GetExpressionType() != ExpressionType::VALUE_TUPLE &&
-          child_expr_map.count(child_expr)) {
-        // EvaluateExpression({child_expr_map}, child_expr);
-        expr->SetChild(i,
-                       new TupleValueExpression(child_expr->GetValueType(), 0,
-                                                child_expr_map[child_expr]));
-      } else {
-        ConvertToTvExpr(child_expr, child_expr_map);
+      for (size_t tuple_idx = 0; tuple_idx < child_expr_maps.size();
+           ++tuple_idx) {
+        if (child_expr->GetExpressionType() != ExpressionType::VALUE_TUPLE &&
+            child_expr_maps[tuple_idx].count(child_expr)) {
+          expr->SetChild(i, new TupleValueExpression(
+                                child_expr->GetValueType(), tuple_idx,
+                                child_expr_maps[tuple_idx][child_expr]));
+          break;
+        }
       }
+      ConvertToTvExpr(expr->GetModifiableChild(i), child_expr_maps);
     }
   }
 
@@ -443,6 +451,42 @@ class ExpressionUtil {
     return ordered_expr;
   }
 
+  /**
+   * Walks an expression trees and find all AggregationExprs subtrees.
+   */
+  static void GetTupleAndAggregateExprs(ExprSet &expr_set,
+                                        AbstractExpression *expr) {
+    std::vector<TupleValueExpression *> tv_exprs;
+    std::vector<AggregateExpression *> aggr_exprs;
+    GetAggregateExprs(aggr_exprs, tv_exprs, expr);
+    for (auto &tv_expr : tv_exprs) {
+      expr_set.insert(tv_expr);
+    }
+    for (auto &aggr_expr : aggr_exprs) {
+      expr_set.insert(aggr_expr);
+    }
+  }
+
+  /**
+   * Walks an expression trees and find all AggregationExprs subtrees.
+   */
+  static void GetTupleAndAggregateExprs(ExprMap &expr_map,
+                                        AbstractExpression *expr) {
+    std::vector<TupleValueExpression *> tv_exprs;
+    std::vector<AggregateExpression *> aggr_exprs;
+    GetAggregateExprs(aggr_exprs, tv_exprs, expr);
+    for (auto &tv_expr : tv_exprs) {
+      if (!expr_map.count(tv_expr)) {
+        expr_map.emplace(tv_expr, expr_map.size());
+      }
+    }
+    for (auto &aggr_expr : aggr_exprs) {
+      if (!expr_map.count(aggr_expr)) {
+        expr_map.emplace(aggr_expr, expr_map.size());
+      }
+    }
+  }
+
   /**
    * Walks an expression trees and find all AggregationExprs subtrees.
    */
@@ -511,8 +555,9 @@ class ExpressionUtil {
     // To evaluate the return type, we need a bottom up approach.
     if (expr == nullptr) return;
     size_t children_size = expr->GetChildrenSize();
-    for (size_t i = 0; i < children_size; i++)
+    for (size_t i = 0; i < children_size; i++) {
       EvaluateExpression(expr_maps, expr->GetModifiableChild(i));
+    }
 
     if (expr->GetExpressionType() == ExpressionType::VALUE_TUPLE) {
       // Point to the correct column returned in the logical tuple underneath
@@ -530,9 +575,12 @@ class ExpressionUtil {
       }
     } else if (IsAggregateExpression(expr->GetExpressionType())) {
       auto aggr_expr = (AggregateExpression *)expr;
-      auto &expr_map = expr_maps[0];
-      auto iter = expr_map.find(expr);
-      if (iter != expr_map.end()) aggr_expr->SetValueIdx(iter->second);
+      for (auto &expr_map : expr_maps) {
+        auto iter = expr_map.find(expr);
+        if (iter != expr_map.end()) {
+          aggr_expr->SetValueIdx(iter->second);
+        }
+      }
     } else if (expr->GetExpressionType() == ExpressionType::FUNCTION) {
       auto func_expr = (expression::FunctionExpression *)expr;
       std::vector<type::TypeId> argtypes;

src/include/optimizer/operators.h

@@ -226,6 +226,9 @@ class LogicalAggregateAndGroupBy
  public:
   static Operator make();
 
+  static Operator make(
+      std::vector<std::shared_ptr<expression::AbstractExpression>> &columns);
+
   static Operator make(
       std::vector<std::shared_ptr<expression::AbstractExpression>> &columns,
       std::vector<AnnotatedExpression> &having);

src/include/optimizer/plan_generator.h

@@ -126,7 +126,8 @@ class PlanGenerator : public OperatorVisitor {
    */
   std::unique_ptr<expression::AbstractExpression> GeneratePredicateForScan(
       const std::shared_ptr<expression::AbstractExpression> predicate_expr,
-      const std::string &alias, std::shared_ptr<catalog::TableCatalogObject> table);
+      const std::string &alias,
+      std::shared_ptr<catalog::TableCatalogObject> table);
 
   /**
    * @brief Generate projection info and projection schema for join
@@ -148,7 +149,7 @@ class PlanGenerator : public OperatorVisitor {
       AggregateType aggr_type,
       const std::vector<std::shared_ptr<expression::AbstractExpression>>
           *groupby_cols,
-      expression::AbstractExpression *having);
+      std::unique_ptr<expression::AbstractExpression> having);
 
   /**
    * @brief The required output property. Note that we have previously enforced

src/include/optimizer/query_to_operator_transformer.h

@@ -66,6 +66,7 @@ class QueryToOperatorTransformer : public SqlNodeVisitor {
   void Visit(expression::ComparisonExpression *expr) override;
   void Visit(expression::OperatorExpression *expr) override;
 
+ private:
   inline oid_t GetAndIncreaseGetId() { return get_id++; }
 
   /**
@@ -80,12 +81,10 @@ class QueryToOperatorTransformer : public SqlNodeVisitor {
    *
    * @param expr The original predicate
    */
-  void CollectPredicates(expression::AbstractExpression *expr);
+  std::vector<AnnotatedExpression> CollectPredicates(
+      expression::AbstractExpression *expr,
+      std::vector<AnnotatedExpression> predicates = {});
 
-  // TODO(boweic): Since we haven't migrated all the functionalities needed to
-  // generate mark-join and single-join to the optimizer, currently this
-  // function has not been tested, and it may be a bit hard to understand. We
-  // may integrate the unnesting functionality in the next PR
   /**
    * @brief Transform a sub-query in an expression to use
    *
@@ -96,14 +95,28 @@ class QueryToOperatorTransformer : public SqlNodeVisitor {
    * @return If the expression could be transformed into sub-query, return true,
    *  return false otherwise
    */
-  bool GenerateSubquerytree(
-      expression::AbstractExpression *expr,
-      std::vector<expression::AbstractExpression *> &select_list,
-      bool single_join = false);
+  bool GenerateSubquerytree(expression::AbstractExpression *expr,
+                            oid_t child_id, bool single_join = false);
 
+  /**
+   * @brief Decide if a conjunctive predicate is supported. We need to extract
+   * conjunction predicate first then call this function to decide if the
+   * predicate is supported by our system
+   *
+   * @param expr The conjunctive predicate provided
+   *
+   * @return True if supported, false otherwise
+   */
+  bool IsSupportedConjunctivePredicate(expression::AbstractExpression *expr);
+  /**
+   * @brief Check if a sub-select statement is supported.
+   *
+   * @param op The select statement
+   *
+   * @return True if supported, false otherwise
+   */
+  bool IsSupportedSubSelect(const parser::SelectStatement *op);
   static bool RequireAggregation(const parser::SelectStatement *op);
-
- private:
   std::shared_ptr<OperatorExpression> output_expr_;
 
   concurrency::TransactionContext *txn_;

src/include/optimizer/rule.h

@@ -108,6 +108,7 @@ struct RuleWithPromise {
   int promise;
 
   bool operator<(const RuleWithPromise &r) const { return promise < r.promise; }
+  bool operator>(const RuleWithPromise &r) const { return promise > r.promise; }
 };
 
 enum class RewriteRuleSetName : uint32_t {