Merge pull request #2320 from kuzudb/refactor-scan-node-table

Refactor node table scan to reuse ScanNodeTableInfo
kuzudb · Nov 1, 2023 · 94a611c · 94a611c
2 parents 4a9fe94 + 8a14755
commit 94a611c
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 29 deletions.
diff --git a/src/include/processor/operator/scan/scan_node_table.h b/src/include/processor/operator/scan/scan_node_table.h
@@ -6,50 +6,64 @@
 namespace kuzu {
 namespace processor {
 
+struct ScanNodeTableInfo {
+    storage::NodeTable* table;
+    std::vector<common::column_id_t> columnIDs;
+
+    ScanNodeTableInfo(storage::NodeTable* table, std::vector<common::column_id_t> columnIDs)
+        : table{table}, columnIDs{std::move(columnIDs)} {}
+    ScanNodeTableInfo(const ScanNodeTableInfo& other)
+        : table{other.table}, columnIDs{other.columnIDs} {}
+
+    inline std::unique_ptr<ScanNodeTableInfo> copy() const {
+        return std::make_unique<ScanNodeTableInfo>(*this);
+    }
+};
+
 class ScanSingleNodeTable : public ScanColumns {
 public:
-    ScanSingleNodeTable(const DataPos& inVectorPos, std::vector<DataPos> outVectorsPos,
-        storage::NodeTable* table, std::vector<common::column_id_t> columnIDs,
-        std::unique_ptr<PhysicalOperator> prevOperator, uint32_t id,
-        const std::string& paramsString)
+    ScanSingleNodeTable(std::unique_ptr<ScanNodeTableInfo> info, const DataPos& inVectorPos,
+        std::vector<DataPos> outVectorsPos, std::unique_ptr<PhysicalOperator> prevOperator,
+        uint32_t id, const std::string& paramsString)
         : ScanColumns{inVectorPos, std::move(outVectorsPos), std::move(prevOperator), id,
               paramsString},
-          table{table}, columnIDs{std::move(columnIDs)} {}
+          info{std::move(info)} {}
 
     bool getNextTuplesInternal(ExecutionContext* context) override;
 
     inline std::unique_ptr<PhysicalOperator> clone() override {
-        return make_unique<ScanSingleNodeTable>(inputNodeIDVectorPos, outPropertyVectorsPos, table,
-            columnIDs, children[0]->clone(), id, paramsString);
+        return make_unique<ScanSingleNodeTable>(info->copy(), inputNodeIDVectorPos,
+            outPropertyVectorsPos, children[0]->clone(), id, paramsString);
     }
 
 private:
-    storage::NodeTable* table;
-    std::vector<common::column_id_t> columnIDs;
+    std::unique_ptr<ScanNodeTableInfo> info;
 };
 
 class ScanMultiNodeTables : public ScanColumns {
 public:
     ScanMultiNodeTables(const DataPos& inVectorPos, std::vector<DataPos> outVectorsPos,
-        std::unordered_map<common::table_id_t, storage::NodeTable*> tables,
-        std::unordered_map<common::table_id_t, std::vector<common::column_id_t>>
-            tableIDToScanColumnIds,
+        std::unordered_map<common::table_id_t, std::unique_ptr<ScanNodeTableInfo>> tables,
         std::unique_ptr<PhysicalOperator> prevOperator, uint32_t id,
         const std::string& paramsString)
         : ScanColumns{inVectorPos, std::move(outVectorsPos), std::move(prevOperator), id,
               paramsString},
-          tables{std::move(tables)}, tableIDToScanColumnIds{std::move(tableIDToScanColumnIds)} {}
+          tables{std::move(tables)} {}
 
     bool getNextTuplesInternal(ExecutionContext* context) override;
 
     inline std::unique_ptr<PhysicalOperator> clone() override {
-        return make_unique<ScanMultiNodeTables>(inputNodeIDVectorPos, outPropertyVectorsPos, tables,
-            tableIDToScanColumnIds, children[0]->clone(), id, paramsString);
+        auto clonedTables =
+            std::unordered_map<common::table_id_t, std::unique_ptr<ScanNodeTableInfo>>{};
+        for (const auto& table : tables) {
+            clonedTables.emplace(table.first, table.second->copy());
+        }
+        return make_unique<ScanMultiNodeTables>(inputNodeIDVectorPos, outPropertyVectorsPos,
+            std::move(clonedTables), children[0]->clone(), id, paramsString);
     }
 
 private:
-    std::unordered_map<common::table_id_t, storage::NodeTable*> tables;
-    std::unordered_map<common::table_id_t, std::vector<common::column_id_t>> tableIDToScanColumnIds;
+    std::unordered_map<common::table_id_t, std::unique_ptr<ScanNodeTableInfo>> tables;
 };
 
 } // namespace processor

diff --git a/src/processor/map/map_scan_node_property.cpp b/src/processor/map/map_scan_node_property.cpp
@@ -23,10 +23,8 @@ std::unique_ptr<PhysicalOperator> PlanMapper::mapScanNodeProperty(
     }
     auto tableIDs = scanProperty.getTableIDs();
     if (tableIDs.size() > 1) {
-        std::unordered_map<table_id_t, std::vector<column_id_t>> tableIDToColumns;
-        std::unordered_map<table_id_t, storage::NodeTable*> tables;
+        std::unordered_map<table_id_t, std::unique_ptr<ScanNodeTableInfo>> tables;
         for (auto& tableID : tableIDs) {
-            tables.insert({tableID, nodeStore.getNodeTable(tableID)});
             std::vector<uint32_t> columns;
             for (auto& expression : scanProperty.getProperties()) {
                 auto property = static_pointer_cast<PropertyExpression>(expression);
@@ -38,11 +36,12 @@ std::unique_ptr<PhysicalOperator> PlanMapper::mapScanNodeProperty(
                             property->getPropertyID(tableID)));
                 }
             }
-            tableIDToColumns.insert({tableID, std::move(columns)});
+            tables.insert({tableID, std::make_unique<ScanNodeTableInfo>(
+                                        nodeStore.getNodeTable(tableID), std::move(columns))});
         }
         return std::make_unique<ScanMultiNodeTables>(inputNodeIDVectorPos, std::move(outVectorsPos),
-            std::move(tables), std::move(tableIDToColumns), std::move(prevOperator),
-            getOperatorID(), scanProperty.getExpressionsForPrinting());
+            std::move(tables), std::move(prevOperator), getOperatorID(),
+            scanProperty.getExpressionsForPrinting());
     } else {
         auto tableID = tableIDs[0];
         auto tableSchema = catalog->getReadOnlyVersion()->getTableSchema(tableID);
@@ -55,9 +54,11 @@ std::unique_ptr<PhysicalOperator> PlanMapper::mapScanNodeProperty(
                 columnIDs.push_back(UINT32_MAX);
             }
         }
-        return std::make_unique<ScanSingleNodeTable>(inputNodeIDVectorPos, std::move(outVectorsPos),
-            nodeStore.getNodeTable(tableID), std::move(columnIDs), std::move(prevOperator),
-            getOperatorID(), scanProperty.getExpressionsForPrinting());
+        auto info = std::make_unique<ScanNodeTableInfo>(
+            nodeStore.getNodeTable(tableID), std::move(columnIDs));
+        return std::make_unique<ScanSingleNodeTable>(std::move(info), inputNodeIDVectorPos,
+            std::move(outVectorsPos), std::move(prevOperator), getOperatorID(),
+            scanProperty.getExpressionsForPrinting());
     }
 }
 

diff --git a/src/processor/operator/scan/scan_node_table.cpp b/src/processor/operator/scan/scan_node_table.cpp
@@ -12,7 +12,7 @@ bool ScanSingleNodeTable::getNextTuplesInternal(ExecutionContext* context) {
     for (auto& outputVector : outPropertyVectors) {
         outputVector->resetAuxiliaryBuffer();
     }
-    table->read(transaction, inputNodeIDVector, columnIDs, outPropertyVectors);
+    info->table->read(transaction, inputNodeIDVector, info->columnIDs, outPropertyVectors);
     return true;
 }
 
@@ -24,8 +24,9 @@ bool ScanMultiNodeTables::getNextTuplesInternal(ExecutionContext* context) {
         inputNodeIDVector
             ->getValue<nodeID_t>(inputNodeIDVector->state->selVector->selectedPositions[0])
             .tableID;
-    tables.at(tableID)->read(
-        transaction, inputNodeIDVector, tableIDToScanColumnIds.at(tableID), outPropertyVectors);
+    auto tableScanInfo = tables.at(tableID).get();
+    tableScanInfo->table->read(
+        transaction, inputNodeIDVector, tableScanInfo->columnIDs, outPropertyVectors);
     return true;
 }