WIP: [feature] Allow parallel execution of SparkShuffleWriter

bolt/common/time/Timer.h

@@ -49,7 +49,7 @@ class MicrosecondTimer {
   }
 
   ~MicrosecondTimer() {
-    if (FOLLY_LIKELY(timer_)) {
+    if (FOLLY_LIKELY(timer_ != nullptr)) {
       auto duration = std::chrono::duration_cast<std::chrono::microseconds>(
           std::chrono::steady_clock::now() - start_);
 
@@ -72,7 +72,7 @@ class NanosecondTimer {
   }
 
   ~NanosecondTimer() {
-    if (FOLLY_LIKELY(timer_)) {
+    if (FOLLY_LIKELY(timer_ != nullptr)) {
       auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(
           std::chrono::steady_clock::now() - start_);
 

bolt/exec/Driver.cpp

@@ -1219,6 +1219,8 @@ std::string blockingReasonToString(BlockingReason reason) {
       return "kWaitForScanScaleUp";
     case BlockingReason::kWaitForIndexLookup:
       return "kWaitForIndexLookup";
+    case BlockingReason::kWaitForShuffle:
+      return "kWaitForShuffle";
     default:
       BOLT_UNREACHABLE(
           fmt::format("Unknown blocking reason {}", static_cast<int>(reason)));

bolt/exec/Driver.h

@@ -220,6 +220,9 @@ enum class BlockingReason {
   /// Used by IndexLookupJoin operator, indicating that it was blocked by the
   /// async index lookup.
   kWaitForIndexLookup,
+  /// Used by SparkShuffleWriter to wait for peer SparkShuffleWriter operators
+  /// in the same pipeline to finish finalization.
+  kWaitForShuffle,
 };
 
 std::string blockingReasonToString(BlockingReason reason);

bolt/exec/LocalPlanner.cpp

@@ -376,11 +376,6 @@ uint32_t maxDrivers(
     // single-threaded for now. This assumption might not hold in the future.
     else if (node->name() == "SparkShuffleReader") {
       return 1;
-    }
-    // multi-threaded spark: SparkShuffleWriter is designed to be
-    // single-threaded for now. This assumption might not hold in the future.
-    else if (node->name() == "SparkShuffleWriter") {
-      return 1;
     } else {
       auto result = Operator::maxDrivers(node);
       if (result) {

bolt/shuffle/sparksql/BoltRowBasedSortShuffleWriter.cpp

@@ -209,11 +209,24 @@ arrow::Status BoltRowBasedSortShuffleWriter::reclaimFixedSize(
 }
 
 arrow::Status BoltRowBasedSortShuffleWriter::stop() {
+  return stopInternal(true);
+}
+
+arrow::Status BoltRowBasedSortShuffleWriter::localStop() {
+  return stopInternal(false);
+}
+
+arrow::Status BoltRowBasedSortShuffleWriter::stopInternal(
+    bool stopPartitionWriter) {
   bytedance::bolt::NanosecondTimer stopTimer(&stopTime_);
   setSplitState(SplitState::kStop);
   RETURN_NOT_OK(tryEvict());
   {
-    RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+    if (stopPartitionWriter) {
+      RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+    } else {
+      RETURN_NOT_OK(partitionWriter_->populateMetrics(&metrics_));
+    }
     metrics_.useRowBased = 1;
     combinedVectorNumber_ = combineVectorTimes_ > 0
         ? (combinedVectorNumber_ / combineVectorTimes_)

bolt/shuffle/sparksql/BoltRowBasedSortShuffleWriter.h

@@ -43,6 +43,8 @@ class BoltRowBasedSortShuffleWriter final : public BoltShuffleWriter {
 
   arrow::Status stop() override;
 
+  arrow::Status localStop() override;
+
   arrow::Status reclaimFixedSize(int64_t size, int64_t* actual) override;
 
   BoltRowBasedSortShuffleWriter(
@@ -52,6 +54,8 @@ class BoltRowBasedSortShuffleWriter final : public BoltShuffleWriter {
       : BoltShuffleWriter(std::move(options), boltPool, pool) {}
 
  private:
+  arrow::Status stopInternal(bool stopPartitionWriter);
+
   arrow::Status init() override;
 
   arrow::Status initFromRowVector(

bolt/shuffle/sparksql/BoltShuffleWriter.cpp

@@ -684,6 +684,14 @@ arrow::Status BoltShuffleWriter::split(
 }
 
 arrow::Status BoltShuffleWriter::stop() {
+  return stopInternal(true);
+}
+
+arrow::Status BoltShuffleWriter::localStop() {
+  return stopInternal(false);
+}
+
+arrow::Status BoltShuffleWriter::stopInternal(bool stopPartitionWriter) {
   bytedance::bolt::NanosecondTimer stopTimer(&stopTime_);
   if (vectorLayout_ != RowVectorLayout::kComposite) {
     partitionWriter_->setRowFormat(false);
@@ -720,9 +728,15 @@ arrow::Status BoltShuffleWriter::stop() {
     {
       SCOPED_TIMER(cpuWallTimingList_[CpuWallTimingStop]);
       setSplitState(SplitState::kStop);
-      RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      if (stopPartitionWriter) {
+        RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      } else {
+        RETURN_NOT_OK(partitionWriter_->populateMetrics(&metrics_));
+      }
       metrics_.avgPreallocSize =
           (preallocCount_ == 0) ? 0 : totalPreallocSize_ / preallocCount_;
+      metrics_.totalPreallocSize = totalPreallocSize_;
+      metrics_.totalPreallocCount = preallocCount_;
       partitionBuffers_.clear();
     }
 
@@ -732,7 +746,11 @@ arrow::Status BoltShuffleWriter::stop() {
     RETURN_NOT_OK(tryEvict());
     {
       SCOPED_TIMER(cpuWallTimingList_[CpuWallTimingStop]);
-      RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      if (stopPartitionWriter) {
+        RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      } else {
+        RETURN_NOT_OK(partitionWriter_->populateMetrics(&metrics_));
+      }
     }
   }
   metrics_.useV2 = 0;

bolt/shuffle/sparksql/BoltShuffleWriter.h

@@ -189,6 +189,9 @@ class BoltShuffleWriter : public ShuffleWriter {
 
   arrow::Status stop() override;
 
+  // stop partition writer in local thread
+  arrow::Status localStop() override;
+
   arrow::Status reclaimFixedSize(int64_t size, int64_t* actual) override;
 
   const uint64_t cachedPayloadSize() const override;
@@ -512,6 +515,8 @@ class BoltShuffleWriter : public ShuffleWriter {
   // for CompositeRowVector
   virtual arrow::Status tryEvict(
       int64_t memLimit = std::numeric_limits<int64_t>::max());
+
+  arrow::Status stopInternal(bool stopPartitionWriter);
   arrow::Status initFromRowVectorForComposite(
       const bytedance::bolt::RowVector& rv);
   virtual arrow::Status splitCompositeVector(

bolt/shuffle/sparksql/BoltShuffleWriterV2.cpp

@@ -262,6 +262,14 @@ arrow::Status BoltShuffleWriterV2::initPartitions() {
 }
 
 arrow::Status BoltShuffleWriterV2::stop() {
+  return stopInternal(true);
+}
+
+arrow::Status BoltShuffleWriterV2::localStop() {
+  return stopInternal(false);
+}
+
+arrow::Status BoltShuffleWriterV2::stopInternal(bool stopPartitionWriter) {
   bytedance::bolt::NanosecondTimer stopTimer(&stopTime_);
   if (vectorLayout_ == RowVectorLayout::kColumnar) {
     partitionWriter_->setRowFormat(false);
@@ -275,7 +283,11 @@ arrow::Status BoltShuffleWriterV2::stop() {
     {
       SCOPED_TIMER(cpuWallTimingList_[CpuWallTimingStop]);
       setSplitState(SplitState::kStop);
-      RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      if (stopPartitionWriter) {
+        RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      } else {
+        RETURN_NOT_OK(partitionWriter_->populateMetrics(&metrics_));
+      }
       metrics_.rowVectorModeCompress = rowVectorModeCompress_;
       releaseBufferPoolMemory();
     }
@@ -286,7 +298,11 @@ arrow::Status BoltShuffleWriterV2::stop() {
     RETURN_NOT_OK(tryEvict());
     {
       SCOPED_TIMER(cpuWallTimingList_[CpuWallTimingStop]);
-      RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      if (stopPartitionWriter) {
+        RETURN_NOT_OK(partitionWriter_->stop(&metrics_));
+      } else {
+        RETURN_NOT_OK(partitionWriter_->populateMetrics(&metrics_));
+      }
     }
   }
   metrics_.useV2 = 1;

bolt/shuffle/sparksql/BoltShuffleWriterV2.h

@@ -206,6 +206,8 @@ class BoltShuffleWriterV2 final : public BoltShuffleWriter {
 
   arrow::Status stop() override;
 
+  arrow::Status localStop() override;
+
   arrow::Status reclaimFixedSize(int64_t size, int64_t* actual) override;
 
   BoltShuffleWriterV2(
@@ -219,6 +221,8 @@ class BoltShuffleWriterV2 final : public BoltShuffleWriter {
   }
 
  private:
+  arrow::Status stopInternal(bool stopPartitionWriter);
+
   void checkLengthBuffer(uint32_t col, uint32_t pid);
 
   void checkNullValue(int32_t col, uint32_t pid, uint32_t numRows);

bolt/shuffle/sparksql/Options.h

@@ -132,6 +132,9 @@ struct PartitionWriterOptions {
   // for LocalPartitionWriter
   std::string dataFile;
   std::vector<std::string> configuredDirs;
+  // For parallel SparkShuffleWriter + LocalPartitionWriter execution, partial
+  // files use a per-driver suffix when part >= 0.
+  int32_t part = -1;
 
   // for CelebornPartitionWriter
   std::shared_ptr<RssClient> rssClient;
@@ -168,6 +171,8 @@ struct ShuffleWriterMetrics {
   int64_t totalEvictTime{0};
   int64_t totalCompressTime{0};
   int64_t maxPartitionBufferSize{0};
+  int64_t totalPreallocSize{0};
+  int64_t totalPreallocCount{0};
   int64_t avgPreallocSize{0};
   int64_t useV2{0};
   int64_t rowVectorModeCompress{0};
@@ -184,6 +189,57 @@ struct ShuffleWriterMetrics {
   int64_t dataSize{0};
   std::vector<int64_t> partitionLengths{};
   std::vector<int64_t> rawPartitionLengths{}; // Uncompressed size.
+  std::string dataFile{};
+
+  ShuffleWriterMetrics& operator+=(const ShuffleWriterMetrics& rhs) {
+    totalInputRowNumber += rhs.totalInputRowNumber;
+    totalInputBatches += rhs.totalInputBatches;
+    totalBytesWritten += rhs.totalBytesWritten;
+    totalBytesEvicted += rhs.totalBytesEvicted;
+    totalWriteTime += rhs.totalWriteTime;
+    totalEvictTime += rhs.totalEvictTime;
+    totalCompressTime += rhs.totalCompressTime;
+    maxPartitionBufferSize =
+        std::max(maxPartitionBufferSize, rhs.maxPartitionBufferSize);
+    totalPreallocSize += rhs.totalPreallocSize;
+    totalPreallocCount += rhs.totalPreallocCount;
+    avgPreallocSize =
+        totalPreallocCount == 0 ? 0 : totalPreallocSize / totalPreallocCount;
+    useV2 = rhs.useV2;
+    rowVectorModeCompress += rhs.rowVectorModeCompress;
+    combinedVectorNumber += rhs.combinedVectorNumber;
+    combineVectorTimes += rhs.combineVectorTimes;
+    combineVectorCost += rhs.combineVectorCost;
+    useRowBased += rhs.useRowBased;
+    splitTime += rhs.splitTime;
+    convertTime += rhs.convertTime;
+    flattenTime += rhs.flattenTime;
+    computePidTime += rhs.computePidTime;
+    shuffleWriteTime += rhs.shuffleWriteTime;
+    dataSize += rhs.dataSize;
+
+    if (!rhs.partitionLengths.empty()) {
+      if (partitionLengths.empty()) {
+        partitionLengths.resize(rhs.partitionLengths.size(), 0);
+      }
+      BOLT_CHECK_EQ(partitionLengths.size(), rhs.partitionLengths.size());
+      for (size_t i = 0; i < partitionLengths.size(); ++i) {
+        partitionLengths[i] += rhs.partitionLengths[i];
+      }
+    }
+
+    if (!rhs.rawPartitionLengths.empty()) {
+      if (rawPartitionLengths.empty()) {
+        rawPartitionLengths.resize(rhs.rawPartitionLengths.size(), 0);
+      }
+      BOLT_CHECK_EQ(rawPartitionLengths.size(), rhs.rawPartitionLengths.size());
+      for (size_t i = 0; i < rawPartitionLengths.size(); ++i) {
+        rawPartitionLengths[i] += rhs.rawPartitionLengths[i];
+      }
+    }
+
+    return *this;
+  }
 };
 
 // Only partitioning that has pid support adaptive shuffle writer, otherwise

bolt/shuffle/sparksql/ShuffleWriter.h

@@ -55,6 +55,8 @@ class ShuffleWriter {
 
   virtual arrow::Status stop() = 0;
 
+  virtual arrow::Status localStop() = 0;
+
   int32_t numPartitions() const {
     return numPartitions_;
   }