replace int with long in RadixSort.java

Author: kiszk

core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java

@@ -176,9 +176,9 @@ public ShuffleSorterIterator getSortedIterator() {
     int offset = 0;
     if (useRadixSort) {
       offset = RadixSort.sort(
-        array, pos,
-        PackedRecordPointer.PARTITION_ID_START_BYTE_INDEX,
-        PackedRecordPointer.PARTITION_ID_END_BYTE_INDEX, false, false);
+        array, (long)pos,
+        (long)PackedRecordPointer.PARTITION_ID_START_BYTE_INDEX,
+        (long)PackedRecordPointer.PARTITION_ID_END_BYTE_INDEX, false, false);
     } else {
       MemoryBlock unused = new MemoryBlock(
         array.getBaseObject(),

core/src/main/java/org/apache/spark/util/collection/unsafe/sort/RadixSort.java

@@ -40,28 +40,28 @@ public class RadixSort {
    *         of always copying the data back to position zero for efficiency.
    */
   public static int sort(
-      LongArray array, int numRecords, int startByteIndex, int endByteIndex,
+      LongArray array, long numRecords, long startByteIndex, long endByteIndex,
       boolean desc, boolean signed) {
     assert startByteIndex >= 0 : "startByteIndex (" + startByteIndex + ") should >= 0";
     assert endByteIndex <= 7 : "endByteIndex (" + endByteIndex + ") should <= 7";
     assert endByteIndex > startByteIndex;
     assert numRecords * 2 <= array.size();
-    int inIndex = 0;
-    int outIndex = numRecords;
+    long inIndex = 0;
+    long outIndex = numRecords;
     if (numRecords > 0) {
       long[][] counts = getCounts(array, numRecords, startByteIndex, endByteIndex);
-      for (int i = startByteIndex; i <= endByteIndex; i++) {
-        if (counts[i] != null) {
+      for (long i = startByteIndex; i <= endByteIndex; i++) {
+        if (counts[(int)i] != null) {
           sortAtByte(
-            array, numRecords, counts[i], i, inIndex, outIndex,
+            array, numRecords, counts[(int)i], i, inIndex, outIndex,
             desc, signed && i == endByteIndex);
-          int tmp = inIndex;
+          long tmp = inIndex;
           inIndex = outIndex;
           outIndex = tmp;
         }
       }
     }
-    return inIndex;
+    return (int)inIndex;
   }
 
   /**
@@ -78,7 +78,7 @@ public static int sort(
    * @param signed whether this is a signed (two's complement) sort (only applies to last byte).
    */
   private static void sortAtByte(
-      LongArray array, int numRecords, long[] counts, int byteIdx, int inIndex, int outIndex,
+      LongArray array, long numRecords, long[] counts, long byteIdx, long inIndex, long outIndex,
       boolean desc, boolean signed) {
     assert counts.length == 256;
     long[] offsets = transformCountsToOffsets(
@@ -106,7 +106,7 @@ private static void sortAtByte(
    *         significant byte. If the byte does not need sorting the array will be null.
    */
   private static long[][] getCounts(
-      LongArray array, int numRecords, int startByteIndex, int endByteIndex) {
+      LongArray array, long numRecords, long startByteIndex, long endByteIndex) {
     long[][] counts = new long[8][];
     // Optimization: do a fast pre-pass to determine which byte indices we can skip for sorting.
     // If all the byte values at a particular index are the same we don't need to count it.
@@ -121,12 +121,12 @@ private static long[][] getCounts(
     }
     long bitsChanged = bitwiseMin ^ bitwiseMax;
     // Compute counts for each byte index.
-    for (int i = startByteIndex; i <= endByteIndex; i++) {
+    for (long i = startByteIndex; i <= endByteIndex; i++) {
       if (((bitsChanged >>> (i * 8)) & 0xff) != 0) {
-        counts[i] = new long[256];
+        counts[(int)i] = new long[256];
         // TODO(ekl) consider computing all the counts in one pass.
         for (long offset = array.getBaseOffset(); offset < maxOffset; offset += 8) {
-          counts[i][(int)((Platform.getLong(baseObject, offset) >>> (i * 8)) & 0xff)]++;
+          counts[(int)i][(int)((Platform.getLong(baseObject, offset) >>> (i * 8)) & 0xff)]++;
         }
       }
     }
@@ -146,7 +146,7 @@ private static long[][] getCounts(
    * @return the input counts array.
    */
   private static long[] transformCountsToOffsets(
-      long[] counts, int numRecords, long outputOffset, int bytesPerRecord,
+      long[] counts, long numRecords, long outputOffset, long bytesPerRecord,
       boolean desc, boolean signed) {
     assert counts.length == 256;
     int start = signed ? 128 : 0;  // output the negative records first (values 129-255).
@@ -176,41 +176,41 @@ private static long[] transformCountsToOffsets(
    */
   public static int sortKeyPrefixArray(
       LongArray array,
-      int startIndex,
-      int numRecords,
-      int startByteIndex,
-      int endByteIndex,
+      long startIndex,
+      long numRecords,
+      long startByteIndex,
+      long endByteIndex,
       boolean desc,
       boolean signed) {
     assert startByteIndex >= 0 : "startByteIndex (" + startByteIndex + ") should >= 0";
     assert endByteIndex <= 7 : "endByteIndex (" + endByteIndex + ") should <= 7";
     assert endByteIndex > startByteIndex;
     assert numRecords * 4 <= array.size();
-    int inIndex = startIndex;
-    int outIndex = startIndex + numRecords * 2;
+    long inIndex = startIndex;
+    long outIndex = startIndex + numRecords * 2L;
     if (numRecords > 0) {
       long[][] counts = getKeyPrefixArrayCounts(
         array, startIndex, numRecords, startByteIndex, endByteIndex);
-      for (int i = startByteIndex; i <= endByteIndex; i++) {
-        if (counts[i] != null) {
+      for (long i = startByteIndex; i <= endByteIndex; i++) {
+        if (counts[(int)i] != null) {
           sortKeyPrefixArrayAtByte(
-            array, numRecords, counts[i], i, inIndex, outIndex,
+            array, numRecords, counts[(int)i], i, inIndex, outIndex,
             desc, signed && i == endByteIndex);
-          int tmp = inIndex;
+          long tmp = inIndex;
           inIndex = outIndex;
           outIndex = tmp;
         }
       }
     }
-    return inIndex;
+    return (int)inIndex;
   }
 
   /**
    * Specialization of getCounts() for key-prefix arrays. We could probably combine this with
    * getCounts with some added parameters but that seems to hurt in benchmarks.
    */
   private static long[][] getKeyPrefixArrayCounts(
-      LongArray array, int startIndex, int numRecords, int startByteIndex, int endByteIndex) {
+      LongArray array, long startIndex, long numRecords, long startByteIndex, long endByteIndex) {
     long[][] counts = new long[8][];
     long bitwiseMax = 0;
     long bitwiseMin = -1L;
@@ -223,11 +223,11 @@ private static long[][] getKeyPrefixArrayCounts(
       bitwiseMin &= value;
     }
     long bitsChanged = bitwiseMin ^ bitwiseMax;
-    for (int i = startByteIndex; i <= endByteIndex; i++) {
+    for (long i = startByteIndex; i <= endByteIndex; i++) {
       if (((bitsChanged >>> (i * 8)) & 0xff) != 0) {
-        counts[i] = new long[256];
+        counts[(int)i] = new long[256];
         for (long offset = baseOffset; offset < limit; offset += 16) {
-          counts[i][(int)((Platform.getLong(baseObject, offset + 8) >>> (i * 8)) & 0xff)]++;
+          counts[(int)i][(int)((Platform.getLong(baseObject, offset + 8) >>> (i * 8)) & 0xff)]++;
         }
       }
     }
@@ -238,7 +238,7 @@ private static long[][] getKeyPrefixArrayCounts(
    * Specialization of sortAtByte() for key-prefix arrays.
    */
   private static void sortKeyPrefixArrayAtByte(
-      LongArray array, int numRecords, long[] counts, int byteIdx, int inIndex, int outIndex,
+      LongArray array, long numRecords, long[] counts, long byteIdx, long inIndex, long outIndex,
       boolean desc, boolean signed) {
     assert counts.length == 256;
     long[] offsets = transformCountsToOffsets(

core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java

@@ -322,7 +322,7 @@ public UnsafeSorterIterator getSortedIterator() {
     if (sortComparator != null) {
       if (this.radixSortSupport != null) {
         offset = RadixSort.sortKeyPrefixArray(
-          array, nullBoundaryPos, (pos - nullBoundaryPos) / 2, 0, 7,
+          array, (long)nullBoundaryPos, (pos - nullBoundaryPos) / 2L, 0L, 7L,
           radixSortSupport.sortDescending(), radixSortSupport.sortSigned());
       } else {
         MemoryBlock unused = new MemoryBlock(

core/src/test/scala/org/apache/spark/util/collection/unsafe/sort/RadixSortSuite.scala

@@ -30,7 +30,7 @@ import org.apache.spark.util.collection.Sorter
 import org.apache.spark.util.random.XORShiftRandom
 
 class RadixSortSuite extends SparkFunSuite with Logging {
-  private val N = 10000  // scale this down for more readable results
+  private val N = 10000.toLong  // scale this down for more readable results
 
   /**
    * Describes a type of sort to test, e.g. two's complement descending. Each sort type has
@@ -40,7 +40,7 @@ class RadixSortSuite extends SparkFunSuite with Logging {
   case class RadixSortType(
     name: String,
     referenceComparator: PrefixComparator,
-    startByteIdx: Int, endByteIdx: Int, descending: Boolean, signed: Boolean, nullsFirst: Boolean)
+    startByteIdx: Long, endByteIdx: Long, descending: Boolean, signed: Boolean, nullsFirst: Boolean)
 
   val SORT_TYPES_TO_TEST = Seq(
     RadixSortType("unsigned binary data asc nulls first",
@@ -73,22 +73,22 @@ class RadixSortSuite extends SparkFunSuite with Logging {
       },
       2, 4, false, false, true))
 
-  private def generateTestData(size: Int, rand: => Long): (Array[JLong], LongArray) = {
-    val ref = Array.tabulate[Long](size) { i => rand }
-    val extended = ref ++ Array.fill[Long](size)(0)
+  private def generateTestData(size: Long, rand: => Long): (Array[JLong], LongArray) = {
+    val ref = Array.tabulate[Long](size.toInt) { i => rand }
+    val extended = ref ++ Array.fill[Long](size.toInt)(0)
     (ref.map(i => new JLong(i)), new LongArray(MemoryBlock.fromLongArray(extended)))
   }
 
-  private def generateKeyPrefixTestData(size: Int, rand: => Long): (LongArray, LongArray) = {
-    val ref = Array.tabulate[Long](size * 2) { i => rand }
-    val extended = ref ++ Array.fill[Long](size * 2)(0)
+  private def generateKeyPrefixTestData(size: Long, rand: => Long): (LongArray, LongArray) = {
+    val ref = Array.tabulate[Long]((size * 2).toInt) { i => rand }
+    val extended = ref ++ Array.fill[Long]((size * 2).toInt)(0)
     (new LongArray(MemoryBlock.fromLongArray(ref)),
      new LongArray(MemoryBlock.fromLongArray(extended)))
   }
 
-  private def collectToArray(array: LongArray, offset: Int, length: Int): Array[Long] = {
+  private def collectToArray(array: LongArray, offset: Int, length: Long): Array[Long] = {
     var i = 0
-    val out = new Array[Long](length)
+    val out = new Array[Long](length.toInt)
     while (i < length) {
       out(i) = array.get(offset + i)
       i += 1
@@ -107,10 +107,10 @@ class RadixSortSuite extends SparkFunSuite with Logging {
     }
   }
 
-  private def referenceKeyPrefixSort(buf: LongArray, lo: Int, hi: Int, refCmp: PrefixComparator) {
+  private def referenceKeyPrefixSort(buf: LongArray, lo: Long, hi: Long, refCmp: PrefixComparator) {
     val sortBuffer = new LongArray(MemoryBlock.fromLongArray(new Array[Long](buf.size().toInt)))
     new Sorter(new UnsafeSortDataFormat(sortBuffer)).sort(
-      buf, lo, hi, new Comparator[RecordPointerAndKeyPrefix] {
+      buf, lo.toInt, hi.toInt, new Comparator[RecordPointerAndKeyPrefix] {
         override def compare(
             r1: RecordPointerAndKeyPrefix,
             r2: RecordPointerAndKeyPrefix): Int = {
@@ -156,7 +156,7 @@ class RadixSortSuite extends SparkFunSuite with Logging {
       val (ref, buffer) = generateTestData(N, rand.nextLong)
       Arrays.sort(ref, toJavaComparator(sortType.referenceComparator))
       val outOffset = RadixSort.sort(
-        buffer, N, sortType.startByteIdx, sortType.endByteIdx,
+        buffer, N.toLong, sortType.startByteIdx, sortType.endByteIdx,
         sortType.descending, sortType.signed)
       val result = collectToArray(buffer, outOffset, N)
       assert(ref.view == result.view)

sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/SortBenchmark.scala

@@ -80,7 +80,7 @@ class SortBenchmark extends BenchmarkBase {
       }
       val buf = new LongArray(MemoryBlock.fromLongArray(array))
       timer.startTiming()
-      RadixSort.sort(buf, size, 0, 7, false, false)
+      RadixSort.sort(buf, size.toLong, 0.toLong, 7.toLong, false, false)
       timer.stopTiming()
     }
     benchmark.addTimerCase("radix sort two bytes") { timer =>
@@ -92,7 +92,7 @@ class SortBenchmark extends BenchmarkBase {
       }
       val buf = new LongArray(MemoryBlock.fromLongArray(array))
       timer.startTiming()
-      RadixSort.sort(buf, size, 0, 7, false, false)
+      RadixSort.sort(buf, size.toLong, 0.toLong, 7.toLong, false, false)
       timer.stopTiming()
     }
     benchmark.addTimerCase("radix sort eight bytes") { timer =>
@@ -104,13 +104,13 @@ class SortBenchmark extends BenchmarkBase {
       }
       val buf = new LongArray(MemoryBlock.fromLongArray(array))
       timer.startTiming()
-      RadixSort.sort(buf, size, 0, 7, false, false)
+      RadixSort.sort(buf, size.toLong, 0.toLong, 7.toLong, false, false)
       timer.stopTiming()
     }
     benchmark.addTimerCase("radix sort key prefix array") { timer =>
       val (_, buf2) = generateKeyPrefixTestData(size, rand.nextLong)
       timer.startTiming()
-      RadixSort.sortKeyPrefixArray(buf2, 0, size, 0, 7, false, false)
+      RadixSort.sortKeyPrefixArray(buf2, 0.toLong, size.toLong, 0.toLong, 7.toLong, false, false)
       timer.stopTiming()
     }
     benchmark.run()