Replace construct iterator callback with virtual class

src/include/storage/store/chunked_group_undo_iterator.h

@@ -13,26 +13,26 @@ class Transaction;
 namespace storage {
 class ChunkedNodeGroup;
 class NodeGroupCollection;
-class ChunkedGroupUndoIterator;
+class VersionRecordHandler;
 
-using chunked_group_undo_op_t = void (
+using version_record_handler_op_t = void (
     ChunkedNodeGroup::*)(common::row_idx_t, common::row_idx_t, common::transaction_t);
 
-using chunked_group_iterator_construct_t = std::function<std::unique_ptr<ChunkedGroupUndoIterator>(
+using version_record_handler_construct_t = std::function<std::unique_ptr<VersionRecordHandler>(
     common::row_idx_t, common::row_idx_t, common::node_group_idx_t, common::transaction_t)>;
 
 // Note: these iterators are not necessarily thread-safe when used on their own
-class ChunkedGroupUndoIterator {
+class VersionRecordHandler {
 public:
-    ChunkedGroupUndoIterator(NodeGroupCollection* nodeGroups, common::row_idx_t startRow,
+    VersionRecordHandler(NodeGroupCollection* nodeGroups, common::row_idx_t startRow,
         common::row_idx_t numRows, common::transaction_t commitTS)
         : startRow(startRow), numRows(numRows), commitTS(commitTS), nodeGroups(nodeGroups) {}
 
-    virtual ~ChunkedGroupUndoIterator() = default;
+    virtual ~VersionRecordHandler() = default;
 
     virtual void initRollbackInsert(const transaction::Transaction* /*transaction*/) {}
-    virtual void finalizeRollbackInsert() {};
-    virtual void iterate(chunked_group_undo_op_t undoFunc) = 0;
+    virtual void finalizeRollbackInsert(){};
+    virtual void applyFuncToChunkedGroups(version_record_handler_op_t func) = 0;
 
 protected:
     common::row_idx_t startRow;
@@ -42,5 +42,14 @@ class ChunkedGroupUndoIterator {
     NodeGroupCollection* nodeGroups;
 };
 
+class VersionRecordHandlerData {
+public:
+    virtual ~VersionRecordHandlerData() = default;
+
+    virtual std::unique_ptr<VersionRecordHandler> constructVersionRecordHandler(
+        common::row_idx_t startRow, common::row_idx_t numRows, common::transaction_t commitTS,
+        common::node_group_idx_t nodeGroupIdx) const = 0;
+};
+
 } // namespace storage
 } // namespace kuzu

src/include/storage/store/csr_node_group.h

@@ -165,12 +165,12 @@ static constexpr common::column_id_t REL_ID_COLUMN_ID = 1;
 struct RelTableScanState;
 class CSRNodeGroup final : public NodeGroup {
 public:
-    class PersistentIterator : public ChunkedGroupUndoIterator {
+    class PersistentIterator : public VersionRecordHandler {
     public:
         PersistentIterator(NodeGroupCollection* nodeGroups, common::node_group_idx_t nodeGroupIdx,
             common::row_idx_t startRow, common::row_idx_t numRows, common::transaction_t commitTS);
 
-        void iterate(chunked_group_undo_op_t undoFunc) override;
+        void applyFuncToChunkedGroups(version_record_handler_op_t func) override;
         void finalizeRollbackInsert() override;
 
     private:

src/include/storage/store/node_group.h

@@ -82,11 +82,11 @@ static auto NODE_GROUP_SCAN_EMMPTY_RESULT = NodeGroupScanResult{};
 struct TableScanState;
 class NodeGroup {
 public:
-    class ChunkedGroupIterator : public ChunkedGroupUndoIterator {
+    class ChunkedGroupIterator : public VersionRecordHandler {
     public:
         ChunkedGroupIterator(NodeGroupCollection* nodeGroups, common::node_group_idx_t nodeGroupIdx,
             common::row_idx_t startRow, common::row_idx_t numRows, common::transaction_t commitTS);
-        void iterate(chunked_group_undo_op_t undoFunc) override;
+        void applyFuncToChunkedGroups(version_record_handler_op_t func) override;
         void finalizeRollbackInsert() override;
 
     protected:

src/include/storage/store/node_group_collection.h

@@ -16,7 +16,7 @@ class NodeGroupCollection {
 public:
     NodeGroupCollection(MemoryManager& memoryManager, const std::vector<common::LogicalType>& types,
         bool enableCompression, FileHandle* dataFH = nullptr, common::Deserializer* deSer = nullptr,
-        const chunked_group_iterator_construct_t* iteratorConstructFunc = nullptr);
+        const VersionRecordHandlerData* versionRecordHandlerData = nullptr);
 
     void append(const transaction::Transaction* transaction,
         const std::vector<common::ValueVector*>& vectors);
@@ -51,7 +51,7 @@ class NodeGroupCollection {
     }
     NodeGroup* getOrCreateNodeGroup(transaction::Transaction* transaction,
         common::node_group_idx_t groupIdx, NodeGroupDataFormat format,
-        const chunked_group_iterator_construct_t* constructIteratorFunc_);
+        const VersionRecordHandlerData* versionRecordHandlerData);
 
     void setNodeGroup(const common::node_group_idx_t nodeGroupIdx,
         std::unique_ptr<NodeGroup> group) {
@@ -82,12 +82,12 @@ class NodeGroupCollection {
     void pushInsertInfo(const transaction::Transaction* transaction,
         common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow,
         common::row_idx_t numRows,
-        const chunked_group_iterator_construct_t* constructIteratorFunc_);
+        const VersionRecordHandlerData* overridedVersionRecordHandlerData);
 
 private:
     void pushInsertInfo(const transaction::Transaction* transaction, NodeGroup* nodeGroup,
         common::row_idx_t numRows,
-        const chunked_group_iterator_construct_t* constructIteratorOverrideFunc = nullptr);
+        const VersionRecordHandlerData* overridedVersionRecordHandlerData = nullptr);
 
     bool enableCompression;
     // Num rows in the collection regardless of deletions.
@@ -96,7 +96,7 @@ class NodeGroupCollection {
     GroupCollection<NodeGroup> nodeGroups;
     FileHandle* dataFH;
     TableStats stats;
-    const chunked_group_iterator_construct_t* iteratorConstructFunc;
+    const VersionRecordHandlerData* versionRecordHandlerData;
 };
 
 } // namespace storage

src/include/storage/store/node_table.h

@@ -21,6 +21,7 @@ class Transaction;
 } // namespace transaction
 
 namespace storage {
+class NodeTable;
 
 struct NodeTableScanState final : TableScanState {
     // Scan state for un-committed data.
@@ -95,6 +96,18 @@ struct PKColumnScanHelper {
     PrimaryKeyIndex* pkIndex;
 };
 
+class NodeTableVersionRecordHandlerData : public VersionRecordHandlerData {
+public:
+    explicit NodeTableVersionRecordHandlerData(NodeTable* nodeTable) : nodeTable(nodeTable) {}
+
+    std::unique_ptr<VersionRecordHandler> constructVersionRecordHandler(common::row_idx_t startRow,
+        common::row_idx_t numRows, common::transaction_t commitTS,
+        common::node_group_idx_t nodeGroupIdx) const override;
+
+private:
+    NodeTable* nodeTable;
+};
+
 class StorageManager;
 class NodeTable final : public Table {
 public:
@@ -201,10 +214,6 @@ class NodeTable final : public Table {
 
     TableStats getStats(const transaction::Transaction* transaction) const;
 
-    const chunked_group_iterator_construct_t& getIteratorConstructFunc() const {
-        return iteratorConstructFunc;
-    }
-
 private:
     void validatePkNotExists(const transaction::Transaction* transaction,
         common::ValueVector* pkVector);
@@ -221,7 +230,7 @@ class NodeTable final : public Table {
     std::unique_ptr<NodeGroupCollection> nodeGroups;
     common::column_id_t pkColumnID;
     std::unique_ptr<PrimaryKeyIndex> pkIndex;
-    chunked_group_iterator_construct_t iteratorConstructFunc;
+    NodeTableVersionRecordHandlerData versionRecordHandlerData;
 };
 
 } // namespace storage

src/include/storage/store/rel_table_data.h

@@ -14,12 +14,27 @@ class Transaction;
 }
 namespace storage {
 class MemoryManager;
+class RelTableData;
 
 struct CSRHeaderColumns {
     std::unique_ptr<Column> offset;
     std::unique_ptr<Column> length;
 };
 
+class RelTableVersionRecordHandlerData : public VersionRecordHandlerData {
+public:
+    RelTableVersionRecordHandlerData(RelTableData* relTableData, CSRNodeGroupScanSource source)
+        : relTableData(relTableData), source(source) {}
+
+    std::unique_ptr<VersionRecordHandler> constructVersionRecordHandler(common::row_idx_t startRow,
+        common::row_idx_t numRows, common::transaction_t commitTS,
+        common::node_group_idx_t nodeGroupIdx) const override;
+
+private:
+    RelTableData* relTableData;
+    CSRNodeGroupScanSource source;
+};
+
 class RelTableData {
 public:
     RelTableData(FileHandle* dataFH, MemoryManager* mm, ShadowFile* shadowFile,
@@ -56,7 +71,7 @@ class RelTableData {
     NodeGroup* getOrCreateNodeGroup(transaction::Transaction* transaction,
         common::node_group_idx_t nodeGroupIdx) const {
         return nodeGroups->getOrCreateNodeGroup(transaction, nodeGroupIdx, NodeGroupDataFormat::CSR,
-            &persistentIteratorConstructFunc);
+            &persistentVersionRecordHandlerData);
     }
 
     common::RelMultiplicity getMultiplicity() const { return multiplicity; }
@@ -70,6 +85,11 @@ class RelTableData {
 
     void serialize(common::Serializer& serializer) const;
 
+    std::unique_ptr<VersionRecordHandler> constructVersionRecordHandler(
+        CSRNodeGroupScanSource source, common::node_group_idx_t nodeGroupIdx,
+        common::row_idx_t startRow, common::row_idx_t numRows,
+        common::transaction_t commitTS) const;
+
 private:
     void initCSRHeaderColumns();
     void initPropertyColumns(const catalog::TableCatalogEntry* tableEntry);
@@ -98,6 +118,9 @@ class RelTableData {
         return types;
     }
 
+    const RelTableVersionRecordHandlerData* getVersionRecordHandlerData(
+        CSRNodeGroupScanSource source);
+
 private:
     FileHandle* dataFH;
     common::table_id_t tableID;
@@ -114,8 +137,8 @@ class RelTableData {
     CSRHeaderColumns csrHeaderColumns;
     std::vector<std::unique_ptr<Column>> columns;
 
-    chunked_group_iterator_construct_t inMemIteratorConstructFunc;
-    chunked_group_iterator_construct_t persistentIteratorConstructFunc;
+    RelTableVersionRecordHandlerData persistentVersionRecordHandlerData;
+    RelTableVersionRecordHandlerData inMemoryVersionRecordHandlerData;
 };
 
 } // namespace storage

src/include/storage/undo_buffer.h

@@ -88,9 +88,11 @@ class UndoBuffer {
     void createSequenceChange(catalog::SequenceCatalogEntry& sequenceEntry,
         const catalog::SequenceRollbackData& data);
     void createInsertInfo(common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow,
-        common::row_idx_t numRows, const chunked_group_iterator_construct_t* iteratorConstructFunc);
+        common::row_idx_t numRows,
+        const storage::VersionRecordHandlerData* versionRecordHandlerData);
     void createDeleteInfo(common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow,
-        common::row_idx_t numRows, const chunked_group_iterator_construct_t* iteratorConstructFunc);
+        common::row_idx_t numRows,
+        const storage::VersionRecordHandlerData* versionRecordHandlerData);
     void createVectorUpdateInfo(UpdateInfo* updateInfo, common::idx_t vectorIdx,
         VectorUpdateInfo* vectorUpdateInfo);
 
@@ -103,7 +105,8 @@ class UndoBuffer {
     uint8_t* createUndoRecord(uint64_t size);
 
     void createVersionInfo(UndoRecordType recordType, common::row_idx_t startRow,
-        common::row_idx_t numRows, const chunked_group_iterator_construct_t* iteratorConstructFunc,
+        common::row_idx_t numRows,
+        const storage::VersionRecordHandlerData* versionRecordHandlerData,
         common::node_group_idx_t nodeGroupIdx = 0);
 
     void commitRecord(UndoRecordType recordType, const uint8_t* record,

src/include/transaction/transaction.h

@@ -23,16 +23,13 @@ class VersionInfo;
 class UpdateInfo;
 struct VectorUpdateInfo;
 class ChunkedNodeGroup;
-class ChunkedGroupUndoIterator;
+class VersionRecordHandler;
+class VersionRecordHandlerData;
 } // namespace storage
 namespace transaction {
 class TransactionManager;
 class Transaction;
 
-using chunked_group_iterator_construct_t =
-    std::function<std::unique_ptr<storage::ChunkedGroupUndoIterator>(common::row_idx_t,
-        common::row_idx_t, common::node_group_idx_t, common::transaction_t commitTS)>;
-
 enum class TransactionType : uint8_t { READ_ONLY, WRITE, CHECKPOINT, DUMMY, RECOVERY };
 
 class Transaction {
@@ -123,10 +120,10 @@ class Transaction {
         const catalog::SequenceRollbackData& data) const;
     void pushInsertInfo(common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow,
         common::row_idx_t numRows,
-        const chunked_group_iterator_construct_t* constructIteratorFunc = nullptr) const;
+        const storage::VersionRecordHandlerData* versionRecordHandlerData) const;
     void pushDeleteInfo(common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow,
         common::row_idx_t numRows,
-        const chunked_group_iterator_construct_t* constructIteratorFunc) const;
+        const storage::VersionRecordHandlerData* versionRecordHandlerData) const;
     void pushVectorUpdateInfo(storage::UpdateInfo& updateInfo, common::idx_t vectorIdx,
         storage::VectorUpdateInfo& vectorUpdateInfo) const;
 

src/storage/store/csr_node_group.cpp

@@ -15,15 +15,15 @@ namespace storage {
 CSRNodeGroup::PersistentIterator::PersistentIterator(NodeGroupCollection* nodeGroups,
     common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow, common::row_idx_t numRows,
     common::transaction_t commitTS)
-    : ChunkedGroupUndoIterator(nodeGroups, startRow, numRows, commitTS), nodeGroup(nullptr) {
+    : VersionRecordHandler(nodeGroups, startRow, numRows, commitTS), nodeGroup(nullptr) {
     if (nodeGroupIdx < nodeGroups->getNumNodeGroups()) {
         nodeGroup = ku_dynamic_cast<CSRNodeGroup*>(nodeGroups->getNodeGroupNoLock(nodeGroupIdx));
     }
 }
 
-void CSRNodeGroup::PersistentIterator::iterate(chunked_group_undo_op_t undoFunc) {
+void CSRNodeGroup::PersistentIterator::applyFuncToChunkedGroups(version_record_handler_op_t func) {
     if (nodeGroup && nodeGroup->persistentChunkGroup) {
-        std::invoke(undoFunc, *nodeGroup->persistentChunkGroup, startRow, numRows, commitTS);
+        std::invoke(func, *nodeGroup->persistentChunkGroup, startRow, numRows, commitTS);
     }
 }
 

src/storage/store/node_group.cpp

@@ -24,13 +24,13 @@ namespace storage {
 NodeGroup::ChunkedGroupIterator::ChunkedGroupIterator(NodeGroupCollection* nodeGroups,
     common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow, common::row_idx_t numRows,
     transaction_t commitTS)
-    : ChunkedGroupUndoIterator(nodeGroups, startRow, numRows, commitTS),
+    : VersionRecordHandler(nodeGroups, startRow, numRows, commitTS),
       nodeGroup(nodeGroups->getNodeGroupNoLock(nodeGroupIdx)),
       numRowsToRollback(std::min(numRows, nodeGroup->getNumRows() - startRow)) {
     KU_ASSERT(startRow <= nodeGroup->getNumRows());
 }
 
-void NodeGroup::ChunkedGroupIterator::iterate(chunked_group_undo_op_t undoFunc) {
+void NodeGroup::ChunkedGroupIterator::applyFuncToChunkedGroups(version_record_handler_op_t func) {
     auto lock = nodeGroup->chunkedGroups.lock();
     const auto [chunkedGroupIdx, startRowInChunkedGroup] =
         nodeGroup->findChunkedGroupIdxFromRowIdxNoLock(startRow);
@@ -44,7 +44,7 @@ void NodeGroup::ChunkedGroupIterator::iterate(chunked_group_undo_op_t undoFunc)
             auto* chunkedGroup = nodeGroup->chunkedGroups.getGroup(lock, curChunkedGroupIdx);
             const auto numRowsForGroup =
                 std::min(numRowsLeft, chunkedGroup->getNumRows() - curStartRowIdxInChunk);
-            std::invoke(undoFunc, *chunkedGroup, curStartRowIdxInChunk, numRowsForGroup, commitTS);
+            std::invoke(func, *chunkedGroup, curStartRowIdxInChunk, numRowsForGroup, commitTS);
 
             ++curChunkedGroupIdx;
             numRowsLeft -= numRowsForGroup;

src/storage/store/node_group_collection.cpp

@@ -14,9 +14,9 @@ namespace storage {
 
 NodeGroupCollection::NodeGroupCollection(MemoryManager& memoryManager,
     const std::vector<LogicalType>& types, const bool enableCompression, FileHandle* dataFH,
-    Deserializer* deSer, const chunked_group_iterator_construct_t* iteratorConstructFunc)
+    Deserializer* deSer, const VersionRecordHandlerData* versionRecordHandlerData)
     : enableCompression{enableCompression}, numTotalRows{0}, types{LogicalType::copy(types)},
-      dataFH{dataFH}, iteratorConstructFunc(iteratorConstructFunc) {
+      dataFH{dataFH}, versionRecordHandlerData(versionRecordHandlerData) {
     if (deSer) {
         deserialize(*deSer, memoryManager);
     }
@@ -155,7 +155,7 @@ row_idx_t NodeGroupCollection::getNumTotalRows() {
 
 NodeGroup* NodeGroupCollection::getOrCreateNodeGroup(transaction::Transaction* transaction,
     node_group_idx_t groupIdx, NodeGroupDataFormat format,
-    const chunked_group_iterator_construct_t* constructIteratorFunc_) {
+    const VersionRecordHandlerData* versionRecordHandlerData) {
     const auto lock = nodeGroups.lock();
     while (groupIdx >= nodeGroups.getNumGroups(lock)) {
         const auto currentGroupIdx = nodeGroups.getNumGroups(lock);
@@ -166,7 +166,7 @@ NodeGroup* NodeGroupCollection::getOrCreateNodeGroup(transaction::Transaction* t
                                              enableCompression, LogicalType::copy(types)));
         // push an insert of size 0 so that we can rollback the creation of this node group if
         // needed
-        pushInsertInfo(transaction, nodeGroups.getLastGroup(lock), 0, constructIteratorFunc_);
+        pushInsertInfo(transaction, nodeGroups.getLastGroup(lock), 0, versionRecordHandlerData);
     }
     KU_ASSERT(groupIdx < nodeGroups.getNumGroups(lock));
     return nodeGroups.getGroup(lock, groupIdx);
@@ -213,17 +213,19 @@ void NodeGroupCollection::rollbackInsert(common::row_idx_t numRows_, bool update
 
 void NodeGroupCollection::pushInsertInfo(const transaction::Transaction* transaction,
     NodeGroup* nodeGroup, common::row_idx_t numRows,
-    const chunked_group_iterator_construct_t* constructIteratorOverrideFunc) {
+    const VersionRecordHandlerData* overridedVersionRecordHandlerData) {
     pushInsertInfo(transaction, nodeGroup->getNodeGroupIdx(), nodeGroup->getNumRows(), numRows,
-        constructIteratorOverrideFunc ? constructIteratorOverrideFunc : iteratorConstructFunc);
+        overridedVersionRecordHandlerData ? overridedVersionRecordHandlerData :
+                                            versionRecordHandlerData);
 };
 
 void NodeGroupCollection::pushInsertInfo(const transaction::Transaction* transaction,
     common::node_group_idx_t nodeGroupIdx, common::row_idx_t startRow, common::row_idx_t numRows,
-    const chunked_group_iterator_construct_t* constructIteratorFunc_) {
+    const VersionRecordHandlerData* overridedVersionRecordHandlerData) {
     // we only append to the undo buffer if the node group collection is persistent
     if (dataFH && transaction->shouldAppendToUndoBuffer()) {
-        transaction->pushInsertInfo(nodeGroupIdx, startRow, numRows, constructIteratorFunc_);
+        transaction->pushInsertInfo(nodeGroupIdx, startRow, numRows,
+            overridedVersionRecordHandlerData);
     }
 }