[SPARK-29503][SQL] Copy result row from RowWriter in GenerateUnsafeProjection when expr is lambdaFunction in MapObject

Author: HeartSaVioR

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala

@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.expressions.codegen
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.catalyst.expressions.objects.LambdaVariable
 import org.apache.spark.sql.types._
 
 /**
@@ -301,6 +302,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     // Evaluate all the subexpression.
     val evalSubexpr = ctx.subexprFunctionsCode
 
+
     val writeExpressions = writeExpressionsToBuffer(
       ctx, ctx.INPUT_ROW, exprEvals, exprSchemas, rowWriter, isTopLevel = true)
 
@@ -310,8 +312,19 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
          |$evalSubexpr
          |$writeExpressions
        """.stripMargin
+
+    val runInsideLoop = expressions.exists {
+      case e: LambdaVariable => true
+      case _ => false
+    }
+    val extractValueCode = if (runInsideLoop) {
+      s"$rowWriter.getRow().copy()"
+    } else {
+      s"$rowWriter.getRow()"
+    }
+
     // `rowWriter` is declared as a class field, so we can access it directly in methods.
-    ExprCode(code, FalseLiteral, JavaCode.expression(s"$rowWriter.getRow()", classOf[UnsafeRow]))
+    ExprCode(code, FalseLiteral, JavaCode.expression(extractValueCode, classOf[UnsafeRow]))
   }
 
   protected def canonicalize(in: Seq[Expression]): Seq[Expression] =

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

@@ -17,9 +17,15 @@
 
 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 +70,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(