Fix bugs in string prefix comparison.

Author: JoshRosen

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

@@ -19,6 +19,7 @@
 
 import com.google.common.base.Charsets;
 import com.google.common.primitives.Longs;
+import com.google.common.primitives.UnsignedBytes;
 
 import org.apache.spark.annotation.Private;
 import org.apache.spark.unsafe.types.UTF8String;
@@ -35,36 +36,33 @@ private PrefixComparators() {}
   public static final class StringPrefixComparator extends PrefixComparator {
     @Override
     public int compare(long aPrefix, long bPrefix) {
-      // TODO: this can certainly be done more efficiently
+      // TODO: can done more efficiently
       byte[] a = Longs.toByteArray(aPrefix);
       byte[] b = Longs.toByteArray(bPrefix);
       for (int i = 0; i < 8; i++) {
-        if (a[i] == b[i]) continue;
-        if (a[i] > b[i]) return -1;
-        else if (a[i] < b[i]) return 1;
+        int c = UnsignedBytes.compare(a[i], b[i]);
+        if (c != 0) return c;
       }
       return 0;
     }
 
-    public long computePrefix(UTF8String value) {
-      // TODO: this can certainly be done more efficiently
-      return value == null ? 0L : computePrefix(value.toString());
-    }
-
-    public long computePrefix(String value) {
-      // TODO: this can certainly be done more efficiently
-      if (value == null || value.length() == 0) {
+    public long computePrefix(byte[] bytes) {
+      if (bytes == null) {
         return 0L;
       } else {
-        String first4Chars = value.substring(0, Math.min(3, value.length() - 1));
-        byte[] utf16Bytes = first4Chars.getBytes(Charsets.UTF_16);
         byte[] padded = new byte[8];
-        if (utf16Bytes.length < 8) {
-          System.arraycopy(utf16Bytes, 0, padded, 0, utf16Bytes.length);
-        }
+        System.arraycopy(bytes, 0, padded, 0, Math.min(bytes.length, 8));
         return Longs.fromByteArray(padded);
       }
     }
+
+    public long computePrefix(String value) {
+      return value == null ? 0L : computePrefix(value.getBytes(Charsets.UTF_8));
+    }
+
+    public long computePrefix(UTF8String value) {
+      return value == null ? 0L : computePrefix(value.getBytes());
+    }
   }
 
   /**

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

@@ -0,0 +1,31 @@
+package org.apache.spark.util.collection.unsafe.sort
+
+import org.scalatest.prop.PropertyChecks
+
+import org.apache.spark.SparkFunSuite
+
+class PrefixComparatorsSuite extends SparkFunSuite with PropertyChecks {
+
+  test("String prefix comparator") {
+
+    def testPrefixComparison(s1: String, s2: String): Unit = {
+      val s1Prefix = PrefixComparators.STRING.computePrefix(s1)
+      val s2Prefix = PrefixComparators.STRING.computePrefix(s2)
+      val prefixComparisonResult = PrefixComparators.STRING.compare(s1Prefix, s2Prefix)
+      assert(
+        (prefixComparisonResult == 0) ||
+        (prefixComparisonResult < 0 && s1 < s2) ||
+        (prefixComparisonResult > 0 && s1 > s2))
+    }
+
+    val regressionTests = Table(
+      ("s1", "s2"),
+      ("abc", "世界"),
+      ("你好", "世界"),
+      ("你好123", "你好122")
+    )
+
+    forAll (regressionTests) { (s1: String, s2: String) => testPrefixComparison(s1, s2) }
+    forAll { (s1: String, s2: String) => testPrefixComparison(s1, s2) }
+  }
+}