comphead commented on code in PR #7256:
URL: https://github.com/apache/datafusion-comet/pull/7256#discussion_r4417654002
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala:
##########
@@ -81,10 +81,38 @@ object CometWindowExec extends
CometOperatorSerde[WindowExec] {
if (windowExprProto.forall(_.isDefined) &&
partitionExprs.forall(_.isDefined)
&& sortOrders.forall(_.isDefined)) {
+ // Spark plans no sort below the window when its input reports the order
the window needs,
+ // a sorted cache for instance, but rows from the JVM arrive through a
ScanExec that
+ // declares none, and the native window rejects them. Send the order for
the planner to
+ // declare (see SortedInputExec), unless a native sort below declares it
or the reported
+ // order may be false. A key Comet cannot express or a collated string
key, which the
+ // native side would order by raw bytes, sends the window back to Spark.
+ val requiredOrdering = op.requiredChildOrdering.head
+ val inputSortOrders =
+ if (requiredOrdering.nonEmpty && !op.child.isInstanceOf[CometSortExec]
&&
+ SortOrder.orderingSatisfies(op.child.outputOrdering,
requiredOrdering) &&
+ !CometHashJoinExec.mayReportFalseInputOrder(op.child)) {
+ val orders = requiredOrdering.map { order =>
Review Comment:
Thanks, good catch. The lower window now sends the required keys and then
the rest of the order its child reports, up to the first key that cannot be
expressed. In your example it declares `(k, v)` on the cache scan, and the
upper window's input carries that order through DataFusion's window equivalence
properties, so `with_declared_order` can return early there. Fixed in
16b1add97e.
I extended the order at the lower window, not at the upper one. A key left
out at the lower window is a collated string or cannot be expressed. The upper
window would be checking the same attribute, so I expect it to fall back to
Spark rather than fail natively.
`stacked native windows over a sorted cache keep the cache's whole order`
runs your shape with the native cache on and off: `sum(v) OVER (PARTITION BY
k)` under `sum(s) OVER (PARTITION BY k, v)`, over a cache sorted by `(k, v)`
with a nullable `k`. It checks the rows, that both windows send `k ASC NULLS
FIRST, v ASC NULLS FIRST`, and that no sort is planned.
##########
native/core/src/execution/planner/window.rs:
##########
@@ -131,17 +131,21 @@ impl OperatorBuilder for WindowBuilder {
// trigger a retract call.
let window_expr = window_expr?;
let all_bounded = window_expr.iter().all(|e| e.uses_bounded_memory());
+
+ // Both window operators reject an input that does not declare its
partition keys' order.
+ let input =
+ planner.with_declared_order(Arc::clone(&child.native_plan),
&wnd.input_sort_orders)?;
Review Comment:
Agreed, an empty input satisfies any order. `keeps_source_order` now accepts
`EmptyExec`, so an empty bucket gets the same declaration as a bucket with a
file. Fixed in 16b1add97e.
`declared_order_reaches_an_empty_bucket` in `planner.rs` checks the
declaration over an `EmptyExec`. `window over a sorted bucketed table with
empty buckets runs natively` in `CometWindowExecSuite` writes your table shape:
one writer, `bucketBy(8, "k").sortBy("k", "v")` and two keys. It reads the
table with the legacy bucketed scan ordering on, then checks `dense_rank() OVER
(PARTITION BY k ORDER BY v)` against Spark. It also checks that the window sits
directly on the bucketed `CometNativeScanExec` and that no sort is planned. The
hash join's probe side uses the same walk. An empty probe input gets the
declaration there too, which changes nothing for it.
##########
spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala:
##########
@@ -3692,6 +3692,91 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ 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, and the
native window
+ // rejected the scan, which declares no order, with "Expects PARTITION BY
expression to be
+ // ordered" (#7223). The partition key is nullable and the order key
descending, so the
+ // declared order has to match the cache's direction and null placement.
dense_rank runs in
+ // DataFusion's BoundedWindowAggExec, and the frame up to UNBOUNDED
FOLLOWING 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) =>
+ val df = cached.selectExpr("k", "v", s"$window AS w")
+ val rows = (0L until 2000L).map(v => Row(key(v), v, expected(v)))
+ QueryTest.sameRows(rows, df.collect().toSeq).foreach(fail(_))
+ val plan = df.queryExecution.executedPlan
+ withClue(s"nativeCache=$nativeCache, $window\n$plan\n") {
+ val windows = collect(plan) { case w: CometWindowExec => w }
+ assert(windows.size == 1)
+ assert(windows.head.nativeOp.getWindow.getInputSortOrdersCount
== 2)
Review Comment:
Done in 16b1add97e. The test now compares the declared keys with `0
Ascending NullsFirst, 1 Descending NullsLast`. It uses a small
`windowInputOrder` helper, which the stacked-window test uses as well.
--
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]
