[SPARK-29918][SQL] RecordBinaryComparator should check endianness when compared by long

sql/core/src/main/java/org/apache/spark/sql/execution/RecordBinaryComparator.java

@@ -22,6 +22,12 @@
 
 public final class RecordBinaryComparator extends RecordComparator {
 
+  boolean isLittlenEndian;
+
+  public RecordBinaryComparator(boolean isLittlenEndian) {
+    this.isLittlenEndian = isLittlenEndian;
+  }
+
   @Override
   public int compare(
       Object leftObj, long leftOff, int leftLen, Object rightObj, long rightOff, int rightLen) {
@@ -38,32 +44,39 @@ public int compare(
     // check if stars align and we can get both offsets to be aligned
     if ((leftOff % 8) == (rightOff % 8)) {
       while ((leftOff + i) % 8 != 0 && i < leftLen) {
-        final int v1 = Platform.getByte(leftObj, leftOff + i) & 0xff;
-        final int v2 = Platform.getByte(rightObj, rightOff + i) & 0xff;
+        final int v1 = Platform.getByte(leftObj, leftOff + i);
+        final int v2 = Platform.getByte(rightObj, rightOff + i);
         if (v1 != v2) {
-          return v1 > v2 ? 1 : -1;
+          return (v1 & 0xff) > (v2 & 0xff) ? 1 : -1;
         }
         i += 1;
       }
     }
     // for architectures that support unaligned accesses, chew it up 8 bytes at a time
     if (Platform.unaligned() || (((leftOff + i) % 8 == 0) && ((rightOff + i) % 8 == 0))) {
       while (i <= leftLen - 8) {
-        final long v1 = Platform.getLong(leftObj, leftOff + i);
-        final long v2 = Platform.getLong(rightObj, rightOff + i);
+        long v1 = Platform.getLong(leftObj, leftOff + i);
+        long v2 = Platform.getLong(rightObj, rightOff + i);
         if (v1 != v2) {
-          return v1 > v2 ? 1 : -1;
+          if (isLittlenEndian) {
+            // if read as little-endian, we have to reverse bytes so that the long comparison result
+            // is equivalent to byte-by-byte comparison result.
+            // See discussion in https://github.com/apache/spark/pull/26548#issuecomment-554645859
+            v1 = Long.reverseBytes(v1);
+            v2 = Long.reverseBytes(v2);
+          }
+          return Long.compareUnsigned(v1, v2);
         }
         i += 8;
       }
     }
     // this will finish off the unaligned comparisons, or do the entire aligned comparison
     // whichever is needed.
     while (i < leftLen) {
-      final int v1 = Platform.getByte(leftObj, leftOff + i) & 0xff;
-      final int v2 = Platform.getByte(rightObj, rightOff + i) & 0xff;
+      final int v1 = Platform.getByte(leftObj, leftOff + i);
+      final int v2 = Platform.getByte(rightObj, rightOff + i);
       if (v1 != v2) {
-        return v1 > v2 ? 1 : -1;
+        return (v1 & 0xff) > (v2 & 0xff) ? 1 : -1;
       }
       i += 1;
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.execution.exchange
 
+import java.nio.ByteOrder
 import java.util.Random
 import java.util.function.Supplier
 
@@ -243,7 +244,8 @@ object ShuffleExchangeExec {
       val newRdd = if (isRoundRobin && SQLConf.get.sortBeforeRepartition) {
         rdd.mapPartitionsInternal { iter =>
           val recordComparatorSupplier = new Supplier[RecordComparator] {
-            override def get: RecordComparator = new RecordBinaryComparator()
+            val isLittlenEndian = ByteOrder.nativeOrder.equals(ByteOrder.LITTLE_ENDIAN)
+            override def get: RecordComparator = new RecordBinaryComparator(isLittlenEndian)
           }
           // The comparator for comparing row hashcode, which should always be Integer.
           val prefixComparator = PrefixComparators.LONG

sql/core/src/test/java/test/org/apache/spark/sql/execution/sort/RecordBinaryComparatorSuite.java

@@ -56,8 +56,8 @@ public class RecordBinaryComparatorSuite {
 
   @Before
   public void beforeEach() {
-    // Only compare between two input rows.
-    array = consumer.allocateArray(2);
+    // At most three input rows
+    array = consumer.allocateArray(3);
     pos = 0;
 
     dataPage = memoryManager.allocatePage(4096, consumer);
@@ -88,7 +88,7 @@ private void insertRow(UnsafeRow row) {
     Platform.copyMemory(recordBase, recordOffset, baseObject, pageCursor, recordLength);
     pageCursor += recordLength;
 
-    assert(pos < 2);
+    assert(pos < 3);
     array.set(pos, recordAddress);
     pos++;
   }
@@ -108,7 +108,7 @@ private int compare(int index1, int index2) {
         baseOffset2, recordLength2);
   }
 
-  private final RecordComparator binaryComparator = new RecordBinaryComparator();
+  private final RecordComparator binaryComparator = new RecordBinaryComparator(true);
 
   // Compute the most compact size for UnsafeRow's backing data.
   private int computeSizeInBytes(int originalSize) {
@@ -273,7 +273,7 @@ public void testBinaryComparatorWhenSubtractionIsDivisibleByMaxIntValue() throws
     insertRow(row1);
     insertRow(row2);
 
-    assert(compare(0, 1) < 0);
+    assert(compare(0, 1) > 0);
   }
 
   @Test
@@ -321,4 +321,69 @@ public void testBinaryComparatorWhenOnlyTheLastColumnDiffers() throws Exception
 
     assert(compare(0, 1) < 0);
   }
+
+  @Test
+  public void testBinaryComparatorGiveSameResultWhenComparedByteByByteAndComparedByLong()
+      throws Exception {
+    int numFields = 1;
+
+    UnsafeRow row1 = new UnsafeRow(numFields);
+    byte[] data1 = new byte[100];
+    row1.pointTo(data1, computeSizeInBytes(numFields * 8));
+    row1.setLong(0, 0x0800000000000000L);
+
+    UnsafeRow row2 = new UnsafeRow(numFields);
+    byte[] data2 = new byte[100];
+    row2.pointTo(data2, computeSizeInBytes(numFields * 8));
+    row2.setLong(0, 0x0000008000000000L);
+
+    UnsafeRow row3 = new UnsafeRow(numFields);
+    byte[] data3 = new byte[100];
+    row3.pointTo(data3, computeSizeInBytes(numFields * 8));
+    row3.setLong(0, 0x0000008000000000L);
+
+    insertRow(row1);
+    insertRow(row2);
+    insertRow(row3);
+
+    // the bytes in row2 and row3 are the same.
+    // the base offset of row1 is 20, row2 is 40, row3 is 60
+    // so the RecordBinaryComparator will compare row1 and row2 with two long comparison directly,
+    // while the comparison between row1 and row3 is started with 4 bytes byte-by-byte comparison,
+    // followed by a long comparison, and lastly 4 bytes byte-by-byte comparison.
+    assert(compare(0, 1) < 0);
+    assert(compare(0, 2) < 0);
+  }
+
+  @Test
+  public void testBinaryComparatorShouldComparedWithUnsignedLong() throws Exception {
+    int numFields = 1;
+
+    UnsafeRow row1 = new UnsafeRow(numFields);
+    byte[] data1 = new byte[100];
+    row1.pointTo(data1, computeSizeInBytes(numFields * 8));
+    row1.setLong(0, 0xa000000000000000L);
+
+    UnsafeRow row2 = new UnsafeRow(numFields);
+    byte[] data2 = new byte[100];
+    row2.pointTo(data2, computeSizeInBytes(numFields * 8));
+    row2.setLong(0, 0x0000000000000000L);
+
+    UnsafeRow row3 = new UnsafeRow(numFields);
+    byte[] data3 = new byte[100];
+    row3.pointTo(data3, computeSizeInBytes(numFields * 8));
+    row3.setLong(0, 0x0000000000000000L);
+
+    insertRow(row1);
+    insertRow(row2);
+    insertRow(row3);
+
+    // the bytes in row2 and row3 are the same.
+    // the base offset of row1 is 20, row2 is 40, row3 is 60
+    // so the RecordBinaryComparator will compare row1 and row2 with two long comparison directly,
+    // while the comparison between row1 and row3 is started with 4 bytes byte-by-byte comparison,
+    // followed by a long comparison, and lastly 4 bytes byte-by-byte comparison.
+    assert(compare(0, 1) > 0);
+    assert(compare(0, 2) > 0);
+  }
 }