diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriterBuffer.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriterBuffer.java
index 6fb41975c0..a1427c711d 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriterBuffer.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriterBuffer.java
@@ -45,13 +45,9 @@ public void addRecord(byte[] recordBuffer, int length) {
         buffers.add(new WrappedBuffer(buffer, nextOffset));
         nextOffset = 0;
       }
-      if (length > bufferSize) {
-        buffer = new byte[length];
-        memoryUsed += length;
-      } else {
-        buffer = new byte[bufferSize];
-        memoryUsed += bufferSize;
-      }
+      int newBufferSize = Math.max(length, bufferSize);
+      buffer = new byte[newBufferSize];
+      memoryUsed += newBufferSize;
     }
 
     try {