Add UTF8StringBuilder

Author: maropu

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

@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.codegen;
+
+import java.nio.charset.StandardCharsets;
+
+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.
+ */
+public class UTF8StringBuilder {
+
+  private StringBuilder buffer;
+
+  public UTF8StringBuilder() {
+    this.buffer = new StringBuilder();
+  }
+
+  public void append(UTF8String value) {
+    buffer.append(value);
+  }
+
+  public void append(String value) {
+    buffer.append(value);
+  }
+
+  public void append(byte[] value) {
+    buffer.append(new String(value, StandardCharsets.UTF_8));
+  }
+
+  @Override
+  public String toString() {
+    return buffer.toString();
+  }
+}

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

@@ -199,30 +199,59 @@ 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 {
     case BinaryType => buildCast[Array[Byte]](_, UTF8String.fromBytes)
-    case StringType => buildCast[UTF8String](_, identity)
     case DateType => buildCast[Int](_, d => UTF8String.fromString(DateTimeUtils.dateToString(d)))
     case TimestampType => buildCast[Long](_,
       t => UTF8String.fromString(DateTimeUtils.timestampToString(t, timeZone)))
     case ar: ArrayType =>
       buildCast[ArrayData](_, array => {
-        val res = new StringBuilder
+        val res = new UTF8StringBuilder
         res.append("[")
         if (array.numElements > 0) {
-          val toStringFunc = castToString(ar.elementType)
-          res.append(toStringFunc(array.get(0, ar.elementType)))
+          val writeElemToBuffer = buildElemWriter(ar.elementType)
+          writeElemToBuffer(array.get(0, ar.elementType), res)
           var i = 1
           while (i < array.numElements) {
             res.append(", ")
-            res.append(toStringFunc(array.get(i, ar.elementType)))
+            writeElemToBuffer(array.get(i, ar.elementType), res)
             i += 1
           }
         }
         res.append("]")
-        UTF8String.fromString(res.toString())
+        UTF8String.fromString(res.toString)
       })
     case _ => buildCast[Any](_, o => UTF8String.fromString(o.toString))
   }
@@ -620,21 +649,20 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
       buffer: String,
       elemTerm: String,
       ctx: CodegenContext): String = dataType match {
-    case BinaryType => s"$buffer.append(new String($elemTerm))"
-    case StringType => s"$buffer.append(new String($elemTerm.getBytes()))"
+    case BinaryType | StringType => s"$buffer.append($elemTerm)"
     case DateType => s"""$buffer.append(
       org.apache.spark.sql.catalyst.util.DateTimeUtils.dateToString($elemTerm))"""
     case TimestampType => s"""$buffer.append(
       org.apache.spark.sql.catalyst.util.DateTimeUtils.timestampToString($elemTerm))"""
     case ar: ArrayType => s"${codegenWriteArrayToBuffer(ar, ctx)}($elemTerm, $buffer)"
-    case _ => s"$buffer.append($elemTerm)"
+    case _ => s"$buffer.append(String.valueOf($elemTerm))"
   }
 
   private[this] def codegenWriteArrayToBuffer(ar: ArrayType, ctx: CodegenContext): String = {
     val loopIndex = ctx.freshName("loopIndex")
     val writeArrayToBuffer = ctx.freshName("writeArrayToBuffer")
     val arTerm = ctx.freshName("arTerm")
-    val bufferClass = "java.lang.StringBuilder"
+    val bufferClass = classOf[UTF8StringBuilder].getName
     val bufferTerm = ctx.freshName("bufferTerm")
     def writeElemCode(elemTerm: String) = {
       writeElemToBufferCode(ar.elementType, bufferTerm, elemTerm, ctx)
@@ -676,7 +704,7 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
       case ar: ArrayType =>
         (c, evPrim, evNull) => {
           val bufferTerm = ctx.freshName("bufferTerm")
-          val bufferClass = "java.lang.StringBuilder"
+          val bufferClass = classOf[UTF8StringBuilder].getName
           val writeArrayToBuffer = codegenWriteArrayToBuffer(ar, ctx)
           s"""
              |$bufferClass $bufferTerm = new $bufferClass();