update metrics

apache · Jul 24, 2024 · 4655d31 · 4655d31
1 parent d892f09
commit 4655d31
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/cpp/velox/shuffle/VeloxSortShuffleWriter.cc b/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
@@ -235,6 +235,13 @@ arrow::Status VeloxSortShuffleWriter::evictAllPartitions() {
 }
 
 arrow::Status VeloxSortShuffleWriter::evictPartition(uint32_t partitionId, size_t begin, size_t end) {
+  auto payload = prepareToEvict(begin, end);
+  RETURN_NOT_OK(partitionWriter_->evict(partitionId, std::move(payload), Evict::type::kSortSpill, false, false, stopped_));
+  return arrow::Status::OK();
+}
+
+std::unique_ptr<InMemoryPayload> VeloxSortShuffleWriter::prepareToEvict(size_t begin, size_t end) {
+  ScopedTimer timer(&sortTime_);
   // Serialize [begin, end)
   uint32_t numRows = end - begin;
   uint64_t rawSize = numRows * sizeof(RowSizeType);
@@ -264,10 +271,7 @@ arrow::Status VeloxSortShuffleWriter::evictPartition(uint32_t partitionId, size_
   std::vector<std::shared_ptr<arrow::Buffer>> buffers;
   buffers.push_back(std::make_shared<arrow::Buffer>(rawData, rawSize));
 
-  auto payload = std::make_unique<InMemoryPayload>(numRows, nullptr, std::move(buffers));
-  RETURN_NOT_OK(
-      partitionWriter_->evict(partitionId, std::move(payload), Evict::type::kSortSpill, false, false, stopped_));
-  return arrow::Status::OK();
+  return std::make_unique<InMemoryPayload>(numRows, nullptr, std::move(buffers));
 }
 
 uint32_t VeloxSortShuffleWriter::maxRowsToInsert(uint32_t offset, uint32_t rows) {

diff --git a/cpp/velox/shuffle/VeloxSortShuffleWriter.h b/cpp/velox/shuffle/VeloxSortShuffleWriter.h
@@ -77,6 +77,8 @@ class VeloxSortShuffleWriter final : public VeloxShuffleWriter {
 
   arrow::Status evictPartition(uint32_t partitionId, size_t begin, size_t end);
 
+  std::unique_ptr<InMemoryPayload> prepareToEvict(size_t begin, size_t end);
+
   uint32_t maxRowsToInsert(uint32_t offset, uint32_t rows);
 
   void acquireNewBuffer(int64_t memLimit, uint64_t minSizeRequired);