update

kuzudb · Sep 5, 2023 · 6ed048c · 6ed048c
1 parent cf7134f
commit 6ed048c
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 21 deletions.
diff --git a/src/include/processor/operator/persistent/reader.h b/src/include/processor/operator/persistent/reader.h
@@ -50,6 +50,7 @@ class Reader : public PhysicalOperator {
 private:
     void getNextNodeGroupInSerial();
     void getNextNodeGroupInParallel();
+    void readNextNodeGroupInParallel();
 
 private:
     std::unique_ptr<ReaderInfo> info;

diff --git a/src/include/storage/copier/reader_state.h b/src/include/storage/copier/reader_state.h
@@ -163,6 +163,7 @@ class ReaderSharedState {
 
 private:
     std::unique_ptr<ReaderMorsel> getMorselOfNextBlock();
+    void readNextBlock(common::DataChunk* dataChunk);
 
 public:
     std::mutex mtx;

diff --git a/src/processor/operator/persistent/reader.cpp b/src/processor/operator/persistent/reader.cpp
@@ -42,10 +42,21 @@ void Reader::getNextNodeGroupInSerial() {
 }
 
 void Reader::getNextNodeGroupInParallel() {
-    while (leftArrowArrays.getLeftNumRows() < DEFAULT_VECTOR_CAPACITY) {
+    readNextNodeGroupInParallel();
+    if (leftArrowArrays.getLeftNumRows() == 0) {
+        dataChunk->state->selVector->selectedSize = 0;
+    } else {
+        int64_t numRowsToReturn =
+            std::min(leftArrowArrays.getLeftNumRows(), DEFAULT_VECTOR_CAPACITY);
+        leftArrowArrays.appendToDataChunk(dataChunk.get(), numRowsToReturn);
+    }
+}
+
+void Reader::readNextNodeGroupInParallel() {
+    if (leftArrowArrays.getLeftNumRows() == 0) {
         auto morsel = sharedState->getParallelMorsel();
         if (morsel->fileIdx == INVALID_VECTOR_IDX) {
-            break;
+            return;
         }
         if (!readFuncData || morsel->fileIdx != readFuncData->fileIdx) {
             readFuncData = initFunc(sharedState->copyDescription->filePaths, morsel->fileIdx,
@@ -54,13 +65,6 @@ void Reader::getNextNodeGroupInParallel() {
         readFunc(*readFuncData, morsel->blockIdx, dataChunk.get());
         leftArrowArrays.appendFromDataChunk(dataChunk.get());
     }
-    if (leftArrowArrays.getLeftNumRows() == 0) {
-        dataChunk->state->selVector->selectedSize = 0;
-    } else {
-        int64_t numRowsToReturn =
-            std::min(leftArrowArrays.getLeftNumRows(), DEFAULT_VECTOR_CAPACITY);
-        leftArrowArrays.appendToDataChunk(dataChunk.get(), numRowsToReturn);
-    }
 }
 
 } // namespace processor

diff --git a/src/storage/copier/reader_state.cpp b/src/storage/copier/reader_state.cpp
@@ -251,11 +251,25 @@ void ReaderSharedState::countBlocks() {
 
 std::unique_ptr<ReaderMorsel> ReaderSharedState::getSerialMorsel(DataChunk* dataChunk) {
     std::unique_lock xLck{mtx};
-    while (leftArrowArrays.getLeftNumRows() < DEFAULT_VECTOR_CAPACITY) {
+    readNextBlock(dataChunk);
+    if (leftArrowArrays.getLeftNumRows() == 0) {
+        dataChunk->state->selVector->selectedSize = 0;
+        return std::make_unique<ReaderMorsel>();
+    } else {
+        auto numRowsToReturn = std::min(leftArrowArrays.getLeftNumRows(), DEFAULT_VECTOR_CAPACITY);
+        leftArrowArrays.appendToDataChunk(dataChunk, numRowsToReturn);
+        auto result = std::make_unique<ReaderMorsel>(currFileIdx, currBlockIdx, currRowIdx);
+        currRowIdx += numRowsToReturn;
+        return result;
+    }
+}
+
+void ReaderSharedState::readNextBlock(common::DataChunk* dataChunk) {
+    if (leftArrowArrays.getLeftNumRows() == 0) {
         auto morsel = getMorselOfNextBlock();
         if (morsel->fileIdx >= copyDescription->filePaths.size()) {
             // No more blocks.
-            break;
+            return;
         }
         if (morsel->fileIdx != readFuncData->fileIdx) {
             readFuncData = initFunc(copyDescription->filePaths, morsel->fileIdx,
@@ -264,16 +278,6 @@ std::unique_ptr<ReaderMorsel> ReaderSharedState::getSerialMorsel(DataChunk* data
         readFunc(*readFuncData, morsel->blockIdx, dataChunk);
         leftArrowArrays.appendFromDataChunk(dataChunk);
     }
-    if (leftArrowArrays.getLeftNumRows() == 0) {
-        dataChunk->state->selVector->selectedSize = 0;
-        return std::make_unique<ReaderMorsel>();
-    } else {
-        auto numRowsToReturn = std::min(leftArrowArrays.getLeftNumRows(), DEFAULT_VECTOR_CAPACITY);
-        leftArrowArrays.appendToDataChunk(dataChunk, numRowsToReturn);
-        auto result = std::make_unique<ReaderMorsel>(currFileIdx, currBlockIdx, currRowIdx);
-        currRowIdx += numRowsToReturn;
-        return result;
-    }
 }
 
 std::unique_ptr<ReaderMorsel> ReaderSharedState::getParallelMorsel() {