fix in stop step and format

apache · Jun 3, 2024 · f669b6a · f669b6a
1 parent 82708df
commit f669b6a
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git a/cpp/velox/shuffle/VeloxHashBasedShuffleWriter.cc b/cpp/velox/shuffle/VeloxHashBasedShuffleWriter.cc
@@ -294,19 +294,19 @@ arrow::Status VeloxHashBasedShuffleWriter::write(std::shared_ptr<ColumnarBatch>
         offset += length;
         numRows -= length;
       } while (numRows);
-    } else if (accumulateBatches_ + rv->size() < maxBatchSize_) {
+    } else if (accumulateRows_ + rv->size() < maxBatchSize_) {
+      accumulateRows_ += rv->size();
       accumulateDataset_.emplace_back(rv);
-      accumulateBatches_ += rv->size();
     } else {
       if (accumulateDataset_.empty()) {
         RETURN_NOT_OK(partitioningAndDoSplit(std::move(rv), memLimit));
       } else {
-        for(auto v : accumulateDataset_) {
+        for (auto v : accumulateDataset_) {
           rv->append(v.get());
         }
-        RETURN_NOT_OK(partitioningAndDoSplit(std::move(rv), memLimit));
         accumulateDataset_.clear();
-        accumulateBatches_ = 0;
+        accumulateRows_ = 0;
+        RETURN_NOT_OK(partitioningAndDoSplit(std::move(rv), memLimit));
       }
     }
   }
@@ -333,6 +333,15 @@ arrow::Status VeloxHashBasedShuffleWriter::partitioningAndDoSplit(facebook::velo
 }
 
 arrow::Status VeloxHashBasedShuffleWriter::stop() {
+  if (!accumulateDataset_.empty()) {
+    auto rv = accumulateDataset_[0];
+    for (int i = 1; i < accumulateDataset_.size(); ++i) {
+      rv->append(accumulateDataset_[i].get());
+    }
+    accumulateDataset_.clear();
+    accumulateRows_ = 0;
+    RETURN_NOT_OK(partitioningAndDoSplit(std::move(rv), kMinMemLimit));
+  }
   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
@@ -249,7 +249,6 @@ class VeloxHashBasedShuffleWriter : public VeloxShuffleWriter {
 
   template <typename T>
   arrow::Status splitFixedType(const uint8_t* srcAddr, const std::vector<uint8_t*>& dstAddrs) {
-    ScopedTimer timer(&splitFixedTypeTimer_);
     for (auto& pid : partitionUsed_) {
       auto dstPidBase = (T*)(dstAddrs[pid] + partitionBufferBase_[pid] * sizeof(T));
       auto pos = partition2RowOffsetBase_[pid];
@@ -402,8 +401,6 @@ class VeloxHashBasedShuffleWriter : public VeloxShuffleWriter {
   facebook::velox::serializer::presto::PrestoVectorSerde serde_;
 
   SplitState splitState_{kInit};
-
-  int64_t splitFixedTypeTimer_{0};
 }; // class VeloxHashBasedShuffleWriter
 
 } // namespace gluten
diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.h b/cpp/velox/shuffle/VeloxShuffleWriter.h
@@ -124,7 +124,7 @@ class VeloxShuffleWriter : public ShuffleWriter {
 
   int32_t maxBatchSize_{0};
 
-  int32_t accumulateBatches_{0};
+  uint32_t accumulateRows_{0};
 
   std::vector<facebook::velox::RowVectorPtr> accumulateDataset_;