Add hooks for post processing generated plans

Reviewers: mtomic, llugovic

Reviewed By: mtomic

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D1824

src/query/plan/planner.hpp

@@ -21,6 +21,48 @@ class SymbolTable;
 
 namespace plan {
 
+class PostProcessor final {
+  Parameters parameters_;
+
+ public:
+  using ProcessedPlan = std::unique_ptr<LogicalOperator>;
+
+  explicit PostProcessor(const Parameters &parameters)
+      : parameters_(parameters) {}
+
+  template <class TPlanningContext>
+  std::unique_ptr<LogicalOperator> Rewrite(
+      std::unique_ptr<LogicalOperator> plan, TPlanningContext *context) {
+    return plan;
+  }
+
+  template <class TVertexCounts>
+  double EstimatePlanCost(const std::unique_ptr<LogicalOperator> &plan,
+                          TVertexCounts *vertex_counts) {
+    return ::query::plan::EstimatePlanCost(vertex_counts, parameters_, *plan);
+  }
+
+  template <class TPlanningContext>
+  std::unique_ptr<LogicalOperator> MergeWithCombinator(
+      std::unique_ptr<LogicalOperator> curr_op,
+      std::unique_ptr<LogicalOperator> last_op, const Tree &combinator,
+      TPlanningContext *context) {
+    if (const auto *union_ = dynamic_cast<const CypherUnion *>(&combinator)) {
+      return std::unique_ptr<LogicalOperator>(
+          impl::GenUnion(*union_, std::move(last_op), std::move(curr_op),
+                         *context->symbol_table));
+    }
+    throw utils::NotYetImplemented("query combinator");
+  }
+
+  template <class TPlanningContext>
+  std::unique_ptr<LogicalOperator> MakeDistinct(
+      std::unique_ptr<LogicalOperator> last_op, TPlanningContext *context) {
+    auto output_symbols = last_op->OutputSymbols(*context->symbol_table);
+    return std::make_unique<Distinct>(std::move(last_op), output_symbols);
+  }
+};
+
 /// @brief Generates the LogicalOperator tree for a single query and returns the
 /// resulting plan.
 ///
@@ -45,63 +87,73 @@ auto MakeLogicalPlanForSingleQuery(
 /// Generates the LogicalOperator tree and returns the resulting plan.
 ///
 /// @tparam TPlanningContext Type of the context used.
+/// @tparam TPlanPostProcess Type of the plan post processor used.
+///
 /// @param context PlanningContext used for generating plans.
-/// @param parameters Parameters used in query .
-/// @param boolean flag use_variable_planner to choose which planner to use
-/// @return pair consisting of the plan's first logical operator @c
-/// LogicalOperator and the estimated cost of that plan
-template <class TPlanningContext>
-auto MakeLogicalPlan(TPlanningContext *context, const Parameters &parameters,
+/// @param post_process performs plan rewrites and cost estimation.
+/// @param use_variable_planner boolean flag to choose which planner to use.
+///
+/// @return pair consisting of the final `TPlanPostProcess::ProcessedPlan` and
+/// the estimated cost of that plan as a `double`.
+template <class TPlanningContext, class TPlanPostProcess>
+auto MakeLogicalPlan(TPlanningContext *context, TPlanPostProcess *post_process,
                      bool use_variable_planner) {
   auto query_parts = CollectQueryParts(*context->symbol_table,
                                        *context->ast_storage, context->query);
   auto &vertex_counts = *context->db;
   double total_cost = 0;
-  std::unique_ptr<LogicalOperator> last_op;
+
+  using ProcessedPlan = typename TPlanPostProcess::ProcessedPlan;
+  ProcessedPlan last_plan;
 
   for (const auto &query_part : query_parts.query_parts) {
-    std::unique_ptr<LogicalOperator> op;
+    std::experimental::optional<ProcessedPlan> curr_plan;
     double min_cost = std::numeric_limits<double>::max();
 
     if (use_variable_planner) {
       auto plans = MakeLogicalPlanForSingleQuery<VariableStartPlanner>(
           query_part.single_query_parts, context);
       for (auto plan : plans) {
-        auto cost = EstimatePlanCost(&vertex_counts, parameters, *plan);
-        if (!op || cost < min_cost) {
-          // Plans are generated lazily and the current plan will disappear, so
-          // it's ok to move it.
-          op = std::move(plan);
+        // Plans are generated lazily and the current plan will disappear, so
+        // it's ok to move it.
+        auto rewritten_plan = post_process->Rewrite(std::move(plan), context);
+        double cost =
+            post_process->EstimatePlanCost(rewritten_plan, &vertex_counts);
+        if (!curr_plan || cost < min_cost) {
+          curr_plan.emplace(std::move(rewritten_plan));
           min_cost = cost;
         }
       }
     } else {
-      op = MakeLogicalPlanForSingleQuery<RuleBasedPlanner>(
+      auto plan = MakeLogicalPlanForSingleQuery<RuleBasedPlanner>(
           query_part.single_query_parts, context);
-      min_cost = EstimatePlanCost(&vertex_counts, parameters, *op);
+      auto rewritten_plan = post_process->Rewrite(std::move(plan), context);
+      min_cost = post_process->EstimatePlanCost(rewritten_plan, &vertex_counts);
+      curr_plan.emplace(std::move(rewritten_plan));
     }
 
     total_cost += min_cost;
-    if (auto *union_ =
-            dynamic_cast<CypherUnion *>(query_part.query_combinator)) {
-      std::shared_ptr<LogicalOperator> curr_op(std::move(op));
-      std::shared_ptr<LogicalOperator> prev_op(std::move(last_op));
-      last_op = std::unique_ptr<LogicalOperator>(
-          impl::GenUnion(*union_, prev_op, curr_op, *context->symbol_table));
-    } else if (query_part.query_combinator) {
-      throw utils::NotYetImplemented("query combinator");
+    if (query_part.query_combinator) {
+      last_plan = post_process->MergeWithCombinator(
+          std::move(*curr_plan), std::move(last_plan),
+          *query_part.query_combinator, context);
     } else {
-      last_op = std::move(op);
+      last_plan = std::move(*curr_plan);
     }
   }
 
   if (query_parts.distinct) {
-    std::shared_ptr<LogicalOperator> prev_op(std::move(last_op));
-    last_op = std::make_unique<Distinct>(
-        prev_op, prev_op->OutputSymbols(*context->symbol_table));
+    last_plan = post_process->MakeDistinct(std::move(last_plan), context);
   }
 
-  return std::make_pair(std::move(last_op), total_cost);
+  return std::make_pair(std::move(last_plan), total_cost);
+}
+
+template <class TPlanningContext>
+auto MakeLogicalPlan(TPlanningContext *context, const Parameters &parameters,
+                     bool use_variable_planner) {
+  PostProcessor post_processor(parameters);
+  return MakeLogicalPlan(context, &post_processor, use_variable_planner);
 }
 
 }  // namespace plan

src/query/plan/rule_based_planner.cpp

@@ -584,7 +584,7 @@ std::unique_ptr<LogicalOperator> GenWith(
 }
 
 std::unique_ptr<LogicalOperator> GenUnion(
-    CypherUnion &cypher_union, std::shared_ptr<LogicalOperator> left_op,
+    const CypherUnion &cypher_union, std::shared_ptr<LogicalOperator> left_op,
     std::shared_ptr<LogicalOperator> right_op, SymbolTable &symbol_table) {
   return std::make_unique<Union>(left_op, right_op, cypher_union.union_symbols_,
                                  left_op->OutputSymbols(symbol_table),

src/query/plan/rule_based_planner.hpp

@@ -135,7 +135,7 @@ std::unique_ptr<LogicalOperator> GenWith(
     std::unordered_set<Symbol> &bound_symbols, AstStorage &storage);
 
 std::unique_ptr<LogicalOperator> GenUnion(
-    CypherUnion &cypher_union, std::shared_ptr<LogicalOperator> left_op,
+    const CypherUnion &cypher_union, std::shared_ptr<LogicalOperator> left_op,
     std::shared_ptr<LogicalOperator> right_op, SymbolTable &symbol_table);
 
 template <class TBoolOperator>