comphead commented on code in PR #7256:
URL: https://github.com/apache/datafusion-comet/pull/7256#discussion_r4437688992
##########
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:
Thanks. 1bb942a33a applies the collation fallback only when the window's
input order can be trusted. That is the same check that decides whether the
order is declared. Over a doubted order the window stays native, as on `main`.
Your fixture fails with the ordering error again instead of handing the
reordered rows to Spark's window. When the order is trusted, Spark's window
gets the rows in the order it relies on and applies the collation.
#6785 also makes Spark's window correct in your fixture, since it keeps the
nested loop join in Spark.
##########
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:
Thanks, applied in 1bb942a33a. The last window now partitions by `k` and
expects `0 Ascending NullsFirst`. It needs its order declared next to the
join's `a`, and the test comment says why.
--
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]
