feat: support WindowGroupLimitExec

[comphead]

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 for ROW_NUMBER, RANK, and DENSE_RANK.

• New native 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's
  SimpleLimitIterator / RankLimitIterator tie semantics exactly.
• Planner (native/core/src/execution/planner.rs) routes ROW_NUMBER without PARTITION BY to LocalLimitExec (the sorted child makes first-K equivalent to top-K); every other combination goes to
  PartitionedRankLimitExec.
• New proto message WindowGroupLimit and enum RankLikeFunction (native/proto/src/proto/operator.proto).
• Scala serde CometWindowGroupLimitExec plus per-version ShimCometWindowGroupLimit (no-op on 3.4, real on 3.5 / 4.0 / 4.1). Registration in CometExecRule.nativeExecs is gated on the shim so the 3.4
  build stays clean.
• New config spark.comet.exec.windowGroupLimit.enabled (default true) via CometConf.COMET_EXEC_WINDOW_GROUP_LIMIT_ENABLED.
• Spark diff patches for 3.5.9 / 4.0.4 / 4.1.3 teach RemoveRedundantWindowGroupLimitsSuite to also count CometWindowGroupLimitExec.
• SQL file tests under spark/src/test/resources/sql-tests/expressions/window/ cover ROW_NUMBER / RANK / DENSE_RANK, datatypes, edge cases, and scalar-subquery-shaped rewrites.
• TPC-DS golden plans updated for q44, q67, q70 (v1_4 3.5 / 4.0) and q67a, q70a (v2_7 3.5).

[comphead]

Interesting, we dont support Subquery? 🤔

[andygrove]

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 rank_limit.rs against Spark's RankLimitIterator, DenseRankLimitIterator, and SimpleLimitIterator and I think it is right. Spark checks rank < limit before calling increaseRank(), so Spark over-emits by one tie group (RANK over [10,10,10,20] with limit 2 emits the 20). Comet computes the exact rank and drops it, which is fine because InsertWindowGroupLimit only fires when there is a Filter on the rank above the Window. The degenerate cases all fall out correctly too. The test coverage is also much better than I expected, especially running every query under windowGroupLimitThreshold=-1,1000 so the pushdown and non-pushdown shapes are both checked, and covering the SPARK-46526 PARTITION BY t2c ORDER BY t2c shape. q44 going from 6 transitions down to 2 is great.

A few things I would like to see addressed.

1. equals and hashCode on CometWindowGroupLimitExec (lines 140 and 151) leave out the rank-like function. The case class carries partitionSpec, orderSpec, and limit but not the function, and nativeOp is not compared either, so two nodes differing only in RowNumber versus Rank compare equal. sameResult and canonicalized bottom out in equals, and those drive ReuseSubquery and the cache manager, so I think two otherwise-identical subqueries using different rank functions could be wrongly reused. CometWindowExec includes windowExpression in its equals for the same reason. Could rankLikeFunction go into the case class and both overrides? Spark's mode is dropped too, and adding it to stringArgs would help, since right now the Partial and Final nodes render identically and that makes the q44 golden file hard to read.

2. A duplicated partition key becomes a hard native error rather than a fallback. LexOrdering::new dedups by PhysicalExpr keeping the first occurrence, so for PARTITION BY a, a ORDER BY a all three collapse to one entry, partition_prefix_len is 2 while expr.len() is 1, and try_new returns DataFusionError::Internal (rank_limit.rs:80). tryConvertToComet has no try/catch, so that fails the query instead of falling back. It is contrived SQL but generated SQL does produce shapes like it. Rather than guarding, would it be cleaner to pass the partition-by and order-by lists to the operator separately and build the two RowConverters from those, using the LexOrdering only for compute_properties and required_input_ordering? That would remove the guard, the failure mode, and the three comment blocks at planner.rs:2400-2432 that exist purely to reason about the dedup. The partition_prefix_len slicing is the only thing coupling correctness to a DataFusion implementation detail and I do not think it has to.

3. The -0.0 versus 0.0 divergence needs more than a test comment. window_group_limit_rank.sql says it outright: the FP section keeps the cutoff above the zero values to avoid exercising it. Spark's SQLOrderingUtil.compareDoubles is if (x == y) 0 else Double.compare(x, y), so -0.0 == 0.0 and RankLimitIterator ties them, while the Arrow row encoder compares bytes and splits them. With rank() = 1 over [0.0, -0.0, 1.0] in one partition Spark keeps two rows and Comet keeps one. compatibility/floating-point.md already has a general positive/negative-zero caveat so this fits an existing family, but nothing a user reads about WindowGroupLimitExec points at it. Could we get a tracking issue, a line in floating-point.md, and the test the file currently declines to write, marked ignore(<issue>) if it fails? Leaving the only record of a known wrong-answer case inside a comment in the test that avoids it means nobody finds it later.

4. ShimCometWindowGroupLimit.extract throws IllegalStateException on an unrecognized rank-like function, in both the 3.5 and 4.x copies. It is unreachable today since InsertWindowGroupLimit only produces the three, but tryConvertToComet does not catch, so a future Spark adding a fourth would fail a query that used to work. Could it return None and let convert record a fallback reason instead?

5. Three hand-edited doc pages still say this is unsupported and need updating here, since GenerateDocs does not touch them. user-guide/latest/operators.md:106 still marks WindowGroupLimitExec as 🔜 with "falls back today", user-guide/latest/compatibility/operators.md:74 still says "not yet supported and falls back to Spark", and contributor-guide/roadmap.md:33 lists it as pending.

6. rank_limit.rs has no unit tests. It is 378 lines with cross-batch state in prev_partition, prev_order, rank, and count, and the SQL test tables top out around 16 rows so everything lands in a single batch. The only real multi-batch coverage is indirect via Verify TPC-DS. Could we add tests for a partition boundary landing exactly on a batch boundary, a limit reached mid-batch with a new partition starting in the next one, a batch that filters to zero rows, and a RANK tie run spanning batches?

7. PartitionedRankLimitExec never implements metrics(), so the numOutputRows that Spark's WindowGroupLimitExec reports shows up as zero even though CometWindowGroupLimitExec inherits baselineMetrics. scan.rs:218 wires up BaselineMetrics the same way, so it should be easy to match.

8. There is no early exit once a partition's limit is reached (rank_limit.rs:324). this_rank is monotonic within a partition, so once keep goes false every remaining row in that partition is dropped, yet the loop still encodes the ORDER BY key and does a row comparison for each. Spark's GroupedLimitIterator.skipRemainingRows skips that work explicitly. For the skewed-partition small-K case this operator exists to optimize, that is the hot path, and only the partition-key compare is needed once the limit is hit.

9. Two comments describe things that do not exist. PartitionedTopKExec in the planner.rs:2426 comment is not anywhere in the repo, so a reader cannot check the comparison. And the required_input_ordering comment at rank_limit.rs:156 says declaring the requirement lets EnforceSorting insert a SortExec if the sort got dropped, but we do not run any DataFusion physical optimizer rules on plans built from the serialized Spark plan, so there is no safety net there. Correctness rests entirely on the JVM-side plan, and I would rather the comment say that, since it is exactly what a reviewer needs to trust.

10. The six test files are under sql-tests/expressions/window/, which holds window expression tests like lag_lead.sql. sql-tests/windows/ is where operator-level window tests live (window_functions.sql), so I think they belong there.

11. The PR description is still all template comments. We generate changelogs from these, so could you fill in the rationale and summary? And since this operator adds per-row Arrow row-encoding on top of what Spark's UnsafeProjection grouping already does, some before/after numbers on q44, q67, and q70 would confirm the win the plan-stability diff implies.

[comphead]

Thanks @andygrove addressed in 7ef79eb

[sunchao]

[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)

[sunchao]

[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 rk

The 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.

[sunchao]

One additional correctness issue at 77702dba, reproduced against the unchanged native operator. Details and the minimal regression case are inline.

[sunchao]

[P1] Record skipped rows before taking the prefix shortcut

Could we set first_dropped_at when this branch drops a row? If a batch starts inside a partition exhausted by the previous batch, its leading drops are not recorded. With fetch = 1 and sorted (partition, order) rows:

Batch 0: (1,10), (1,20), (1,30), (1,40)
Batch 1: (1,50), (2,10), (2,20), (3,10)

Expected output: (1,10), (2,10), (3,10)
Actual output:   (1,10), (1,50), (2,10)

Batch 1 builds the correct mask [false, true, false, true], but records first_dropped_at = 2 and kept = 2. The shortcut at lines 419–420 therefore returns batch.slice(0, 2) instead of applying that mask. This permanently loses partition 3's top row, which a subsequent Final limiter or window/filter cannot recover.

I reproduced this for ROW_NUMBER, RANK, and DENSE_RANK by compiling and executing the unchanged rank_limit.rs from 77702dba in an isolated harness with DataFusion 54.1.0 and Arrow 58.4.0. The six existing tests pass and all three regression cases fail. Recording the first dropped position in this branch makes all nine pass. Could we include this mixed-partition batch case in the native tests?

[comphead]

Checking this, not sure if we can have also SQL test for this, but let me try

[pingz-oai]

Reviewed exact commit 77702dba39c8472094bf55d596f8b0365684c024.

The existing P1 batch-prefix row-loss report remains reproducible. I independently executed the unchanged rank_limit.rs with DataFusion 54.1.0, Arrow 58.4.0, and Futures 0.3.33: all six existing tests pass, while the reported regression fails for ROW_NUMBER, RANK, and DENSE_RANK. A separate Spark 4.1.3 baseline returns the expected three groups. Recording the first dropped position in the exhausted-partition branch in a diagnostic copy makes all nine native tests pass.

The merge-base version leaves WindowGroupLimit in Spark; this PR enables the faulty native path by default. I am withholding approval for that P1. Its existing inline thread already contains the triggering batches, incorrect result, and affected lines, so I have not duplicated the inline finding. No additional reproducible P1/P2 regression found.

[ziting-openai]

Re-reviewed the current head: exhausted cross-batch partition prefixes now record every dropped row, regression coverage spans ROW_NUMBER/RANK/DENSE_RANK, and non-default collated partition/order keys correctly fall back.

[pingz-oai]

Reviewed exact commit 8d2ea9830d1b5e0898ac43f349cd835343b7501d.

The previously reported cross-batch row-loss regression is fixed. The unchanged native source passes all 7 current unit tests, all 3 prior focused regression probes, and a 6,912-case matrix with zero mismatches using the pinned Arrow 58.4.0 / DataFusion 54.1.0 dependencies. I also traced Spark version shims, ordering/distribution, plan identity, collation fallback, and optimizer safety conditions. No remaining P1/P2 findings.

Validation included 22 Spark 4.1.3 plan/semantics probes and isolated native execution; a full Comet/Spark integration build was not run.

Reviewed by Codex on behalf of pingz-oai.

[sunchao]

Re-reviewed commit 8d2ea9830d1b5e0898ac43f349cd835343b7501d with five fresh independent reviewers. The earlier compilation, collation fallback, plan-identity, and cross-batch row-loss issues are addressed. No new, unique P1/P2 findings.

Validation included all seven current native unit tests, three earlier regression probes, and a 30,720-case native matrix. Full Comet integration was not run locally. The existing floating-point peer-equality limitation remains.

CI is still incomplete: the macOS Spark 4.0 scans job hit a native thread-cleanup SIGSEGV after a passing test. The available evidence does not establish whether this PR caused that crash.

[comphead]

Thanks @sunchao for the review!