fix: Fix the core correctness issues with windowed aggregate queries

[andygrove]

Summary

This PR enables a subset of window aggregate functions (COUNT, SUM, MIN, MAX) for native execution in Comet. Window functions were previously disabled with "Native WindowExec has known correctness issues" - this PR addresses the core issues for simple cases.

Tracking Issue: #2721

What's Now Supported

• Window aggregates: COUNT, SUM, MIN, MAX
• PARTITION BY only (no ORDER BY)
• ORDER BY only (no PARTITION BY)
• PARTITION BY with ORDER BY when partition columns are a subset of order columns
  (e.g., PARTITION BY a ORDER BY a, b works)

Example Supported Queries

-- PARTITION BY only
SELECT a, COUNT(*) OVER (PARTITION BY a) as cnt FROM table;
SELECT a, SUM(c) OVER (PARTITION BY a) as sum_c FROM table;

-- ORDER BY only
SELECT a, MIN(c) OVER (ORDER BY b) as min_c FROM table;

-- Multiple aggregates
SELECT a, b, c,
  COUNT(*) OVER (PARTITION BY a) as cnt,
  SUM(c) OVER (PARTITION BY a) as sum_c,
  MIN(c) OVER (PARTITION BY a) as min_c,
  MAX(c) OVER (PARTITION BY a) as max_c
FROM table;

What's NOT Yet Supported (Falls Back to Spark)

• AVG window aggregate (native implementation has known issues)
• PARTITION BY a ORDER BY b where partition columns differ from order columns
• ROWS BETWEEN frames with PARTITION BY and ORDER BY on different columns
• Ranking functions: ROW_NUMBER, RANK, DENSE_RANK, PERCENT_RANK, NTILE, CUME_DIST
• Offset functions: LAG, LEAD
• Value functions: FIRST_VALUE, LAST_VALUE, NTH_VALUE
• RANGE BETWEEN with numeric/temporal expressions (Invalid argument error: Invalid arithmetic operation: Int32 - Int64 #1246)

Changes

1. Change Support Level from Blanket Incompatible

File: spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala

Changed getSupportLevel from blanket Incompatible("Native WindowExec has known correctness issues") to:

• Compatible() when partition expressions are a subset of order expressions (or when one is empty)
• Unsupported() when DataFusion's constraint is violated

2. Remove Overly Restrictive Validation

File: spark/src/main/scala/org/apache/spark/sql/comet/CometWindowExec.scala

Removed validatePartitionAndSortSpecsForWindowFunc which required partition columns to exactly match order columns. This was more restrictive than necessary.

3. Add Explicit Sort for ORDER BY

File: native/core/src/execution/planner.rs

Added explicit SortExec before BoundedWindowAggExec when ORDER BY is present to ensure correct input ordering.

4. Update Tests

File: spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala

• Enabled tests for COUNT, SUM, MIN, MAX with PARTITION BY only
• Kept original failing tests as ignore to document what still doesn't work
• Changed ranking function tests to verify fallback behavior

5. Update Documentation

File: docs/source/user-guide/latest/compatibility.md

Updated Window Functions section to accurately document supported and unsupported functionality.

Related Issues

• Partially addresses [EPIC] Improve native execution for window functions #2721 (window function correctness - simple cases now work)
• Does NOT address Invalid argument error: Invalid arithmetic operation: Int32 - Int64 #1246 (RANGE BETWEEN)

[andygrove]

This seems to be the main fix

[codecov-commenter]

Codecov Report

❌ Patch coverage is 81.81818% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 60.67%. Comparing base (f09f8af) to head (7e985ab).
⚠️ Report is 858 commits behind head on main.

Files with missing lines	Patch %	Lines
...a/org/apache/spark/sql/comet/CometWindowExec.scala	80.00%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #3191      +/-   ##
============================================
+ Coverage     56.12%   60.67%   +4.54%     
- Complexity      976     1437     +461     
============================================
  Files           119      170      +51     
  Lines         11743    15733    +3990     
  Branches       2251     2596     +345     
============================================
+ Hits           6591     9546    +2955     
- Misses         4012     4842     +830     
- Partials       1140     1345     +205     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[andygrove]

I need to update the golden files

[andygrove]

@comphead this is ready for another look. I am not 100% happy with this but it does fix some tests and the feature is still disabled by default, so it seems like some progress. Some of the tests got updated so we need to review and make sure that these are valid changes and not ignoring issues.