sunchao commented on code in PR #7256:
URL: https://github.com/apache/datafusion-comet/pull/7256#discussion_r4429697946
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala:
##########
@@ -75,16 +75,54 @@ object CometWindowExec extends
CometOperatorSerde[WindowExec] {
}
}
+ // The native window finds partitions and peers by comparing keys byte by
byte, so it would
+ // tell apart collated strings that Spark treats as equal, as
CometWindowGroupLimitExec notes.
+ val collated = (op.partitionSpec ++ op.orderSpec.map(_.child))
+ .filter(key => hasNonDefaultStringCollation(key.dataType))
+ if (collated.nonEmpty) {
+ withFallbackReason(
+ op,
+ collated
+ .map(_.sql)
+ .mkString("Window: non-default string collation on key(s): ", ", ",
""))
+ return None
Review Comment:
This fallback can still give Spark a window input whose rows are out of
order. I cached 12 rows in one partition, sorted by integer `k = v % 3` and `s
COLLATE UTF8_LCASE`, with `s` cycling through `A, a, B, b`. I then broadcast
two rows (`z = 0, 1`) and joined on `v % 2 <= z`. The native nested-loop join
emits each `k` in two separate runs. The collated order key sends `count(*)
OVER (PARTITION BY k ORDER BY s ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED
FOLLOWING)` to Spark, which has already omitted the sort.
On Spark 4.1.3 the fallback returned 2 on six rows and 4 on twelve. Vanilla
Spark and a physical sort above the same native join both returned 6 on all 18.
Caching the join output also reproduced it. The base window serializer rejects
this unordered input in the same fixture. Restoring the full `(k, s)` ordering
before the fallback window, or preventing the join from advertising it, would
let Spark apply the collation correctly.
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -3692,6 +3692,175 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ // The order a native window declares on its input, one "index direction
nulls" per key.
+ private def windowInputOrder(window: CometWindowExec): Seq[String] =
+ window.nativeOp.getWindow.getInputSortOrdersList.asScala.toSeq.map { expr
=>
+ val order = expr.getSortOrder
+ s"${order.getChild.getBound.getIndex} ${order.getDirection}
${order.getNullOrdering}"
+ }
+
+ // Checks `df`'s rows, the order each native window declares, the top one
first, and that no
+ // sort is planned.
+ private def checkWindows(df: DataFrame, rows: Seq[Row], orders:
Seq[Seq[String]]) = {
+ QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_))
+ val plan = df.queryExecution.executedPlan
+ withClue(s"$plan\n") {
+ assert(collect(plan) { case w: CometWindowExec => w
}.map(windowInputOrder) == orders)
+ assert(collect(plan) { case s @ (_: CometSortExec | _: SortExec) => s
}.isEmpty)
+ }
+ }
+
+ test("a native window over a sorted cache declares the cache's order") {
+ // Spark plans no sort below a window over the cache's sort keys, so the
scan has to declare
+ // the partition key's order (#7223). The key is nullable, so the null
placement has to match,
+ // and the descending order key is not declared. dense_rank runs in
BoundedWindowAggExec and
+ // the UNBOUNDED FOLLOWING frame in WindowAggExec. Within a key, v = k +
20 * j for j in
+ // 0 until 100, and a NULL key holds the multiples of 20.
+ val key = (v: Long) => if (v % 20 == 0) null else v % 20
+ val queries = Seq(
+ "dense_rank() OVER (PARTITION BY k ORDER BY v DESC)" ->
+ ((v: Long) => (100 - v / 20).toInt),
+ "sum(v) OVER (PARTITION BY k ORDER BY v DESC " +
+ "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)" -> { (v: Long) =>
+ val j = v / 20
+ (j + 1) * (v % 20) + 10 * j * (j + 1)
+ })
+ Seq(true, false).foreach { nativeCache =>
+ withConversions(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key ->
nativeCache.toString) {
+ try {
+ val cached = spark
+ .range(2000)
+ .selectExpr("CASE WHEN id % 20 = 0 THEN NULL ELSE id % 20 END AS
k", "id AS v")
+ .repartition(2, col("k"))
+ .sortWithinPartitions(col("k"), col("v").desc)
+ .cache()
+ cached.count()
+
+ queries.foreach { case (window, expected) =>
+ withClue(s"nativeCache=$nativeCache, $window: ") {
+ checkWindows(
+ cached.selectExpr("k", "v", s"$window AS w"),
+ (0L until 2000L).map(v => Row(key(v), v, expected(v))),
+ Seq(Seq("0 Ascending NullsFirst")))
+ }
+ }
+ } finally {
+ spark.catalog.clearCache()
+ }
+ }
+ }
+ }
+
+ test("a native window declares its order next to the order its input
declares") {
+ // Spark plans no sort above an operator that keeps the cache's order, so
a window's input may
+ // declare only the order that operator needed: an upper window's input
the lower window's,
+ // and an outer join's output its streamed side's. The window's order is
declared next to it
+ // (#7223). Alias a of sort key k shares its order. One partition, so no
exchange comes in
+ // between.
+ val ids = 0L until 24L
+ val key = (id: Long) => if (id % 4 == 0) None else Some(id % 3)
+ val nullable = (k: Option[Long]) => k.map(Long.box).orNull
+ val keySums = ids.groupBy(key).map { case (k, group) => k -> group.sum }
+ val keyCounts = ids.groupBy(key).map { case (k, group) => k -> group.size }
+ val pairCounts =
+ ids.groupBy(id => (key(id), id % 2)).map { case (p, group) => p ->
group.size }
+ withConversions(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ try {
+ val cached = spark
+ .range(0, 24, 1, 1)
+ .selectExpr(
+ "CASE WHEN id % 4 = 0 THEN NULL ELSE id % 3 END AS k",
+ "id % 2 AS x",
+ "id AS v")
+ .sortWithinPartitions("k", "x")
+ .selectExpr("k", "k AS a", "x", "v")
+ .cache()
+ cached.count()
+
+ checkWindows(
+ cached
+ .selectExpr("k", "x", "v", "sum(v) OVER (PARTITION BY k) AS s")
+ .selectExpr("k", "x", "v", "sum(s) OVER (PARTITION BY k, x) AS t"),
+ ids.map { id =>
+ val k = key(id)
+ Row(nullable(k), id % 2, id, keySums(k) * pairCounts((k, id % 2)))
+ },
+ Seq(
+ Seq("0 Ascending NullsFirst", "1 Ascending NullsFirst"),
+ Seq("0 Ascending NullsFirst")))
+ checkWindows(
+ cached
+ .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY k) AS s")
+ .selectExpr("k", "a", "v", "sum(s) OVER (PARTITION BY a) AS t"),
+ ids.map { id =>
+ val k = key(id)
+ Row(nullable(k), nullable(k), id, keySums(k) * keyCounts(k))
+ },
+ Seq(Seq("1 Ascending NullsFirst"), Seq("0 Ascending NullsFirst")))
+ // The residual condition leaves rows of key 0 unmatched among matched
ones, so the key
+ // stays together only if the join keeps its unmatched probe rows in
place.
+ val one = spark.range(1).toDF("j").hint("broadcast")
+ checkWindows(
+ cached
+ .join(one, col("k") === col("j") && col("v") % 2 === 0,
"left_outer")
+ .selectExpr("k", "a", "v", "sum(v) OVER (PARTITION BY a) AS s"),
Review Comment:
This last query passes with the pre-fix native library. On Spark 4.1.3 the
join declares `a` (ordinal 1), and `PARTITION BY a` asks for the same column,
so the window never needs the new right-outer-join branch. I ran the repository
test with that library and it passed.
Changing only this final window to `PARTITION BY k` and its expected
declaration below to `0 Ascending NullsFirst` preserves the expected rows. That
version fails with the old library (`All partition by columns should have an
ordering`) and passes with this commit’s library. It covers the end-to-end case
where the window needs an equivalent key that the join did not declare.
--
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]
