[SPARK-29503][SQL] Copy nested unsafe data in return value of lambda function of MapObject

HeartSaVioR · HeartSaVioR · commit 07514a149e9b · 2019-10-21T23:07:00.000+09:00
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeArrayData.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeArrayData.java
@@ -365,6 +365,11 @@ public UnsafeArrayData copy() {
     return arrayCopy;
   }
 
+  @Override
+  public ArrayData copyUnsafeData(ArrayType dataType) {
+    return copy();
+  }
+
   @Override
   public boolean[] toBooleanArray() {
     boolean[] values = new boolean[numElements];
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeMapData.java
@@ -29,6 +29,7 @@
 import com.esotericsoftware.kryo.io.Output;
 
 import org.apache.spark.sql.catalyst.util.MapData;
+import org.apache.spark.sql.types.MapType;
 import org.apache.spark.unsafe.Platform;
 
 import static org.apache.spark.unsafe.Platform.BYTE_ARRAY_OFFSET;
@@ -135,6 +136,11 @@ public UnsafeMapData copy() {
     return mapCopy;
   }
 
+  @Override
+  public MapData copyUnsafeData(MapType dataType) {
+    return copy();
+  }
+
   @Override
   public void writeExternal(ObjectOutput out) throws IOException {
     byte[] bytes = UnsafeDataUtils.getBytes(baseObject, baseOffset, sizeInBytes);
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
@@ -468,6 +468,11 @@ public UnsafeRow copy() {
     return rowCopy;
   }
 
+  @Override
+  public InternalRow copyUnsafeData(StructType dataType) {
+    return copy();
+  }
+
   /**
    * Creates an empty UnsafeRow from a byte array with specified numBytes and numFields.
    * The returned row is invalid until we call copyFrom on it.
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarMap.java b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarMap.java
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.vectorized;
 
 import org.apache.spark.sql.catalyst.util.MapData;
+import org.apache.spark.sql.types.MapType;
 
 /**
  * Map abstraction in {@link ColumnVector}.
@@ -50,4 +51,9 @@ public ColumnarArray valueArray() {
   public ColumnarMap copy() {
     throw new UnsupportedOperationException();
   }
+
+  @Override
+  public MapData copyUnsafeData(MapType dataType) {
+    throw new UnsupportedOperationException();
+  }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/InternalRow.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/InternalRow.scala
@@ -63,6 +63,53 @@ abstract class InternalRow extends SpecializedGetters with Serializable {
    */
   def copy(): InternalRow
 
+  def copyUnsafeData(dataTypeJson: String): InternalRow = {
+    val dataType = DataType.fromJson(dataTypeJson)
+    assert(dataType.isInstanceOf[StructType])
+    copyUnsafeData(dataType.asInstanceOf[StructType])
+  }
+
+  def copyUnsafeData(dataType: StructType): InternalRow = {
+    def updateRetIfNecessary(
+        ret: InternalRow,
+        field: AnyRef,
+        newField: AnyRef,
+        idx: Int): InternalRow = {
+      var newRet: InternalRow = ret
+      if (field.ne(newField)) {
+        if (newRet == null) newRet = this.copy()
+        if (newField != null) {
+          newRet.update(idx, newField)
+        } else {
+          newRet.setNullAt(idx)
+        }
+      }
+      newRet
+    }
+
+    var ret: InternalRow = null
+    dataType.map(_.dataType).zipWithIndex.foreach {
+      case (ty: StructType, idx) =>
+        val field = getStruct(idx, ty.size)
+        val newField = field.copyUnsafeData(ty)
+        ret = updateRetIfNecessary(ret, field, newField, idx)
+
+      case (ty: ArrayType, idx) =>
+        val field = getArray(idx)
+        val newField = field.copyUnsafeData(ty)
+        ret = updateRetIfNecessary(ret, field, newField, idx)
+
+      case (ty: MapType, idx) =>
+        val field = getMap(idx)
+        val newField = field.copyUnsafeData(ty)
+        ret = updateRetIfNecessary(ret, field, newField, idx)
+
+      case _ =>
+    }
+
+    if (ret != null) ret else this
+  }
+
   /** Returns true if there are any NULL values in this row. */
   def anyNull: Boolean = {
     val len = numFields
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala
@@ -24,6 +24,8 @@ import scala.collection.mutable.Builder
 import scala.reflect.ClassTag
 import scala.util.Try
 
+import org.apache.commons.text.StringEscapeUtils
+
 import org.apache.spark.{SparkConf, SparkEnv}
 import org.apache.spark.serializer._
 import org.apache.spark.sql.Row
@@ -885,13 +887,16 @@ case class MapObjects private(
         )
     }
 
-    // Make a copy of the data if it's unsafe-backed
-    def makeCopyIfInstanceOf(clazz: Class[_ <: Any], value: String) =
-      s"$value instanceof ${clazz.getSimpleName}? ${value}.copy() : $value"
+    // Make a copy of the unsafe data if the result contains any
+    def makeCopyUnsafeData(dataType: DataType, value: String) = {
+      val typeToJson = StringEscapeUtils.escapeJava(StringEscapeUtils.escapeJson(dataType.json))
+      s"""${value}.copyUnsafeData("${typeToJson}")"""
+    }
+
     val genFunctionValue: String = lambdaFunction.dataType match {
-      case StructType(_) => makeCopyIfInstanceOf(classOf[UnsafeRow], genFunction.value)
-      case ArrayType(_, _) => makeCopyIfInstanceOf(classOf[UnsafeArrayData], genFunction.value)
-      case MapType(_, _, _) => makeCopyIfInstanceOf(classOf[UnsafeMapData], genFunction.value)
+      case ty: StructType => makeCopyUnsafeData(ty, genFunction.value)
+      case ty: ArrayType => makeCopyUnsafeData(ty, genFunction.value)
+      case ty: MapType => makeCopyUnsafeData(ty, genFunction.value)
       case _ => genFunction.value
     }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ArrayBasedMapData.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ArrayBasedMapData.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.catalyst.util
 
 import java.util.{Map => JavaMap}
 
+import org.apache.spark.sql.types.{ArrayType, MapType}
+
 /**
  * A simple `MapData` implementation which is backed by 2 arrays.
  *
@@ -35,6 +37,22 @@ class ArrayBasedMapData(val keyArray: ArrayData, val valueArray: ArrayData) exte
   override def toString: String = {
     s"keys: $keyArray, values: $valueArray"
   }
+
+  override def copyUnsafeData(dataType: MapType): MapData = {
+    val keyType = dataType.keyType
+    val valueType = dataType.valueType
+
+    val keyArr = keyArray
+    val valueArr = valueArray
+    val newKeyArray = keyArr.copyUnsafeData(ArrayType(keyType))
+    val newValueArray = valueArr.copyUnsafeData(ArrayType(valueType))
+
+    if (keyArr.ne(newKeyArray) || valueArr.ne(newValueArray)) {
+      new ArrayBasedMapData(newKeyArray, newValueArray)
+    } else {
+      this
+    }
+  }
 }
 
 object ArrayBasedMapData {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ArrayData.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ArrayData.scala
@@ -68,6 +68,69 @@ abstract class ArrayData extends SpecializedGetters with Serializable {
 
   def copy(): ArrayData
 
+  def copyUnsafeData(dataTypeJson: String): ArrayData = {
+    val dataType = DataType.fromJson(dataTypeJson)
+    assert(dataType.isInstanceOf[ArrayType])
+    copyUnsafeData(dataType.asInstanceOf[ArrayType])
+  }
+
+  def copyUnsafeData(dataType: ArrayType): ArrayData = {
+    def updateRetIfNecessary(
+        ret: ArrayData,
+        field: AnyRef,
+        newField: AnyRef,
+        idx: Int): ArrayData = {
+      var newRet: ArrayData = ret
+      if (field.ne(newField)) {
+        if (newRet == null) newRet = this.copy()
+        if (newField != null) {
+          newRet.update(idx, newField)
+        } else {
+          newRet.setNullAt(idx)
+        }
+      }
+      newRet
+    }
+
+    var ret: ArrayData = null
+    dataType.elementType match {
+      case ty: StructType =>
+        val len = numElements()
+        var i = 0
+        while (i < len) {
+          val field = getStruct(i, ty.size)
+          val newField = field.copyUnsafeData(ty)
+          ret = updateRetIfNecessary(ret, field, newField, i)
+          i += 1
+        }
+        if (ret != null) ret else this
+
+      case ty: ArrayType =>
+        val len = numElements()
+        var i = 0
+        while (i < len) {
+          val field = getArray(i)
+          val newField = field.copyUnsafeData(ty)
+          ret = updateRetIfNecessary(ret, field, newField, i)
+          i += 1
+        }
+        if (ret != null) ret else this
+
+      case ty: MapType =>
+        val len = numElements()
+        var i = 0
+        while (i < len) {
+          val field = getMap(i)
+          val newField = field.copyUnsafeData(ty)
+          ret = updateRetIfNecessary(ret, field, newField, i)
+          i += 1
+        }
+        if (ret != null) ret else this
+
+      case _ => this
+    }
+  }
+
   def array: Array[Any]
 
   def toSeq[T](dataType: DataType): IndexedSeq[T] =
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/MapData.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/MapData.scala
@@ -17,7 +17,8 @@
 
 package org.apache.spark.sql.catalyst.util
 
-import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType}
 
 /**
  * This is an internal data representation for map type in Spark SQL. This should not implement
@@ -34,6 +35,14 @@ abstract class MapData extends Serializable {
 
   def copy(): MapData
 
+  def copyUnsafeData(dataTypeJson: String): MapData = {
+    val dataType = DataType.fromJson(dataTypeJson)
+    assert(dataType.isInstanceOf[MapType])
+    copyUnsafeData(dataType.asInstanceOf[MapType])
+  }
+
+  def copyUnsafeData(dataType: MapType): MapData
+
   def foreach(keyType: DataType, valueType: DataType, f: (Any, Any) => Unit): Unit = {
     val length = numElements()
     val keys = keyArray()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameComplexTypeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameComplexTypeSuite.scala
@@ -17,9 +17,16 @@
 
 package org.apache.spark.sql
 
+import scala.collection.mutable
+
 import org.apache.spark.sql.catalyst.DefinedByConstructorParams
+import org.apache.spark.sql.catalyst.expressions.{Expression, GenericRowWithSchema}
+import org.apache.spark.sql.catalyst.expressions.objects.MapObjects
 import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.ArrayType
+
 
 /**
  * A test suite to test DataFrame/SQL functionalities with complex types (i.e. array, struct, map).
@@ -64,6 +71,33 @@ class DataFrameComplexTypeSuite extends QueryTest with SharedSparkSession {
     val ds100_5 = Seq(S100_5()).toDS()
     ds100_5.rdd.count
   }
+
+  test("SPARK-29503 nest unsafe struct inside safe array") {
+    withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") {
+      val exampleDS = spark.sparkContext.parallelize(Seq(Seq(1, 2, 3))).toDF("items")
+
+      // items: Seq[Int] => items.map { item => Seq(Struct(item)) }
+      val result = exampleDS.select(
+        new Column(MapObjects(
+          (item: Expression) => array(struct(new Column(item))).expr,
+          $"items".expr,
+          exampleDS.schema("items").dataType.asInstanceOf[ArrayType].elementType
+        )) as "items"
+      ).collect()
+
+      def getValueInsideDepth(result: Row, index: Int): Int = {
+        // expected output:
+        // WrappedArray([WrappedArray(WrappedArray([1]), WrappedArray([2]), WrappedArray([3]))])
+        result.getSeq[mutable.WrappedArray[_]](0)(index)(0)
+          .asInstanceOf[GenericRowWithSchema].getInt(0)
+      }
+
+      assert(result.size === 1)
+      assert(getValueInsideDepth(result.head, 0) === 1)
+      assert(getValueInsideDepth(result.head, 1) === 2)
+      assert(getValueInsideDepth(result.head, 2) === 3)
+    }
+  }
 }
 
 class S100(
