Partial aggregation stops skipping after early emits under memory pressure #27350

Describe the bug

PartialHashAggregateStream gives the skip probe all rows seen so far, but only the groups currently in its table (hash_stream.rs#L401):

self.update_skip_aggregation_probe(input_rows, hash_table.building_group_count());

Under memory pressure the stream emits all groups early and continues with an empty table (hash_stream.rs#L306-L313). After that, num_groups counts only the groups since the last emit, while input_rows keeps growing (skip_partial.rs#L78-L91). The ratio stays below skip_partial_aggregation_probe_ratio_threshold, so the probe fires late or never, even for unique keys. The legacy GroupedHashAggregateStream has the same pattern (grouped_hash_stream.rs#L1399).

To Reproduce

SET datafusion.execution.target_partitions = 2;
SET datafusion.runtime.memory_limit = '4M';
EXPLAIN ANALYZE
SELECT (value * 1103515245 + 12345) % 2147483648 AS k, SUM(value)
FROM generate_series(1, 2000000)
GROUP BY k;

All 2M keys are unique. The affine map only makes the input unordered, so the hash path is used instead of the clustered one.

AggregateExec: mode=Partial on 51960dbad:

The limited runs send the same 2M rows downstream, at about 4x the partial CPU. The fair pool (datafusion-cli -m 16M --mem-pool-type fair) behaves the same: 48 emits, 0 skipped.

Expected behavior

The probe should still skip after skip_partial_aggregation_probe_rows_threshold rows. Counting the groups already emitted would fix it:

self.update_skip_aggregation_probe(
    input_rows,
    self.emitted_groups + hash_table.building_group_count(),
);

This counts output rows, the same total reduction_factor reports, which is what the skip decision needs. Resetting input_rows on each emit would not work, because frequent emits would keep the probe below its rows threshold.

Additional context

Repeated keys are affected too. 20M rows over about 200K keys, 8 partitions and a 32 MB fair pool gave 590 emits, 0 skipped rows, and 360 ms instead of 70 ms of partial CPU.

Related: #26124 (proactive flushing would hit this on every flush), #22405, #24996.

Willingness to contribute

I plan to contribute a fix for this bug independently