[PATCH] Consider window functions when reordering GROUP BY items

Hi hackers,

Currently, preprocess_groupclause() reorders the GROUP BY items to
match the query's ORDER BY clause so that a single sort can serve both
the grouping step and the final ordering. However, when the query
contains window functions, the sort performed directly above the
grouping step is the one required by the first window clause, not the
one for ORDER BY, which is only performed above the WindowAgg nodes.
Unless the window's sort requirements happen to coincide with a prefix
of the ORDER BY clause, an additional sort is needed to provide the
first WindowAgg with correctly ordered input:

CREATE TABLE t(a int, b int);
SET enable_hashagg = off;

-- master
EXPLAIN (COSTS OFF)
SELECT a, b, count(*) OVER (PARTITION BY b) FROM t GROUP BY a, b;
QUERY PLAN
------------------------------------------
WindowAgg
Window: w1 AS (PARTITION BY b)
-> Sort
Sort Key: b
-> Group
Group Key: a, b
-> Sort
Sort Key: a, b
-> Seq Scan on t

The ordering of the GROUP BY items is semantically insignificant, so
the attached patch teaches preprocess_groupclause() to match them
against the first active window's PARTITION BY and ORDER BY keys when
there are window functions, then against the query's ORDER BY clause
for any remaining items. With the patch, the query above needs a
single sort:

-- patched
QUERY PLAN
------------------------------------
WindowAgg
Window: w1 AS (PARTITION BY b)
-> Group
Group Key: b, a
-> Sort
Sort Key: b, a
-> Seq Scan on t

We match the first window only, because it alone can share a sort with
the grouping step: select_active_windows() has already fixed the
windows' evaluation order at this point, placing the window with the
strongest sort requirements first, directly above the grouping step.
If the first window requires no sort at all then no active window
does, and we fall back to matching the ORDER BY clause as before.

Appending the query's ORDER BY keys after the window's keys orders the
otherwise-arbitrary tail of the GROUP BY items. Since WindowAgg
preserves its input ordering, this can also save the final sort (e.g.
PARTITION BY c ... GROUP BY a, b, c ORDER BY c, b now sorts once
instead of the two sorts you'd get from matching the window alone),
and it retains master's behavior for queries where the window's keys
form a prefix of the ORDER BY clause, which the existing ORDER BY
matching already handles well. When the ORDER BY does not lead with
the window's keys, the final sort remains but one sort is still saved:

-- same query plus ORDER BY a: master needs 3 sorts, patched 2
master: Sort(a) <- WindowAgg <- Sort(b) <- Group <- Sort(a, b)
patched: Sort(a) <- WindowAgg <- Group <- Sort(b, a)

Benchmark (3M rows, ~200k groups, enable_hashagg = off,
work_mem = 512MB, jit = off, no parallelism, median of 3):

CREATE TABLE t_hi AS
SELECT (i % 2)::int a, (random() * 100000)::int b
FROM generate_series(1, 3000000) i;

SELECT a, b, count(*) OVER (PARTITION BY b)
FROM t_hi GROUP BY a, b;

master: 5215 ms patched: 2340 ms (2.2x faster)

One honest caveat. Like the long-standing ORDER BY matching, this
heuristic is insensitive to column cardinalities, and so is
cost_sort(), so the planner cannot correct a bad choice. Moving a
low-cardinality window key to the front of a large grouping sort
defeats tuplesort's leading-key specialization and abbreviated keys,
and can cost more than the eliminated sort saves when the number of
groups is small relative to the input. Swapping the cardinalities in
the table above shows the worst case:

CREATE TABLE t_lo AS
SELECT (random() * 100000)::int a, (i % 2)::int b
FROM generate_series(1, 3000000) i;

SELECT a, b, count(*) OVER (PARTITION BY b)
FROM t_lo GROUP BY a, b;

master: 2379 ms patched: 5175 ms (2.2x slower)

Here the patch trades a cheap 200k-row sort for making the 3M-row sort
lead with a 2-value column. Note that master shows the same failure
mode today without any windows involved (e.g. GROUP BY a, b ORDER BY b
with a low-cardinality b), so the patch extends an already-accepted
trade-off rather than introducing a new class of problem. It is also
mitigated in practice: a small number of groups is exactly when hash
aggregation tends to win (with default settings the planner picks
HashAggregate for the query above and the issue disappears), whereas
sorted aggregation tends to be chosen when groups are numerous, which
is when the eliminated sort is large. Choosing the ordering based on
cost would require a column-order-aware sort cost model and belongs in
get_useful_group_keys_orderings(); I left that as future work, but I'm
happy to explore it if there's interest.

The patch moves the preprocessing of a plain GROUP BY clause until
after select_active_windows(); nothing in between examines
root->processed_groupClause, and the grouping-sets path is unaffected.
New regression tests cover the plan shapes above. make check passes.

Patch attached.

Regards,
Zhenglong Li