fix: Fix the core correctness issues with windowed aggregate queries by andygrove · Pull Request #3191 · apache/datafusion-comet

andygrove · 2026-01-15T16:26:58Z

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 · 2026-01-15T16:43:00Z

native/core/src/execution/planner.rs

+                // Ensure input is properly sorted when ORDER BY is present
+                // BoundedWindowAggExec requires InputOrderMode::Sorted
+                let needs_explicit_sort = !sort_exprs.is_empty();
+                let sorted_child: Arc<dyn ExecutionPlan> = if needs_explicit_sort {
+                    // Insert explicit sort to ensure data ordering
+                    Arc::new(SortExec::new(
+                        LexOrdering::new(sort_exprs.to_vec()).unwrap(),
+                        Arc::clone(&child.native_plan),
+                    ))
+                } else {
+                    Arc::clone(&child.native_plan)
+                };


This seems to be the main fix

…omet into fix-windowed-aggs

codecov-commenter · 2026-01-15T19:26:07Z

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 · 2026-01-15T20:43:06Z

I need to update the golden files

- Restore original ignored tests for ROWS BETWEEN with PARTITION BY + ORDER BY and convert them to use checkSparkAnswerAndFallbackReason to verify the correct fallback message: "Partition expressions must be a subset of order expressions for native window" - Fix documentation to accurately reflect what's supported: - Remove AVG from supported list (has native implementation issues) - Clarify PARTITION BY/ORDER BY restriction (partition must be subset of order) - Clarify ROWS BETWEEN limitations - Fix misleading test names ("COUNT with ROWS frame" -> "COUNT with PARTITION BY only") - Convert several ignored tests to active tests that verify fallback behavior: - ROWS BETWEEN tests (COUNT, SUM, AVG, MAX) - ORDER BY DESC test - Multiple PARTITION BY/ORDER BY columns tests - RANGE BETWEEN tests Co-Authored-By: Claude Opus 4.5 <[email protected]>

…c results Window function tests with ROWS BETWEEN frames were failing due to non-deterministic sort order when using ORDER BY with non-unique values. When CometSort encounters rows with identical sort keys, the relative order of those rows is undefined, causing window function results to vary between Spark and Comet executions. Updated all window fallback tests to include column 'c' (which contains unique values) as a tiebreaker in ORDER BY clauses, ensuring deterministic ordering and consistent test results. Tests updated: - ROWS BETWEEN tests (COUNT, SUM, AVG, MAX) - ORDER BY DESC test - Multiple PARTITION BY/ORDER BY columns tests - RANGE BETWEEN tests Co-Authored-By: Claude Sonnet 4.5 <[email protected]>

RANGE window frames with value boundaries cannot have multiple ORDER BY columns in Spark SQL. Reverted the two RANGE BETWEEN tests back to single ORDER BY column. These tests use b=i (unique values) instead of b=i%5, so ordering is already deterministic. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>

andygrove · 2026-01-20T20:08:08Z

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

# Conflicts: # docs/source/user-guide/latest/configs.md

andygrove · 2026-02-04T19:02:53Z

I am going to break this down into some smaller PRs

fix window aggregates

851ad54

andygrove changed the title ~~fix: Fix some bugs in windowed aggregate implementation~~ fix: Fix the core correctness issues with windowed aggregate queries Jan 15, 2026

docs

24a4e4f

andygrove requested a review from comphead January 15, 2026 16:40

andygrove commented Jan 15, 2026

View reviewed changes

andygrove added 2 commits January 15, 2026 09:51

prettier

0a71a26

scalastyle

32ea9e9

andygrove marked this pull request as ready for review January 15, 2026 18:02

andygrove added 2 commits January 15, 2026 12:02

update docs

4a36882

Merge branch 'fix-windowed-aggs' of github.com:andygrove/datafusion-c…

7a1f47d

…omet into fix-windowed-aggs

andygrove added 5 commits January 15, 2026 15:01

remove feature flags

f92396d

update golden files

d20cc50

update docs

092f6b4

update docs

17cec95

update docs

a83b73f

andygrove marked this pull request as draft January 16, 2026 01:06

andygrove added 3 commits January 15, 2026 19:17

fix

964e2f9

docs

2710ff8

update golden files

d50e34e

andygrove marked this pull request as ready for review January 16, 2026 18:32

andygrove marked this pull request as draft January 16, 2026 18:48

andygrove force-pushed the fix-windowed-aggs branch from 52ab85c to 11b615d Compare January 16, 2026 19:01

andygrove and others added 2 commits January 20, 2026 09:17

andygrove marked this pull request as ready for review January 20, 2026 20:06

Merge remote-tracking branch 'origin/main' into fix-windowed-aggs

256a21d

# Conflicts: # docs/source/user-guide/latest/configs.md

andygrove closed this Feb 4, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: Fix the core correctness issues with windowed aggregate queries#3191

fix: Fix the core correctness issues with windowed aggregate queries#3191
andygrove wants to merge 18 commits intoapache:mainfrom
andygrove:fix-windowed-aggs

andygrove commented Jan 15, 2026 •

edited

Loading

Uh oh!

andygrove Jan 15, 2026

Uh oh!

codecov-commenter commented Jan 15, 2026 •

edited

Loading

Uh oh!

andygrove commented Jan 15, 2026

Uh oh!

andygrove commented Jan 20, 2026

Uh oh!

andygrove commented Feb 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

andygrove commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What's Now Supported

Example Supported Queries

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

Changes

1. Change Support Level from Blanket Incompatible

2. Remove Overly Restrictive Validation

3. Add Explicit Sort for ORDER BY

4. Update Tests

5. Update Documentation

Related Issues

Uh oh!

andygrove Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

andygrove commented Jan 15, 2026

Uh oh!

andygrove commented Jan 20, 2026

Uh oh!

andygrove commented Feb 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

andygrove commented Jan 15, 2026 •

edited

Loading

codecov-commenter commented Jan 15, 2026 •

edited

Loading