review comments

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

@@ -523,14 +523,14 @@ case class TransformKeys(
     MapType(function.dataType, map.valueType, map.valueContainsNull)
   }
 
-  @transient val MapType(keyType, valueType, valueContainsNull) = argument.dataType
+  @transient lazy val MapType(keyType, valueType, valueContainsNull) = argument.dataType
 
   override def bind(f: (Expression, Seq[(DataType, Boolean)]) => LambdaFunction): TransformKeys = {
     copy(function = f(function, (keyType, false) :: (valueType, valueContainsNull) :: Nil))
   }
 
   @transient lazy val (keyVar, valueVar) = {
-    val LambdaFunction(
+    @transient lazy val LambdaFunction(
     _, (keyVar: NamedLambdaVariable) :: (valueVar: NamedLambdaVariable) :: Nil, _) = function
     (keyVar, valueVar)
   }
@@ -544,7 +544,7 @@ case class TransformKeys(
       keyVar.value.set(map.keyArray().get(i, keyVar.dataType))
       valueVar.value.set(map.valueArray().get(i, valueVar.dataType))
       val result = f.eval(inputRow)
-      if (result ==  null) {
+      if (result == null) {
         throw new RuntimeException("Cannot use null as map key!")
       }
       resultKeys.update(i, result)
@@ -554,7 +554,7 @@ case class TransformKeys(
   }
 
   override def prettyName: String = "transform_keys"
-  }
+}
 
 /**
  * Merges two given maps into a single map by applying function to the pair of values with

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

@@ -328,7 +328,7 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper
       MapType(StringType, StringType, valueContainsNull = true))
     val as2 = Literal.create(null,
       MapType(StringType, StringType, valueContainsNull = false))
-    val asn = Literal.create(Map.empty[StringType, StringType],
+    val as3 = Literal.create(Map.empty[StringType, StringType],
       MapType(StringType, StringType, valueContainsNull = true))
 
     val concatValue: (Expression, Expression) => Expression = (k, v) => Concat(Seq(k, v))
@@ -340,16 +340,16 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper
     checkEvaluation(
       transformKeys(transformKeys(as0, concatValue), concatValue),
       Map("axyxy" -> "xy", "bbyzyz" -> "yz", "ccczxzx" -> "zx"))
-    checkEvaluation(transformKeys(asn, concatValue), Map.empty[String, String])
+    checkEvaluation(transformKeys(as3, concatValue), Map.empty[String, String])
     checkEvaluation(
-      transformKeys(transformKeys(asn, concatValue), convertKeyToKeyLength),
+      transformKeys(transformKeys(as3, concatValue), convertKeyToKeyLength),
       Map.empty[Int, String])
     checkEvaluation(transformKeys(as0, convertKeyToKeyLength),
       Map(2 -> "xy", 3 -> "yz", 4 -> "zx"))
     checkEvaluation(transformKeys(as1, convertKeyToKeyLength),
       Map(2 -> "xy", 3 -> "yz", 4 -> null))
     checkEvaluation(transformKeys(as2, convertKeyToKeyLength), null)
-    checkEvaluation(transformKeys(asn, convertKeyToKeyLength), Map.empty[Int, String])
+    checkEvaluation(transformKeys(as3, convertKeyToKeyLength), Map.empty[Int, String])
 
     val ax0 = Literal.create(
       Map(1 -> "x", 2 -> "y", 3 -> "z"),

sql/core/src/test/resources/sql-tests/inputs/higher-order-functions.sql

@@ -53,8 +53,8 @@ select exists(ys, y -> y > 30) as v from nested;
 select exists(cast(null as array<int>), y -> y > 30) as v;
 
 create or replace temporary view nested as values
-  (1, map(1,1,2,2,3,3)),
-  (2, map(4,4,5,5,6,6))
+  (1, map(1, 1, 2, 2, 3, 3)),
+  (2, map(4, 4, 5, 5, 6, 6))
   as t(x, ys);
 
 -- Identity Transform Keys in a map

sql/core/src/test/resources/sql-tests/results/higher-order-functions.sql.out

@@ -167,8 +167,8 @@ NULL
 
 -- !query 17
 create or replace temporary view nested as values
-  (1, map(1,1,2,2,3,3)),
-  (2, map(4,4,5,5,6,6))
+  (1, map(1, 1, 2, 2, 3, 3)),
+  (2, map(4, 4, 5, 5, 6, 6))
   as t(x, ys)
 -- !query 17 schema
 struct<>

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

@@ -2302,13 +2302,13 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
     assert(ex5.getMessage.contains("function map_zip_with does not support ordering on type map"))
   }
 
-  test("transform keys function - test various primitive data types combinations") {
+  test("transform keys function - primitive data types") {
     val dfExample1 = Seq(
       Map[Int, Int](1 -> 1, 9 -> 9, 8 -> 8, 7 -> 7)
     ).toDF("i")
 
     val dfExample2 = Seq(
-      Map[Int, Double](1 -> 1.0E0, 2 -> 1.4E0, 3 -> 1.7E0)
+      Map[Int, Double](1 -> 1.0, 2 -> 1.40, 3 -> 1.70)
     ).toDF("j")
 
     val dfExample3 = Seq(
@@ -2357,34 +2357,37 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
   }
 
   test("transform keys function - Invalid lambda functions and exceptions") {
+
     val dfExample1 = Seq(
-      Map[Int, Int](1 -> 1, 9 -> 9, 8 -> 8, 7 -> 7)
+      Map[String, String]("a" -> null)
     ).toDF("i")
 
     val dfExample2 = Seq(
-      Map[String, String]("a" -> "b")
+      Seq(1, 2, 3, 4)
     ).toDF("j")
 
-    val dfExample3 = Seq(
-      Map[String, String]("a" -> null)
-    ).toDF("x")
-
     def testInvalidLambdaFunctions(): Unit = {
       val ex1 = intercept[AnalysisException] {
-        dfExample1.selectExpr("transform_keys(i, k -> k )")
+        dfExample1.selectExpr("transform_keys(i, k -> k)")
       }
       assert(ex1.getMessage.contains("The number of lambda function arguments '1' does not match"))
 
       val ex2 = intercept[AnalysisException] {
-        dfExample2.selectExpr("transform_keys(j, (k, v, x) -> k + 1)")
+        dfExample1.selectExpr("transform_keys(i, (k, v, x) -> k + 1)")
       }
       assert(ex2.getMessage.contains(
-      "The number of lambda function arguments '3' does not match"))
+        "The number of lambda function arguments '3' does not match"))
 
       val ex3 = intercept[RuntimeException] {
-        dfExample3.selectExpr("transform_keys(x, (k, v) -> v)").show()
+        dfExample1.selectExpr("transform_keys(i, (k, v) -> v)").show()
       }
       assert(ex3.getMessage.contains("Cannot use null as map key!"))
+
+      val ex4 = intercept[AnalysisException] {
+        dfExample2.selectExpr("transform_keys(j, (k, v) -> k + 1)")
+      }
+      assert(ex4.getMessage.contains(
+        "data type mismatch: argument 1 requires map type"))
     }
 
     testInvalidLambdaFunctions()