Play around with Unsafe implementations of more string methods.

Author: JoshRosen

sql/catalyst/src/main/scala/org/apache/spark/sql/types/UTF8String.scala

@@ -19,8 +19,9 @@ package org.apache.spark.sql.types
 
 import java.util.Arrays
 
-import org.apache.spark.unsafe.PlatformDependent
-import org.apache.spark.unsafe.string.{UTF8StringPointer, UTF8StringMethods}
+import org.apache.spark.unsafe.PlatformDependent.BYTE_ARRAY_OFFSET
+import org.apache.spark.unsafe.array.ByteArrayMethods
+import org.apache.spark.unsafe.string.UTF8StringMethods
 
 /**
  *  A UTF-8 String, as internal representation of StringType in SparkSQL
@@ -35,8 +36,6 @@ final class UTF8String extends Ordered[UTF8String] with Serializable {
 
   private[this] var bytes: Array[Byte] = _
 
-  private val pointer: UTF8StringPointer = new UTF8StringPointer
-
   /**
    * Update the UTF8String with String.
    */
@@ -49,7 +48,6 @@ final class UTF8String extends Ordered[UTF8String] with Serializable {
    */
   def set(bytes: Array[Byte]): UTF8String = {
     this.bytes = bytes
-    pointer.set(bytes, PlatformDependent.BYTE_ARRAY_OFFSET, bytes.length)
     this
   }
 
@@ -59,7 +57,7 @@ final class UTF8String extends Ordered[UTF8String] with Serializable {
    * This is only used by Substring() when `start` is negative.
    */
   def length(): Int = {
-    pointer.getLengthInCodePoints
+    UTF8StringMethods.getLengthInCodePoints(bytes, BYTE_ARRAY_OFFSET, bytes.length)
   }
 
   def getBytes: Array[Byte] = {
@@ -107,19 +105,27 @@ final class UTF8String extends Ordered[UTF8String] with Serializable {
   }
 
   def startsWith(prefix: UTF8String): Boolean = {
-    val b = prefix.getBytes
-    if (b.length > bytes.length) {
-      return false
-    }
-    Arrays.equals(Arrays.copyOfRange(bytes, 0, b.length), b)
+    val prefixBytes = prefix.getBytes
+    UTF8StringMethods.startsWith(
+      bytes,
+      BYTE_ARRAY_OFFSET,
+      bytes.length,
+      prefixBytes,
+      BYTE_ARRAY_OFFSET,
+      prefixBytes.length
+    )
   }
 
   def endsWith(suffix: UTF8String): Boolean = {
-    val b = suffix.getBytes
-    if (b.length > bytes.length) {
-      return false
-    }
-    Arrays.equals(Arrays.copyOfRange(bytes, bytes.length - b.length, bytes.length), b)
+    val suffixBytes = suffix.getBytes
+    UTF8StringMethods.endsWith(
+      bytes,
+      BYTE_ARRAY_OFFSET,
+      bytes.length,
+      suffixBytes,
+      BYTE_ARRAY_OFFSET,
+      suffixBytes.length
+    )
   }
 
   def toUpperCase(): UTF8String = {
@@ -139,13 +145,14 @@ final class UTF8String extends Ordered[UTF8String] with Serializable {
   override def clone(): UTF8String = new UTF8String().set(this.bytes)
 
   override def compare(other: UTF8String): Int = {
+    val otherBytes = other.getBytes
     UTF8StringMethods.compare(
-      pointer.getBaseObject,
-      pointer.getBaseOffset,
-      pointer.getLengthInBytes,
-      other.pointer.getBaseObject,
-      other.pointer.getBaseOffset,
-      other.pointer.getLengthInBytes
+      bytes,
+      BYTE_ARRAY_OFFSET,
+      bytes.length,
+      otherBytes,
+      BYTE_ARRAY_OFFSET,
+      otherBytes.length
     )
   }
 
@@ -155,7 +162,14 @@ final class UTF8String extends Ordered[UTF8String] with Serializable {
 
   override def equals(other: Any): Boolean = other match {
     case s: UTF8String =>
-      Arrays.equals(bytes, s.getBytes)
+      val otherBytes = s.getBytes
+      otherBytes.length == bytes.length && ByteArrayMethods.arrayEquals(
+        bytes,
+        BYTE_ARRAY_OFFSET,
+        otherBytes,
+        BYTE_ARRAY_OFFSET,
+        otherBytes.length
+      )
     case s: String =>
       // This is only used for Catalyst unit tests
       // fail fast

unsafe/src/main/java/org/apache/spark/unsafe/array/ByteArrayMethods.java

@@ -23,10 +23,34 @@
 
 public class ByteArrayMethods {
 
+  // TODO: there are substantial opportunities for optimization here and we should investigate them.
+  // See the fast comparisions in Guava's UnsignedBytes, for instance:
+  // https://code.google.com/p/guava-libraries/source/browse/guava/src/com/google/common/primitives/UnsignedBytes.java
+
   private ByteArrayMethods() {
     // Private constructor, since this class only contains static methods.
   }
 
+  /**
+   * Optimized  equality check for equal-length byte arrays.
+   * @return true if the arrays are equal, false otherwise
+   */
+  public static boolean arrayEquals(
+      Object leftBaseObject,
+      long leftBaseOffset,
+      Object rightBaseObject,
+      long rightBaseOffset,
+      long arrayLengthInBytes) {
+    for (int i = 0; i < arrayLengthInBytes; i++) {
+      final byte left =
+        PlatformDependent.UNSAFE.getByte(leftBaseObject, leftBaseOffset + i);
+      final byte right =
+        PlatformDependent.UNSAFE.getByte(rightBaseObject, rightBaseOffset + i);
+      if (left != right) return false;
+    }
+    return true;
+  }
+
   /**
    * Optimized byte array equality check for 8-byte-word-aligned byte arrays.
    * @return true if the arrays are equal, false otherwise

unsafe/src/main/java/org/apache/spark/unsafe/string/UTF8StringMethods.java

@@ -18,6 +18,7 @@
 package org.apache.spark.unsafe.string;
 
 import org.apache.spark.unsafe.PlatformDependent;
+import org.apache.spark.unsafe.array.ByteArrayMethods;
 
 import java.io.UnsupportedEncodingException;import java.lang.Object;import java.lang.String;
 
@@ -43,19 +44,57 @@ static long getLengthInBytes(Object baseObject, long baseOffset) {
   public static int compare(
       Object leftBaseObject,
       long leftBaseOffset,
-      int leftBaseLengthInBytes,
+      int leftLengthInBytes,
       Object rightBaseObject,
       long rightBaseOffset,
-      int rightBaseLengthInBytes) {
+      int rightLengthInBytes) {
     int i = 0;
-    while (i < leftBaseLengthInBytes && i < rightBaseLengthInBytes) {
+    while (i < leftLengthInBytes && i < rightLengthInBytes) {
       final byte leftByte = PlatformDependent.UNSAFE.getByte(leftBaseObject, leftBaseOffset + i);
       final byte rightByte = PlatformDependent.UNSAFE.getByte(rightBaseObject, rightBaseOffset + i);
       final int res = leftByte - rightByte;
       if (res != 0) return res;
       i += 1;
     }
-    return leftBaseLengthInBytes - rightBaseLengthInBytes;
+    return leftLengthInBytes - rightLengthInBytes;
+  }
+
+  public static boolean startsWith(
+      Object strBaseObject,
+      long strBaseOffset,
+      int strLengthInBytes,
+      Object prefixBaseObject,
+      long prefixBaseOffset,
+      int prefixLengthInBytes) {
+    if (prefixLengthInBytes > strLengthInBytes) {
+      return false;
+    } {
+      return ByteArrayMethods.arrayEquals(
+        strBaseObject,
+        strBaseOffset,
+        prefixBaseObject,
+        prefixBaseOffset,
+        prefixLengthInBytes);
+    }
+  }
+
+  public static boolean endsWith(
+      Object strBaseObject,
+      long strBaseOffset,
+      int strLengthInBytes,
+      Object suffixBaseObject,
+      long suffixBaseOffset,
+      int suffixLengthInBytes) {
+    if (suffixLengthInBytes > strLengthInBytes) {
+      return false;
+    } {
+      return ByteArrayMethods.arrayEquals(
+        strBaseObject,
+        strBaseOffset + strLengthInBytes - suffixLengthInBytes,
+        suffixBaseObject,
+        suffixBaseOffset,
+        suffixLengthInBytes);
+    }
   }
 
   /**