From 1fae9d927208f304a03eb7611ccbde707b16432c Mon Sep 17 00:00:00 2001
From: GideonPotok <g.potok4@gmail.com>
Date: Wed, 22 May 2024 14:53:59 -0400
Subject: [PATCH] review feedback

---
 .../catalyst/expressions/aggregate/Mode.scala   | 17 ++++++++---------
 .../sql/CollationSQLExpressionsSuite.scala      | 12 ++++++------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala
index 908e533410709..ada20934def93 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.trees.UnaryLike
 import org.apache.spark.sql.catalyst.types.PhysicalDataType
 import org.apache.spark.sql.catalyst.util.{CollationFactory, GenericArrayData, UnsafeRowUtils}
 import org.apache.spark.sql.errors.QueryCompilationErrors
-import org.apache.spark.sql.types.{AbstractDataType, AnyDataType, ArrayType, BooleanType, DataType, StringType, StructType}
+import org.apache.spark.sql.types.{AbstractDataType, AnyDataType, ArrayType, BooleanType, DataType, StringType}
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.collection.OpenHashMap
 
@@ -53,15 +53,14 @@ case class Mode(
     val defaultCheck = super.checkInputDataTypes()
     if (defaultCheck.isFailure) {
       defaultCheck
+    } else if (UnsafeRowUtils.isBinaryStable(child.dataType) ||
+      child.dataType.isInstanceOf[StringType]) {
+      defaultCheck
     } else {
-      child.dataType match {
-        case _: StructType | _: ArrayType if !UnsafeRowUtils.isBinaryStable(child.dataType) =>
-          TypeCheckResult.TypeCheckFailure(
-            s"Input to function mode was a complex type" +
-              s" with non-binary collated fields," +
-              s" which is not yet supported by mode.")
-        case _ => TypeCheckResult.TypeCheckSuccess
-      }
+      TypeCheckResult.TypeCheckFailure(
+        "The input to the function 'mode' was a complex type with non-binary collated fields," +
+          " which are currently not supported by 'mode'."
+      )
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala
index 7b6653e0c1e67..bdda678761972 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala
@@ -1510,8 +1510,8 @@ class CollationSQLExpressionsSuite
           // Input to function mode was a complex type with strings collated on non-binary
           // collations, which is not yet supported.. SQLSTATE: 42K09; line 1 pos 13;
           val params = Seq(("sqlExpr", "\"mode(i)\""),
-            ("msg", "Input to function mode was a complex type with non-binary collated fields" +
-              ", which is not yet supported by mode."),
+            ("msg", "The input to the function 'mode' was a complex type with non-binary collated" +
+              " fields, which are currently not supported by 'mode'."),
             ("hint", "")).toMap
           checkError(
             exception = intercept[AnalysisException] {
@@ -1559,8 +1559,8 @@ class CollationSQLExpressionsSuite
           // Input to function mode was a complex type with strings collated on non-binary
           // collations, which is not yet supported.. SQLSTATE: 42K09; line 1 pos 13;
           val params = Seq(("sqlExpr", "\"mode(i)\""),
-            ("msg", "Input to function mode was a complex type with non-binary collated fields" +
-              ", which is not yet supported by mode."),
+            ("msg", "The input to the function 'mode' was a complex type with non-binary collated" +
+              " fields, which are currently not supported by 'mode'."),
             ("hint", "")).toMap
           checkError(
             exception = intercept[AnalysisException] {
@@ -1606,8 +1606,8 @@ class CollationSQLExpressionsSuite
         val query = s"SELECT lower(element_at(element_at(mode(i), 1).s1.a2, 1)) FROM ${tableName}"
         if(t.collationId == "utf8_binary_lcase" || t.collationId == "unicode_ci") {
           val params = Seq(("sqlExpr", "\"mode(i)\""),
-            ("msg", "Input to function mode was a complex type with non-binary collated fields" +
-              ", which is not yet supported by mode."),
+            ("msg", "The input to the function 'mode' was a complex type with non-binary collated" +
+              " fields, which are currently not supported by 'mode'."),
             ("hint", "")).toMap
           checkError(
             exception = intercept[AnalysisException] {