From 470afce23ae397fe216ff0a4537019ea236d7095 Mon Sep 17 00:00:00 2001
From: Haibo Yan <tristan.yim@gmail.com>
Date: Tue, 11 Aug 2026 11:30:00 -0700
Subject: [PATCH v1] Optimize redundant ORDER BY in COUNT aggregates

COUNT is insensitive to input order, but an aggregate-local ORDER BY
currently causes it to be treated as an ordered aggregate, which may
require sorting and can prevent partial or hash aggregation.

Teach the existing COUNT Aggref simplification support to remove ORDER BY
when doing so only discards bare Var or Const sort expressions.  More
complex expressions are left unchanged since removing them could suppress
side effects or errors.

The simplification happens before aggregate preprocessing, allowing the
existing planner machinery to treat the result as an ordinary COUNT.
---
 src/backend/utils/adt/int8.c             | 124 ++++++++---
 src/test/regress/expected/aggregates.out | 259 ++++++++++++++++++++++-
 src/test/regress/sql/aggregates.sql      | 108 +++++++++-
 3 files changed, 455 insertions(+), 36 deletions(-)

diff --git a/src/backend/utils/adt/int8.c b/src/backend/utils/adt/int8.c
index 1f59d831600..7a508e4f1ba 100644
--- a/src/backend/utils/adt/int8.c
+++ b/src/backend/utils/adt/int8.c
@@ -20,6 +20,7 @@
 #include "common/int.h"
 #include "funcapi.h"
 #include "libpq/pqformat.h"
+#include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "nodes/supportnodes.h"
 #include "optimizer/optimizer.h"
@@ -780,6 +781,34 @@ int8dec_any(PG_FUNCTION_ARGS)
 	return int8dec(fcinfo);
 }
 
+/*
+ * count_orderby_is_removable
+ *		Can the aggregate-local ORDER BY of a COUNT(ANY) Aggref be removed?
+ *
+ * Removing aggregate ORDER BY also removes evaluation of resjunk sort
+ * expressions.  Only discard bare Vars and Consts: even an immutable
+ * expression such as "1 / b" can still throw.
+ */
+static bool
+count_orderby_is_removable(Aggref *agg)
+{
+	ListCell   *lc;
+
+	foreach(lc, agg->aggorder)
+	{
+		SortGroupClause *sgc = lfirst_node(SortGroupClause, lc);
+		TargetEntry *tle = get_sortgroupclause_tle(sgc, agg->args);
+
+		if (!tle->resjunk)
+			continue;			/* real argument; evaluation survives */
+
+		if (!IsA(tle->expr, Var) && !IsA(tle->expr, Const))
+			return false;
+	}
+
+	return true;
+}
+
 /*
  * int8inc_support
  *		prosupport function for int8inc() and int8inc_any()
@@ -856,44 +885,87 @@ int8inc_support(PG_FUNCTION_ARGS)
 		Aggref	   *agg = req->aggref;
 
 		/*
-		 * Check for COUNT(ANY) and try to convert to COUNT(*). The input
-		 * argument cannot be NULL, we can't have an ORDER BY / DISTINCT in
-		 * the aggregate, and agglevelsup must be 0.
-		 *
-		 * Technically COUNT(ANY) must have 1 arg, but be paranoid and check.
+		 * COUNT is insensitive to input order, but an aggregate-local ORDER
+		 * BY still needs to be removed explicitly, else it's treated as an
+		 * ordered aggregate.  While at it, also convert to COUNT(*) when the
+		 * argument is provably non-NULL.  Both are handled in one pass, since
+		 * simplify_aggref() does not re-invoke us on the node we return.
 		 */
-		if (agg->aggfnoid == F_COUNT_ANY && list_length(agg->args) == 1)
+		if (agg->aggfnoid == F_COUNT_ANY)
 		{
-			TargetEntry *tle = (TargetEntry *) linitial(agg->args);
-			Expr	   *arg = tle->expr;
+			TargetEntry *arg_tle;
+			Expr	   *arg;
+			bool		nonnullable;
+			Aggref	   *newagg;
+			ListCell   *lc;
+
+			if (agg->aggdistinct != NIL || agg->agglevelsup != 0)
+				PG_RETURN_POINTER(NULL);
+
+			/*
+			 * COUNT(ANY) has one real argument, in the first position of
+			 * agg->args; any further entries are resjunk ORDER-BY-only
+			 * expressions.  The parser guarantees this shape, but an
+			 * unexpected one is simply not an optimization opportunity here,
+			 * not a reason to fail.
+			 */
+			if (agg->args == NIL || list_length(agg->aggargtypes) != 1)
+				PG_RETURN_POINTER(NULL);
 
-			/* Check for unsupported cases */
-			if (agg->aggdistinct != NIL || agg->aggorder != NIL ||
-				agg->agglevelsup != 0)
+			arg_tle = linitial_node(TargetEntry, agg->args);
+			if (arg_tle->resjunk)
 				PG_RETURN_POINTER(NULL);
 
-			/* If the arg isn't NULLable, do the conversion */
-			if (expr_is_nonnullable(req->root, arg, NOTNULL_SOURCE_HASHTABLE))
+			for_each_from(lc, agg->args, 1)
 			{
-				Aggref	   *newagg;
+				TargetEntry *tle = lfirst_node(TargetEntry, lc);
 
-				/* We don't expect these to have been set yet */
-				Assert(agg->aggtransno == -1);
-				Assert(agg->aggtranstype == InvalidOid);
+				if (!tle->resjunk)
+					PG_RETURN_POINTER(NULL);
+			}
 
-				/* Convert COUNT(ANY) to COUNT(*) by making a new Aggref */
-				newagg = makeNode(Aggref);
-				memcpy(newagg, agg, sizeof(Aggref));
-				newagg->aggfnoid = F_COUNT_;
+			arg = arg_tle->expr;
+
+			/* An ORDER BY that isn't removable rules out any simplification. */
+			if (agg->aggorder != NIL && !count_orderby_is_removable(agg))
+				PG_RETURN_POINTER(NULL);
+
+			nonnullable = expr_is_nonnullable(req->root, arg,
+											  NOTNULL_SOURCE_HASHTABLE);
+
+			/* Nothing to do when there's no ORDER BY and the arg may be NULL. */
+			if (agg->aggorder == NIL && !nonnullable)
+				PG_RETURN_POINTER(NULL);
+
+			/* We don't expect these to have been set yet */
+			Assert(agg->aggtransno == -1);
+			Assert(agg->aggtranstype == InvalidOid);
 
-				/* count(*) has no args */
-				newagg->aggargtypes = NULL;
-				newagg->args = NULL;
+			newagg = makeNode(Aggref);
+			memcpy(newagg, agg, sizeof(Aggref));
+			newagg->aggorder = NIL;
+
+			if (nonnullable)
+			{
+				/* Convert COUNT(ANY) to COUNT(*); this also drops the args. */
+				newagg->aggfnoid = F_COUNT_;
+				newagg->aggargtypes = NIL;
+				newagg->args = NIL;
 				newagg->aggstar = true;
 				newagg->location = -1;
-
-				PG_RETURN_POINTER(newagg);
 			}
+			else
+			{
+				/*
+				 * Rebuild with a fresh TargetEntry, canonicalizing to plain
+				 * COUNT(arg) so it can share transition state with one
+				 * elsewhere in the query.
+				 */
+				Assert(agg->aggorder != NIL);
+				newagg->args = list_make1(makeTargetEntry(arg, 1, NULL, false));
+			}
+
+			PG_RETURN_POINTER(newagg);
 		}
 	}
 
diff --git a/src/test/regress/expected/aggregates.out b/src/test/regress/expected/aggregates.out
index 5f0668382ed..01b4cdc4abe 100644
--- a/src/test/regress/expected/aggregates.out
+++ b/src/test/regress/expected/aggregates.out
@@ -2933,6 +2933,12 @@ select pg_typeof(cleast_agg(variadic array[4.5,f1])) from int4_tbl;
 --
 begin;
 create table agg_simplify (a int, not_null_col int not null, nullable_col int);
+insert into agg_simplify values
+    (1, 10, 100),
+    (2, 20, null),
+    (3, 30, 300),
+    (4, 40, null),
+    (5, 50, 0);
 -- Ensure count(not_null_col) uses count(*)
 explain (costs off, verbose)
 select count(not_null_col) from agg_simplify;
@@ -2977,7 +2983,8 @@ select count(nullable_col) from agg_simplify;
          Output: a, not_null_col, nullable_col
 (4 rows)
 
--- Ensure there's no optimization with DISTINCT aggs
+-- Ensure there's no optimization with DISTINCT aggs, with or without an
+-- aggregate-local ORDER BY
 explain (costs off, verbose)
 select count(distinct not_null_col) from agg_simplify;
                  QUERY PLAN                  
@@ -2991,20 +2998,171 @@ select count(distinct not_null_col) from agg_simplify;
                Output: not_null_col
 (7 rows)
 
--- Ensure there's no optimization with ORDER BY aggs
+explain (costs off, verbose)
+select count(distinct a order by a) from agg_simplify;
+                 QUERY PLAN                  
+---------------------------------------------
+ Aggregate
+   Output: count(DISTINCT a ORDER BY a)
+   ->  Sort
+         Output: a
+         Sort Key: agg_simplify.a
+         ->  Seq Scan on public.agg_simplify
+               Output: a
+(7 rows)
+
+-- ORDER BY is redundant for COUNT, so it's removed when every ORDER-BY-only
+-- expression is a bare Var or Const.  A non-nullable arg then reaches count(*).
 explain (costs off, verbose)
 select count(not_null_col order by not_null_col) from agg_simplify;
-                     QUERY PLAN                      
------------------------------------------------------
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(*)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- Same, but a nullable arg is only simplified to count(<arg>)
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
  Aggregate
-   Output: count(not_null_col ORDER BY not_null_col)
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- A separate resjunk ORDER BY key that's a bare Var is likewise removable
+explain (costs off, verbose)
+select count(nullable_col order by not_null_col) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- A mix of the real argument and a resjunk Var as ORDER BY keys
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col, not_null_col) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- A constant ORDER BY key is likewise removable
+explain (costs off, verbose)
+select count(nullable_col order by 1) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- Sort direction and NULLS ordering don't affect removability
+explain (costs off, verbose)
+select count(nullable_col order by not_null_col desc) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+explain (costs off, verbose)
+select count(nullable_col order by not_null_col nulls first) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- FILTER doesn't block removal of a safe ORDER BY
+explain (costs off, verbose)
+select count(a order by not_null_col) filter (where a <> 2) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(a) FILTER (WHERE (a <> 2))
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- A complex expression used as the real (non-resjunk) argument is not
+-- affected by the resjunk restriction: its evaluation survives regardless,
+-- so a matching ORDER BY on it is still redundant and removable
+explain (costs off, verbose)
+select count(not_null_col + 1 order by not_null_col + 1) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count((not_null_col + 1))
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- Nontrivial ORDER-BY-only expressions are never removed, since dropping
+-- them could suppress side effects or errors that would otherwise surface
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col + 1) from agg_simplify;
+                         QUERY PLAN                          
+-------------------------------------------------------------
+ Aggregate
+   Output: count(nullable_col ORDER BY ((nullable_col + 1)))
    ->  Sort
-         Output: not_null_col
-         Sort Key: agg_simplify.not_null_col
+         Output: nullable_col, ((nullable_col + 1))
+         Sort Key: ((agg_simplify.nullable_col + 1))
          ->  Seq Scan on public.agg_simplify
-               Output: not_null_col
+               Output: nullable_col, (nullable_col + 1)
+(7 rows)
+
+explain (costs off, verbose)
+select count(nullable_col order by random()) from agg_simplify;
+                    QUERY PLAN                     
+---------------------------------------------------
+ Aggregate
+   Output: count(nullable_col ORDER BY (random()))
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
+-- The rule is all-or-nothing across multiple sort keys: one nontrivial key
+-- blocks removal even when another key is a bare Var ...
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col, not_null_col + 1) from agg_simplify;
+                                   QUERY PLAN                                   
+--------------------------------------------------------------------------------
+ Aggregate
+   Output: count(nullable_col ORDER BY nullable_col, ((not_null_col + 1)))
+   ->  Sort
+         Output: nullable_col, not_null_col, ((not_null_col + 1))
+         Sort Key: agg_simplify.nullable_col, ((agg_simplify.not_null_col + 1))
+         ->  Seq Scan on public.agg_simplify
+               Output: nullable_col, not_null_col, (not_null_col + 1)
 (7 rows)
 
+-- ... while an all-Var/Const set of keys still simplifies
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col, not_null_col) from agg_simplify;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Aggregate
+   Output: count(nullable_col)
+   ->  Seq Scan on public.agg_simplify
+         Output: a, not_null_col, nullable_col
+(4 rows)
+
 -- Ensure we don't optimize to count(*) with agglevelsup > 0
 explain (costs off, verbose)
 select a from agg_simplify a group by a
@@ -3022,6 +3180,91 @@ having exists (select 1 from onek b where count(a.not_null_col) = b.four);
            Filter: (count(a.not_null_col) = b.four)
 (9 rows)
 
+-- A volatile ORDER-BY-only expression is not removed, so it still runs once
+-- per row that reaches the aggregate
+create sequence count_order_seq;
+select count(not_null_col order by nextval('count_order_seq')) from agg_simplify;
+ count 
+-------
+     5
+(1 row)
+
+select currval('count_order_seq') as expect_5;
+ expect_5 
+----------
+        5
+(1 row)
+
+select count(not_null_col order by nextval('count_order_seq')) filter (where false) from agg_simplify;
+ count 
+-------
+     0
+(1 row)
+
+select currval('count_order_seq') as expect_5_still;
+ expect_5_still 
+----------------
+              5
+(1 row)
+
+select count(not_null_col order by nextval('count_order_seq')) filter (where a <> 2) from agg_simplify;
+ count 
+-------
+     4
+(1 row)
+
+select currval('count_order_seq') as expect_9;
+ expect_9 
+----------
+        9
+(1 row)
+
+-- Result equality: count(arg ORDER BY <removable keys>) matches count(arg),
+-- covering nullable/non-nullable args, sort direction, an all-NULL filtered
+-- input, and an empty filtered input
+select count(nullable_col order by not_null_col) as ordered,
+       count(nullable_col) as plain
+from agg_simplify;
+ ordered | plain 
+---------+-------
+       3 |     3
+(1 row)
+
+select count(nullable_col order by nullable_col desc) as ordered,
+       count(nullable_col) as plain
+from agg_simplify;
+ ordered | plain 
+---------+-------
+       3 |     3
+(1 row)
+
+select count(nullable_col order by not_null_col) filter (where nullable_col is null) as ordered,
+       count(nullable_col) filter (where nullable_col is null) as plain
+from agg_simplify;
+ ordered | plain 
+---------+-------
+       0 |     0
+(1 row)
+
+select count(nullable_col order by not_null_col) filter (where false) as ordered,
+       count(nullable_col) filter (where false) as plain
+from agg_simplify;
+ ordered | plain 
+---------+-------
+       0 |     0
+(1 row)
+
+-- Immutable expressions can still throw errors.  The strict Var/Const rule
+-- keeps this resjunk expression from disappearing, so the error remains ...
+savepoint sp_div0;
+select count(a order by 1 / nullable_col) from agg_simplify;
+ERROR:  division by zero
+rollback to savepoint sp_div0;
+-- ... but when the same expression is the real argument instead, its
+-- evaluation was never conditional on the ORDER BY, so the error remains
+-- even after the (now redundant) ORDER BY is removed
+select count(1 / nullable_col order by 1 / nullable_col) from agg_simplify;
+ERROR:  division by zero
 rollback;
 -- test aggregates with common transition functions share the same states
 begin work;
diff --git a/src/test/regress/sql/aggregates.sql b/src/test/regress/sql/aggregates.sql
index b788152f0c2..9ab13d13ad8 100644
--- a/src/test/regress/sql/aggregates.sql
+++ b/src/test/regress/sql/aggregates.sql
@@ -1156,6 +1156,12 @@ select pg_typeof(cleast_agg(variadic array[4.5,f1])) from int4_tbl;
 --
 begin;
 create table agg_simplify (a int, not_null_col int not null, nullable_col int);
+insert into agg_simplify values
+    (1, 10, 100),
+    (2, 20, null),
+    (3, 30, 300),
+    (4, 40, null),
+    (5, 50, 0);
 
 -- Ensure count(not_null_col) uses count(*)
 explain (costs off, verbose)
@@ -1173,19 +1179,117 @@ select count(null) from agg_simplify;
 explain (costs off, verbose)
 select count(nullable_col) from agg_simplify;
 
--- Ensure there's no optimization with DISTINCT aggs
+-- Ensure there's no optimization with DISTINCT aggs, with or without an
+-- aggregate-local ORDER BY
 explain (costs off, verbose)
 select count(distinct not_null_col) from agg_simplify;
 
--- Ensure there's no optimization with ORDER BY aggs
+explain (costs off, verbose)
+select count(distinct a order by a) from agg_simplify;
+
+-- ORDER BY is redundant for COUNT, so it's removed when every ORDER-BY-only
+-- expression is a bare Var or Const.  A non-nullable arg then reaches count(*).
 explain (costs off, verbose)
 select count(not_null_col order by not_null_col) from agg_simplify;
 
+-- Same, but a nullable arg is only simplified to count(<arg>)
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col) from agg_simplify;
+
+-- A separate resjunk ORDER BY key that's a bare Var is likewise removable
+explain (costs off, verbose)
+select count(nullable_col order by not_null_col) from agg_simplify;
+
+-- A mix of the real argument and a resjunk Var as ORDER BY keys
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col, not_null_col) from agg_simplify;
+
+-- A constant ORDER BY key is likewise removable
+explain (costs off, verbose)
+select count(nullable_col order by 1) from agg_simplify;
+
+-- Sort direction and NULLS ordering don't affect removability
+explain (costs off, verbose)
+select count(nullable_col order by not_null_col desc) from agg_simplify;
+
+explain (costs off, verbose)
+select count(nullable_col order by not_null_col nulls first) from agg_simplify;
+
+-- FILTER doesn't block removal of a safe ORDER BY
+explain (costs off, verbose)
+select count(a order by not_null_col) filter (where a <> 2) from agg_simplify;
+
+-- A complex expression used as the real (non-resjunk) argument is not
+-- affected by the resjunk restriction: its evaluation survives regardless,
+-- so a matching ORDER BY on it is still redundant and removable
+explain (costs off, verbose)
+select count(not_null_col + 1 order by not_null_col + 1) from agg_simplify;
+
+-- Nontrivial ORDER-BY-only expressions are never removed, since dropping
+-- them could suppress side effects or errors that would otherwise surface
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col + 1) from agg_simplify;
+
+explain (costs off, verbose)
+select count(nullable_col order by random()) from agg_simplify;
+
+-- The rule is all-or-nothing across multiple sort keys: one nontrivial key
+-- blocks removal even when another key is a bare Var ...
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col, not_null_col + 1) from agg_simplify;
+
+-- ... while an all-Var/Const set of keys still simplifies
+explain (costs off, verbose)
+select count(nullable_col order by nullable_col, not_null_col) from agg_simplify;
+
 -- Ensure we don't optimize to count(*) with agglevelsup > 0
 explain (costs off, verbose)
 select a from agg_simplify a group by a
 having exists (select 1 from onek b where count(a.not_null_col) = b.four);
 
+-- A volatile ORDER-BY-only expression is not removed, so it still runs once
+-- per row that reaches the aggregate
+create sequence count_order_seq;
+
+select count(not_null_col order by nextval('count_order_seq')) from agg_simplify;
+select currval('count_order_seq') as expect_5;
+
+select count(not_null_col order by nextval('count_order_seq')) filter (where false) from agg_simplify;
+select currval('count_order_seq') as expect_5_still;
+
+select count(not_null_col order by nextval('count_order_seq')) filter (where a <> 2) from agg_simplify;
+select currval('count_order_seq') as expect_9;
+
+-- Result equality: count(arg ORDER BY <removable keys>) matches count(arg),
+-- covering nullable/non-nullable args, sort direction, an all-NULL filtered
+-- input, and an empty filtered input
+select count(nullable_col order by not_null_col) as ordered,
+       count(nullable_col) as plain
+from agg_simplify;
+
+select count(nullable_col order by nullable_col desc) as ordered,
+       count(nullable_col) as plain
+from agg_simplify;
+
+select count(nullable_col order by not_null_col) filter (where nullable_col is null) as ordered,
+       count(nullable_col) filter (where nullable_col is null) as plain
+from agg_simplify;
+
+select count(nullable_col order by not_null_col) filter (where false) as ordered,
+       count(nullable_col) filter (where false) as plain
+from agg_simplify;
+
+-- Immutable expressions can still throw errors.  The strict Var/Const rule
+-- keeps this resjunk expression from disappearing, so the error remains ...
+savepoint sp_div0;
+select count(a order by 1 / nullable_col) from agg_simplify;
+rollback to savepoint sp_div0;
+
+-- ... but when the same expression is the real argument instead, its
+-- evaluation was never conditional on the ORDER BY, so the error remains
+-- even after the (now redundant) ORDER BY is removed
+select count(1 / nullable_col order by 1 / nullable_col) from agg_simplify;
+
 rollback;
 
 -- test aggregates with common transition functions share the same states
-- 
2.54.0

