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 b87c0ede1ab81..d119b8c276f86 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 b894abe614761..4c56e18cb33ee 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,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) 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) 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")) + } + } }