Clean up prefix comparators for integral types

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 org.apache.spark.annotation.Private;
 import org.apache.spark.unsafe.types.UTF8String;
 
@@ -27,8 +28,7 @@ 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 IntegralPrefixComparator INTEGRAL = new IntegralPrefixComparator();
   public static final FloatPrefixComparator FLOAT = new FloatPrefixComparator();
   public static final DoublePrefixComparator DOUBLE = new DoublePrefixComparator();
 
@@ -67,22 +67,10 @@ public long computePrefix(String value) {
     }
   }
 
-  public static final class IntPrefixComparator extends PrefixComparator {
-    @Override
-    public int compare(long aPrefix, long bPrefix) {
-      int a = (int) aPrefix;
-      int b = (int) bPrefix;
-      return (a < b) ? -1 : (a > b) ? 1 : 0;
-    }
-
-    public final long NULL_PREFIX = computePrefix(Integer.MIN_VALUE);
-
-    public long computePrefix(int value) {
-      return value & 0xffffffffL;
-    }
-  }
-
-  public static final class LongPrefixComparator extends PrefixComparator {
+  /**
+   * Prefix comparator for all integral types (boolean, byte, short, int, long).
+   */
+  public static final class IntegralPrefixComparator extends PrefixComparator {
     @Override
     public int compare(long a, long b) {
       return (a < b) ? -1 : (a > b) ? 1 : 0;

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

@@ -38,8 +38,7 @@ object SortPrefixUtils {
   def getPrefixComparator(sortOrder: SortOrder): PrefixComparator = {
     sortOrder.dataType match {
       case StringType => PrefixComparators.STRING
-      case IntegerType => PrefixComparators.INTEGER
-      case LongType => PrefixComparators.LONG
+      case BooleanType | ByteType | ShortType | IntegerType | LongType => PrefixComparators.INTEGRAL
       case FloatType => PrefixComparators.FLOAT
       case DoubleType => PrefixComparators.DOUBLE
       case _ => NoOpPrefixComparator
@@ -51,16 +50,37 @@ object SortPrefixUtils {
       case StringType => (row: InternalRow) => {
         PrefixComparators.STRING.computePrefix(sortOrder.child.eval(row).asInstanceOf[UTF8String])
       }
-      case IntegerType => (row: InternalRow) => {
-        val exprVal = sortOrder.child.eval(row)
-        if (exprVal == null) PrefixComparators.INTEGER.NULL_PREFIX
-        else PrefixComparators.INTEGER.computePrefix(sortOrder.child.eval(row).asInstanceOf[Int])
-      }
-      case LongType => (row: InternalRow) => {
-        val exprVal = sortOrder.child.eval(row)
-        if (exprVal == null) PrefixComparators.LONG.NULL_PREFIX
-        else sortOrder.child.eval(row).asInstanceOf[Long]
-      }
+      case BooleanType =>
+        (row: InternalRow) => {
+          val exprVal = sortOrder.child.eval(row)
+          if (exprVal == null) PrefixComparators.INTEGRAL.NULL_PREFIX
+          else if (sortOrder.child.eval(row).asInstanceOf[Boolean]) 1
+          else 0
+        }
+      case ByteType =>
+        (row: InternalRow) => {
+          val exprVal = sortOrder.child.eval(row)
+          if (exprVal == null) PrefixComparators.INTEGRAL.NULL_PREFIX
+          else sortOrder.child.eval(row).asInstanceOf[Byte]
+        }
+      case ShortType =>
+        (row: InternalRow) => {
+          val exprVal = sortOrder.child.eval(row)
+          if (exprVal == null) PrefixComparators.INTEGRAL.NULL_PREFIX
+          else sortOrder.child.eval(row).asInstanceOf[Short]
+        }
+      case IntegerType =>
+        (row: InternalRow) => {
+          val exprVal = sortOrder.child.eval(row)
+          if (exprVal == null) PrefixComparators.INTEGRAL.NULL_PREFIX
+          else sortOrder.child.eval(row).asInstanceOf[Int]
+        }
+      case LongType =>
+        (row: InternalRow) => {
+          val exprVal = sortOrder.child.eval(row)
+          if (exprVal == null) PrefixComparators.INTEGRAL.NULL_PREFIX
+          else sortOrder.child.eval(row).asInstanceOf[Long]
+        }
       case FloatType => (row: InternalRow) => {
         val exprVal = sortOrder.child.eval(row)
         if (exprVal == null) PrefixComparators.FLOAT.NULL_PREFIX