Add abstraction of LocalTableData to LocalStorage

src/include/storage/local_storage/local_storage.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <unordered_map>
+
+#include "storage/local_storage/local_table.h"
+
+namespace kuzu {
+namespace storage {
+
+class Column;
+class MemoryManager;
+
+// Data structures in LocalStorage are not thread-safe.
+// For now, we only support single thread insertions and updates. Once we optimize them with
+// multiple threads, LocalStorage and its related data structures should be reworked to be
+// thread-safe.
+class LocalStorage {
+public:
+    LocalStorage(storage::MemoryManager* mm);
+
+    // This function will create the local table data if not exists.
+    LocalTableData* getOrCreateLocalTableData(
+        common::table_id_t tableID, const std::vector<std::unique_ptr<Column>>& columns);
+    // This function will return nullptr if the local table data does not exist.
+    LocalTableData* getLocalTableData(common::table_id_t tableID);
+    std::unordered_set<common::table_id_t> getTableIDsWithUpdates();
+
+private:
+    std::unordered_map<common::table_id_t, std::unique_ptr<LocalTable>> tables;
+    storage::MemoryManager* mm;
+};
+
+} // namespace storage
+} // namespace kuzu

src/include/storage/local_table.h → src/include/storage/local_storage/local_table.h

@@ -2,6 +2,7 @@
 
 #include <map>
 
+#include "common/enums/table_type.h"
 #include "common/vector/value_vector.h"
 
 namespace kuzu {
@@ -79,15 +80,7 @@ class LocalNodeGroup {
     friend class NodeTableData;
 
 public:
-    LocalNodeGroup(
-        const std::vector<std::unique_ptr<common::LogicalType>>& dataTypes, MemoryManager* mm) {
-        columns.resize(dataTypes.size());
-        for (auto i = 0u; i < dataTypes.size(); ++i) {
-            // To avoid unnecessary memory consumption, we chunk local changes of each column in the
-            // node group into chunks of size DEFAULT_VECTOR_CAPACITY.
-            columns[i] = std::make_unique<LocalVectorCollection>(dataTypes[i].get(), mm);
-        }
-    }
+    LocalNodeGroup(std::vector<common::LogicalType*> dataTypes, MemoryManager* mm);
 
     void scan(common::ValueVector* nodeIDVector, const std::vector<common::column_id_t>& columnIDs,
         const std::vector<common::ValueVector*>& outputVectors);
@@ -113,13 +106,12 @@ class LocalNodeGroup {
     std::vector<std::unique_ptr<LocalVectorCollection>> columns;
 };
 
-class LocalTable {
+class LocalTableData {
     friend class NodeTableData;
 
 public:
-    explicit LocalTable(common::table_id_t tableID,
-        std::vector<std::unique_ptr<common::LogicalType>> dataTypes, MemoryManager* mm)
-        : tableID{tableID}, dataTypes{std::move(dataTypes)}, mm{mm} {};
+    LocalTableData(std::vector<common::LogicalType*> dataTypes, MemoryManager* mm)
+        : dataTypes{std::move(dataTypes)}, mm{mm} {}
 
     void scan(common::ValueVector* nodeIDVector, const std::vector<common::column_id_t>& columnIDs,
         const std::vector<common::ValueVector*>& outputVectors);
@@ -139,11 +131,26 @@ class LocalTable {
     common::node_group_idx_t initializeLocalNodeGroup(common::offset_t nodeOffset);
 
 private:
-    common::table_id_t tableID;
-    std::vector<std::unique_ptr<common::LogicalType>> dataTypes;
+    std::vector<common::LogicalType*> dataTypes;
     MemoryManager* mm;
     std::unordered_map<common::node_group_idx_t, std::unique_ptr<LocalNodeGroup>> nodeGroups;
 };
 
+class Column;
+class LocalTable {
+public:
+    LocalTable(common::table_id_t tableID, common::TableType tableType)
+        : tableID{tableID}, tableType{tableType} {};
+
+    LocalTableData* getOrCreateLocalTableData(
+        const std::vector<std::unique_ptr<Column>>& columns, MemoryManager* mm);
+    inline LocalTableData* getLocalTableData() { return localTableData.get(); }
+
+private:
+    common::table_id_t tableID;
+    common::TableType tableType;
+    std::unique_ptr<LocalTableData> localTableData;
+};
+
 } // namespace storage
 } // namespace kuzu

src/include/storage/store/column.h

@@ -25,6 +25,7 @@ using batch_lookup_func_t = read_values_to_page_func_t;
 
 class NullColumn;
 class StructColumn;
+class LocalVectorCollection;
 class Column {
     friend class LocalColumn;
     friend class StringLocalColumn;

src/include/storage/store/node_table.h

@@ -67,8 +67,8 @@ class NodeTable : public Table {
         common::ValueVector* defaultValueVector) final;
     inline void dropColumn(common::column_id_t columnID) final { tableData->dropColumn(columnID); }
 
-    void prepareCommit(LocalTable* localTable) final;
-    void prepareRollback(LocalTable* localTable) final;
+    void prepareCommit(LocalTableData* localTable) final;
+    void prepareRollback(LocalTableData* localTable) final;
     void checkpointInMemory() final;
     void rollbackInMemory() final;
 

src/include/storage/store/node_table_data.h

@@ -5,6 +5,8 @@
 namespace kuzu {
 namespace storage {
 
+class LocalTableData;
+
 class NodeTableData : public TableData {
 public:
     NodeTableData(BMFileHandle* dataFH, BMFileHandle* metadataFH, common::table_id_t tableID,
@@ -33,7 +35,7 @@ class NodeTableData : public TableData {
 
     void append(NodeGroup* nodeGroup) final;
 
-    void prepareLocalTableToCommit(LocalTable* localTable);
+    void prepareLocalTableToCommit(LocalTableData* localTable);
 };
 
 } // namespace storage

src/include/storage/store/rel_table.h

@@ -44,8 +44,8 @@ class RelTable : public Table {
                                                      bwdRelTableData->append(nodeGroup);
     }
 
-    void prepareCommit(LocalTable* localTable) final;
-    void prepareRollback(LocalTable* localTable) final;
+    void prepareCommit(LocalTableData* localTable) final;
+    void prepareRollback(LocalTableData* localTable) final;
     void checkpointInMemory() final;
     void rollbackInMemory() final;
 

src/include/storage/store/rel_table_data.h

@@ -78,7 +78,7 @@ class RelTableData : public TableData {
                    common::ColumnDataFormat::CSR;
     }
 
-    void prepareLocalTableToCommit(LocalTable* localTable);
+    void prepareLocalTableToCommit(LocalTableData* localTable);
 
 private:
     std::unique_ptr<Column> adjColumn;

src/include/storage/store/table.h

@@ -31,8 +31,8 @@ class Table {
         common::ValueVector* defaultValueVector) = 0;
     virtual void dropColumn(common::column_id_t columnID) = 0;
 
-    virtual void prepareCommit(LocalTable* localTable) = 0;
-    virtual void prepareRollback(LocalTable* localTable) = 0;
+    virtual void prepareCommit(LocalTableData* localTable) = 0;
+    virtual void prepareRollback(LocalTableData* localTable) = 0;
     virtual void checkpointInMemory() = 0;
     virtual void rollbackInMemory() = 0;
 

src/include/storage/store/table_data.h

@@ -12,6 +12,7 @@ struct TableReadState {
     std::vector<common::column_id_t> columnIDs;
 };
 
+class LocalTableData;
 class TableData {
 public:
     virtual ~TableData() = default;
@@ -43,7 +44,7 @@ class TableData {
         return columns[0]->getNumNodeGroups(transaction);
     }
 
-    virtual void prepareLocalTableToCommit(LocalTable* localTable) = 0;
+    virtual void prepareLocalTableToCommit(LocalTableData* localTable) = 0;
     virtual void checkpointInMemory();
     virtual void rollbackInMemory();
 

src/include/transaction/transaction.h

@@ -2,11 +2,14 @@
 
 #include <memory>
 
-#include "storage/local_storage.h"
+#include "storage/local_storage/local_storage.h"
 
 namespace kuzu {
+namespace storage {
+class LocalStorage;
+class MemoryManager;
+} // namespace storage
 namespace transaction {
-
 class TransactionManager;
 
 enum class TransactionType : uint8_t { READ_ONLY, WRITE };

src/storage/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_subdirectory(buffer_manager)
 add_subdirectory(compression)
 add_subdirectory(index)
+add_subdirectory(local_storage)
 add_subdirectory(stats)
 add_subdirectory(storage_structure)
 add_subdirectory(store)
@@ -9,8 +10,6 @@ add_subdirectory(wal)
 add_library(kuzu_storage
         OBJECT
         file_handle.cpp
-        local_storage.cpp
-        local_table.cpp
         storage_info.cpp
         storage_manager.cpp
         storage_utils.cpp

src/storage/local_storage/CMakeLists.txt

@@ -0,0 +1,8 @@
+add_library(kuzu_storage_local_storage
+        OBJECT
+        local_table.cpp
+        local_storage.cpp)
+
+set(ALL_OBJECT_FILES
+        ${ALL_OBJECT_FILES} $<TARGET_OBJECTS:kuzu_storage_local_storage>
+        PARENT_SCOPE)

src/storage/local_storage/local_storage.cpp

@@ -0,0 +1,38 @@
+#include "storage/local_storage/local_storage.h"
+
+#include "storage/local_storage/local_table.h"
+#include "storage/store/column.h"
+
+using namespace kuzu::common;
+using namespace kuzu::transaction;
+
+namespace kuzu {
+namespace storage {
+
+LocalStorage::LocalStorage(MemoryManager* mm) : mm{mm} {}
+
+LocalTableData* LocalStorage::getOrCreateLocalTableData(
+    common::table_id_t tableID, const std::vector<std::unique_ptr<Column>>& columns) {
+    if (!tables.contains(tableID)) {
+        tables.emplace(tableID, std::make_unique<LocalTable>(tableID, TableType::NODE));
+    }
+    return tables.at(tableID)->getOrCreateLocalTableData(columns, mm);
+}
+
+LocalTableData* LocalStorage::getLocalTableData(common::table_id_t tableID) {
+    if (!tables.contains(tableID)) {
+        return nullptr;
+    }
+    return tables.at(tableID)->getLocalTableData();
+}
+
+std::unordered_set<table_id_t> LocalStorage::getTableIDsWithUpdates() {
+    std::unordered_set<table_id_t> tableSetToUpdate;
+    for (auto& [tableID, _] : tables) {
+        tableSetToUpdate.insert(tableID);
+    }
+    return tableSetToUpdate;
+}
+
+} // namespace storage
+} // namespace kuzu