[SPARK-58069][SQL] approx_top_k returns incorrect/garbage sort-key bytes instead of actual values for collated strings

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java

@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions;
+
+import java.util.Arrays;
+
+import org.apache.datasketches.common.ArrayOfItemsSerDe;
+import org.apache.datasketches.common.ArrayOfStringsSerDe;
+import org.apache.datasketches.memory.Memory;
+
+import org.apache.spark.sql.catalyst.util.CollationFactory;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/**
+ * SerDe for {@link CollatedString} items used by {@code approx_top_k} over non-binary collated
+ * strings (SPARK-58069).
+ * <p>
+ * Only the {@code original} value is written to the wire, reusing the plain-string format of
+ * {@link ArrayOfStringsSerDe}; the collation key is recomputed on read from {@code collationId}.
+ * As a result the serialized bytes are exactly a string array (the extra per-item key is derived,
+ * not persisted), and the on-wire layout stays identical to the plain-string sketch.
+ */
+public class ArrayOfCollatedStringsSerDe extends ArrayOfItemsSerDe<CollatedString> {
+
+    private final int collationId;
+    private final ArrayOfStringsSerDe stringSerDe = new ArrayOfStringsSerDe();
+
+    public ArrayOfCollatedStringsSerDe(int collationId) {
+        this.collationId = collationId;
+    }
+
+    private CollatedString wrap(String original) {
+        String key = CollationFactory.getCollationKey(
+            UTF8String.fromString(original), collationId).toString();
+        return new CollatedString(key, original);
+    }
+
+    @Override
+    public byte[] serializeToByteArray(CollatedString item) {
+        return stringSerDe.serializeToByteArray(item.original());
+    }
+
+    @Override
+    public byte[] serializeToByteArray(CollatedString[] items) {
+        String[] originals = new String[items.length];
+        for (int i = 0; i < items.length; i++) {
+            originals[i] = items[i].original();
+        }
+        return stringSerDe.serializeToByteArray(originals);
+    }
+
+    @Override
+    public CollatedString[] deserializeFromMemory(Memory mem, long offsetBytes, int numItems) {
+        String[] originals = stringSerDe.deserializeFromMemory(mem, offsetBytes, numItems);
+        return Arrays.stream(originals).map(this::wrap).toArray(CollatedString[]::new);
+    }
+
+    @Override
+    public int sizeOf(CollatedString item) {
+        return stringSerDe.sizeOf(item.original());
+    }
+
+    @Override
+    public int sizeOf(Memory mem, long offsetBytes, int numItems) {
+        return stringSerDe.sizeOf(mem, offsetBytes, numItems);
+    }
+
+    @Override
+    public String toString(CollatedString item) {
+        return item.original();
+    }
+
+    @Override
+    public Class<CollatedString> getClassOfT() {
+        return CollatedString.class;
+    }
+}

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java

@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions;
+
+/**
+ * A DataSketches ItemsSketch item for non-binary collated strings (SPARK-58069).
+ * <p>
+ * Equality and hashing are driven solely by the collation {@code key}, so that collation-equal
+ * strings (e.g. {@code 'HELLO'} and {@code 'hello'} under {@code UTF8_LCASE}) are counted as a
+ * single item. The {@code original} field retains an actual input value to return in the result,
+ * mirroring how {@code mode()} returns a real value rather than the normalized collation key.
+ */
+public class CollatedString {
+    private final String key;
+    private final String original;
+
+    public CollatedString(String key, String original) {
+        this.key = key;
+        this.original = original;
+    }
+
+    public String key() {
+        return key;
+    }
+
+    public String original() {
+        return original;
+    }
+
+    @Override
+    public int hashCode() {
+        return key.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (!(obj instanceof CollatedString)) {
+            return false;
+        }
+        return key.equals(((CollatedString) obj).key);
+    }
+}

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

@@ -24,12 +24,13 @@ import org.apache.datasketches.common._
 import org.apache.datasketches.frequencies.{ErrorType, ItemsSketch}
 import org.apache.datasketches.memory.Memory
 
+import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, TypeCheckResult}
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
-import org.apache.spark.sql.catalyst.expressions.{ArrayOfDecimalsSerDe, Expression, ExpressionDescription, ImplicitCastInputTypes, Literal}
+import org.apache.spark.sql.catalyst.expressions.{ArrayOfCollatedStringsSerDe, ArrayOfDecimalsSerDe, CollatedString, Expression, ExpressionDescription, ImplicitCastInputTypes, Literal}
 import org.apache.spark.sql.catalyst.trees.{BinaryLike, TernaryLike}
-import org.apache.spark.sql.catalyst.util.{CollationFactory, GenericArrayData}
+import org.apache.spark.sql.catalyst.util.{CollationFactory, GenericArrayData, UnsafeRowUtils}
 import org.apache.spark.sql.errors.QueryExecutionErrors
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
@@ -70,6 +71,11 @@ import org.apache.spark.unsafe.types.UTF8String
       > SELECT _FUNC_(expr, 10, 100) FROM VALUES (0), (1), (1), (2), (2), (2) AS tab(expr);
        [{"item":2,"count":3},{"item":1,"count":2},{"item":0,"count":1}]
   """,
+  note = """
+    When `expr` is a string with a non-UTF8_BINARY collation, values that are equal under the
+    collation are counted as one item, and the returned item is one of the actual input values of
+    that group; which one is returned is not deterministic (as with the `mode` function).
+  """,
   group = "agg_funcs",
   since = "4.1.0")
 // scalastyle:on line.size.limit
@@ -253,8 +259,12 @@ object ApproxTopK {
         new ItemsSketch[Long](maxMapSize).asInstanceOf[ItemsSketch[Any]]
       case _: DoubleType =>
         new ItemsSketch[Double](maxMapSize).asInstanceOf[ItemsSketch[Any]]
-      case _: StringType =>
-        new ItemsSketch[String](maxMapSize).asInstanceOf[ItemsSketch[Any]]
+      case st: StringType =>
+        if (UnsafeRowUtils.isBinaryStable(st)) {
+          new ItemsSketch[String](maxMapSize).asInstanceOf[ItemsSketch[Any]]
+        } else {
+          new ItemsSketch[CollatedString](maxMapSize).asInstanceOf[ItemsSketch[Any]]
+        }
       case _: DecimalType =>
         new ItemsSketch[Decimal](maxMapSize).asInstanceOf[ItemsSketch[Any]]
     }
@@ -269,8 +279,12 @@ object ApproxTopK {
         new ArrayOfLongsSerDe().asInstanceOf[ArrayOfItemsSerDe[Any]]
       case _: DoubleType =>
         new ArrayOfDoublesSerDe().asInstanceOf[ArrayOfItemsSerDe[Any]]
-      case _: StringType =>
-        new ArrayOfStringsSerDe().asInstanceOf[ArrayOfItemsSerDe[Any]]
+      case st: StringType =>
+        if (UnsafeRowUtils.isBinaryStable(st)) {
+          new ArrayOfStringsSerDe().asInstanceOf[ArrayOfItemsSerDe[Any]]
+        } else {
+          new ArrayOfCollatedStringsSerDe(st.collationId).asInstanceOf[ArrayOfItemsSerDe[Any]]
+        }
       case dt: DecimalType =>
         new ArrayOfDecimalsSerDe(dt).asInstanceOf[ArrayOfItemsSerDe[Any]]
     }
@@ -285,7 +299,9 @@ object ApproxTopK {
 
   def dataTypeToDDL(dataType: DataType): String = dataType match {
     case _: StringType =>
-      // Hide collation information in DDL format, otherwise CollationExpressionWalkerSuite fails
+      // Strip collation from the user-facing state DDL to keep the persisted format stable across
+      // collations. Collation is recovered from the state struct's static field-2 type (see
+      // withCollationOf in ApproxTopKCombine.update and the JSON encoding in CombineInternal).
       s"item string not null"
     case other =>
       StructField("item", other, nullable = false).toDDL
@@ -295,6 +311,11 @@ object ApproxTopK {
     StructType.fromDDL(ddl).fields.head.dataType
   }
 
+  def withCollationOf(base: DataType, source: DataType): DataType = (base, source) match {
+    case (_: StringType, st: StringType) => st
+    case _ => base
+  }
+
   def checkStateFieldAndType(state: Expression): TypeCheckResult = {
     val stateStructType = state.dataType.asInstanceOf[StructType]
     if (stateStructType.length != 4) {
@@ -358,8 +379,14 @@ class ApproxTopKAggregateBuffer[T](val sketch: ItemsSketch[T], private var nullC
         case _: TimestampNTZType =>
           sketch.asInstanceOf[ItemsSketch[Long]].update(v.asInstanceOf[Long])
         case st: StringType =>
-          val cKey = CollationFactory.getCollationKey(v.asInstanceOf[UTF8String], st.collationId)
-          sketch.asInstanceOf[ItemsSketch[String]].update(cKey.toString)
+          val orig = v.asInstanceOf[UTF8String]
+          if (UnsafeRowUtils.isBinaryStable(st)) {
+            sketch.asInstanceOf[ItemsSketch[String]].update(orig.toString)
+          } else {
+            val cKey = CollationFactory.getCollationKey(orig, st.collationId).toString
+            sketch.asInstanceOf[ItemsSketch[CollatedString]]
+              .update(new CollatedString(cKey, orig.toString))
+          }
         case _: DecimalType =>
           sketch.asInstanceOf[ItemsSketch[Decimal]].update(v.asInstanceOf[Decimal])
       }
@@ -429,7 +456,12 @@ class ApproxTopKAggregateBuffer[T](val sketch: ItemsSketch[T], private var nullC
                _: DateType | _: TimestampType | _: TimestampNTZType =>
             curFrequentItem.getItem
           case _: StringType =>
-            UTF8String.fromString(curFrequentItem.getItem.asInstanceOf[String])
+            curFrequentItem.getItem match {
+              case cs: CollatedString => UTF8String.fromString(cs.original)
+              case s: String => UTF8String.fromString(s)
+              case other => throw SparkException.internalError(
+                s"Unexpected sketch item type for a string column: ${other.getClass.getName}")
+            }
         }
         fiIndex += 1 // move to next frequent item
         (item, itemEstimate)
@@ -653,22 +685,25 @@ class CombineInternal[T](
    * Serialize the CombineInternal instance to a byte array.
    * Serialization format:
    *     maxItemsTracked (4 bytes int) +
-   *     itemDataTypeDDL length n in byte  (4 bytes int) +
-   *     itemDataTypeDDL (n bytes) +
+   *     itemDataType JSON length n in byte  (4 bytes int) +
+   *     itemDataType JSON (n bytes) +
    *     sketchBytes
+   *
+   * The item data type is encoded as collation-preserving JSON (not the collation-stripped DDL)
+   * so that a collated sketch is deserialized and merged by collation key across shuffle
+   * boundaries (SPARK-58069).
    */
   def serialize(): Array[Byte] = {
     val sketchWithNullCountBytes = sketchWithNullCount.serialize(
       ApproxTopK.genSketchSerDe(itemDataType).asInstanceOf[ArrayOfItemsSerDe[T]])
-    val itemDataTypeDDL = ApproxTopK.dataTypeToDDL(itemDataType)
-    val ddlBytes: Array[Byte] = itemDataTypeDDL.getBytes(StandardCharsets.UTF_8)
+    val typeBytes: Array[Byte] = itemDataType.json.getBytes(StandardCharsets.UTF_8)
     val byteArray = new Array[Byte](
-      sketchWithNullCountBytes.length + Integer.BYTES + Integer.BYTES + ddlBytes.length)
+      sketchWithNullCountBytes.length + Integer.BYTES + Integer.BYTES + typeBytes.length)
 
     val byteBuffer = ByteBuffer.wrap(byteArray)
     byteBuffer.putInt(maxItemsTracked)
-    byteBuffer.putInt(ddlBytes.length)
-    byteBuffer.put(ddlBytes)
+    byteBuffer.putInt(typeBytes.length)
+    byteBuffer.put(typeBytes)
     byteBuffer.put(sketchWithNullCountBytes)
     byteArray
   }
@@ -679,22 +714,21 @@ object CombineInternal {
    * Deserialize a byte array to a CombineInternal instance.
    * Serialization format:
    *     maxItemsTracked (4 bytes int) +
-   *     itemDataTypeDDL length n in byte  (4 bytes int) +
-   *     itemDataTypeDDL (n bytes) +
+   *     itemDataType JSON length n in byte  (4 bytes int) +
+   *     itemDataType JSON (n bytes) +
    *     sketchBytes
    */
   def deserialize(buffer: Array[Byte]): CombineInternal[Any] = {
     val byteBuffer = ByteBuffer.wrap(buffer)
     // read maxItemsTracked
     val maxItemsTracked = byteBuffer.getInt
-    // read itemDataTypeDDL
-    val ddlLength = byteBuffer.getInt
-    val ddlBytes = new Array[Byte](ddlLength)
-    byteBuffer.get(ddlBytes)
-    val itemDataTypeDDL = new String(ddlBytes, StandardCharsets.UTF_8)
-    val itemDataType = ApproxTopK.DDLToDataType(itemDataTypeDDL)
+    // read itemDataType JSON
+    val typeLength = byteBuffer.getInt
+    val typeBytes = new Array[Byte](typeLength)
+    byteBuffer.get(typeBytes)
+    val itemDataType = DataType.fromJson(new String(typeBytes, StandardCharsets.UTF_8))
     // read sketchBytes
-    val sketchBytes = new Array[Byte](buffer.length - Integer.BYTES - Integer.BYTES - ddlLength)
+    val sketchBytes = new Array[Byte](buffer.length - Integer.BYTES - Integer.BYTES - typeLength)
     byteBuffer.get(sketchBytes)
     val sketchWithNullCount = ApproxTopKAggregateBuffer.deserialize(
       sketchBytes, ApproxTopK.genSketchSerDe(itemDataType))
@@ -817,7 +851,8 @@ case class ApproxTopKCombine(
     val inputSketchBytes = inputState.getBinary(0)
     val inputMaxItemsTracked = inputState.getInt(1)
     val inputItemDataTypeDDL = inputState.getUTF8String(3).toString
-    val inputItemDataType = ApproxTopK.DDLToDataType(inputItemDataTypeDDL)
+    val inputItemDataType = ApproxTopK.withCollationOf(
+      ApproxTopK.DDLToDataType(inputItemDataTypeDDL), uncheckedItemDataType)
     // update maxItemsTracked (throw error if not match)
     buffer.updateMaxItemsTracked(combineSizeSpecified, inputMaxItemsTracked)
     // update itemDataType (throw error if not match)

sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala

@@ -269,6 +269,50 @@ class ApproxTopKSuite extends SharedSparkSession {
     checkAnswer(res, Row(Seq(Row("c", 4), Row("d", 2))))
   }
 
+  Seq("UTF8_LCASE", "UNICODE_CI").foreach { collation =>
+    test(s"SPARK-58069: approx_top_k returns an actual value, not the collation key ($collation)") {
+      val res = sql(
+        s"""SELECT approx_top_k(c, 2)
+           |FROM (SELECT CAST(col AS STRING COLLATE $collation) AS c
+           |      FROM VALUES ('HELLO'), ('HELLO'), ('HELLO'), ('world') AS t(col))
+           |""".stripMargin)
+      checkAnswer(res, Row(Seq(Row("HELLO", 3), Row("world", 1))))
+    }
+
+    test("SPARK-58069: approx_top_k_accumulate/estimate returns an actual value, " +
+      s"not the collation key ($collation)") {
+      val res = sql(
+        s"""SELECT approx_top_k_estimate(approx_top_k_accumulate(c), 2)
+           |FROM (SELECT CAST(col AS STRING COLLATE $collation) AS c
+           |      FROM VALUES ('HELLO'), ('HELLO'), ('HELLO'), ('world') AS t(col))
+           |""".stripMargin)
+      checkAnswer(res, Row(Seq(Row("HELLO", 3), Row("world", 1))))
+    }
+
+    test("SPARK-58069: approx_top_k_combine merges collation-equal values across sketches " +
+      s"and a shuffle ($collation)") {
+      withSQLConf("spark.sql.shuffle.partitions" -> "2") {
+        val sketches = sql(
+          s"""SELECT approx_top_k_accumulate(CAST(col AS STRING COLLATE $collation)) AS sketch
+             |  FROM VALUES ('HELLO'), ('HELLO') AS t(col)
+             |UNION ALL
+             |SELECT approx_top_k_accumulate(CAST(col AS STRING COLLATE $collation)) AS sketch
+             |  FROM VALUES ('hello'), ('WORLD') AS t(col)
+             |""".stripMargin).repartition(2)
+        sketches.createOrReplaceTempView("approx_top_k_sketches")
+        val res = sql(
+          "SELECT approx_top_k_estimate(approx_top_k_combine(sketch, 100), 2) " +
+            "FROM approx_top_k_sketches")
+        // 'HELLO' x2 and 'hello' x1 are collation-equal, so they merge to count 3; 'WORLD' has 1.
+        // Assert on lowercased items so the test is independent of which actual value survives as
+        // the (non-deterministic) representative.
+        val items = res.collect()(0).getSeq[Row](0)
+          .map(r => (r.getString(0).toLowerCase(java.util.Locale.ROOT), r.getLong(1))).toSet
+        assert(items === Set(("hello", 3L), ("world", 1L)))
+      }
+    }
+  }
+
   test("SPARK-52588: accumulate and estimate of Decimal(4, 1)") {
     val res = sql("SELECT approx_top_k_estimate(approx_top_k_accumulate(expr, 10)) " +
       "FROM VALUES CAST(0.0 AS DECIMAL(4, 1)), CAST(0.0 AS DECIMAL(4, 1)), " +