Make initialSize configurable in UnsafeSorter

Author: JoshRosen

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 =

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++) {