[CELEBORN] Fix push small data (apache#3766)

zhanglistar · Nov 20, 2023 · cb18fc1 · cb18fc1
1 parent dc75cca
commit cb18fc1
Showing 1 changed file with 3 additions and 2 deletions.
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
@@ -35,8 +35,8 @@ class CelebornPartitionPusher(
 
   @throws[IOException]
   def pushPartitionData(partitionId: Int, buffer: Array[Byte], length: Int): Int = {
-    logDebug(s"Push record, size ${buffer.length}.")
-    if (buffer.length > clientPushBufferMaxSize) {
+    if (length > clientPushBufferMaxSize) {
+      logDebug(s"Push Data, size $length.")
       client.pushData(
         shuffleId,
         mapId,
@@ -48,6 +48,7 @@ class CelebornPartitionPusher(
         numMappers,
         numPartitions)
     } else {
+      logDebug(s"Merge Data, size $length.")
       client.mergeData(
         shuffleId,
         mapId,