[SPARK-22825][SQL] Fix incorrect results of Casting Array to String

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

@@ -0,0 +1,89 @@
+/*
+ * 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.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 byte buffer
+ * and get written data as `UTF8String`.
+ */
+public class UTF8StringBuilder {
+
+  private static final int ARRAY_MAX = ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH;
+
+  private byte[] buffer;
+  private int cursor = Platform.BYTE_ARRAY_OFFSET;
+
+  public UTF8StringBuilder() {
+    // 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) {
+    grow(value.numBytes());
+    value.writeToMemory(buffer, cursor);
+    cursor += value.numBytes();
+  }
+
+  public void append(String value) {
+    append(value.getBytes(StandardCharsets.UTF_8));
+  }
+
+  public void append(byte[] value) {
+    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);
+  }
+
+  public int totalSize() {
+    return cursor - Platform.BYTE_ARRAY_OFFSET;
+  }
+}

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

@@ -206,6 +206,23 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
     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 UTF8StringBuilder
+        res.append("[")
+        if (array.numElements > 0) {
+          val toUTF8String = castToString(ar.elementType)
+          res.append(toUTF8String(array.get(0, ar.elementType)).asInstanceOf[UTF8String])
+          var i = 1
+          while (i < array.numElements) {
+            res.append(", ")
+            res.append(toUTF8String(array.get(i, ar.elementType)).asInstanceOf[UTF8String])
+            i += 1
+          }
+        }
+        res.append("]")
+        res.toUTF8String
+      })
     case _ => buildCast[Any](_, o => UTF8String.fromString(o.toString))
   }
 
@@ -597,6 +614,52 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
     """
   }
 
+  private[this] def writeElemToBufferCode(
+      dataType: DataType,
+      buffer: String,
+      elemTerm: String,
+      ctx: CodegenContext): String = dataType match {
+    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(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 = classOf[UTF8StringBuilder].getName
+    val bufferTerm = ctx.freshName("bufferTerm")
+    def writeElemCode(elemTerm: String) = {
+      writeElemToBufferCode(ar.elementType, bufferTerm, elemTerm, ctx)
+    }
+    def writeToBufferCode(i: String) = {
+      val elemTerm = ctx.freshName("elemTerm")
+      s"""
+         |${ctx.javaType(ar.elementType)} $elemTerm = ${ctx.getValue(arTerm, ar.elementType, i)};
+         |${writeElemCode(elemTerm)};
+       """.stripMargin
+    }
+    ctx.addNewFunction(writeArrayToBuffer,
+      s"""
+         |private void $writeArrayToBuffer(ArrayData $arTerm, $bufferClass $bufferTerm) {
+         |  $bufferTerm.append("[");
+         |  if ($arTerm.numElements() > 0) {
+         |    ${writeToBufferCode("0")}
+         |    for (int $loopIndex = 1; $loopIndex < $arTerm.numElements(); $loopIndex++) {
+         |      $bufferTerm.append(", ");
+         |      ${writeToBufferCode(loopIndex)}
+         |    }
+         |  }
+         |  $bufferTerm.append("]");
+         |}
+       """.stripMargin)
+  }
+
   private[this] def castToStringCode(from: DataType, ctx: CodegenContext): CastFunction = {
     from match {
       case BinaryType =>
@@ -608,6 +671,17 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
         val tz = ctx.addReferenceObj("timeZone", timeZone)
         (c, evPrim, evNull) => s"""$evPrim = UTF8String.fromString(
           org.apache.spark.sql.catalyst.util.DateTimeUtils.timestampToString($c, $tz));"""
+      case ar: ArrayType =>
+        (c, evPrim, evNull) => {
+          val bufferTerm = ctx.freshName("bufferTerm")
+          val bufferClass = classOf[UTF8StringBuilder].getName
+          val writeArrayToBuffer = codegenWriteArrayToBuffer(ar, ctx)
+          s"""
+             |$bufferClass $bufferTerm = new $bufferClass();
+             |$writeArrayToBuffer($c, $bufferTerm);
+             |$evPrim = $bufferTerm.toUTF8String();
+           """.stripMargin
+        }
       case _ =>
         (c, evPrim, evNull) => s"$evPrim = UTF8String.fromString(String.valueOf($c));"
     }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala

@@ -853,4 +853,27 @@ class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
     cast("2", LongType).genCode(ctx)
     assert(ctx.inlinedMutableStates.length == 0)
   }
+
+  test("SPARK-22825 Cast array to string") {
+    val ret1 = cast(Literal.create(Array(1, 2, 3, 4, 5)), StringType)
+    checkEvaluation(ret1, "[1, 2, 3, 4, 5]")
+    val ret2 = cast(Literal.create(Array("ab", "cde", "f")), StringType)
+    checkEvaluation(ret2, "[ab, cde, f]")
+    val ret3 = cast(Literal.create(Array("ab".getBytes, "cde".getBytes, "f".getBytes)), StringType)
+    checkEvaluation(ret3, "[ab, cde, f]")
+    val ret4 = cast(
+      Literal.create(Array("2014-12-03", "2014-12-04", "2014-12-06").map(Date.valueOf)),
+      StringType)
+    checkEvaluation(ret4, "[2014-12-03, 2014-12-04, 2014-12-06]")
+    val ret5 = cast(
+      Literal.create(Array("2014-12-03 13:01:00", "2014-12-04 15:05:00").map(Timestamp.valueOf)),
+      StringType)
+    checkEvaluation(ret5, "[2014-12-03 13:01:00, 2014-12-04 15:05:00]")
+    val ret6 = cast(Literal.create(Array(Array(1, 2, 3), Array(4, 5))), StringType)
+    checkEvaluation(ret6, "[[1, 2, 3], [4, 5]]")
+    val ret7 = cast(
+      Literal.create(Array(Array(Array("a"), Array("b", "c")), Array(Array("d")))),
+      StringType)
+    checkEvaluation(ret7, "[[[a], [b, c]], [[d]]]")
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -20,16 +20,14 @@ package org.apache.spark.sql
 import java.io.File
 import java.math.MathContext
 import java.net.{MalformedURLException, URL}
-import java.sql.Timestamp
+import java.sql.{Date, Timestamp}
 import java.util.concurrent.atomic.AtomicBoolean
 
 import org.apache.spark.{AccumulatorSuite, SparkException}
 import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart}
 import org.apache.spark.sql.catalyst.util.StringUtils
 import org.apache.spark.sql.execution.aggregate
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec}
-import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation}
-import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, CartesianProductExec, SortMergeJoinExec}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
@@ -2775,4 +2773,48 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       }
     }
   }
+
+  test("SPARK-22825 Cast array to string") {
+    Seq("true", "false").foreach { codegen =>
+      withSQLConf("spark.sql.codegen.wholeStage" -> codegen) {
+        withTable("t") {
+          Seq(Seq(0, 1, 2, 3, 4)).toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[0, 1, 2, 3, 4]"))
+        }
+        withTable("t") {
+          Seq(Seq("ab", "cde", "f")).toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[ab, cde, f]"))
+        }
+        withTable("t") {
+          Seq(Seq("ab".getBytes, "cde".getBytes, "f".getBytes)).toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[ab, cde, f]"))
+        }
+        withTable("t") {
+          Seq(Seq("2014-12-03", "2014-12-04", "2014-12-06").map(Date.valueOf))
+            .toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[2014-12-03, 2014-12-04, 2014-12-06]"))
+        }
+        withTable("t") {
+          Seq(Seq("2014-12-03 13:01:00", "2014-12-04 15:05:00").map(Timestamp.valueOf))
+            .toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[2014-12-03 13:01:00, 2014-12-04 15:05:00]"))
+        }
+        withTable("t") {
+          Seq(Seq(Seq(1, 2), Seq(3), Seq(4, 5, 6))).toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[[1, 2], [3], [4, 5, 6]]"))
+        }
+        withTable("t") {
+          Seq(Seq(Seq(Seq("a"), Seq("b", "c")), Seq(Seq("d")))).toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[[[a], [b, c]], [[d]]]"))
+        }
+      }
+    }
+  }
 }