[VL] Reduce Velox hash shuffle partition buffer memory by evicting large partitions after split

backends-velox/src-celeborn/main/scala/org/apache/spark/shuffle/VeloxCelebornColumnarShuffleWriter.scala

@@ -150,6 +150,7 @@ class VeloxCelebornColumnarShuffleWriter[K, V](
           GlutenShuffleUtils.getStartPartitionId(dep.nativePartitioning, context.partitionId),
           nativeBufferSize,
           GlutenConfig.get.columnarShuffleReallocThreshold,
+          GlutenConfig.get.columnarShuffleEvictPartitionSize,
           partitionWriterHandle
         )
       case SortShuffleWriterType =>

backends-velox/src-uniffle/main/java/org/apache/spark/shuffle/writer/VeloxUniffleColumnarShuffleWriter.java

@@ -185,6 +185,7 @@ protected void writeImpl(Iterator<Product2<K, V>> records) {
                         columnarDep.nativePartitioning(), partitionId),
                     nativeBufferSize,
                     reallocThreshold,
+                    GlutenConfig.get().columnarShuffleEvictPartitionSize(),
                     partitionWriterHandle);
           }
 

backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala

@@ -192,6 +192,7 @@ class ColumnarShuffleWriter[K, V](
                 taskContext.partitionId),
               nativeBufferSize,
               reallocThreshold,
+              GlutenConfig.get.columnarShuffleEvictPartitionSize,
               partitionWriterHandle
             )
           }

cpp/core/jni/JniWrapper.cc

@@ -990,6 +990,7 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe
     jint startPartitionId,
     jint splitBufferSize,
     jdouble splitBufferReallocThreshold,
+    jint evictPartitionSize,
     jlong partitionWriterHandle) {
   JNI_METHOD_START
   const auto ctx = getRuntime(env, wrapper);
@@ -1004,7 +1005,8 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe
       toPartitioning(jStringToCString(env, partitioningNameJstr)),
       startPartitionId,
       splitBufferSize,
-      splitBufferReallocThreshold);
+      splitBufferReallocThreshold,
+      evictPartitionSize);
 
   return ctx->saveObject(ctx->createShuffleWriter(numPartitions, partitionWriter, shuffleWriterOptions));
   JNI_METHOD_END(kInvalidObjectHandle)

cpp/core/shuffle/Options.h

@@ -27,6 +27,7 @@
 namespace gluten {
 
 static constexpr int16_t kDefaultBatchSize = 4096;
+static constexpr int32_t kDefaultEvictPartitionSize = 256 * 1024;
 static constexpr int32_t kDefaultShuffleWriterBufferSize = 4096;
 static constexpr int64_t kDefaultSortBufferThreshold = 64 << 20;
 static constexpr int64_t kDefaultPushMemoryThreshold = 4096;
@@ -81,17 +82,20 @@ struct ShuffleWriterOptions {
 struct HashShuffleWriterOptions : ShuffleWriterOptions {
   int32_t splitBufferSize = kDefaultShuffleWriterBufferSize;
   double splitBufferReallocThreshold = kDefaultSplitBufferReallocThreshold;
+  int32_t evictPartitionSize = kDefaultEvictPartitionSize;
 
   HashShuffleWriterOptions() : ShuffleWriterOptions(ShuffleWriterType::kHashShuffle) {}
 
   HashShuffleWriterOptions(
       Partitioning partitioning,
       int32_t startPartitionId,
       int32_t partitionBufferSize,
-      double partitionBufferReallocThreshold)
+      double partitionBufferReallocThreshold,
+      int32_t partitionEvictPartitionSize = kDefaultEvictPartitionSize)
       : ShuffleWriterOptions(ShuffleWriterType::kHashShuffle, partitioning, startPartitionId),
         splitBufferSize(partitionBufferSize),
-        splitBufferReallocThreshold(partitionBufferReallocThreshold) {}
+        splitBufferReallocThreshold(partitionBufferReallocThreshold),
+        evictPartitionSize(partitionEvictPartitionSize) {}
 
  protected:
   HashShuffleWriterOptions(ShuffleWriterType shuffleWriterType) : ShuffleWriterOptions(shuffleWriterType) {}
@@ -101,10 +105,12 @@ struct HashShuffleWriterOptions : ShuffleWriterOptions {
       Partitioning partitioning,
       int32_t startPartitionId,
       int32_t partitionBufferSize,
-      double partitionBufferReallocThreshold)
+      double partitionBufferReallocThreshold,
+      int32_t partitionEvictPartitionSize = kDefaultEvictPartitionSize)
       : ShuffleWriterOptions(shuffleWriterType, partitioning, startPartitionId),
         splitBufferSize(partitionBufferSize),
-        splitBufferReallocThreshold(partitionBufferReallocThreshold) {}
+        splitBufferReallocThreshold(partitionBufferReallocThreshold),
+        evictPartitionSize(partitionEvictPartitionSize) {}
 };
 
 struct SortShuffleWriterOptions : ShuffleWriterOptions {

cpp/core/shuffle/Payload.cc

@@ -60,7 +60,7 @@ arrow::Result<uint8_t> readPayloadType(arrow::io::InputStream* is) {
 }
 
 arrow::Result<int64_t> compressBuffer(
-    const std::shared_ptr<arrow::Buffer>& buffer,
+    const std::shared_ptr<arrow::Buffer> buffer,
     uint8_t* output,
     int64_t outputLength,
     arrow::util::Codec* codec) {

cpp/velox/shuffle/VeloxHashShuffleWriter.cc

@@ -440,9 +440,41 @@ arrow::Status VeloxHashShuffleWriter::doSplit(const facebook::velox::RowVector&
   printPartitionBuffer();
 
   setSplitState(SplitState::kInit);
+  if (evictPartitionSize_ > 0) {
+    // After split, evict large partition buffers to free up memory for the next input RowVector.
+    const auto partitionBytes = estimatePartitionBufferBytes();
+    for (uint32_t pid = 0; pid < partitionBytes.size(); ++pid) {
+      if (partitionBufferBase_[pid] > 0 && partitionBytes[pid] >= evictPartitionSize_) {
+        RETURN_NOT_OK(evictPartitionBuffers(pid, false));
+      }
+    }
+  }
   return arrow::Status::OK();
 }
 
+std::vector<int64_t> VeloxHashShuffleWriter::estimatePartitionBufferBytes() const {
+  std::vector<int64_t> partitionBytes(numPartitions_, 0);
+
+  for (const auto& columnBuffers : partitionBuffers_) {
+    for (uint32_t pid = 0; pid < columnBuffers.size(); ++pid) {
+      for (const auto& buffer : columnBuffers[pid]) {
+        if (buffer) {
+          partitionBytes[pid] += buffer->capacity();
+        }
+      }
+    }
+  }
+
+  for (uint32_t pid = 0; pid < complexTypeFlushBuffer_.size(); ++pid) {
+    const auto& buffer = complexTypeFlushBuffer_[pid];
+    if (buffer) {
+      partitionBytes[pid] += buffer->capacity();
+    }
+  }
+
+  return partitionBytes;
+}
+
 arrow::Status VeloxHashShuffleWriter::splitRowVector(const facebook::velox::RowVector& rv) {
   SCOPED_TIMER(cpuWallTimingList_[CpuWallTimingSplitRV]);
 

cpp/velox/shuffle/VeloxHashShuffleWriter.h

@@ -214,7 +214,8 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter {
       MemoryManager* memoryManager)
       : VeloxShuffleWriter(numPartitions, partitionWriter, options, memoryManager),
         splitBufferSize_(options->splitBufferSize),
-        splitBufferReallocThreshold_(options->splitBufferReallocThreshold) {
+        splitBufferReallocThreshold_(options->splitBufferReallocThreshold),
+        evictPartitionSize_(options->evictPartitionSize) {
     arenas_.resize(numPartitions);
   }
 
@@ -223,6 +224,8 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter {
 
   arrow::Status initColumnTypes(const facebook::velox::RowVector& rv);
 
+  std::vector<int64_t> estimatePartitionBufferBytes() const;
+
   arrow::Status splitRowVector(const facebook::velox::RowVector& rv);
 
   arrow::Status initFromRowVector(const facebook::velox::RowVector& rv);
@@ -326,6 +329,7 @@ class VeloxHashShuffleWriter : public VeloxShuffleWriter {
  protected:
   int32_t splitBufferSize_;
   double splitBufferReallocThreshold_;
+  int32_t evictPartitionSize_;
 
   std::shared_ptr<arrow::Schema> schema_;