From cc8f88adf08a93f8d2cc07695569345bee16dae8 Mon Sep 17 00:00:00 2001 From: Matheus Alcantara Date: Mon, 31 Aug 2026 10:49:50 -0300 Subject: [PATCH v8 4/4] Enable full partitionwise aggregate for partition keys wrapped by RelabelType group_by_has_partkey() decides whether a partitioned aggregate can be computed independently per partition (PARTITIONWISE_AGGREGATE_FULL) or whether partial aggregates have to be combined afterwards (PARTITIONWISE_AGGREGATE_PARTIAL). It does so by matching each partition key expression against the GROUP BY expressions with equal(). The grouping expressions have their RelabelType decorations stripped before that comparison, but the partition key expressions do not. When the partition key is an expression involving a binary-compatible cast, for example PARTITION BY LIST ((col::text)) on a varchar column, the stored partition key expression is itself a RelabelType, so no GROUP BY expression can ever match it: grouping by the column compares a bare Var against a RelabelType, and spelling out the cast strips the grouping side down to a bare Var while the partition key side stays wrapped. Full partitionwise aggregation is therefore never chosen for such a key, and the plan falls back to partial aggregation with a finalize step on top. Fix by stripping RelabelType from the partition key expression as well. Collation correctness is unaffected: the partition collation comes from the PartitionScheme rather than from the expression, and the grouping expression's collation is captured before it is stripped, so the existing comparison of the two is unchanged. When the partition key is stripped in the function group_by_has_partkey(), grouping employs the equality operator of the argument type, whereas partitioning uses the result type, and a binary-coercible cast is not required to preserve the equality: citext is binary-coercible to text but is compared in a case-insensitive manner. As a result, a group might span across partitions and yield an incorrect answer. Therefore, require the grouping clause's equality operator to be a member of the partitioning operator family, as have_partkey_equi_join() already does for the clause operator. Discussion: https://postgr.es/m/DEYVEFEFMSDC.23KLLYQ3F81R9@gmail.com --- src/backend/optimizer/plan/planner.c | 23 ++++- .../regress/expected/partition_aggregate.out | 99 +++++++++++++++++++ src/test/regress/sql/partition_aggregate.sql | 47 +++++++++ 3 files changed, 168 insertions(+), 1 deletion(-) diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c index c3c158a253d..c7b0154a261 100644 --- a/src/backend/optimizer/plan/planner.c +++ b/src/backend/optimizer/plan/planner.c @@ -8531,12 +8531,26 @@ group_by_has_partkey(RelOptInfo *input_rel, foreach(lc, partexprs) { ListCell *lg; + ListCell *lgc; Expr *partexpr = lfirst(lc); Oid partcoll = input_rel->part_scheme->partcollation[cnt]; + Oid partopfamily = input_rel->part_scheme->partopfamily[cnt]; - foreach(lg, groupexprs) + /* + * Strip any RelabelType decorations, to match the stripping done + * on the grouping expressions below. A partition key involving a + * binary-compatible cast, such as ((col::text)) on a varchar + * column, is itself stored wrapped in a RelabelType. The + * collation is not lost, since partcoll and groupcoll are + * compared separately below. + */ + while (partexpr && IsA(partexpr, RelabelType)) + partexpr = ((RelabelType *) partexpr)->arg; + + forboth(lg, groupexprs, lgc, groupClause) { Expr *groupexpr = lfirst(lg); + SortGroupClause *sgc = lfirst_node(SortGroupClause, lgc); Oid groupcoll = exprCollation((Node *) groupexpr); /* @@ -8557,6 +8571,13 @@ group_by_has_partkey(RelOptInfo *input_rel, partcoll != groupcoll) return false; + /* + * Reject a match if the operator that grouping will use + * is not part of the partitioning operator family. + */ + if (!op_in_opfamily(sgc->eqop, partopfamily)) + return false; + found = true; break; } diff --git a/src/test/regress/expected/partition_aggregate.out b/src/test/regress/expected/partition_aggregate.out index c30304b99c7..4ce1154466b 100644 --- a/src/test/regress/expected/partition_aggregate.out +++ b/src/test/regress/expected/partition_aggregate.out @@ -961,6 +961,105 @@ SELECT a, c, sum(b), avg(c), count(*) FROM pagg_tab_m GROUP BY (a+b)/2, 2, 1 HAV 20 | 40 | 50 | 40.0000000000000000 | 5 (6 rows) +-- Partition by an expression that is a binary-compatible cast, so that the +-- stored partition key expression is itself wrapped in a RelabelType +CREATE TABLE pagg_tab_v (a int, c varchar(40)) PARTITION BY LIST ((c::text)); +CREATE TABLE pagg_tab_v_p1 PARTITION OF pagg_tab_v FOR VALUES IN ('0000', '0001'); +CREATE TABLE pagg_tab_v_p2 PARTITION OF pagg_tab_v FOR VALUES IN ('0002'); +INSERT INTO pagg_tab_v SELECT i, to_char(i % 3, 'FM0000') FROM generate_series(0, 2999) i; +ANALYZE pagg_tab_v; +-- Full aggregation as GROUP BY clause matches with PARTITION KEY +EXPLAIN (COSTS OFF) +SELECT c, sum(a), count(*) FROM pagg_tab_v GROUP BY c ORDER BY 1; + QUERY PLAN +---------------------------------------------------------- + Sort + Sort Key: pagg_tab_v.c + -> Append + -> HashAggregate + Group Key: pagg_tab_v.c + -> Seq Scan on pagg_tab_v_p1 pagg_tab_v + -> HashAggregate + Group Key: pagg_tab_v_1.c + -> Seq Scan on pagg_tab_v_p2 pagg_tab_v_1 +(9 rows) + +SELECT c, sum(a), count(*) FROM pagg_tab_v GROUP BY c ORDER BY 1; + c | sum | count +------+---------+------- + 0000 | 1498500 | 1000 + 0001 | 1499500 | 1000 + 0002 | 1500500 | 1000 +(3 rows) + +-- Full aggregation also when the GROUP BY clause spells out the cast +EXPLAIN (COSTS OFF) +SELECT c::text, sum(a), count(*) FROM pagg_tab_v GROUP BY c::text ORDER BY 1; + QUERY PLAN +---------------------------------------------------------- + Sort + Sort Key: ((pagg_tab_v.c)::text) + -> Append + -> HashAggregate + Group Key: (pagg_tab_v.c)::text + -> Seq Scan on pagg_tab_v_p1 pagg_tab_v + -> HashAggregate + Group Key: (pagg_tab_v_1.c)::text + -> Seq Scan on pagg_tab_v_p2 pagg_tab_v_1 +(9 rows) + +SELECT c::text, sum(a), count(*) FROM pagg_tab_v GROUP BY c::text ORDER BY 1; + c | sum | count +------+---------+------- + 0000 | 1498500 | 1000 + 0001 | 1499500 | 1000 + 0002 | 1500500 | 1000 +(3 rows) + +-- A binary-compatible cast need not preserve equality semantics, so matching +-- the stripped partition key is not enough on its own. Build a type that is +-- binary-coercible to text but compares case insensitively. +CREATE TYPE pagg_ci; +CREATE FUNCTION pagg_ci_in(cstring) RETURNS pagg_ci STRICT IMMUTABLE LANGUAGE internal AS 'textin'; +NOTICE: return type pagg_ci is only a shell +CREATE FUNCTION pagg_ci_out(pagg_ci) RETURNS cstring STRICT IMMUTABLE LANGUAGE internal AS 'textout'; +NOTICE: argument type pagg_ci is only a shell +LINE 1: CREATE FUNCTION pagg_ci_out(pagg_ci) RETURNS cstring STRICT ... + ^ +CREATE TYPE pagg_ci (input = pagg_ci_in, output = pagg_ci_out, like = text); +CREATE CAST (pagg_ci AS text) WITHOUT FUNCTION; +CREATE FUNCTION pagg_ci_eq(pagg_ci, pagg_ci) RETURNS bool + STRICT IMMUTABLE LANGUAGE sql AS $$SELECT lower($1::text) = lower($2::text)$$; +CREATE OPERATOR = (leftarg = pagg_ci, rightarg = pagg_ci, procedure = pagg_ci_eq); +CREATE FUNCTION pagg_ci_hash(pagg_ci) RETURNS int4 STRICT IMMUTABLE LANGUAGE sql AS $$SELECT hashtext(lower($1::text))$$; +CREATE OPERATOR CLASS pagg_ci_ops DEFAULT FOR TYPE pagg_ci USING hash AS OPERATOR 1 =, FUNCTION 1 pagg_ci_hash(pagg_ci); +CREATE TABLE pagg_tab_ci (a int, c pagg_ci) PARTITION BY LIST ((c::text)); +CREATE TABLE pagg_tab_ci_p1 PARTITION OF pagg_tab_ci FOR VALUES IN ('A'); +CREATE TABLE pagg_tab_ci_p2 PARTITION OF pagg_tab_ci FOR VALUES IN ('a'); +INSERT INTO pagg_tab_ci SELECT i, (CASE WHEN i % 3 = 0 THEN 'A' ELSE 'a' END)::pagg_ci FROM generate_series(1, 3000) i; +ANALYZE pagg_tab_ci; +-- Partial aggregation only +EXPLAIN (COSTS OFF) +SELECT c, count(*) FROM pagg_tab_ci GROUP BY c; + QUERY PLAN +------------------------------------------------------------ + Finalize HashAggregate + Group Key: pagg_tab_ci.c + -> Append + -> Partial HashAggregate + Group Key: pagg_tab_ci.c + -> Seq Scan on pagg_tab_ci_p1 pagg_tab_ci + -> Partial HashAggregate + Group Key: pagg_tab_ci_1.c + -> Seq Scan on pagg_tab_ci_p2 pagg_tab_ci_1 +(9 rows) + +SELECT c, count(*) FROM pagg_tab_ci GROUP BY c; + c | count +---+------- + A | 3000 +(1 row) + -- Test with multi-level partitioning scheme CREATE TABLE pagg_tab_ml (a int, b int, c text) PARTITION BY RANGE(a); CREATE TABLE pagg_tab_ml_p1 PARTITION OF pagg_tab_ml FOR VALUES FROM (0) TO (12); diff --git a/src/test/regress/sql/partition_aggregate.sql b/src/test/regress/sql/partition_aggregate.sql index 7c725e2663a..8befdca71fc 100644 --- a/src/test/regress/sql/partition_aggregate.sql +++ b/src/test/regress/sql/partition_aggregate.sql @@ -208,6 +208,53 @@ SELECT a, c, sum(b), avg(c), count(*) FROM pagg_tab_m GROUP BY (a+b)/2, 2, 1 HAV SELECT a, c, sum(b), avg(c), count(*) FROM pagg_tab_m GROUP BY (a+b)/2, 2, 1 HAVING sum(b) = 50 AND avg(c) > 25 ORDER BY 1, 2, 3; +-- Partition by an expression that is a binary-compatible cast, so that the +-- stored partition key expression is itself wrapped in a RelabelType + +CREATE TABLE pagg_tab_v (a int, c varchar(40)) PARTITION BY LIST ((c::text)); +CREATE TABLE pagg_tab_v_p1 PARTITION OF pagg_tab_v FOR VALUES IN ('0000', '0001'); +CREATE TABLE pagg_tab_v_p2 PARTITION OF pagg_tab_v FOR VALUES IN ('0002'); +INSERT INTO pagg_tab_v SELECT i, to_char(i % 3, 'FM0000') FROM generate_series(0, 2999) i; +ANALYZE pagg_tab_v; + +-- Full aggregation as GROUP BY clause matches with PARTITION KEY +EXPLAIN (COSTS OFF) +SELECT c, sum(a), count(*) FROM pagg_tab_v GROUP BY c ORDER BY 1; +SELECT c, sum(a), count(*) FROM pagg_tab_v GROUP BY c ORDER BY 1; + +-- Full aggregation also when the GROUP BY clause spells out the cast +EXPLAIN (COSTS OFF) +SELECT c::text, sum(a), count(*) FROM pagg_tab_v GROUP BY c::text ORDER BY 1; +SELECT c::text, sum(a), count(*) FROM pagg_tab_v GROUP BY c::text ORDER BY 1; + +-- A binary-compatible cast need not preserve equality semantics, so matching +-- the stripped partition key is not enough on its own. Build a type that is +-- binary-coercible to text but compares case insensitively. + +CREATE TYPE pagg_ci; +CREATE FUNCTION pagg_ci_in(cstring) RETURNS pagg_ci STRICT IMMUTABLE LANGUAGE internal AS 'textin'; +CREATE FUNCTION pagg_ci_out(pagg_ci) RETURNS cstring STRICT IMMUTABLE LANGUAGE internal AS 'textout'; +CREATE TYPE pagg_ci (input = pagg_ci_in, output = pagg_ci_out, like = text); +CREATE CAST (pagg_ci AS text) WITHOUT FUNCTION; + +CREATE FUNCTION pagg_ci_eq(pagg_ci, pagg_ci) RETURNS bool + STRICT IMMUTABLE LANGUAGE sql AS $$SELECT lower($1::text) = lower($2::text)$$; +CREATE OPERATOR = (leftarg = pagg_ci, rightarg = pagg_ci, procedure = pagg_ci_eq); +CREATE FUNCTION pagg_ci_hash(pagg_ci) RETURNS int4 STRICT IMMUTABLE LANGUAGE sql AS $$SELECT hashtext(lower($1::text))$$; +CREATE OPERATOR CLASS pagg_ci_ops DEFAULT FOR TYPE pagg_ci USING hash AS OPERATOR 1 =, FUNCTION 1 pagg_ci_hash(pagg_ci); + +CREATE TABLE pagg_tab_ci (a int, c pagg_ci) PARTITION BY LIST ((c::text)); +CREATE TABLE pagg_tab_ci_p1 PARTITION OF pagg_tab_ci FOR VALUES IN ('A'); +CREATE TABLE pagg_tab_ci_p2 PARTITION OF pagg_tab_ci FOR VALUES IN ('a'); +INSERT INTO pagg_tab_ci SELECT i, (CASE WHEN i % 3 = 0 THEN 'A' ELSE 'a' END)::pagg_ci FROM generate_series(1, 3000) i; +ANALYZE pagg_tab_ci; + +-- Partial aggregation only +EXPLAIN (COSTS OFF) +SELECT c, count(*) FROM pagg_tab_ci GROUP BY c; +SELECT c, count(*) FROM pagg_tab_ci GROUP BY c; + + -- Test with multi-level partitioning scheme CREATE TABLE pagg_tab_ml (a int, b int, c text) PARTITION BY RANGE(a); -- 2.47.3