Brush up UTF8StringBuilder

Author: maropu

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UTF8StringBuilder.java

@@ -19,34 +19,71 @@
 
 import java.nio.charset.StandardCharsets;
 
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.ByteArrayMethods;
 import org.apache.spark.unsafe.types.UTF8String;
 
 /**
- * A helper class to write `UTF8String`, `String`, and `byte[]` data into an internal buffer
- * and get a final concatenated string.
+ * A helper class to write `UTF8String`, `String`, and `byte[]` data into an internal byte buffer
+ * and get written data as `UTF8String`.
  */
 public class UTF8StringBuilder {
 
-  private StringBuilder buffer;
+  private static final int ARRAY_MAX = ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH;
+
+  private byte[] buffer;
+  private int cursor = Platform.BYTE_ARRAY_OFFSET;
 
   public UTF8StringBuilder() {
-    this.buffer = new StringBuilder();
+    // Since initial buffer size is 16 in `StringBuilder`, we set the same size here
+    this.buffer = new byte[16];
+  }
+
+  // Grows the buffer by at least `neededSize`
+  private void grow(int neededSize) {
+    if (neededSize > ARRAY_MAX - totalSize()) {
+      throw new UnsupportedOperationException(
+        "Cannot grow internal buffer by size " + neededSize + " because the size after growing " +
+          "exceeds size limitation " + ARRAY_MAX);
+    }
+    final int length = totalSize() + neededSize;
+    if (buffer.length < length) {
+      int newLength = length < ARRAY_MAX / 2 ? length * 2 : ARRAY_MAX;
+      final byte[] tmp = new byte[newLength];
+      Platform.copyMemory(
+        buffer,
+        Platform.BYTE_ARRAY_OFFSET,
+        tmp,
+        Platform.BYTE_ARRAY_OFFSET,
+        totalSize());
+      buffer = tmp;
+    }
   }
 
   public void append(UTF8String value) {
-    buffer.append(value);
+    grow(value.numBytes());
+    value.writeToMemory(buffer, cursor);
+    cursor += value.numBytes();
   }
 
   public void append(String value) {
-    buffer.append(value);
+    append(value.getBytes(StandardCharsets.UTF_8));
   }
 
   public void append(byte[] value) {
-    buffer.append(new String(value, StandardCharsets.UTF_8));
+    grow(value.length);
+    Platform.copyMemory(value, Platform.BYTE_ARRAY_OFFSET, buffer, cursor, value.length);
+    cursor += value.length;
+  }
+
+  public UTF8String toUTF8String() {
+    final int len = totalSize();
+    final byte[] bytes = new byte[len];
+    Platform.copyMemory(buffer, Platform.BYTE_ARRAY_OFFSET, bytes, Platform.BYTE_ARRAY_OFFSET, len);
+    return UTF8String.fromBytes(bytes);
   }
 
-  @Override
-  public String toString() {
-    return buffer.toString();
+  public int totalSize() {
+    return cursor - Platform.BYTE_ARRAY_OFFSET;
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala

@@ -199,36 +199,6 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
 
   // [[func]] assumes the input is no longer null because eval already does the null check.
   @inline private[this] def buildCast[T](a: Any, func: T => Any): Any = func(a.asInstanceOf[T])
-  @inline private[this] def buildWriter[T](
-      a: Any, buffer: UTF8StringBuilder, writer: (T, UTF8StringBuilder) => Unit): Unit = {
-    writer(a.asInstanceOf[T], buffer)
-  }
-
-  private[this] def buildElemWriter(
-      from: DataType): (Any, UTF8StringBuilder) => Unit = from match {
-    case BinaryType => buildWriter[Array[Byte]](_, _, (b, buf) => buf.append(b))
-    case StringType => buildWriter[UTF8String](_, _, (b, buf) => buf.append(b))
-    case DateType => buildWriter[Int](_, _,
-      (d, buf) => buf.append(DateTimeUtils.dateToString(d)))
-    case TimestampType => buildWriter[Long](_, _,
-      (t, buf) => buf.append(DateTimeUtils.timestampToString(t)))
-    case ar: ArrayType =>
-      buildWriter[ArrayData](_, _, (array, buf) => {
-         buf.append("[")
-        if (array.numElements > 0) {
-          val writeElemToBuffer = buildElemWriter(ar.elementType)
-          writeElemToBuffer(array.get(0, ar.elementType), buf)
-          var i = 1
-          while (i < array.numElements) {
-            buf.append(", ")
-            writeElemToBuffer(array.get(i, ar.elementType), buf)
-            i += 1
-          }
-        }
-        buf.append("]")
-      })
-    case _ => buildWriter[Any](_, _, (o, buf) => buf.append(String.valueOf(o)))
-  }
 
   // UDFToString
   private[this] def castToString(from: DataType): Any => Any = from match {
@@ -241,17 +211,17 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
         val res = new UTF8StringBuilder
         res.append("[")
         if (array.numElements > 0) {
-          val writeElemToBuffer = buildElemWriter(ar.elementType)
-          writeElemToBuffer(array.get(0, ar.elementType), res)
+          val toUTF8String = castToString(ar.elementType)
+          res.append(toUTF8String(array.get(0, ar.elementType)).asInstanceOf[UTF8String])
           var i = 1
           while (i < array.numElements) {
             res.append(", ")
-            writeElemToBuffer(array.get(i, ar.elementType), res)
+            res.append(toUTF8String(array.get(i, ar.elementType)).asInstanceOf[UTF8String])
             i += 1
           }
         }
         res.append("]")
-        UTF8String.fromString(res.toString)
+        res.toUTF8String
       })
     case _ => buildCast[Any](_, o => UTF8String.fromString(o.toString))
   }
@@ -709,7 +679,7 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
           s"""
              |$bufferClass $bufferTerm = new $bufferClass();
              |$writeArrayToBuffer($c, $bufferTerm);
-             |$evPrim = UTF8String.fromString($bufferTerm.toString());
+             |$evPrim = $bufferTerm.toUTF8String();
            """.stripMargin
         }
       case _ =>