perf: drain a large ordered aggregate range once instead of in batch-size steps

[msafonov]

Problem

push_ranking_results (prod thrift dbt) runs ~5x slower on Comet than on Spark, and slower every night as its input grows: 1473 s on 1.1-8 vs a vanilla median of 305 s. Reproduced on GCP with the same input (10 executors, c4a-lssd):

	query time
vanilla	89 s
Comet 1.1-8	1290 s
Comet 1.1-7	1288 s
Comet 1.1-8, window.partitionAggregate.enabled=false	1243 s

So it is not a 1.1-8 regression. Per-operator metrics: the time is in the COUNT(DISTINCT ranking_id) hash aggregates: Partial ~829 s and PartialMerge ~813 s, both in one task holding the largest window partition (~30M groups).

Root cause

The aggregates follow MAX(...) OVER (PARTITION BY partition_date, winner_type), so their input is sorted on a prefix of the grouping keys and DataFusion runs them in the ordered partial / ordered final streams. Groups become ready only when the sorted prefix changes or input ends. The ready range was then emitted with EmitTo::First(batch_size) per output batch. Each such emit walks the whole group hash table (map.retain) and copies every remaining key and accumulator state, so draining N groups costs O(N² / batch_size).

Change

aggregate_hash_table/common_ordered.rs in vendored datafusion-physical-plan: when the ready range is larger than batch_size and the stream outputs states (Partial, PartialReduce), emit the range once into a materialized batch and hand it out in batch_size slices, as the unordered hash streams already do (MaterializedAggregateOutput). The undelivered slices count towards the table's memory and emptiness. Under memory pressure they are prepended to the state batch. Final-value output (is_final) keeps the old path.

Tests

• ordered_drain_tests.rs (new):
  • 300k groups behind one prefix value: ordered Partial and ordered PartialReduce must stay within 5x + 0.5 s of the unordered stream. Before the fix: 19.0 s vs 31 ms and 19.4 s vs 37 ms. After: pass.
  • Memory-starved ordered Partial (6 MiB pool, 4 prefixes): all groups and counts are kept, and the pool is released.
• CometAggregateSuite: COUNT(DISTINCT) over a window-sorted input with one huge ordered key range, checkSparkAnswerAndOperator with batch size 64.

GCP measurement with a test image and a data compare follow in a comment.

🤖 Generated with Claude Code

[msafonov]

GCP measurement (push_ranking_results query, same input 2026-10-10..11, 10 executors x 8 cores, c4a-lssd, prod thrift Comet config)

variant	original query	deterministic variant	rows (det)	sha256 of sorted rows (det)
vanilla	89 s	88.8 s	4347	3476ff8c…
Comet 1.1-8	1290 s	—	—	—
Comet 1.1-7	1288 s	—	—	—
this branch (test image agg-emit-test-ee5d84128)	60 s	62.9 s	4347	3476ff8c… (identical to vanilla)

The original model is nondeterministic: FIRST_VALUE(requestType) ORDER BY weight DESC has ~75k rankings with tied top weights, and four runs gave four different row counts. The data compare therefore uses the same query with the window order completed by requestType, requestId. Its output is identical to vanilla.

Local: vendor datafusion-physical-plan 1769/0; rust workspace 1929/0; CometAggregateSuite + CometWindowExecSuite + CometPartitionAggregateWindowSuite 301/0 (2 canceled, pre-existing); agg fuzz full 31/31; fmt/clippy clean.

🤖 Generated with Claude Code