test(p3): add nlj init check (#526)

Author: christopherlim98

src/common/bustub_instance.cpp

@@ -20,6 +20,7 @@
 #include "common/util/string_util.h"
 #include "concurrency/lock_manager.h"
 #include "concurrency/transaction.h"
+#include "execution/check_options.h"
 #include "execution/execution_engine.h"
 #include "execution/executor_context.h"
 #include "execution/executors/mock_scan_executor.h"
@@ -172,10 +173,11 @@ see the execution plan of your query.
   WriteOneCell(help, writer);
 }
 
-auto BustubInstance::ExecuteSql(const std::string &sql, ResultWriter &writer) -> bool {
+auto BustubInstance::ExecuteSql(const std::string &sql, ResultWriter &writer,
+                                std::shared_ptr<CheckOptions> check_options) -> bool {
   auto txn = txn_manager_->Begin();
   try {
-    auto result = ExecuteSqlTxn(sql, writer, txn);
+    auto result = ExecuteSqlTxn(sql, writer, txn, std::move(check_options));
     txn_manager_->Commit(txn);
     delete txn;
     return result;
@@ -186,7 +188,8 @@ auto BustubInstance::ExecuteSql(const std::string &sql, ResultWriter &writer) ->
   }
 }
 
-auto BustubInstance::ExecuteSqlTxn(const std::string &sql, ResultWriter &writer, Transaction *txn) -> bool {
+auto BustubInstance::ExecuteSqlTxn(const std::string &sql, ResultWriter &writer, Transaction *txn,
+                                   std::shared_ptr<CheckOptions> check_options) -> bool {
   if (!sql.empty() && sql[0] == '\\') {
     // Internal meta-commands, like in `psql`.
     if (sql == "\\dt") {
@@ -257,6 +260,9 @@ auto BustubInstance::ExecuteSqlTxn(const std::string &sql, ResultWriter &writer,
 
     // Execute the query.
     auto exec_ctx = MakeExecutorContext(txn);
+    if (check_options != nullptr) {
+      exec_ctx->InitCheckOptions(std::move(check_options));
+    }
     std::vector<Tuple> result_set{};
     is_successful &= execution_engine_->Execute(optimized_plan, &result_set, txn, exec_ctx.get());
 

src/execution/CMakeLists.txt

@@ -8,6 +8,7 @@ add_library(
         fmt_impl.cpp
         hash_join_executor.cpp
         index_scan_executor.cpp
+        init_check_executor.cpp
         insert_executor.cpp
         limit_executor.cpp
         mock_scan_executor.cpp

src/execution/executor_factory.cpp

@@ -21,6 +21,7 @@
 #include "execution/executors/filter_executor.h"
 #include "execution/executors/hash_join_executor.h"
 #include "execution/executors/index_scan_executor.h"
+#include "execution/executors/init_check_executor.h"
 #include "execution/executors/insert_executor.h"
 #include "execution/executors/limit_executor.h"
 #include "execution/executors/mock_scan_executor.h"
@@ -33,6 +34,7 @@
 #include "execution/executors/update_executor.h"
 #include "execution/executors/values_executor.h"
 #include "execution/plans/filter_plan.h"
+#include "execution/plans/init_check_plan.h"
 #include "execution/plans/mock_scan_plan.h"
 #include "execution/plans/projection_plan.h"
 #include "execution/plans/sort_plan.h"
@@ -95,6 +97,17 @@ auto ExecutorFactory::CreateExecutor(ExecutorContext *exec_ctx, const AbstractPl
       auto nested_loop_join_plan = dynamic_cast<const NestedLoopJoinPlanNode *>(plan.get());
       auto left = ExecutorFactory::CreateExecutor(exec_ctx, nested_loop_join_plan->GetLeftPlan());
       auto right = ExecutorFactory::CreateExecutor(exec_ctx, nested_loop_join_plan->GetRightPlan());
+      auto check_options_set = exec_ctx->GetCheckOptions()->check_options_set_;
+      if (check_options_set.find(CheckOption::ENABLE_NLJ_CHECK) != check_options_set.end()) {
+        auto left_init_check_plan = dynamic_cast<const InitCheckPlanNode *>(nested_loop_join_plan->GetLeftPlan().get());
+        auto right_init_check_plan =
+            dynamic_cast<const InitCheckPlanNode *>(nested_loop_join_plan->GetRightPlan().get());
+        auto left_check = std::make_unique<InitCheckExecutor>(exec_ctx, left_init_check_plan, std::move(left));
+        auto right_check = std::make_unique<InitCheckExecutor>(exec_ctx, right_init_check_plan, std::move(right));
+        exec_ctx->AddCheckExecutor(left_check.get(), right_check.get());
+        return std::make_unique<NestedLoopJoinExecutor>(exec_ctx, nested_loop_join_plan, std::move(left_check),
+                                                        std::move(right_check));
+      }
       return std::make_unique<NestedLoopJoinExecutor>(exec_ctx, nested_loop_join_plan, std::move(left),
                                                       std::move(right));
     }

src/execution/init_check_executor.cpp

@@ -0,0 +1,40 @@
+//===----------------------------------------------------------------------===//
+//
+//                         BusTub
+//
+// init_check_executor.cpp
+//
+// Identification: src/execution/init_check_executor.cpp
+//
+// Copyright (c) 2015-2021, Carnegie Mellon University Database Group
+//
+//===----------------------------------------------------------------------===//
+
+#include "execution/executors/init_check_executor.h"
+
+namespace bustub {
+
+InitCheckExecutor::InitCheckExecutor(ExecutorContext *exec_ctx, const InitCheckPlanNode *plan,
+                                     std::unique_ptr<AbstractExecutor> &&child_executor)
+    : AbstractExecutor{exec_ctx}, child_executor_{std::move(child_executor)} {}
+
+void InitCheckExecutor::Init() {
+  if (!child_executor_) {
+    return;
+  }
+  n_init_++;
+  // Initialize the child executor
+  child_executor_->Init();
+}
+
+auto InitCheckExecutor::Next(Tuple *tuple, RID *rid) -> bool {
+  if (!child_executor_) {
+    return EXECUTOR_EXHAUSTED;
+  }
+
+  // Emit the next tuple
+  n_next_++;
+  return child_executor_->Next(tuple, rid);
+}
+
+}  // namespace bustub

src/include/common/bustub_instance.h

@@ -25,6 +25,7 @@
 #include "catalog/catalog.h"
 #include "common/config.h"
 #include "common/util/string_util.h"
+#include "execution/check_options.h"
 #include "libfort/lib/fort.hpp"
 #include "type/value.h"
 
@@ -226,12 +227,14 @@ class BustubInstance {
   /**
    * Execute a SQL query in the BusTub instance.
    */
-  auto ExecuteSql(const std::string &sql, ResultWriter &writer) -> bool;
+  auto ExecuteSql(const std::string &sql, ResultWriter &writer, std::shared_ptr<CheckOptions> check_options = nullptr)
+      -> bool;
 
   /**
    * Execute a SQL query in the BusTub instance with provided txn.
    */
-  auto ExecuteSqlTxn(const std::string &sql, ResultWriter &writer, Transaction *txn) -> bool;
+  auto ExecuteSqlTxn(const std::string &sql, ResultWriter &writer, Transaction *txn,
+                     std::shared_ptr<CheckOptions> check_options = nullptr) -> bool;
 
   /**
    * FOR TEST ONLY. Generate test tables in this BusTub instance.

src/include/execution/check_options.h

@@ -0,0 +1,36 @@
+// Copyright 2022 RisingLight Project Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <memory>
+#include <unordered_set>
+
+#pragma once
+
+namespace bustub {
+
+enum class CheckOption : uint8_t {
+  ENABLE_NLJ_CHECK = 0,
+  ENABLE_TOPN_CHECK = 1,
+};
+
+/**
+ * The CheckOptions class contains the set of check options used for testing
+ * executor logic.
+ */
+class CheckOptions {
+ public:
+  std::unordered_set<CheckOption> check_options_set_;
+};
+
+};  // namespace bustub

src/include/execution/execution_engine.h

@@ -20,6 +20,7 @@
 #include "concurrency/transaction_manager.h"
 #include "execution/executor_context.h"
 #include "execution/executor_factory.h"
+#include "execution/executors/init_check_executor.h"
 #include "execution/plans/abstract_plan.h"
 #include "storage/table/tuple.h"
 
@@ -63,6 +64,7 @@ class ExecutionEngine {
     try {
       executor->Init();
       PollExecutor(executor.get(), plan, result_set);
+      PerformChecks(exec_ctx);
     } catch (const ExecutionException &ex) {
 #ifndef NDEBUG
       LOG_ERROR("Error Encountered in Executor Execution: %s", ex.what());
@@ -76,6 +78,15 @@ class ExecutionEngine {
     return executor_succeeded;
   }
 
+  void PerformChecks(ExecutorContext *exec_ctx) {
+    for (const auto &[left_executor, right_executor] : exec_ctx->GetNLJCheckExecutorSet()) {
+      auto casted_left_executor = dynamic_cast<const InitCheckExecutor *>(left_executor);
+      auto casted_right_executor = dynamic_cast<const InitCheckExecutor *>(right_executor);
+      BUSTUB_ASSERT(casted_left_executor->GetNextCount() == casted_right_executor->GetInitCount(),
+                    "nlj check failed, are you initialising the right executor correctly?");
+    }
+  }
+
  private:
   /**
    * Poll the executor until exhausted, or exception escapes.

src/include/execution/executor_context.h

@@ -12,15 +12,20 @@
 
 #pragma once
 
+#include <deque>
+#include <memory>
 #include <unordered_set>
 #include <utility>
 #include <vector>
 
 #include "catalog/catalog.h"
 #include "concurrency/transaction.h"
+#include "execution/check_options.h"
+#include "execution/executors/abstract_executor.h"
 #include "storage/page/tmp_tuple_page.h"
 
 namespace bustub {
+class AbstractExecutor;
 /**
  * ExecutorContext stores all the context necessary to run an executor.
  */
@@ -36,7 +41,11 @@ class ExecutorContext {
    */
   ExecutorContext(Transaction *transaction, Catalog *catalog, BufferPoolManager *bpm, TransactionManager *txn_mgr,
                   LockManager *lock_mgr)
-      : transaction_(transaction), catalog_{catalog}, bpm_{bpm}, txn_mgr_(txn_mgr), lock_mgr_(lock_mgr) {}
+      : transaction_(transaction), catalog_{catalog}, bpm_{bpm}, txn_mgr_(txn_mgr), lock_mgr_(lock_mgr) {
+    nlj_check_exec_set_ = std::deque<std::pair<AbstractExecutor *, AbstractExecutor *>>(
+        std::deque<std::pair<AbstractExecutor *, AbstractExecutor *>>{});
+    check_options_ = std::make_shared<CheckOptions>();
+  }
 
   ~ExecutorContext() = default;
 
@@ -60,6 +69,23 @@ class ExecutorContext {
   /** @return the transaction manager */
   auto GetTransactionManager() -> TransactionManager * { return txn_mgr_; }
 
+  /** @return the set of nlj check executors */
+  auto GetNLJCheckExecutorSet() -> std::deque<std::pair<AbstractExecutor *, AbstractExecutor *>> & {
+    return nlj_check_exec_set_;
+  }
+
+  /** @return the check options */
+  auto GetCheckOptions() -> std::shared_ptr<CheckOptions> { return check_options_; }
+
+  void AddCheckExecutor(AbstractExecutor *left_exec, AbstractExecutor *right_exec) {
+    nlj_check_exec_set_.emplace_back(left_exec, right_exec);
+  }
+
+  void InitCheckOptions(std::shared_ptr<CheckOptions> &&check_options) {
+    BUSTUB_ASSERT(check_options, "nullptr");
+    check_options_ = std::move(check_options);
+  }
+
  private:
   /** The transaction context associated with this executor context */
   Transaction *transaction_;
@@ -71,6 +97,10 @@ class ExecutorContext {
   TransactionManager *txn_mgr_;
   /** The lock manager associated with this executor context */
   LockManager *lock_mgr_;
+  /** The set of NLJ check executors associated with this executor context */
+  std::deque<std::pair<AbstractExecutor *, AbstractExecutor *>> nlj_check_exec_set_;
+  /** The set of check options associated with this executor context */
+  std::shared_ptr<CheckOptions> check_options_;
 };
 
 }  // namespace bustub

src/include/execution/executors/abstract_executor.h

@@ -16,6 +16,7 @@
 #include "storage/table/tuple.h"
 
 namespace bustub {
+class ExecutorContext;
 /**
  * The AbstractExecutor implements the Volcano tuple-at-a-time iterator model.
  * This is the base class from which all executors in the BustTub execution

src/include/execution/executors/init_check_executor.h

@@ -0,0 +1,76 @@
+//===----------------------------------------------------------------------===//
+//
+//                         BusTub
+//
+// init_check_executor.h
+//
+// Identification: src/include/execution/executors/init_check_executor.h
+//
+// Copyright (c) 2015-2021, Carnegie Mellon University Database Group
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <memory>
+#include <utility>
+
+#include "execution/executor_context.h"
+#include "execution/executors/abstract_executor.h"
+#include "execution/plans/init_check_plan.h"
+
+namespace bustub {
+
+/**
+ * InitCheckExecutor counts the number of times the child operator calls init.
+ */
+class InitCheckExecutor : public AbstractExecutor {
+ public:
+  /**
+   * Construct a new InitCheckExecutor instance.
+   * @param exec_ctx The executor context
+   * @param plan The init check plan to be executed
+   * @param child_executor The child executor from which init calls are counted
+   */
+  InitCheckExecutor(ExecutorContext *exec_ctx, const InitCheckPlanNode *plan,
+                    std::unique_ptr<AbstractExecutor> &&child_executor);
+
+  /** Initialize the InitCheck */
+  void Init() override;
+
+  /**
+   * Yield the next tuple from the child executor.
+   * @param[out] tuple The next tuple produced by the child executor
+   * @param[out] rid The next tuple RID produced by the child executor
+   * @return `true` if a tuple was produced, `false` if there are no more tuples
+   */
+  auto Next(Tuple *tuple, RID *rid) -> bool override;
+
+  /** @return The output schema for the child executor */
+  auto GetOutputSchema() const -> const Schema & override { return plan_->OutputSchema(); };
+
+  /** @return The number of inits */
+  auto GetInitCount() const -> std::size_t { return n_init_; };
+
+  /** @return The number of nexts */
+  auto GetNextCount() const -> std::size_t { return n_next_; };
+
+ private:
+  /** InitCheckExecutor returns `true` when it should be polled again */
+  constexpr static const bool EXECUTOR_ACTIVE{true};
+
+  /** InitCheckExecutor returns `false` when child executor is exhausted */
+  constexpr static const bool EXECUTOR_EXHAUSTED{false};
+
+  /** The init check plan node to be executed */
+  const InitCheckPlanNode *plan_;
+
+  /** The child executor from which tuples are obtained */
+  std::unique_ptr<AbstractExecutor> child_executor_;
+
+  /** The number of times init was called */
+  std::size_t n_init_{0};
+  std::size_t n_next_{0};
+};
+
+}  // namespace bustub