You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2026/09/06 08:59:06 UTC
(spark) branch branch-4.x updated: [SPARK-48701][SQL][PS] Make PandasMode collation-aware
This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 47a34bb9ccb7 [SPARK-48701][SQL][PS] Make PandasMode collation-aware
47a34bb9ccb7 is described below
commit 47a34bb9ccb78cc2ecdf4536e47e23b8bd83dbb3
Author: Vinod KC <vi...@gmail.com>
AuthorDate: Sun Sep 6 17:58:37 2026 +0900
[SPARK-48701][SQL][PS] Make PandasMode collation-aware
### What changes were proposed in this pull request?
Make the internal `PandasMode` aggregate (`pandas_mode`) collation-aware.
This PR extracts the folding logic into a shared `ModeCollationAware` trait mixed into both `Mode` and `PandasMode`, and calls `getCollationAwareBuffer` in `PandasMode.eval` (with a null-key guard, since `PandasMode` stores a null key when `ignoreNA = false`). Eval-only change; `Mode`'s behavior is unchanged.
### Why are the changes needed?
`pandas_mode` (behind Pandas API on Spark `Series.mode()` / `DataFrame.mode()`) silently returns wrong results on non-binary collated strings: collation-equal values (e.g. `'b'` and `'B'` under `UTF8_LCASE`) are counted as separate keys.
Example : a `UTF8_LCASE` column with values `a, a, b, B, B` (so `'b'` and `'B'` are equal: counts `a`=2, `b/B`=3, mode should be the single `b/B` group):
```python
sdf = spark.sql("""
SELECT CAST(c AS STRING COLLATE UTF8_LCASE) AS s
FROM VALUES ('a'), ('a'), ('b'), ('B'), ('B') AS t(c)
""")
print(sorted(sdf.pandas_api()["s"].mode().to_numpy().tolist()))
```
Before: ['B', 'a'] # wrong: a=2 ties B=2, b counted separately (2 modes)
After: ['b'] # correct: b/B fold to one group of count 3, winning (1 mode)
### Does this PR introduce _any_ user-facing change?
Yes. pandas_mode (Pandas API on Spark Series.mode() / DataFrame.mode()) now returns the correct mode(s) on columns with a non-binary collation, instead of counting collation-equal strings as distinct values. Results on default (UTF8_BINARY) collation and non-string types are unchanged.
### How was this patch tested?
Added unit tests in `CollationAggregationSuite`
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude code.
Closes #58233 from vinodkc/pandas-mode-collation.
Authored-by: Vinod KC <vi...@gmail.com>
Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
(cherry picked from commit f37971cd0add244fcc2a6debd456b8cf85ae67f6)
Signed-off-by: Hyukjin Kwon <hy...@databricks.com>
---
.../sql/catalyst/expressions/aggregate/Mode.scala | 126 ++++++++++++---------
.../sql/collation/CollationAggregationSuite.scala | 117 ++++++++++++++++++-
2 files changed, 190 insertions(+), 53 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 b87c0ede1ab8..d119b8c276f8 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
@@ -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)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala
index d62782f7f6d2..b3e9b482f456 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala
@@ -17,7 +17,9 @@
package org.apache.spark.sql.collation
-import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, Row}
+import java.util.Locale
+
+import org.apache.spark.sql.{AnalysisException, Column, DataFrame, QueryTest, Row}
import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode
import org.apache.spark.sql.execution.WholeStageCodegenExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
@@ -25,6 +27,7 @@ import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAg
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.StringType
import org.apache.spark.unsafe.types.CalendarInterval
class CollationAggregationSuite
@@ -524,6 +527,39 @@ 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("collect_set on UTF8_BINARY strings is unchanged (case-sensitive)") {
val tblName = "collect_set_binary"
withTable(tblName) {
@@ -537,4 +573,83 @@ class CollationAggregationSuite
Seq(Row(Seq("FOO", "bar", "foo"))))
}
}
+
+ 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"))
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org