feat: support WindowGroupLimitExec - #4870
Conversation
14ee3bb to
d781362
Compare
| : : : +- CometWindowGroupLimitExec | ||
| : : : +- CometSort | ||
| : : : +- CometFilter | ||
| : : : : +- Subquery |
There was a problem hiding this comment.
Interesting, we dont support Subquery? 🤔
|
Up front: I used an LLM to help work through this review, so please push back on anything that looks wrong or off base. Thanks for taking this on. I traced the rank arithmetic in A few things I would like to see addressed.
|
|
Thanks @andygrove addressed in 7ef79eb |
| let _timer = self.baseline_metrics.elapsed_compute().timer(); | ||
| self.process_batch(&batch) |
There was a problem hiding this comment.
[P1] Release the metrics borrow before calling process_batch
Could we clone elapsed_compute() into a local before creating this timer? ScopedTimerGuard holds an immutable borrow of self.baseline_metrics until the end of this block, while self.process_batch(&batch) requires a mutable borrow of self. Compiling this file at 7ef79eb0 produces E0502: cannot borrow self as mutable because it is also borrowed as immutable on the process_batch call. This prevents native builds even when WindowGroupLimit is disabled. The preceding add9758e version compiles against the same dependencies.
This small change compiles while preserving the timing:
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
let _timer = elapsed_compute.timer();
self.process_batch(&batch)| val partitionProtos = fields.partitionSpec.map(e => e -> exprToProto(e, childOutput)) | ||
| val orderProtos = fields.orderSpec.map(e => e -> exprToProto(e, childOutput)) |
There was a problem hiding this comment.
[P2] Reject collated keys before native window-group pruning
Could we reject non-default string collations on these keys before lowering the operator? For an ordinary, uncollated Parquet table containing (grp, s) = (1, 'A'), (1, 'a'), (1, 'b'), this window expression filtered to rk <= 1 must retain both A and a:
RANK() OVER (
PARTITION BY grp
ORDER BY CAST(s AS STRING COLLATE UTF8_LCASE)
) AS rkThe collated key is serialized as ordinary UTF8, and the new limiter compares its row-encoded bytes, so it drops a. DENSE_RANK has the same problem.
The collated-scan fallback does not protect this case because the stored column is plain STRING. The cast stays in a native projection through the default codegen dispatcher, and the required two-key [grp, cast-result] sort passes supportedSortType, which only checks collations in its single-key branch. This particular A, a, b sequence is correctly sorted under both binary and case-insensitive ordering, so the lost row is not explained by an existing sort-order mismatch. The conversion path also permits this with spark.comet.exec.window.enabled=false, and a later Spark Window cannot restore a row already discarded by native Partial WGL.
I verified that Spark 4.0.1 returns both peers and that the isolated native Rank and DenseRank operators each retain only A. The native probes used a diagnostic copy with only the timer-borrow compilation issue corrected, not an end-to-end run of the unchanged head. Please gate collated key types, including nested collated types, until peer equality preserves Spark's collation semantics.
Which issue does this PR close?
Closes #4837 .
Rationale for this change
Adds native support for Spark's
WindowGroupLimitExec(SPARK-37099, Spark 3.5+), which computes per-partition top-K forROW_NUMBER,RANK, andDENSE_RANK.PartitionedRankLimitExec(native/core/src/execution/operators/rank_limit.rs) does a single streaming pass over the child sorted by[partition_keys..., order_keys...], matching Spark'sSimpleLimitIterator/RankLimitIteratortie semantics exactly.native/core/src/execution/planner.rs) routesROW_NUMBERwithoutPARTITION BYtoLocalLimitExec(the sorted child makes first-K equivalent to top-K); every other combination goes toPartitionedRankLimitExec.WindowGroupLimitand enumRankLikeFunction(native/proto/src/proto/operator.proto).CometWindowGroupLimitExecplus per-versionShimCometWindowGroupLimit(no-op on 3.4, real on 3.5 / 4.0 / 4.1). Registration inCometExecRule.nativeExecsis gated on the shim so the 3.4build stays clean.
spark.comet.exec.windowGroupLimit.enabled(defaulttrue) viaCometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.RemoveRedundantWindowGroupLimitsSuiteto also countCometWindowGroupLimitExec.spark/src/test/resources/sql-tests/expressions/window/coverROW_NUMBER/RANK/DENSE_RANK, datatypes, edge cases, and scalar-subquery-shaped rewrites.