[SPARK-24216][SQL] Spark TypedAggregateExpression uses getSimpleName that is not safe in scala

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -2715,6 +2715,66 @@ private[spark] object Utils extends Logging {
     HashCodes.fromBytes(secretBytes).toString()
   }
 
+  /**
+   * A safer version than scala obj's getClass.getSimpleName and Utils.getFormattedClassName
+   * which may throw Malformed class name error.
+   * This method mimicks scalatest's getSimpleNameOfAnObjectsClass.
+   */
+  def getSimpleName(fullyQualifiedName: String): String = {
+    stripDollars(parseSimpleName(fullyQualifiedName))
+  }
+
+  /**
+   * Remove the packages from full qualified class name
+   */
+  private def parseSimpleName(fullyQualifiedName: String): String = {
+    // Find last dot position
+    val dotPos = fullyQualifiedName.lastIndexOf('.')
+    // Need to check the dotPos != fullyQualifiedName.length
+    if (dotPos != -1 && dotPos != fullyQualifiedName.length) {
+      fullyQualifiedName.substring(dotPos + 1)
+    } else {
+      fullyQualifiedName
+    }
+  }
+
+  /**
+   * Remove trailing dollar signs from qualified class name,
+   * and return the trailing part after the last dollar sign in the middle
+   */
+  private def stripDollars(s: String): String = {
+    val lastDollarIndex = s.lastIndexOf('$')
+    if (lastDollarIndex < s.length - 1) {
+      // The last char is not a dollar sign
+      if (lastDollarIndex == -1 || !s.contains("$iw")) {
+        // The name does not have dollar sign or is not an intepreter
+        // generated class, so we should return the full string
+        s
+      } else {
+        // The class name is intepreter generated,
+        // return the part after the last dollar sign
+        // This is the same behavior as getClass.getSimpleName
+        s.substring(lastDollarIndex + 1)
+      }
+    }
+    else {
+      // The last char is a dollar sign
+      // Find last non-dollar char
+      val lastNonDollarChar = s.reverse.find(_ != '$')
+      lastNonDollarChar match {
+        case None => s
+        case Some(c) =>
+          val lastNonDollarIndex = s.lastIndexOf(c)
+          if (lastNonDollarIndex == -1) {
+            s
+          } else {
+            // Strip the trailing dollar signs
+            // Invoke stripDollars again to get the simple name
+            stripDollars(s.substring(0, lastNonDollarIndex + 1))
+          }
+      }
+    }
+  }
 }
 
 private[util] object CallerContext extends Logging {

core/src/test/scala/org/apache/spark/util/UtilsSuite.scala

@@ -1168,6 +1168,32 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties with Logging {
       Utils.checkAndGetK8sMasterUrl("k8s://foo://host:port")
     }
   }
+
+  object A {
+    class B
+  }
+
+  test("Safe getSimpleName") {
+    val fullname1 = "org.apache.spark.TestClass$MyClass"
+    assert(Utils.getSimpleName(fullname1) === "TestClass$MyClass")
+
+    val fullname2 = "org.apache.spark.TestClass$MyClass$"
+    assert(Utils.getSimpleName(fullname2) === "TestClass$MyClass")
+
+    val fullname3 = "org.apache.spark.TestClass$MyClass$1$"
+    assert(Utils.getSimpleName(fullname3) === "TestClass$MyClass$1")
+
+    val fullname4 = "TestClass$MyClass$1$"
+    assert(Utils.getSimpleName(fullname4) === "TestClass$MyClass$1")
+
+    val fullname5 = "$iwC$iwC$$iwC$$iwC$TestClass$MyClass$"
+    assert(Utils.getSimpleName(fullname5) === "MyClass")
+
+    intercept[java.lang.InternalError] {
+      classOf[A.B].getSimpleName
+    }
+    assert(Utils.getSimpleName(classOf[A.B].getName) === "UtilsSuite$A$B")
+  }
 }
 
 private class SimpleExtension

sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala

@@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateSafeProjection
 import org.apache.spark.sql.catalyst.expressions.objects.Invoke
 import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.types._
+import org.apache.spark.util.Utils
 
 object TypedAggregateExpression {
   def apply[BUF : Encoder, OUT : Encoder](
@@ -109,7 +110,9 @@ trait TypedAggregateExpression extends AggregateFunction {
     s"$nodeName($input)"
   }
 
-  override def nodeName: String = aggregator.getClass.getSimpleName.stripSuffix("$")
+  // aggregator.getClass.getSimpleName can cause Malformed class name error,
+  // call safer `Utils.getSimpleName` instead
+  override def nodeName: String = Utils.getSimpleName(aggregator.getClass.getName);
 }
 
 // TODO: merge these 2 implementations once we refactor the `AggregateFunction` interface.