From 2d76099497299a756418878564a3a44ddf13f482 Mon Sep 17 00:00:00 2001
From: Dima <dimanowq@gmail.com>
Date: Fri, 13 Dec 2024 14:35:04 +0100
Subject: [PATCH] Reply to comments from @MaxGekk

---
 .../sql/catalyst/analysis/CheckAnalysis.scala |  6 ++---
 .../org/apache/spark/sql/DataFrameSuite.scala | 23 +++++++------------
 2 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
index 9a6b630094ab9..06e8e26c0cebb 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
@@ -868,7 +868,7 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
               summary = j.origin.context.summary)
 
           // TODO: although map type is not orderable, technically map type should be able to be
-          //   used in equality comparison, remove this type check once we support it.
+          // used in equality comparison, remove this type check once we support it.
           case o if mapColumnInSetOperation(o).isDefined =>
             val mapCol = mapColumnInSetOperation(o).get
             o.failAnalysis(
@@ -891,7 +891,7 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
             o.failAnalysis(
               errorClass = "UNSUPPORTED_FEATURE.PARTITION_BY_VARIANT",
               messageParameters = Map(
-                "expr" -> variantExpr.sql,
+                "expr" -> toSQLExpr(variantExpr),
                 "dataType" -> toSQLType(variantExpr.dataType)))
 
           case o if mapExprInPartitionExpression(o).isDefined =>
@@ -899,7 +899,7 @@ trait CheckAnalysis extends PredicateHelper with LookupCatalog with QueryErrorsB
             o.failAnalysis(
               errorClass = "UNSUPPORTED_FEATURE.PARTITION_BY_MAP",
               messageParameters = Map(
-                "expr" -> mapExpr.sql,
+                "expr" -> toSQLExpr(mapExpr),
                 "dataType" -> toSQLType(mapExpr.dataType)))
 
           case o if o.expressions.exists(!_.deterministic) &&
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 4dae686947e49..3e5bb49c59e8c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -316,7 +316,7 @@ class DataFrameSuite extends QueryTest
       exception = intercept[AnalysisException](df.repartition(5, col("v"))),
       condition = "UNSUPPORTED_FEATURE.PARTITION_BY_VARIANT",
       parameters = Map(
-        "expr" -> "v",
+        "expr" -> "\"v\"",
         "dataType" -> "\"VARIANT\"")
     )
     // nested variant column
@@ -324,7 +324,7 @@ class DataFrameSuite extends QueryTest
       exception = intercept[AnalysisException](df.repartition(5, col("s"))),
       condition = "UNSUPPORTED_FEATURE.PARTITION_BY_VARIANT",
       parameters = Map(
-        "expr" -> "s",
+        "expr" -> "\"s\"",
         "dataType" -> "\"STRUCT<v: VARIANT NOT NULL>\"")
     )
     // variant producing expression
@@ -333,7 +333,7 @@ class DataFrameSuite extends QueryTest
         intercept[AnalysisException](df.repartition(5, parse_json(col("id").cast("string")))),
       condition = "UNSUPPORTED_FEATURE.PARTITION_BY_VARIANT",
       parameters = Map(
-        "expr" -> "parse_json(CAST(id AS STRING))",
+        "expr" -> "\"parse_json(CAST(id AS STRING))\"",
         "dataType" -> "\"VARIANT\"")
     )
     // Partitioning by non-variant column works
@@ -350,7 +350,7 @@ class DataFrameSuite extends QueryTest
         exception = intercept[AnalysisException](sql("SELECT * FROM tv DISTRIBUTE BY v")),
         condition = "UNSUPPORTED_FEATURE.PARTITION_BY_VARIANT",
         parameters = Map(
-          "expr" -> "tv.v",
+          "expr" -> "\"v\"",
           "dataType" -> "\"VARIANT\""),
         context = ExpectedContext(
           fragment = "DISTRIBUTE BY v",
@@ -361,7 +361,7 @@ class DataFrameSuite extends QueryTest
         exception = intercept[AnalysisException](sql("SELECT * FROM tv DISTRIBUTE BY s")),
         condition = "UNSUPPORTED_FEATURE.PARTITION_BY_VARIANT",
         parameters = Map(
-          "expr" -> "tv.s",
+          "expr" -> "\"s\"",
           "dataType" -> "\"STRUCT<v: VARIANT NOT NULL>\""),
         context = ExpectedContext(
           fragment = "DISTRIBUTE BY s",
@@ -378,7 +378,7 @@ class DataFrameSuite extends QueryTest
       exception = intercept[AnalysisException](df.repartition(5, col("m"))),
       condition = "UNSUPPORTED_FEATURE.PARTITION_BY_MAP",
       parameters = Map(
-        "expr" -> "m",
+        "expr" -> "\"m\"",
         "dataType" -> "\"MAP<BIGINT, BIGINT>\"")
     )
     // map producing expression
@@ -386,16 +386,9 @@ class DataFrameSuite extends QueryTest
       exception = intercept[AnalysisException](df.repartition(5, map(col("id"), col("id")))),
       condition = "UNSUPPORTED_FEATURE.PARTITION_BY_MAP",
       parameters = Map(
-        "expr" -> "map(id, id)",
+        "expr" -> "\"map(id, id)\"",
         "dataType" -> "\"MAP<BIGINT, BIGINT>\"")
     )
-    // Partitioning by non-map column works
-    try {
-      df.repartition(5, col("id")).collect()
-    } catch {
-      case e: Exception =>
-        fail(s"Expected no exception to be thrown but an exception was thrown: ${e.getMessage}")
-    }
     // SQL
     withTempView("tv") {
       df.createOrReplaceTempView("tv")
@@ -403,7 +396,7 @@ class DataFrameSuite extends QueryTest
         exception = intercept[AnalysisException](sql("SELECT * FROM tv DISTRIBUTE BY m")),
         condition = "UNSUPPORTED_FEATURE.PARTITION_BY_MAP",
         parameters = Map(
-          "expr" -> "tv.m",
+          "expr" -> "\"m\"",
           "dataType" -> "\"MAP<BIGINT, BIGINT>\""),
         context = ExpectedContext(
           fragment = "DISTRIBUTE BY m",