Github user ueshin commented on a diff in the pull request:
https://github.com/apache/spark/pull/22013#discussion_r210161746
--- Diff:
sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala ---
@@ -2302,6 +2302,97 @@ 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") {
+ 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)
+ ).toDF("j")
+
+ val dfExample3 = Seq(
+ Map[Int, Boolean](25 -> true, 26 -> false)
+ ).toDF("x")
+
+ val dfExample4 = Seq(
+ Map[Array[Int], Boolean](Array(1, 2) -> false)
+ ).toDF("y")
+
+
+ def testMapOfPrimitiveTypesCombination(): Unit = {
+ checkAnswer(dfExample1.selectExpr("transform_keys(i, (k, v) -> k +
v)"),
+ Seq(Row(Map(2 -> 1, 18 -> 9, 16 -> 8, 14 -> 7))))
+
+ checkAnswer(dfExample2.selectExpr("transform_keys(j, " +
+ "(k, v) -> map_from_arrays(ARRAY(1, 2, 3), ARRAY('one', 'two',
'three'))[k])"),
+ Seq(Row(Map("one" -> 1.0, "two" -> 1.4, "three" -> 1.7))))
+
+ checkAnswer(dfExample2.selectExpr("transform_keys(j, (k, v) ->
CAST(v * 2 AS BIGINT) + k)"),
+ Seq(Row(Map(3 -> 1.0, 4 -> 1.4, 6 -> 1.7))))
+
+ checkAnswer(dfExample2.selectExpr("transform_keys(j, (k, v) -> k +
v)"),
+ Seq(Row(Map(2.0 -> 1.0, 3.4 -> 1.4, 4.7 -> 1.7))))
+
+ checkAnswer(dfExample3.selectExpr("transform_keys(x, (k, v) -> k %
2 = 0 OR v)"),
+ Seq(Row(Map(true -> true, true -> false))))
+
+ checkAnswer(dfExample3.selectExpr("transform_keys(x, (k, v) -> if(v,
2 * k, 3 * k))"),
+ Seq(Row(Map(50 -> true, 78 -> false))))
+
+ checkAnswer(dfExample3.selectExpr("transform_keys(x, (k, v) -> if(v,
2 * k, 3 * k))"),
+ Seq(Row(Map(50 -> true, 78 -> false))))
+
+ checkAnswer(dfExample4.selectExpr("transform_keys(y, (k, v) ->
array_contains(k, 3) AND v)"),
+ Seq(Row(Map(false -> false))))
+ }
+ // Test with local relation, the Project will be evaluated without
codegen
+ testMapOfPrimitiveTypesCombination()
+ dfExample1.cache()
+ dfExample2.cache()
+ dfExample3.cache()
+ dfExample4.cache()
+ // Test with cached relation, the Project will be evaluated with
codegen
+ testMapOfPrimitiveTypesCombination()
+ }
+
+ test("transform keys function - Invalid lambda functions and
exceptions") {
+ val dfExample1 = Seq(
+ Map[Int, Int](1 -> 1, 9 -> 9, 8 -> 8, 7 -> 7)
+ ).toDF("i")
+
+ val dfExample2 = Seq(
+ Map[String, String]("a" -> "b")
+ ).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 )")
+ }
+ 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)")
+ }
+ assert(ex2.getMessage.contains(
+ "The number of lambda function arguments '3' does not match"))
--- End diff --
nit: indent
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
