accumulate batch when shuffle

apache · Jun 3, 2024 · 34f87f0 · 34f87f0
1 parent 26ff58d
commit 34f87f0
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 1 deletion.
diff --git a/cpp/velox/shuffle/VeloxHashBasedShuffleWriter.cc b/cpp/velox/shuffle/VeloxHashBasedShuffleWriter.cc
@@ -295,7 +295,17 @@ arrow::Status VeloxHashBasedShuffleWriter::write(std::shared_ptr<ColumnarBatch>
         numRows -= length;
       } while (numRows);
     } else {
-      RETURN_NOT_OK(partitioningAndDoSplit(std::move(rv), memLimit));
+      if (accumulateRows_ + rv->size() < maxBatchSize_) {
+        accumulateRows_ += rv->size();
+        initAccumulateDataset(rv);
+        accumulateDataset_->append(rv.get());
+      } else {
+        initAccumulateDataset(rv);
+        accumulateDataset_->append(rv.get());
+        RETURN_NOT_OK(partitioningAndDoSplit(std::move(accumulateDataset_), memLimit));
+        accumulateDataset_ = nullptr;
+        accumulateRows_ = 0;
+      }
     }
   }
   return arrow::Status::OK();
@@ -321,6 +331,10 @@ arrow::Status VeloxHashBasedShuffleWriter::partitioningAndDoSplit(facebook::velo
 }
 
 arrow::Status VeloxHashBasedShuffleWriter::stop() {
+  if (accumulateDataset_ != nullptr) {
+    RETURN_NOT_OK(partitioningAndDoSplit(std::move(accumulateDataset_), kMinMemLimit));
+    accumulateRows_ = 0;
+  }
   if (options_.partitioning != Partitioning::kSingle) {
     for (auto pid = 0; pid < numPartitions_; ++pid) {
       RETURN_NOT_OK(evictPartitionBuffers(pid, false));

diff --git a/cpp/velox/shuffle/VeloxHashBasedShuffleWriter.h b/cpp/velox/shuffle/VeloxHashBasedShuffleWriter.h
@@ -303,6 +303,18 @@ class VeloxHashBasedShuffleWriter : public VeloxShuffleWriter {
 
   arrow::Status partitioningAndDoSplit(facebook::velox::RowVectorPtr rv, int64_t memLimit);
 
+  void initAccumulateDataset(facebook::velox::RowVectorPtr& rv) {
+    if (accumulateDataset_) {
+      return;
+    }
+    std::vector<facebook::velox::VectorPtr> children;
+    for (int i = 0; i < rv->children().size(); ++i) {
+      children.push_back(nullptr);
+    }
+    accumulateDataset_ =
+        std::make_shared<facebook::velox::RowVector>(veloxPool_.get(), rv->type(), nullptr, 0, std::move(children));
+  }
+
   BinaryArrayResizeState binaryArrayResizeState_{};
 
   bool hasComplexType_ = false;

diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.h b/cpp/velox/shuffle/VeloxShuffleWriter.h
@@ -124,6 +124,10 @@ class VeloxShuffleWriter : public ShuffleWriter {
 
   int32_t maxBatchSize_{0};
 
+  uint32_t accumulateRows_{0};
+
+  facebook::velox::RowVectorPtr accumulateDataset_;
+
   enum EvictState { kEvictable, kUnevictable };
 
   // stat