diff --git a/core/src/main/java/org/apache/spark/unsafe/sort/UnsafeSorter.java b/core/src/main/java/org/apache/spark/unsafe/sort/UnsafeSorter.java
index 6da89004d2f53..517d737368864 100644
--- a/core/src/main/java/org/apache/spark/unsafe/sort/UnsafeSorter.java
+++ b/core/src/main/java/org/apache/spark/unsafe/sort/UnsafeSorter.java
@@ -67,7 +67,7 @@ public static abstract class PrefixComparator {
    * Within this buffer, position {@code 2 * i} holds a pointer pointer to the record at
    * index {@code i}, while position {@code 2 * i + 1} in the array holds an 8-byte key prefix.
    */
-  private long[] sortBuffer = new long[1024];
+  private long[] sortBuffer;
 
   private int sortBufferInsertPosition = 0;
 
@@ -82,7 +82,10 @@ public UnsafeSorter(
       final TaskMemoryManager memoryManager,
       final RecordComparator recordComparator,
       PrefixComputer prefixComputer,
-      final PrefixComparator prefixComparator) {
+      final PrefixComparator prefixComparator,
+      int initialSize) {
+    assert (initialSize > 0);
+    this.sortBuffer = new long[initialSize * 2];
     this.memoryManager = memoryManager;
     this.prefixComputer = prefixComputer;
     this.sorter =
diff --git a/core/src/test/java/org/apache/spark/unsafe/sort/UnsafeSorterSuite.java b/core/src/test/java/org/apache/spark/unsafe/sort/UnsafeSorterSuite.java
index c22edfb412e1b..95d54ea79cf40 100644
--- a/core/src/test/java/org/apache/spark/unsafe/sort/UnsafeSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/sort/UnsafeSorterSuite.java
@@ -106,8 +106,8 @@ public int compare(long prefix1, long prefix2) {
         return (int) prefix1 - (int) prefix2;
       }
     };
-    final UnsafeSorter sorter =
-      new UnsafeSorter(memoryManager, recordComparator, prefixComputer, prefixComparator);
+    final UnsafeSorter sorter = new UnsafeSorter(memoryManager, recordComparator, prefixComputer,
+      prefixComparator, dataToSort.length);
     // Given a page of records, insert those records into the sorter one-by-one:
     position = dataPage.getBaseOffset();
     for (int i = 0; i < dataToSort.length; i++) {