Implement prefix sort for strings (albeit inefficiently).

Author: JoshRosen

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

@@ -17,17 +17,56 @@
 
 package org.apache.spark.util.collection.unsafe.sort;
 
+import com.google.common.base.Charsets;
+import com.google.common.primitives.Longs;
 import org.apache.spark.annotation.Private;
+import org.apache.spark.unsafe.types.UTF8String;
 
 @Private
 public class PrefixComparators {
   private PrefixComparators() {}
 
+  public static final StringPrefixComparator STRING = new StringPrefixComparator();
   public static final IntPrefixComparator INTEGER = new IntPrefixComparator();
   public static final LongPrefixComparator LONG = new LongPrefixComparator();
   public static final FloatPrefixComparator FLOAT = new FloatPrefixComparator();
   public static final DoublePrefixComparator DOUBLE = new DoublePrefixComparator();
 
+  public static final class StringPrefixComparator extends PrefixComparator {
+    @Override
+    public int compare(long aPrefix, long bPrefix) {
+      // TODO: this can certainly be 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;
+      }
+      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) {
+        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);
+        }
+        return Longs.fromByteArray(padded);
+      }
+    }
+  }
+
   public static final class IntPrefixComparator extends PrefixComparator {
     @Override
     public int compare(long aPrefix, long bPrefix) {

sql/core/src/main/scala/org/apache/spark/sql/execution/SortPrefixUtils.scala

@@ -20,7 +20,8 @@ package org.apache.spark.sql.execution
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.SortOrder
-import org.apache.spark.sql.types.{DoubleType, FloatType, LongType, IntegerType}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.collection.unsafe.sort.{PrefixComparators, PrefixComparator}
 
 
@@ -36,6 +37,7 @@ object SortPrefixUtils {
 
   def getPrefixComparator(sortOrder: SortOrder): PrefixComparator = {
     sortOrder.dataType match {
+      case StringType => PrefixComparators.STRING
       case IntegerType => PrefixComparators.INTEGER
       case LongType => PrefixComparators.LONG
       case FloatType => PrefixComparators.FLOAT
@@ -46,6 +48,8 @@ object SortPrefixUtils {
 
   def getPrefixComputer(sortOrder: SortOrder): InternalRow => Long = {
     sortOrder.dataType match {
+      case StringType => (row: InternalRow) =>
+        PrefixComparators.STRING.computePrefix(sortOrder.child.eval(row).asInstanceOf[UTF8String])
       case IntegerType => (row: InternalRow) =>
         PrefixComparators.INTEGER.computePrefix(sortOrder.child.eval(row).asInstanceOf[Int])
       case LongType => (row: InternalRow) => sortOrder.child.eval(row).asInstanceOf[Long]