[SPARK-19104][SQL] Lambda variables in ExternalMapToCatalyst should be global

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

@@ -736,7 +736,27 @@ class CodegenContext {
       // Cannot split these expressions because they are not created from a row object.
       return expressions.mkString("\n")
     }
-    splitExpressions(expressions, "apply", ("InternalRow", row) :: Nil)
+    splitExpressions(row, expressions, Seq.empty)
+  }
+
+  /**
+   * Splits the generated code of expressions into multiple functions, because function has
+   * 64kb code size limit in JVM
+   *
+   * @param row the variable name of row that is used by expressions
+   * @param expressions the codes to evaluate expressions.
+   * @param arguments the additional arguments to the functions.
+   */
+  def splitExpressions(
+      row: String,
+      expressions: Seq[String],
+      arguments: Seq[(String, String)]): String = {
+    if (row == null || currentVars != null) {
+      // Cannot split these expressions because they are not created from a row object.
+      return expressions.mkString("\n")
+    }
+    val params = arguments ++ Seq(("InternalRow", row))
+    splitExpressions(expressions, "apply", params)
   }
 
   /**

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

@@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions.codegen._
+import org.apache.spark.sql.catalyst.expressions.objects.LambdaVariable
 import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData, TypeUtils}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.Platform
@@ -342,19 +343,38 @@ case class CreateNamedStruct(children: Seq[Expression]) extends CreateNamedStruc
     val values = ctx.freshName("values")
     ctx.addMutableState("Object[]", values, s"$values = null;")
 
+    // `splitExpressions` might split codes to multiple functions. The local variables of
+    // `LambdaVariable` can't be accessed in the functions. We need to add them into the parameters
+    // of the functions.
+    val (valExprCodes, valExprParams) = valExprs.map { expr =>
+      val exprCode = expr.genCode(ctx)
+      val lambdaVars = expr.collect {
+        case l: LambdaVariable => l
+      }.flatMap { lambda =>
+        val valueParam = ctx.javaType(lambda.dataType) -> lambda.value
+        if (lambda.isNull == "false") {
+          Seq(valueParam)
+        } else {
+          Seq(valueParam, "boolean" -> lambda.isNull)
+        }
+      }
+      (exprCode, lambdaVars)
+    }.unzip
+
+    val splitFuncsParams = valExprParams.flatten.distinct
+
     ev.copy(code = s"""
       $values = new Object[${valExprs.size}];""" +
       ctx.splitExpressions(
         ctx.INPUT_ROW,
-        valExprs.zipWithIndex.map { case (e, i) =>
-          val eval = e.genCode(ctx)
+        valExprCodes.zipWithIndex.map { case (eval, i) =>
           eval.code + s"""
           if (${eval.isNull}) {
             $values[$i] = null;
           } else {
             $values[$i] = ${eval.value};
           }"""
-        }) +
+        }, splitFuncsParams) +
       s"""
         final InternalRow ${ev.value} = new $rowClass($values);
         $values = null;

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

@@ -39,6 +39,9 @@ case class ComplexClass(seq: SeqClass, list: ListClass, queue: QueueClass)
 
 case class ComplexMapClass(map: MapClass, lhmap: LHMapClass)
 
+case class InnerData(name: String, value: Int)
+case class NestedData(id: Int, param: Map[String, InnerData])
+
 package object packageobject {
   case class PackageClass(value: Int)
 }
@@ -354,4 +357,9 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
     checkDataset(Seq(PackageClass(1)).toDS(), PackageClass(1))
   }
 
+  test("SPARK-19104: lambda variables should work when parent expression splits generated codes") {
+    val data = Seq.tabulate(10)(i => NestedData(1, Map("key" -> InnerData("name", i + 100))))
+    val ds = spark.createDataset(data)
+    checkDataset(ds, data: _*)
+  }
 }