diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java new file mode 100644 index 0000000000000..543cc3b5a7271 --- /dev/null +++ b/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). + *

+ * 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 { + + 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 getClassOfT() { + return CollatedString.class; + } +} diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java new file mode 100644 index 0000000000000..27dc7ab5cf9fe --- /dev/null +++ b/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). + *

+ * 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); + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala index 7ae542f190d56..b09f242ccbace 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala +++ b/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) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala index bc8ddd42f2ef8..e579382e24377 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala +++ b/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)), " + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala index b748ae4c0edac..dd22f647ab31d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala @@ -385,7 +385,13 @@ class CollationExpressionWalkerSuite extends SharedSparkSession { "sha", "crc32", "ascii", - "time_trunc" + "time_trunc", + // The result/sketch embeds the original item value, which now preserves the + // input case for collated strings, so it is not comparable across collations. + "approx_top_k", + "approx_top_k_accumulate", + "approx_top_k_combine", + "approx_top_k_estimate" ) logInfo("Total number of expression: " + expressionCounter)