[CELEBORN] Optimize the performance of Celeborn's pushPartitionData (f…

…acebookincubator#3533)
marin-ma · Oct 26, 2023 · 87372db · 87372db
1 parent 5038204
commit 87372db
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/...mmon/src/main/scala/org/apache/spark/shuffle/CelebornHashBasedColumnarShuffleWriter.scala b/...mmon/src/main/scala/org/apache/spark/shuffle/CelebornHashBasedColumnarShuffleWriter.scala
@@ -51,14 +51,16 @@ abstract class CelebornHashBasedColumnarShuffleWriter[K, V](
 
   protected val mapId: Int = context.partitionId()
 
+  protected val clientPushBufferMaxSize: Int = celebornConf.clientPushBufferMaxSize
+
   protected val celebornPartitionPusher = new CelebornPartitionPusher(
     shuffleId,
     numMappers,
     numPartitions,
     context,
     mapId,
     client,
-    celebornConf)
+    clientPushBufferMaxSize)
 
   protected val blockManager: BlockManager = SparkEnv.get.blockManager
 

diff --git a/gluten-celeborn/common/src/main/scala/org/apache/spark/shuffle/CelebornPartitionPusher.scala b/gluten-celeborn/common/src/main/scala/org/apache/spark/shuffle/CelebornPartitionPusher.scala
@@ -20,7 +20,6 @@ import org.apache.spark._
 import org.apache.spark.internal.Logging
 
 import org.apache.celeborn.client.ShuffleClient
-import org.apache.celeborn.common.CelebornConf
 
 import java.io.IOException
 
@@ -31,13 +30,13 @@ class CelebornPartitionPusher(
     val context: TaskContext,
     val mapId: Int,
     val client: ShuffleClient,
-    val celebornConf: CelebornConf)
+    val clientPushBufferMaxSize: Int)
   extends Logging {
 
   @throws[IOException]
   def pushPartitionData(partitionId: Int, buffer: Array[Byte], length: Int): Int = {
     logDebug(s"Push record, size ${buffer.length}.")
-    if (buffer.length > celebornConf.clientPushBufferMaxSize) {
+    if (buffer.length > clientPushBufferMaxSize) {
       client.pushData(
         shuffleId,
         mapId,

diff --git a/...src/main/scala/org/apache/spark/shuffle/VeloxCelebornHashBasedColumnarShuffleWriter.scala b/...src/main/scala/org/apache/spark/shuffle/VeloxCelebornHashBasedColumnarShuffleWriter.scala
@@ -78,7 +78,7 @@ class VeloxCelebornHashBasedColumnarShuffleWriter[K, V](
             customizedCompressionCodec,
             bufferCompressThreshold,
             GlutenConfig.getConf.columnarShuffleCompressionMode,
-            celebornConf.clientPushBufferMaxSize,
+            clientPushBufferMaxSize,
             celebornPartitionPusher,
             NativeMemoryManagers
               .create(