sunchao commented on code in PR #7256:
URL: https://github.com/apache/datafusion-comet/pull/7256#discussion_r4416117731
##########
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:
Could this also declare the order for an empty bucket? I tried a Parquet
table written by one writer with `bucketBy(8, "k").sortBy("k", "v")`, two
distinct `k` values, and legacy bucketed scan output ordering enabled. The head
plans `CometWindowExec` directly over `CometNativeScan` for `dense_rank() OVER
(PARTITION BY k ORDER BY v)` and sends both sort keys. Six of the eight tasks
have no files. `build_parquet_scan_plan` gives those tasks an `EmptyExec`,
which `keeps_source_order` rejects. DataFusion's bounded window checks its
partition ordering before polling the empty stream, so one empty bucket is
enough to fail the query. Vanilla Spark returns all 16 rows. An empty input
satisfies any ordering. Could we handle that case and cover a table with
unpopulated buckets?
##########
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:
Could we preserve the usable part of the cache's ordering for the next
window? On a single-partition cache sorted by `(k, v)`, compute `s = sum(v)
OVER (PARTITION BY k)`, then `sum(s) OVER (PARTITION BY k, v)`. Spark 4.1.3
returns all 24 rows in my example with no intervening sort or exchange. On this
head, both `CometWindowExec`s advertise `(k, v)` to Spark, but the first
serializes only `k`. The second sends both keys and `with_declared_order`
returns early because its input already has an ordering. DataFusion therefore
has no ordering for `v` when the second window starts. Keeping the trusted,
expressible suffix at the first window, or extending a partial declaration
safely at the second, would cover this. Could we add a stacked-window
regression?
##########
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:
Could we assert the columns, directions and null placements in
`input_sort_orders` here? If only the declaration for `v` changed from `DESC
NULLS LAST` to `ASC NULLS FIRST`, the cache would still deliver the same rows
and the window would still use its separate `order_by_list` to calculate both
results. The count would also stay at two. Checking `k ASC NULLS FIRST, v DESC
NULLS LAST` directly would catch that regression.
--
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]
