add primary key support and change tuple meta format (#631)

* add primary key support and change tuple meta format

Signed-off-by: Alex Chi <[email protected]>

* fix

Signed-off-by: Alex Chi <[email protected]>

---------

Signed-off-by: Alex Chi <[email protected]>

src/binder/bind_create.cpp

@@ -68,6 +68,14 @@ auto Binder::BindColumnDefinition(duckdb_libpgquery::PGColumnDef *cdef) -> Colum
     return {colname, TypeId::INTEGER};
   }
 
+  if (name == "double") {
+    return {colname, TypeId::DECIMAL};
+  }
+
+  if (name == "bool") {
+    return {colname, TypeId::BOOLEAN};
+  }
+
   if (name == "varchar") {
     auto exprs = BindExpressionList(cdef->typeName->typmods);
     if (exprs.size() != 1) {
@@ -85,6 +93,7 @@ auto Binder::BindCreate(duckdb_libpgquery::PGCreateStmt *pg_stmt) -> std::unique
   auto table = std::string(pg_stmt->relation->relname);
   auto columns = std::vector<Column>{};
   size_t column_count = 0;
+  std::vector<std::string> pk;
 
   for (auto c = pg_stmt->tableElts->head; c != nullptr; c = lnext(c)) {
     auto node = reinterpret_cast<duckdb_libpgquery::PGNode *>(c->data.ptr_value);
@@ -93,14 +102,44 @@ auto Binder::BindCreate(duckdb_libpgquery::PGCreateStmt *pg_stmt) -> std::unique
         auto cdef = reinterpret_cast<duckdb_libpgquery::PGColumnDef *>(c->data.ptr_value);
         auto centry = BindColumnDefinition(cdef);
         if (cdef->constraints != nullptr) {
-          throw NotImplementedException("constraints not supported");
+          for (auto constr = cdef->constraints->head; constr != nullptr; constr = constr->next) {
+            auto constraint = reinterpret_cast<duckdb_libpgquery::PGConstraint *>(constr->data.ptr_value);
+            switch (constraint->contype) {
+              case duckdb_libpgquery::PG_CONSTR_PRIMARY: {
+                if (!pk.empty()) {
+                  throw NotImplementedException("cannot have two primary keys");
+                }
+                pk = {centry.GetName()};
+                break;
+              }
+              default:
+                throw NotImplementedException("unsupported constraint");
+            }
+          }
         }
         columns.push_back(std::move(centry));
         column_count++;
         break;
       }
       case duckdb_libpgquery::T_PGConstraint: {
-        throw NotImplementedException("constraints not supported");
+        for (auto con = c; con != nullptr; con = con->next) {
+          auto constraint = reinterpret_cast<duckdb_libpgquery::PGConstraint *>(con->data.ptr_value);
+          switch (constraint->contype) {
+            case duckdb_libpgquery::PG_CONSTR_PRIMARY: {
+              std::vector<std::string> columns;
+              for (auto kc = constraint->keys->head; kc != nullptr; kc = kc->next) {
+                columns.emplace_back(reinterpret_cast<duckdb_libpgquery::PGValue *>(kc->data.ptr_value)->val.str);
+              }
+              if (!pk.empty()) {
+                throw NotImplementedException("cannot have two primary keys");
+              }
+              pk = std::move(columns);
+              break;
+            }
+            default:
+              throw NotImplementedException("unsupported constraint");
+          }
+        }
         break;
       }
       default:
@@ -112,7 +151,7 @@ auto Binder::BindCreate(duckdb_libpgquery::PGCreateStmt *pg_stmt) -> std::unique
     throw bustub::Exception("should have at least 1 column");
   }
 
-  return std::make_unique<CreateStatement>(std::move(table), std::move(columns));
+  return std::make_unique<CreateStatement>(std::move(table), std::move(columns), std::move(pk));
 }
 
 auto Binder::BindIndex(duckdb_libpgquery::PGIndexStmt *stmt) -> std::unique_ptr<IndexStatement> {

src/binder/statement/create_statement.cpp

@@ -5,11 +5,14 @@
 
 namespace bustub {
 
-CreateStatement::CreateStatement(std::string table, std::vector<Column> columns)
-    : BoundStatement(StatementType::CREATE_STATEMENT), table_(std::move(table)), columns_(std::move(columns)) {}
+CreateStatement::CreateStatement(std::string table, std::vector<Column> columns, std::vector<std::string> primary_key)
+    : BoundStatement(StatementType::CREATE_STATEMENT),
+      table_(std::move(table)),
+      columns_(std::move(columns)),
+      primary_key_(std::move(primary_key)) {}
 
 auto CreateStatement::ToString() const -> std::string {
-  return fmt::format("BoundCreate {{\n  table={}\n  columns={}\n}}", table_, columns_);
+  return fmt::format("BoundCreate {{\n  table={}\n  columns={}\n  primary_key={}\n}}", table_, columns_, primary_key_);
 }
 
 }  // namespace bustub

src/catalog/table_generator.cpp

@@ -77,8 +77,7 @@ void TableGenerator::FillTable(TableInfo *info, TableInsertMeta *table_meta) {
       for (const auto &col : values) {
         entry.emplace_back(col[i]);
       }
-      auto rid =
-          info->table_->InsertTuple(TupleMeta{INVALID_TXN_ID, INVALID_TXN_ID, false}, Tuple(entry, &info->schema_));
+      auto rid = info->table_->InsertTuple(TupleMeta{0, false}, Tuple(entry, &info->schema_));
       BUSTUB_ENSURE(rid != std::nullopt, "Sequential insertion cannot fail");
       num_inserted++;
     }

src/common/bustub_ddl.cpp

@@ -15,6 +15,7 @@
 #include "binder/statement/select_statement.h"
 #include "binder/statement/set_show_statement.h"
 #include "buffer/buffer_pool_manager.h"
+#include "catalog/catalog.h"
 #include "catalog/schema.h"
 #include "catalog/table_generator.h"
 #include "common/bustub_instance.h"
@@ -43,12 +44,44 @@ namespace bustub {
 void BustubInstance::HandleCreateStatement(Transaction *txn, const CreateStatement &stmt, ResultWriter &writer) {
   std::unique_lock<std::shared_mutex> l(catalog_lock_);
   auto info = catalog_->CreateTable(txn, stmt.table_, Schema(stmt.columns_));
+  IndexInfo *index = nullptr;
+  if (!stmt.primary_key_.empty()) {
+    std::vector<uint32_t> col_ids;
+    for (const auto &col : stmt.primary_key_) {
+      auto idx = info->schema_.GetColIdx(col);
+      col_ids.push_back(idx);
+      if (info->schema_.GetColumn(idx).GetType() != TypeId::INTEGER) {
+        throw NotImplementedException("only support creating index on integer column");
+      }
+    }
+    auto key_schema = Schema::CopySchema(&info->schema_, col_ids);
+
+    // TODO(spring2023): If you want to support composite index key for leaderboard optimization, remove this assertion
+    // and create index with different key type that can hold multiple keys based on number of index columns.
+    //
+    // You can also create clustered index that directly stores value inside the index by modifying the value type.
+
+    if (col_ids.empty() || col_ids.size() > 2) {
+      throw NotImplementedException("only support creating index with exactly one or two columns");
+    }
+
+    index = catalog_->CreateIndex<IntegerKeyType, IntegerValueType, IntegerComparatorType>(
+        txn, stmt.table_ + "_pk", stmt.table_, info->schema_, key_schema, col_ids, TWO_INTEGER_SIZE,
+        IntegerHashFunctionType{}, true);
+  }
   l.unlock();
 
   if (info == nullptr) {
     throw bustub::Exception("Failed to create table");
   }
-  WriteOneCell(fmt::format("Table created with id = {}", info->oid_), writer);
+
+  if (index != nullptr) {
+    WriteOneCell(fmt::format("Table created with id = {}, Primary key index created with id = {}", info->oid_,
+                             index->index_oid_),
+                 writer);
+  } else {
+    WriteOneCell(fmt::format("Table created with id = {}", info->oid_), writer);
+  }
 }
 
 void BustubInstance::HandleIndexStatement(Transaction *txn, const IndexStatement &stmt, ResultWriter &writer) {

src/common/bustub_instance.cpp

@@ -49,8 +49,10 @@ BustubInstance::BustubInstance(const std::string &db_file_name) {
   // Storage related.
   disk_manager_ = std::make_unique<DiskManager>(db_file_name);
 
+#ifndef DISABLE_CHECKPOINT_MANAGER
   // Log related.
   log_manager_ = std::make_unique<LogManager>(disk_manager_.get());
+#endif
 
   // We need more frames for GenerateTestTable to work. Therefore, we use 128 instead of the default
   // buffer pool size specified in `config.h`.
@@ -62,24 +64,34 @@ BustubInstance::BustubInstance(const std::string &db_file_name) {
     buffer_pool_manager_ = nullptr;
   }
 
-  // Transaction (txn) related.
+// Transaction (txn) related.
+#ifndef DISABLE_LOCK_MANAGER
   lock_manager_ = std::make_unique<LockManager>();
+  txn_manager_ = std::make_unique<TransactionManager>(lock_manager_.get());
+#else
+  txn_manager_ = std::make_unique<TransactionManager>();
+#endif
 
-  txn_manager_ = std::make_unique<TransactionManager>(lock_manager_.get(), log_manager_.get());
-
+#ifndef DISABLE_LOCK_MANAGER
   lock_manager_->txn_manager_ = txn_manager_.get();
 
 #ifndef __EMSCRIPTEN__
   lock_manager_->StartDeadlockDetection();
 #endif
 
+#endif
+
+#ifndef DISABLE_CHECKPOINT_MANAGER
   // Checkpoint related.
   checkpoint_manager_ =
       std::make_unique<CheckpointManager>(txn_manager_.get(), log_manager_.get(), buffer_pool_manager_.get());
+#endif
 
   // Catalog related.
   catalog_ = std::make_unique<Catalog>(buffer_pool_manager_.get(), lock_manager_.get(), log_manager_.get());
 
+  txn_manager_->catalog_ = catalog_.get();
+
   // Execution engine related.
   execution_engine_ = std::make_unique<ExecutionEngine>(buffer_pool_manager_.get(), txn_manager_.get(), catalog_.get());
 }
@@ -90,8 +102,10 @@ BustubInstance::BustubInstance() {
   // Storage related.
   disk_manager_ = std::make_unique<DiskManagerUnlimitedMemory>();
 
+#ifndef DISABLE_CHECKPOINT_MANAGER
   // Log related.
   log_manager_ = std::make_unique<LogManager>(disk_manager_.get());
+#endif
 
   // We need more frames for GenerateTestTable to work. Therefore, we use 128 instead of the default
   // buffer pool size specified in `config.h`.
@@ -103,24 +117,32 @@ BustubInstance::BustubInstance() {
     buffer_pool_manager_ = nullptr;
   }
 
-  // Transaction (txn) related.
+#ifndef DISABLE_LOCK_MANAGER
   lock_manager_ = std::make_unique<LockManager>();
+  txn_manager_ = std::make_unique<TransactionManager>(lock_manager_.get());
+#else
+  txn_manager_ = std::make_unique<TransactionManager>();
+#endif
 
-  txn_manager_ = std::make_unique<TransactionManager>(lock_manager_.get(), log_manager_.get());
-
+#ifndef DISABLE_LOCK_MANAGER
   lock_manager_->txn_manager_ = txn_manager_.get();
 
 #ifndef __EMSCRIPTEN__
   lock_manager_->StartDeadlockDetection();
 #endif
+#endif
 
+#ifndef DISABLE_CHECKPOINT_MANAGER
   // Checkpoint related.
   checkpoint_manager_ =
       std::make_unique<CheckpointManager>(txn_manager_.get(), log_manager_.get(), buffer_pool_manager_.get());
+#endif
 
   // Catalog related.
   catalog_ = std::make_unique<Catalog>(buffer_pool_manager_.get(), lock_manager_.get(), log_manager_.get());
 
+  txn_manager_->catalog_ = catalog_.get();
+
   // Execution engine related.
   execution_engine_ = std::make_unique<ExecutionEngine>(buffer_pool_manager_.get(), txn_manager_.get(), catalog_.get());
 }
@@ -188,7 +210,7 @@ see the execution plan of your query.
 
 auto BustubInstance::ExecuteSql(const std::string &sql, ResultWriter &writer,
                                 std::shared_ptr<CheckOptions> check_options) -> bool {
-  auto txn = txn_manager_->Begin();
+  auto *txn = txn_manager_->Begin();
   try {
     auto result = ExecuteSqlTxn(sql, writer, txn, std::move(check_options));
     txn_manager_->Commit(txn);
@@ -316,7 +338,7 @@ auto BustubInstance::ExecuteSqlTxn(const std::string &sql, ResultWriter &writer,
  * create / drop table and insert for now. Should remove it in the future.
  */
 void BustubInstance::GenerateTestTable() {
-  auto txn = txn_manager_->Begin();
+  auto *txn = txn_manager_->Begin();
   auto exec_ctx = MakeExecutorContext(txn, false);
   TableGenerator gen{exec_ctx.get()};
 

src/include/binder/statement/create_statement.h

@@ -21,10 +21,11 @@ namespace bustub {
 
 class CreateStatement : public BoundStatement {
  public:
-  explicit CreateStatement(std::string table, std::vector<Column> columns);
+  explicit CreateStatement(std::string table, std::vector<Column> columns, std::vector<std::string> primary_key);
 
   std::string table_;
   std::vector<Column> columns_;
+  std::vector<std::string> primary_key_;
 
   auto ToString() const -> std::string override;
 };

src/include/catalog/catalog.h

@@ -72,13 +72,14 @@ struct IndexInfo {
    * @param key_size The size of the index key, in bytes
    */
   IndexInfo(Schema key_schema, std::string name, std::unique_ptr<Index> &&index, index_oid_t index_oid,
-            std::string table_name, size_t key_size)
+            std::string table_name, size_t key_size, bool is_primary_key)
       : key_schema_{std::move(key_schema)},
         name_{std::move(name)},
         index_{std::move(index)},
         index_oid_{index_oid},
         table_name_{std::move(table_name)},
-        key_size_{key_size} {}
+        key_size_{key_size},
+        is_primary_key_{is_primary_key} {}
   /** The schema for the index key */
   Schema key_schema_;
   /** The name of the index */
@@ -91,6 +92,8 @@ struct IndexInfo {
   std::string table_name_;
   /** The size of the index key, in bytes */
   const size_t key_size_;
+  /** Is primary key index? */
+  bool is_primary_key_;
 };
 
 /**
@@ -203,7 +206,7 @@ class Catalog {
   template <class KeyType, class ValueType, class KeyComparator>
   auto CreateIndex(Transaction *txn, const std::string &index_name, const std::string &table_name, const Schema &schema,
                    const Schema &key_schema, const std::vector<uint32_t> &key_attrs, std::size_t keysize,
-                   HashFunction<KeyType> hash_function) -> IndexInfo * {
+                   HashFunction<KeyType> hash_function, bool is_primary_key = false) -> IndexInfo * {
     // Reject the creation request for nonexistent table
     if (table_names_.find(table_name) == table_names_.end()) {
       return NULL_INDEX_INFO;
@@ -220,7 +223,7 @@ class Catalog {
     }
 
     // Construct index metdata
-    auto meta = std::make_unique<IndexMetadata>(index_name, table_name, &schema, key_attrs);
+    auto meta = std::make_unique<IndexMetadata>(index_name, table_name, &schema, key_attrs, is_primary_key);
 
     // Construct the index, take ownership of metadata
     // TODO(Kyle): We should update the API for CreateIndex
@@ -241,8 +244,8 @@ class Catalog {
     const auto index_oid = next_index_oid_.fetch_add(1);
 
     // Construct index information; IndexInfo takes ownership of the Index itself
-    auto index_info =
-        std::make_unique<IndexInfo>(key_schema, index_name, std::move(index), index_oid, table_name, keysize);
+    auto index_info = std::make_unique<IndexInfo>(key_schema, index_name, std::move(index), index_oid, table_name,
+                                                  keysize, is_primary_key);
     auto *tmp = index_info.get();
 
     // Update internal tracking

src/include/common/bustub_instance.h

@@ -263,6 +263,7 @@ class BustubInstance {
 
   std::unique_ptr<DiskManager> disk_manager_;
   std::unique_ptr<BufferPoolManager> buffer_pool_manager_;
+
   std::unique_ptr<LockManager> lock_manager_;
   std::unique_ptr<TransactionManager> txn_manager_;
   std::unique_ptr<LogManager> log_manager_;

src/include/common/config.h

@@ -39,11 +39,13 @@ static constexpr int LRUK_REPLACER_K = 10;  // lookback window for lru-k replace
 
 using frame_id_t = int32_t;    // frame id type
 using page_id_t = int32_t;     // page id type
-using txn_id_t = int32_t;      // transaction id type
+using txn_id_t = int64_t;      // transaction id type
 using lsn_t = int32_t;         // log sequence number type
 using slot_offset_t = size_t;  // slot offset type
 using oid_t = uint16_t;
 
+const txn_id_t TXN_START_ID = 1LL << 62;  // first txn id
+
 static constexpr int VARCHAR_DEFAULT_LENGTH = 128;  // default length for varchar when constructing the column
 
 }  // namespace bustub

src/include/concurrency/lock_manager.h

@@ -86,6 +86,7 @@ class LockManager {
   }
 
   ~LockManager() {
+#ifndef DISABLE_LOCK_MANAGER
     UnlockAll();
 
     enable_cycle_detection_ = false;
@@ -94,6 +95,7 @@ class LockManager {
       cycle_detection_thread_->join();
       delete cycle_detection_thread_;
     }
+#endif
   }
 
   /**

src/include/concurrency/transaction_manager.h

@@ -17,6 +17,7 @@
 #include <unordered_map>
 #include <unordered_set>
 
+#include "catalog/catalog.h"
 #include "common/config.h"
 #include "concurrency/lock_manager.h"
 #include "concurrency/transaction.h"
@@ -96,6 +97,8 @@ class TransactionManager {
   /** Resumes all transactions, used for checkpointing. */
   void ResumeTransactions();
 
+  Catalog *catalog_{nullptr};
+
  private:
   /**
    * Releases all the locks held by the given transaction.