[SPARK-48701][SQL] Make PandasMode collation-aware

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala

@@ -44,53 +44,22 @@ private[aggregate] object ModeKeyNormalizer {
   }
 }
 
-case class Mode(
-    child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0,
-    reverseOpt: Option[Boolean] = None)
-  extends TypedAggregateWithHashMapAsBuffer with ImplicitCastInputTypes
-    with SupportsOrderingWithinGroup with UnaryLike[Expression] {
-
-  def this(child: Expression) = this(child, 0, 0)
-
-  def this(child: Expression, reverse: Boolean) = {
-    this(child, 0, 0, Some(reverse))
-  }
-
-  // Returns null for empty inputs
-  override def nullable: Boolean = true
-
-  override def dataType: DataType = child.dataType
-
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
-
-  override def prettyName: String = "mode"
-
-  @transient private lazy val keyNormalizer: Any => Any =
-    ModeKeyNormalizer.forType(child.dataType)
-
-  override def update(
-      buffer: OpenHashMap[AnyRef, Long],
-      input: InternalRow): OpenHashMap[AnyRef, Long] = {
-    val key = child.eval(input)
-
-    if (key != null) {
-      buffer.changeValue(keyNormalizer(key).asInstanceOf[AnyRef], 1L, _ + 1L)
-    }
-    buffer
-  }
-
-  override def merge(
-      buffer: OpenHashMap[AnyRef, Long],
-      other: OpenHashMap[AnyRef, Long]): OpenHashMap[AnyRef, Long] = {
-    other.foreach { case (key, count) =>
-      buffer.changeValue(key, count, _ + count)
-    }
-    buffer
-  }
+/**
+ * Shared collation-aware buffer folding for the `mode` family of aggregates.
+ *
+ * The aggregation buffer keeps the original (float-normalized) values as keys, so under a
+ * non-binary collation two collation-equal strings (e.g. 'b' and 'B' under UTF8_LCASE) land
+ * in separate buffer entries. At eval time [[getCollationAwareBuffer]] folds those entries
+ * into one group keyed on the collation key, summing their counts and keeping the first-seen
+ * original value as the group's representative. Binary-stable types skip the fold entirely
+ * (zero overhead) and complex types (struct/array/map) are handled recursively.
+ *
+ * Mixed into both [[Mode]] and [[PandasMode]] so the two share one implementation.
+ */
+private[aggregate] trait ModeCollationAware { self: Expression =>
+  protected def child: Expression
 
-  private def getCollationAwareBuffer(
+  protected def getCollationAwareBuffer(
       childDataType: DataType,
       buffer: OpenHashMap[AnyRef, Long]): Iterable[(AnyRef, Long)] = {
     def groupAndReduceBuffer(groupingFunction: AnyRef => _): Iterable[(AnyRef, Long)] = {
@@ -101,7 +70,11 @@ case class Mode(
         childDataType: DataType): Option[AnyRef => _] = {
       childDataType match {
         case _ if UnsafeRowUtils.isBinaryStable(child.dataType) => None
-        case _ => Some(collationAwareTransform(_, childDataType))
+        // A null key is kept as its own group: PandasMode may store one when `ignoreNA`
+        // is false, and passing null to collationAwareTransform would throw. Mode never
+        // stores a null key, so its behavior is unchanged.
+        case _ => Some((key: AnyRef) =>
+          if (key == null) null else collationAwareTransform(key, childDataType))
       }
     }
     determineBufferingFunction(childDataType).map(groupAndReduceBuffer).getOrElse(buffer)
@@ -148,6 +121,53 @@ case class Mode(
     }
     transformedKeys.zip(transformedValues).toMap
   }
+}
+
+case class Mode(
+    child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0,
+    reverseOpt: Option[Boolean] = None)
+  extends TypedAggregateWithHashMapAsBuffer with ImplicitCastInputTypes
+    with SupportsOrderingWithinGroup with UnaryLike[Expression] with ModeCollationAware {
+
+  def this(child: Expression) = this(child, 0, 0)
+
+  def this(child: Expression, reverse: Boolean) = {
+    this(child, 0, 0, Some(reverse))
+  }
+
+  // Returns null for empty inputs
+  override def nullable: Boolean = true
+
+  override def dataType: DataType = child.dataType
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+
+  override def prettyName: String = "mode"
+
+  @transient private lazy val keyNormalizer: Any => Any =
+    ModeKeyNormalizer.forType(child.dataType)
+
+  override def update(
+      buffer: OpenHashMap[AnyRef, Long],
+      input: InternalRow): OpenHashMap[AnyRef, Long] = {
+    val key = child.eval(input)
+
+    if (key != null) {
+      buffer.changeValue(keyNormalizer(key).asInstanceOf[AnyRef], 1L, _ + 1L)
+    }
+    buffer
+  }
+
+  override def merge(
+      buffer: OpenHashMap[AnyRef, Long],
+      other: OpenHashMap[AnyRef, Long]): OpenHashMap[AnyRef, Long] = {
+    other.foreach { case (key, count) =>
+      buffer.changeValue(key, count, _ + count)
+    }
+    buffer
+  }
 
   override def eval(buffer: OpenHashMap[AnyRef, Long]): Any = {
     if (buffer.isEmpty) {
@@ -224,7 +244,6 @@ case class Mode(
     copy(child = newChild)
 }
 
-// TODO: SPARK-48701: PandasMode (all collations)
 // scalastyle:off line.size.limit
 @ExpressionDescription(
   usage = """
@@ -307,7 +326,7 @@ case class PandasMode(
     ignoreNA: Boolean = true,
     mutableAggBufferOffset: Int = 0,
     inputAggBufferOffset: Int = 0) extends TypedAggregateWithHashMapAsBuffer
-  with ImplicitCastInputTypes with UnaryLike[Expression] {
+  with ImplicitCastInputTypes with UnaryLike[Expression] with ModeCollationAware {
 
   def this(child: Expression) = this(child, true, 0, 0)
 
@@ -353,11 +372,14 @@ case class PandasMode(
       return new GenericArrayData(Array.empty)
     }
 
+    // Fold collation-equal keys into one group before selecting the mode(s), so that under a
+    // non-binary collation values such as 'b' and 'B' (equal under UTF8_LCASE) are counted
+    // together. Binary-stable types return the buffer unchanged (no overhead). Mirrors Mode.eval.
+    val collationAwareBuffer = getCollationAwareBuffer(child.dataType, buffer)
+
     val modes = collection.mutable.ArrayBuffer.empty[AnyRef]
     var maxCount = -1L
-    val iter = buffer.iterator
-    while (iter.hasNext) {
-      val (key, count) = iter.next()
+    collationAwareBuffer.foreach { case (key, count) =>
       if (maxCount < count) {
         modes.clear()
         modes.append(key)

sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala

@@ -17,10 +17,14 @@
 
 package org.apache.spark.sql.collation
 
-import org.apache.spark.sql.Row
+import java.util.Locale
+
+import org.apache.spark.sql.{Column, DataFrame, Row}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.functions.{col, lit, lower}
 import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.StringType
 
 class CollationAggregationSuite
   extends SharedSparkSession
@@ -86,4 +90,116 @@ class CollationAggregationSuite
       }
     }
   }
+
+  // `pandas_mode` (backing pandas-on-Spark Series.mode / DataFrame.mode) is an internal
+  // expression, so it is invoked here via `Column.internalFn` rather than SQL. Its second
+  // argument is `ignoreNA` (true = drop NULLs, mirroring pandas `dropna`). It returns an
+  // array of all the most frequent values; the array order is unspecified, so the helpers
+  // below normalize before asserting.
+
+  private def pandasMode(df: DataFrame, colName: String, ignoreNA: Boolean): Seq[AnyRef] = {
+    df.select(Column.internalFn("pandas_mode", col(colName), lit(ignoreNA)))
+      .collect().head.getSeq[AnyRef](0)
+  }
+
+  // Case-insensitive, order-insensitive view of a string-valued mode result.
+  private def normalizedStringModes(modes: Seq[AnyRef]): Set[String] =
+    modes.map {
+      case null => null
+      case s: String => s.toLowerCase(Locale.ROOT)
+    }.toSet
+
+  test("SPARK-48701: pandas_mode is collation-aware for non-binary collations") {
+    Seq("UTF8_LCASE", "UNICODE_CI").foreach { collation =>
+      withTable("t") {
+        sql(s"CREATE TABLE t (c STRING COLLATE $collation) USING parquet")
+        // Spread across partitions to also exercise the partial-buffer merge path.
+        sql("INSERT INTO t VALUES ('a'), ('a'), ('b'), ('B')")
+        val modes = pandasMode(spark.table("t").repartition(4), "c", ignoreNA = true)
+        // 'b' and 'B' are collation-equal, so they fold into one group of 2 that ties
+        // 'a' (also 2). Both are modes. Without folding 'a' (2) would win alone.
+        assert(modes.length == 2, s"$collation: expected two modes, got $modes")
+        assert(normalizedStringModes(modes) == Set("a", "b"))
+      }
+    }
+  }
+
+  test("SPARK-48701: pandas_mode keeps binary collation unchanged") {
+    withTable("t") {
+      // Default UTF8_BINARY: 'b' and 'B' are distinct, so 'a' (2) is the sole mode.
+      sql("CREATE TABLE t (c STRING) USING parquet")
+      sql("INSERT INTO t VALUES ('a'), ('a'), ('b'), ('B')")
+      val modes = pandasMode(spark.table("t").repartition(4), "c", ignoreNA = true)
+      assert(modes == Seq("a"))
+    }
+  }
+
+  test("SPARK-48701: pandas_mode collation-aware with ignoreNA controlling NULLs") {
+    withTable("t") {
+      sql("CREATE TABLE t (c STRING COLLATE UTF8_LCASE) USING parquet")
+      sql("INSERT INTO t VALUES ('a'), ('b'), ('B'), (null), (null), (null)")
+      val df = spark.table("t").repartition(4)
+
+      // ignoreNA = true: NULLs dropped. 'b'/'B' fold to 2, outvoting 'a' (1).
+      val dropped = pandasMode(df, "c", ignoreNA = true)
+      assert(dropped.length == 1)
+      assert(normalizedStringModes(dropped) == Set("b"))
+
+      // ignoreNA = false: the null key is preserved as its own group (count 3) and wins.
+      // This also exercises the null guard in getCollationAwareBuffer's folding.
+      val kept = pandasMode(df, "c", ignoreNA = false)
+      assert(kept == Seq(null))
+    }
+  }
+
+  test("SPARK-48701: pandas_mode collation-aware for collated string nested in struct") {
+    withTable("t") {
+      sql("CREATE TABLE t (c STRUCT<f: STRING COLLATE UTF8_LCASE>) USING parquet")
+      sql(
+        """INSERT INTO t VALUES (named_struct('f', 'a')), (named_struct('f', 'a')),
+          |  (named_struct('f', 'b')), (named_struct('f', 'B'))""".stripMargin)
+      val modes = pandasMode(spark.table("t").repartition(4), "c", ignoreNA = true)
+        .map(_.asInstanceOf[Row])
+      // Same fold as the top-level case, applied to the collated struct field.
+      assert(modes.length == 2)
+      assert(modes.map(_.getString(0).toLowerCase(Locale.ROOT)).toSet == Set("a", "b"))
+    }
+  }
+
+  // `mode` (the public aggregate) is already collation-aware; these tests guard that
+  // behavior, which currently has no coverage (the original tests were removed with
+  // CollationSQLExpressionsSuite by SPARK-51067). Unlike pandas_mode, `mode` returns a
+  // single value and shares the same eval-time folding via ModeCollationAware.
+
+  test("SPARK-47353: mode is collation-aware for non-binary collations") {
+    // Buffer counts a=3, b=2, B=2. Under a case-insensitive collation 'b' and 'B' fold
+    // into one group of 4 that outvotes 'a' (3); under binary collation 'a' (3) wins.
+    Seq(
+      ("UTF8_BINARY", "a"),
+      ("UTF8_LCASE", "b"),
+      ("UNICODE", "a"),
+      ("UNICODE_CI", "b")).foreach { case (collation, expected) =>
+      withTable("t") {
+        sql(s"CREATE TABLE t (c STRING COLLATE $collation) USING parquet")
+        sql("INSERT INTO t VALUES ('a'), ('a'), ('a'), ('b'), ('b'), ('B'), ('B')")
+        val df = sql("SELECT mode(c) AS m FROM t")
+        // The result keeps the input's collated string type.
+        assert(df.schema("m").dataType.sameType(StringType(collation)))
+        // Representative case within a folded group is unspecified; normalize with lower.
+        checkAnswer(df.select(lower(col("m"))), Row(expected))
+      }
+    }
+  }
+
+  test("SPARK-47353: mode is collation-aware for collated string nested in struct") {
+    withTable("t") {
+      sql("CREATE TABLE t (c STRUCT<f: STRING COLLATE UTF8_LCASE>) USING parquet")
+      sql(
+        """INSERT INTO t VALUES (named_struct('f', 'a')), (named_struct('f', 'a')),
+          |  (named_struct('f', 'a')), (named_struct('f', 'b')), (named_struct('f', 'b')),
+          |  (named_struct('f', 'B')), (named_struct('f', 'B'))""".stripMargin)
+      // The collated struct field folds: {b, B} (4) outvotes {a} (3).
+      checkAnswer(sql("SELECT lower(mode(c).f) FROM t"), Row("b"))
+    }
+  }
 }