From 3b5af0e4982e7c1353452be5e92d7c7e40cc4600 Mon Sep 17 00:00:00 2001 From: "fqaiser94@gmail.com" Date: Mon, 3 Aug 2020 11:16:29 -0400 Subject: [PATCH 1/3] WithFields Expression is not foldable --- .../catalyst/expressions/complexTypeCreator.scala | 2 -- .../org/apache/spark/sql/ColumnExpressionSuite.scala | 12 ++++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala index cf7cc3a5e16ff..563ce7133a3dc 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala @@ -563,8 +563,6 @@ case class WithFields( override def dataType: StructType = evalExpr.dataType.asInstanceOf[StructType] - override def foldable: Boolean = structExpr.foldable && valExprs.forall(_.foldable) - override def nullable: Boolean = structExpr.nullable override def prettyName: String = "with_fields" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala index 761632e76b165..06bec88b2f028 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala @@ -1420,4 +1420,16 @@ class ColumnExpressionSuite extends QueryTest with SharedSparkSession { }.getMessage should include("No such struct field b in a, B") } } + + test("SPARK-32521: withField with statically evaluable arguments should succeed") { + checkAnswerAndSchema( + sql("SELECT named_struct('a', 1, 'b', 2) a").select($"a".withField("c", lit(3)).as("a")), + Row(Row(1, 2, 3)) :: Nil, + StructType(Seq( + StructField("a", StructType(Seq( + StructField("a", IntegerType, nullable = false), + StructField("b", IntegerType, nullable = false), + StructField("c", IntegerType, nullable = false))), + nullable = false)))) + } } From 0adf6203088ddb8415cac26285fb0c9c38170ca5 Mon Sep 17 00:00:00 2001 From: "fqaiser94@gmail.com" Date: Mon, 3 Aug 2020 12:17:33 -0400 Subject: [PATCH 2/3] add tests for user facing examples --- .../spark/sql/ColumnExpressionSuite.scala | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala index 06bec88b2f028..802e17a43bafc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala @@ -1421,6 +1421,38 @@ class ColumnExpressionSuite extends QueryTest with SharedSparkSession { } } + test("withField user-facing examples") { + checkAnswer( + sql("SELECT named_struct('a', 1, 'b', 2) struct_col") + .select($"struct_col".withField("c", lit(3))), + Row(Row(1, 2, 3))) + + checkAnswer( + sql("SELECT named_struct('a', 1, 'b', 2) struct_col") + .select($"struct_col".withField("b", lit(3))), + Row(Row(1, 3))) + + checkAnswer( + sql("SELECT CAST(NULL AS struct) struct_col") + .select($"struct_col".withField("c", lit(3))), + Row(null)) + + checkAnswer( + sql("SELECT named_struct('a', 1, 'b', 2, 'b', 3) struct_col") + .select($"struct_col".withField("b", lit(100))), + Row(Row(1, 100, 100))) + + checkAnswer( + sql("SELECT named_struct('a', named_struct('a', 1, 'b', 2)) struct_col") + .select($"struct_col".withField("a.c", lit(3))), + Row(Row(Row(1, 2, 3)))) + + intercept[AnalysisException] { + sql("SELECT named_struct('a', named_struct('b', 1), 'a', named_struct('c', 2)) struct_col") + .select($"struct_col".withField("a.c", lit(3))) + }.getMessage should include("Ambiguous reference to fields") + } + test("SPARK-32521: withField with statically evaluable arguments should succeed") { checkAnswerAndSchema( sql("SELECT named_struct('a', 1, 'b', 2) a").select($"a".withField("c", lit(3)).as("a")), From 0e1c02835e989279a887b8fc81956152f99f9f2b Mon Sep 17 00:00:00 2001 From: "fqaiser94@gmail.com" Date: Tue, 4 Aug 2020 07:41:15 -0400 Subject: [PATCH 3/3] remove duplicated test case --- .../org/apache/spark/sql/ColumnExpressionSuite.scala | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala index 802e17a43bafc..36c10874cc183 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala @@ -1452,16 +1452,4 @@ class ColumnExpressionSuite extends QueryTest with SharedSparkSession { .select($"struct_col".withField("a.c", lit(3))) }.getMessage should include("Ambiguous reference to fields") } - - test("SPARK-32521: withField with statically evaluable arguments should succeed") { - checkAnswerAndSchema( - sql("SELECT named_struct('a', 1, 'b', 2) a").select($"a".withField("c", lit(3)).as("a")), - Row(Row(1, 2, 3)) :: Nil, - StructType(Seq( - StructField("a", StructType(Seq( - StructField("a", IntegerType, nullable = false), - StructField("b", IntegerType, nullable = false), - StructField("c", IntegerType, nullable = false))), - nullable = false)))) - } }