From 7b7c5b3fc0b57ade123a2cb7e92ecfdb07637a74 Mon Sep 17 00:00:00 2001
From: Josip Mrden <josip.mrden@memgraph.io>
Date: Wed, 21 Jun 2023 16:35:59 +0200
Subject: [PATCH] Add cost estimator to subqueries

---
 src/query/plan/cost_estimator.hpp | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/query/plan/cost_estimator.hpp b/src/query/plan/cost_estimator.hpp
index 30eceb2d3..e666f0f49 100644
--- a/src/query/plan/cost_estimator.hpp
+++ b/src/query/plan/cost_estimator.hpp
@@ -186,9 +186,7 @@ class CostEstimator : public HierarchicalLogicalOperatorVisitor {
   // TODO: Cost estimate ScanAllById?
 
   bool PostVisit(Expand &expand) override {
-    const auto &scope = scopes_.back();
     auto card_param = CardParam::kExpand;
-
     auto stats = GetStatsFor(expand.input_symbol_);
 
     if (stats.has_value()) {
@@ -289,10 +287,13 @@ class CostEstimator : public HierarchicalLogicalOperatorVisitor {
   }
 
   bool PreVisit(Apply &op) override {
-    double input_cost = EstimateCostOnBranch(&op.input_);
-    double subquery_cost = EstimateCostOnBranch(&op.subquery_);
+    op.input_->Accept(*this);
+
+    auto last_scope = scopes_.back();
+    double subquery_cost = EstimateCostOnBranch(&op.subquery_, last_scope);
 
     // if the query is a unit subquery, we don't want the cost to be zero but 1xN
+    double input_cost = cost();
     input_cost = input_cost == 0 ? 1 : input_cost;
     subquery_cost = subquery_cost == 0 ? 1 : subquery_cost;
 
@@ -302,6 +303,11 @@ class CostEstimator : public HierarchicalLogicalOperatorVisitor {
     return false;
   }
 
+  bool PostVisit(EmptyResult & /*op*/) override {
+    scopes_.emplace_back();
+    return true;
+  }
+
   bool Visit(Once &) override { return true; }
 
   auto cost() const { return cost_; }
@@ -330,6 +336,12 @@ class CostEstimator : public HierarchicalLogicalOperatorVisitor {
     return cost_estimator.cost();
   }
 
+  double EstimateCostOnBranch(std::shared_ptr<LogicalOperator> *branch, Scope scope) {
+    CostEstimator<TDbAccessor> cost_estimator(db_accessor_, parameters, table_, scope);
+    (*branch)->Accept(cost_estimator);
+    return cost_estimator.cost();
+  }
+
   // converts an optional ScanAll range bound into a property value
   // if the bound is present and is a constant expression convertible to
   // a property value. otherwise returns nullopt