From 50411c3b41dd1de97ea61b4799cfd9b42256784a Mon Sep 17 00:00:00 2001 From: LinSimon-901101 Date: Fri, 25 Sep 2026 01:47:10 +0800 Subject: [PATCH] fix: reject non-default collations in all sort keys --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../apache/comet/serde/QueryPlanSerde.scala | 7 +- .../windows/window_group_limit_collation.sql | 19 +- .../spark/sql/CometSortCollationSuite.scala | 202 ++++++++++++++++++ 5 files changed, 216 insertions(+), 14 deletions(-) create mode 100644 spark/src/test/spark-4.x/org/apache/spark/sql/CometSortCollationSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 8abd6c92496..efadc1e21c3 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -572,6 +572,7 @@ jobs: org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite org.apache.spark.sql.CometCollationSuite + org.apache.spark.sql.CometSortCollationSuite org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index af6bbc48c85..372ee5e2719 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -220,6 +220,7 @@ jobs: org.apache.comet.serde.operator.CometNativeScanSuite org.apache.spark.sql.CometToPrettyStringSuite org.apache.spark.sql.CometCollationSuite + org.apache.spark.sql.CometSortCollationSuite org.apache.comet.CometFuzzAggregateSuite org.apache.spark.sql.comet.execution.arrow.CometArrowStreamSuite org.apache.spark.sql.CometSparkInternalFunctionsSuite diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 3944b13f49d..77f4b432b4a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1291,7 +1291,12 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } def supportedSortType(op: SparkPlan, sortOrder: Seq[SortOrder]): Boolean = { - if (sortOrder.length == 1) { + // Both single- and multi-column sorts compare strings by raw bytes. Check nested types + // before the single-column kernel restrictions, since multi-column keys bypass those. + if (sortOrder.exists(order => hasNonDefaultStringCollation(order.dataType))) { + withFallbackReason(op, "Sort does not support non-default string collation") + false + } else if (sortOrder.length == 1) { val canSort = sortOrder.head.dataType match { case ArrayType(elementType, _) => supportedScalarSortElementType(elementType) case MapType(_, valueType, _) => supportedScalarSortElementType(valueType) diff --git a/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql b/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql index cac9e87e501..8bcb422a0bc 100644 --- a/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql +++ b/spark/src/test/resources/sql-tests/windows/window_group_limit_collation.sql @@ -20,12 +20,9 @@ -- collation semantics that make Spark tie e.g. 'A' with 'a' under UTF8_LCASE. Falling -- back to Spark keeps peer equality intact. -- --- Unlike the other window_group_limit_* fixtures this file pins the threshold instead of running --- the -1,1000 matrix. `spark.sql.optimizer.windowGroupLimitThreshold=-1` makes --- `InferWindowGroupLimit` a no-op (see its `apply`), so no `WindowGroupLimit` node is planned and --- the `expect_fallback` reason below - which only `CometWindowGroupLimitExec.convert` can emit - --- is unreachable by construction. The plain `WindowExec` path that the -1 arm would exercise has --- no collation guard of its own and is tracked separately. +-- Pin the threshold so these queries exercise WindowGroupLimit plans. The sort below the +-- window must also reject collated keys (#6158); its fallback can happen before the +-- WindowGroupLimit converter, so either operator may supply the collation fallback reason. -- MinSparkVersion: 4.0 -- Config: spark.sql.optimizer.windowGroupLimitThreshold=1000 @@ -34,14 +31,10 @@ statement CREATE TABLE test_wgl_collation(grp int, s string) USING parquet statement -INSERT INTO test_wgl_collation VALUES (1, 'A'), (1, 'a'), (1, 'b') +INSERT INTO test_wgl_collation VALUES (1, 'A'), (1, 'a'), (1, 'b'), (1, 'B') --- Keep the values byte-order-compatible with UTF8_LCASE order ('A' < 'a' < 'b' holds under both). --- `QueryPlanSerde.supportedSortType` only rejects collated strings for single-column sorts, so the --- two-column `Sort [key ASC, s ASC]` that Spark injects below `WindowGroupLimitExec` still runs on --- Comet and sorts by raw bytes. Values where the two orders disagree (e.g. 'a' before 'B' under --- UTF8_LCASE but after it by byte value) would feed Spark's fallback operator a wrongly ordered --- stream and fail for a reason unrelated to what this file pins. +-- Byte order places 'B' before 'a', while UTF8_LCASE places it after. Spark's fallback window +-- operators must receive collation-aware ordering, including for multi-column sort keys. -- Case-insensitive ORDER BY key: 'A' and 'a' must tie at rank 1 (Spark keeps both). query expect_fallback(non-default string collation) diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/CometSortCollationSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/CometSortCollationSuite.scala new file mode 100644 index 00000000000..6916ae72b3c --- /dev/null +++ b/spark/src/test/spark-4.x/org/apache/spark/sql/CometSortCollationSuite.scala @@ -0,0 +1,202 @@ +/* + * 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 + +import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, SortOrder} +import org.apache.spark.sql.catalyst.plans.physical.RoundRobinPartitioning +import org.apache.spark.sql.comet.{CometSortExec, CometTakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.{LocalTableScanExec, SortExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types._ + +import org.apache.comet.{CometConf, CometExplainInfo} +import org.apache.comet.serde.QueryPlanSerde + +class CometSortCollationSuite extends CometTestBase { + + private val sortCollationReason = "Sort does not support non-default string collation" + + private def assertSparkSort(plan: SparkPlan): Unit = { + assert(collect(plan) { case sort: SortExec => sort }.nonEmpty, plan.toString) + assert(collect(plan) { case sort: CometSortExec => sort }.isEmpty, plan.toString) + } + + private def sortPlan(dataTypes: Seq[DataType]): SortExec = { + val keys = dataTypes.zipWithIndex.map { case (dataType, index) => + AttributeReference(s"key$index", dataType)() + } + SortExec( + keys.map(key => SortOrder(key, Ascending)), + global = false, + child = LocalTableScanExec(keys, Nil, None)) + } + + test("window sort with UTF8_LCASE and multiple keys falls back to Spark (issue #6158)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable(Seq((1, "b"), (2, "A"), (3, "a"), (4, "B")), "tbl") { + // Window ordering uses SinglePartition, so the hash/range shuffle collation guards + // cannot hide a missing sort guard. The row numbers expose incorrect byte ordering. + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + "SELECT _1, c, row_number() OVER (ORDER BY c, _1) AS rn " + + "FROM (SELECT _1, _2 COLLATE UTF8_LCASE AS c FROM tbl)", + sortCollationReason) + assertSparkSort(cometPlan) + } + } + } + + test("sort type check rejects non-default collation in every key and nested type") { + val collated = StringType("UTF8_LCASE") + val collatedTypes = Seq( + collated, + StructType(Seq(StructField("text", collated))), + ArrayType(collated), + MapType(collated, IntegerType), + MapType(IntegerType, collated), + ArrayType(StructType(Seq(StructField("nested", ArrayType(collated))))), + StructType(Seq(StructField("nested", MapType(IntegerType, ArrayType(collated)))))) + + for { + dataType <- collatedTypes + keyTypes <- Seq(Seq(dataType), Seq(dataType, IntegerType), Seq(IntegerType, dataType)) + } { + withClue(s"Sort key types: $keyTypes: ") { + // Maps cannot be ordered in SQL, but the recursive type guard must inspect both + // their keys and values if it receives such a physical sort expression. + val plan = sortPlan(keyTypes) + assert(!QueryPlanSerde.supportedSortType(plan, plan.sortOrder)) + assert( + plan + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .exists(_.contains(sortCollationReason))) + } + } + } + + test("sort type check preserves default collation and existing single-key restrictions") { + val defaultTypes = Seq( + StringType, + ArrayType(StringType), + MapType(StringType, IntegerType), + MapType(IntegerType, StringType)) + for { + dataType <- defaultTypes + keyTypes <- Seq(Seq(dataType), Seq(dataType, IntegerType), Seq(IntegerType, dataType)) + } { + withClue(s"Sort key types: $keyTypes: ") { + val plan = sortPlan(keyTypes) + assert(QueryPlanSerde.supportedSortType(plan, plan.sortOrder)) + } + } + + // The collation guard must not extend the existing single-key restrictions to multi-key + // sorting, which uses a different representation and already supports structs. + val structType = StructType(Seq(StructField("text", StringType))) + val singleKey = sortPlan(Seq(structType)) + assert(!QueryPlanSerde.supportedSortType(singleKey, singleKey.sortOrder)) + val multipleKeys = sortPlan(Seq(structType, IntegerType)) + assert(QueryPlanSerde.supportedSortType(multipleKeys, multipleKeys.sortOrder)) + } + + test("window sort rejects collated keys after another key and inside nested types") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + // All values of the first key tie, so ordering must actually compare the collated key. + withParquetTable(Seq((0, 1, "b"), (0, 2, "A"), (0, 3, "a"), (0, 4, "B")), "tbl") { + for (sortKeys <- Seq("_1, c, _2", "struct(array(c)), _2")) { + withClue(s"Sort keys: $sortKeys: ") { + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + s"SELECT _2, c, row_number() OVER (ORDER BY $sortKeys) AS rn " + + "FROM (SELECT _1, _2, _3 COLLATE UTF8_LCASE AS c FROM tbl)", + sortCollationReason) + assertSparkSort(cometPlan) + } + } + } + } + } + + test("local sort after round-robin shuffle rejects non-default collation") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm") { + val data = (1 to 8).flatMap { n => + Seq((n * 4, "b"), (n * 4 + 1, "A"), (n * 4 + 2, "a"), (n * 4 + 3, "B")) + } + withParquetTable(data, "tbl") { + // checkSparkAnswer preserves collected order when the logical plan contains Sort, + // including sortWithinPartitions; this checks ordering, not only row membership. + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + sql("SELECT _1, _2 COLLATE UTF8_LCASE AS c FROM tbl") + .repartition(2) + .sortWithinPartitions("c", "_1"), + sortCollationReason) + assertSparkSort(cometPlan) + assert( + collect(cometPlan) { + case plan if plan.outputPartitioning.isInstanceOf[RoundRobinPartitioning] => plan + }.nonEmpty, + cometPlan.toString) + } + } + } + + test("top-k with multiple keys rejects non-default collation") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable(Seq((1, "b"), (2, "A"), (3, "a"), (4, "B")), "tbl") { + // LIMIT 2 changes row membership under byte ordering: A,B instead of A,a. + // Keep COLLATE below the exchange so TopK projects attributes and reaches the sort + // guard instead of falling back while serializing a fused COLLATE projection. + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + sql("SELECT _1, _2 COLLATE UTF8_LCASE AS c FROM tbl") + .repartition(1) + .orderBy("c", "_1") + .limit(2), + sortCollationReason) + assert( + collect(cometPlan) { case topK: TakeOrderedAndProjectExec => topK }.nonEmpty, + cometPlan.toString) + assert( + collect(cometPlan) { case topK: CometTakeOrderedAndProjectExec => topK }.isEmpty, + cometPlan.toString) + } + } + } + + test("default UTF8_BINARY string and multi-key struct sorting remain native") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable(Seq((1, "b"), (2, "A"), (3, "a"), (4, "B")), "tbl") { + for (sortKeys <- Seq("_2, _1", "struct(_2), _1")) { + withClue(s"Sort keys: $sortKeys: ") { + val (_, cometPlan) = checkSparkAnswerAndOperator( + s"SELECT _1, _2, row_number() OVER (ORDER BY $sortKeys) AS rn FROM tbl") + assert( + collect(cometPlan) { case sort: CometSortExec => sort }.nonEmpty, + cometPlan.toString) + } + } + val (_, topKPlan) = + checkSparkAnswerAndOperator("SELECT _1, _2 FROM tbl ORDER BY _2, _1 LIMIT 2") + assert( + collect(topKPlan) { case topK: CometTakeOrderedAndProjectExec => topK }.nonEmpty, + topKPlan.toString) + } + } + } +}