peter-toth commented on code in PR #59306:
URL: https://github.com/apache/spark/pull/59306#discussion_r4282262232
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -6793,6 +6797,98 @@ class KeyGroupedPartitioningSuite
}
}
+ /**
+ * The setup of the SPARK-60044 tests. Creates `items`, which reports no
ordering, so the scan
+ * derives [id, name] from its partition keys and drops the transform
between them. The splits
+ * are sorted by the full key, so within id 1 the 'bb' split comes before
the 'aa' one. Returns a
+ * window query that groups by id and orders by name. Under
`kWayMergeWindowConfs` a k-way merge
+ * delivers that ordering, so EnsureRequirements adds no SortExec. The query
selects
+ * `arrive_time`, so the transform's column stays in the scan output.
+ */
+ private def createKWayMergeWindowQuery(): String = {
+ createTable(items, itemsColumns, Array(identity("id"),
days("arrive_time"), identity("name")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'bb', 10.0, cast('2021-01-01' as timestamp)), " +
+ "(1, 'aa', 20.0, cast('2022-01-01' as timestamp)), " +
+ "(2, 'cc', 30.0, cast('2021-01-01' as timestamp))")
+ s"""SELECT id, name, arrive_time, ROW_NUMBER() OVER (PARTITION BY id ORDER
BY name) AS rn
+ |FROM testcat.ns.$items
+ |""".stripMargin
+ }
+
+ private val kWayMergeWindowConfs = Seq(
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true")
+
+ private val kWayMergeWindowRows = Seq(
+ Row(1, "aa", Timestamp.valueOf("2022-01-01 00:00:00"), 1),
+ Row(1, "bb", Timestamp.valueOf("2021-01-01 00:00:00"), 2),
+ Row(2, "cc", Timestamp.valueOf("2021-01-01 00:00:00"), 1))
+
+ test("SPARK-60044: a planned k-way merge keeps its ordering after a later
config change") {
+ // Turning partitionKeyOrdering off after planning empties the scan's
ordering. The planned
+ // merge must still merge by [id, name]. Otherwise the window numbers the
rows of id 1 in split
+ // order.
+ //
+ // Both AQE modes are covered, see the SPARK-59279 test above.
+ val query = createKWayMergeWindowQuery()
+ Seq(true, false).foreach { aqeEnabled =>
+ val confs = (SQLConf.ADAPTIVE_EXECUTION_ENABLED.key ->
aqeEnabled.toString) +:
+ kWayMergeWindowConfs
+ withSQLConf(confs: _*) {
+ val df = sql(query)
+ // Force the plan without executing it, so the k-way merge is decided
under this config.
+ val plan = df.queryExecution.executedPlan
+ val merging =
collectAllGroupPartitions(plan).filter(_.enableSortedMerge)
+ assert(merging.length == 1, "expected one k-way merge")
+ assert(merging.head.outputOrdering.map(_.child.sql) == Seq("id",
"name"))
+ assert(collect(plan) { case s: SortExec => s }.isEmpty,
+ "the k-way merge satisfies the window's ordering, so no SortExec
should be added")
+
+ // Flip only the key ordering config. The plan above is already
committed.
+ withSQLConf(SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key ->
"false") {
+ assert(merging.head.child.outputOrdering.isEmpty,
+ "test setup: the scan no longer derives an ordering from its keys")
+ checkAnswer(df, kWayMergeWindowRows)
Review Comment:
Done in
[`34f2672de31`](https://github.com/apache/spark/pull/59306/commits/34f2672de316aad274f3f555af3c04b9cc745a44),
here and in the SPARK-59279 test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
