[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,80 @@
+/*
+ * 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;
+    }
+  }
+
+  private int totalSize() {
+    return cursor - Platform.BYTE_ARRAY_OFFSET;
+  }
+
+  public void append(UTF8String value) {
+    grow(value.numBytes());
+    value.writeToMemory(buffer, cursor);
+    cursor += value.numBytes();
+  }
+
+  public void append(String value) {
+    append(UTF8String.fromString(value));
+  }
+
+  public UTF8String build() {
+    return UTF8String.fromBytes(buffer, 0, totalSize());
+  }
+}

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

@@ -206,6 +206,28 @@ 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 ArrayType(et, _) =>
+      buildCast[ArrayData](_, array => {
+        val builder = new UTF8StringBuilder
+        builder.append("[")
+        if (array.numElements > 0) {
+          val toUTF8String = castToString(et)
+          if (!array.isNullAt(0)) {
+            builder.append(toUTF8String(array.get(0, et)).asInstanceOf[UTF8String])
+          }
+          var i = 1
+          while (i < array.numElements) {
+            builder.append(",")
+            if (!array.isNullAt(i)) {
+              builder.append(" ")
+              builder.append(toUTF8String(array.get(i, et)).asInstanceOf[UTF8String])
+            }
+            i += 1
+          }
+        }
+        builder.append("]")
+        builder.build()
+      })
     case _ => buildCast[Any](_, o => UTF8String.fromString(o.toString))
   }
 
@@ -597,6 +619,44 @@ case class Cast(child: Expression, dataType: DataType, timeZoneId: Option[String
     """
   }
 
+  private[this] def codegenWriteArrayElemCode(et: DataType, ctx: CodegenContext): String = {
+    val elementToStringCode = castToStringCode(et, ctx)
+    val funcName = ctx.freshName("elementToString")
+    val elementToStringFunc = ctx.addNewFunction(funcName,
+      s"""
+         |private UTF8String $funcName(${ctx.javaType(et)} element) {
+         |  UTF8String elementStr = null;
+         |  ${elementToStringCode("element", "elementStr", null /* resultIsNull won't be used */)}
+         |  return elementStr;
+         |}
+       """.stripMargin)
+
+    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")
+    ctx.addNewFunction(writeArrayToBuffer,
+      s"""
+         |private void $writeArrayToBuffer(ArrayData $arTerm, $bufferClass $bufferTerm) {
+         |  $bufferTerm.append("[");
+         |  if ($arTerm.numElements() > 0) {
+         |    if (!$arTerm.isNullAt(0)) {
+         |      $bufferTerm.append($elementToStringFunc(${ctx.getValue(arTerm, et, "0")}));
+         |    }
+         |    for (int $loopIndex = 1; $loopIndex < $arTerm.numElements(); $loopIndex++) {
+         |      $bufferTerm.append(",");
+         |      if (!$arTerm.isNullAt($loopIndex)) {
+         |        $bufferTerm.append(" ");
+         |        $bufferTerm.append($elementToStringFunc(${ctx.getValue(arTerm, et, loopIndex)}));
+         |      }
+         |    }
+         |  }
+         |  $bufferTerm.append("]");
+         |}
+       """.stripMargin)
+  }
+
   private[this] def castToStringCode(from: DataType, ctx: CodegenContext): CastFunction = {
     from match {
       case BinaryType =>
@@ -608,6 +668,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 ArrayType(et, _) =>
+        (c, evPrim, evNull) => {
+          val bufferTerm = ctx.freshName("bufferTerm")
+          val bufferClass = classOf[UTF8StringBuilder].getName
+          val writeArrayElemCode = codegenWriteArrayElemCode(et, ctx)
+          s"""
+             |$bufferClass $bufferTerm = new $bufferClass();
+             |$writeArrayElemCode($c, $bufferTerm);
+             |$evPrim = $bufferTerm.build();
+           """.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,29 @@ 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", null, "c")), StringType)
+    checkEvaluation(ret3, "[ab,, c]")
+    val ret4 = cast(Literal.create(Array("ab".getBytes, "cde".getBytes, "f".getBytes)), StringType)
+    checkEvaluation(ret4, "[ab, cde, f]")
+    val ret5 = cast(
+      Literal.create(Array("2014-12-03", "2014-12-04", "2014-12-06").map(Date.valueOf)),
+      StringType)
+    checkEvaluation(ret5, "[2014-12-03, 2014-12-04, 2014-12-06]")
+    val ret6 = cast(
+      Literal.create(Array("2014-12-03 13:01:00", "2014-12-04 15:05:00").map(Timestamp.valueOf)),
+      StringType)
+    checkEvaluation(ret6, "[2014-12-03 13:01:00, 2014-12-04 15:05:00]")
+    val ret7 = cast(Literal.create(Array(Array(1, 2, 3), Array(4, 5))), StringType)
+    checkEvaluation(ret7, "[[1, 2, 3], [4, 5]]")
+    val ret8 = cast(
+      Literal.create(Array(Array(Array("a"), Array("b", "c")), Array(Array("d")))),
+      StringType)
+    checkEvaluation(ret8, "[[[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,53 @@ 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", null, "c")).toDF("a").write.saveAsTable("t")
+          val df = sql("SELECT CAST(a AS STRING) FROM t")
+          checkAnswer(df, Row("[ab,, c]"))
+        }
+        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]]]"))
+        }
+      }
+    }
+  }
 }