Fix COUNT bug when pulling up scalar subqueries (fallback planner)

A correlated subquery like

    select ... from t1 where t1.a > (select count(*) from t2 where t2.a = t1.d)

must return COUNT = 0 (not NULL) for outer rows that have no match. The
fallback planner pulled such subqueries into an INNER join, which dropped
every no-match row -- the classic COUNT bug.

PR #376 (df87bf0) turned the join into a LEFT join and wrapped the aggregate
in COALESCE(agg, default) with COUNT -> 0. But the comparison stayed on the
join itself (a "Join Filter"), so a matched row that failed the comparison
was treated as unmatched, null-extended, and let back in by the COALESCE
default -- e.g. count(*) returned 100 instead of 99.

Fix: when the pulled-up join is LEFT, evaluate the comparison ABOVE the join
(wrap it in a FromExpr) instead of as the join condition. The INNER path is
unchanged.

Tests in subselect_gp cover COUNT and mixed COUNT+other aggregate expressions
on both the optimizer=off and ORCA->planner fallback paths.

Co-Authored-By: excaliiibur <excaliiibur@foxmail.com>

Author: Alena0704

src/backend/cdb/cdbsubselect.c

@@ -17,6 +17,7 @@
 #include "access/htup_details.h"
 #include "access/skey.h"
 #include "catalog/pg_operator.h"
+#include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "nodes/makefuncs.h"
 #include "optimizer/clauses.h"
@@ -43,6 +44,10 @@ static JoinExpr *make_join_expr(Node *larg, int r_rtindex, int join_type);
 static Node *make_lasj_quals(PlannerInfo *root, SubLink *sublink, int subquery_indx);
 
 static Node *add_null_match_clause(Node *clause);
+static bool contains_count_agg_expr(Node *expr);
+static Expr *build_count_coalesce_expr(Var *aggVar, Expr *defaultExpr);
+static Expr *build_count_default_expr(Node *expr);
+static Node *replace_agg_with_empty_default_mutator(Node *node, void *context);
 
 typedef struct NonNullableVarsContext
 {
@@ -574,6 +579,7 @@ convert_EXPR_to_join(PlannerInfo *root, OpExpr *opexp)
 		 * targetlist.
 		 */
 		TargetEntry *origSubqueryTLE = (TargetEntry *) list_nth(subselect->targetList, 0);
+		bool		has_count_expr = contains_count_agg_expr((Node *) origSubqueryTLE->expr);
 
 		List	   *subselectTargetList = (List *) copyObject(ctx1.targetList);
 
@@ -624,14 +630,111 @@ convert_EXPR_to_join(PlannerInfo *root, OpExpr *opexp)
 											 exprCollation((Node *) subselectAggTLE->expr),
 											 0);
 
-		list_nth_replace(opexp->args, 1, aggVar);
+		if (has_count_expr)
+		{
+			/*
+			 * Expressions containing COUNT over no matching rows must evaluate
+			 * with COUNT=0 instead of NULL after pull-up. For non-COUNT
+			 * aggregates in the same expression, use NULL empty-input defaults.
+			 * Preserve semantics with LEFT JOIN + COALESCE(agg_expr, default_expr).
+			 */
+			Expr	   *defaultExpr = build_count_default_expr((Node *) origSubqueryTLE->expr);
+
+			join_expr->jointype = JOIN_LEFT;
+			list_nth_replace(opexp->args, 1, build_count_coalesce_expr(aggVar, defaultExpr));
+		}
+		else
+		{
+			list_nth_replace(opexp->args, 1, aggVar);
+		}
 
 		return join_expr;
 	}
 
 	return NULL;
 }
 
+static bool
+contains_count_agg_expr(Node *expr)
+{
+	if (expr == NULL)
+		return false;
+
+	if (IsA(expr, Aggref))
+	{
+		Aggref	   *aggref = (Aggref *) expr;
+		return aggref->aggfnoid == COUNT_ANY_OID ||
+			   aggref->aggfnoid == COUNT_STAR_OID;
+	}
+
+	return expression_tree_walker(expr,
+								  contains_count_agg_expr,
+								  NULL);
+}
+
+static Expr *
+build_count_coalesce_expr(Var *aggVar, Expr *defaultExpr)
+{
+	CoalesceExpr *coalesce;
+
+	Assert(aggVar != NULL);
+	Assert(defaultExpr != NULL);
+
+	coalesce = makeNode(CoalesceExpr);
+	coalesce->coalescetype = exprType((Node *) aggVar);
+	coalesce->coalescecollid = exprCollation((Node *) aggVar);
+	coalesce->args = list_make2(aggVar, defaultExpr);
+	coalesce->location = -1;
+
+	return (Expr *) coalesce;
+}
+
+static Expr *
+build_count_default_expr(Node *expr)
+{
+	Node	   *rewritten;
+
+	rewritten = replace_agg_with_empty_default_mutator(copyObject(expr), NULL);
+	return (Expr *) rewritten;
+}
+
+static Node *
+replace_agg_with_empty_default_mutator(Node *node, void *context)
+{
+	Aggref	   *aggref;
+	Oid			default_type;
+	Oid			default_collation;
+	int16		typlen;
+	bool		typbyval;
+
+	if (node == NULL)
+		return NULL;
+
+	if (IsA(node, Aggref))
+	{
+		aggref = (Aggref *) node;
+		if (aggref->aggfnoid == COUNT_ANY_OID ||
+			aggref->aggfnoid == COUNT_STAR_OID)
+		{
+			default_type = INT8OID;
+			default_collation = InvalidOid;
+		}
+		else
+		{
+			default_type = aggref->aggtype;
+			default_collation = exprCollation((Node *) aggref);
+		}
+
+		get_typlenbyval(default_type, &typlen, &typbyval);
+		return (Node *) makeConst(default_type, -1, default_collation, typlen,
+								  (default_type == INT8OID) ? Int64GetDatum(0) : (Datum) 0,
+								  (default_type != INT8OID), typbyval);
+	}
+
+	return expression_tree_mutator(node, replace_agg_with_empty_default_mutator,
+								   context);
+}
+
 /* NOTIN subquery transformation -start */
 
 /* check if NOT IN conversion to antijoin is possible */

src/backend/optimizer/prep/prepjointree.c

@@ -691,6 +691,26 @@ pull_up_sublinks_qual_recurse(PlannerInfo *root, Node *node,
 				{
 					/* Yes, insert the new join node into the join tree */
 					j->larg = *jtlink1;
+
+					if (j->jointype == JOIN_LEFT)
+					{
+						/*
+						 * COUNT-preserving pull-up (see convert_EXPR_to_join):
+						 * the comparison in opexp must be evaluated ABOVE the
+						 * LEFT JOIN, never as its join qual.  If it stayed as a
+						 * qual of the join's larg it would be pulled up into the
+						 * outer join's condition; then a row that matches the
+						 * correlation but fails the comparison would be treated
+						 * as unmatched, get null-extended, and pass again through
+						 * the COALESCE empty-input default -- resurrecting rows
+						 * that should have been filtered out.  Wrap the join in a
+						 * FromExpr that carries opexp as a post-join qual.
+						 */
+						*jtlink1 = (Node *) makeFromExpr(list_make1(j), node);
+						return NULL;
+					}
+
+					/* Inner-join case: opexp stays as an ordinary qual. */
 					*jtlink1 = (Node *) j;
 				}
 				return node;

src/test/regress/expected/subselect_gp.out

@@ -499,6 +499,120 @@ select count(*) from csq_t1 t1 where a > ( select avg(a)::int from csq_t1 t2 whe
     49
 (1 row)
 
+-- COUNT correlated scalar subquery should keep no-match rows (COUNT = 0)
+-- Scenario 1: optimizer=off (always PostgreSQL planner path)
+set optimizer=off;
+SET
+select count(*) from csq_t1 t1 where a > (select count(*) from csq_t1 t2 where t2.a = t1.d);
+ count 
+-------
+    99
+(1 row)
+
+select min(a), max(a) from csq_t1 t1 where (select count(*) from csq_t1 t2 where t2.a = t1.d) = 0;
+ min | max 
+-----+-----
+ 100 | 100
+(1 row)
+
+reset optimizer;
+RESET
+-- Scenario 2: optimizer=on, but query shape forces ORCA fallback to PostgreSQL planner
+-- (ordered aggregate disabled in ORCA by default)
+set optimizer=on;
+SET
+set optimizer_enable_orderedagg=off;
+SET
+explain select string_agg(t1.a::text, ',' order by t1.a)
+from csq_t1 t1
+where t1.a > (select count(*) from csq_t1 t2 where t2.a = t1.d);
+                                             QUERY PLAN                                              
+-----------------------------------------------------------------------------------------------------
+ Aggregate  (cost=16.29..16.30 rows=1 width=32)
+   ->  Gather Motion 3:1  (slice2; segments: 3)  (cost=7.75..16.04 rows=34 width=4)
+         ->  Hash Left Join  (cost=7.75..14.71 rows=12 width=4)
+               Hash Cond: (t1.d = t2.a)
+               Filter: (t1.a > COALESCE((count(*)), '0'::bigint))
+               ->  Redistribute Motion 3:3  (slice1; segments: 3)  (cost=0.00..6.00 rows=34 width=8)
+                     Hash Key: t1.d
+                     ->  Seq Scan on csq_t1 t1  (cost=0.00..4.00 rows=34 width=8)
+               ->  Hash  (cost=6.50..6.50 rows=34 width=12)
+                     ->  HashAggregate  (cost=4.50..5.50 rows=34 width=12)
+                           Group Key: t2.a
+                           ->  Seq Scan on csq_t1 t2  (cost=0.00..4.00 rows=34 width=4)
+ Optimizer: Postgres query optimizer
+(13 rows)
+
+select string_agg(t1.a::text, ',' order by t1.a)
+from csq_t1 t1
+where t1.a > (select count(*) from csq_t1 t2 where t2.a = t1.d);
+                                                                                                                                             string_agg                                                                                                                                              
+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+ 2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100
+(1 row)
+
+-- Non-plain COUNT expression (COUNT(*) + 1) should also preserve no-match semantics
+-- Scenario 1: optimizer=off (always PostgreSQL planner path)
+set optimizer=off;
+SET
+select min(a), max(a), count(*) from csq_t1 t1
+where (select count(*) + 1 from csq_t1 t2 where t2.a = t1.d) = 1;
+ min | max | count 
+-----+-----+-------
+ 100 | 100 |     1
+(1 row)
+
+select min(a), max(a), count(*) from csq_t1 t1
+where (select (count(*) + 1)::bigint from csq_t1 t2 where t2.a = t1.d) = 1;
+ min | max | count 
+-----+-----+-------
+ 100 | 100 |     1
+(1 row)
+
+select min(a), max(a), count(*) from csq_t1 t1
+where (select case when count(*) > 0 then count(*) + 1 else 1 end
+       from csq_t1 t2 where t2.a = t1.d) = 1;
+ min | max | count 
+-----+-----+-------
+ 100 | 100 |     1
+(1 row)
+
+select min(a), max(a), count(*) from csq_t1 t1
+where (select abs(count(*) - 1) from csq_t1 t2 where t2.a = t1.d) = 1;
+ min | max | count 
+-----+-----+-------
+ 100 | 100 |     1
+(1 row)
+
+select min(a), max(a), count(*) from csq_t1 t1
+where (select count(*) + coalesce(sum(t2.a), 0) from csq_t1 t2 where t2.a = t1.d) = 0;
+ min | max | count 
+-----+-----+-------
+ 100 | 100 |     1
+(1 row)
+
+select min(a), max(a), count(*) from csq_t1 t1
+where (select coalesce(count(*) + sum(t2.a), -1) from csq_t1 t2 where t2.a = t1.d) = -1;
+ min | max | count 
+-----+-----+-------
+ 100 | 100 |     1
+(1 row)
+
+reset optimizer;
+RESET
+-- Scenario 2: optimizer=on, but query shape forces ORCA fallback to PostgreSQL planner
+set optimizer=on;
+SET
+set optimizer_enable_orderedagg=off;
+SET
+select string_agg(t1.a::text, ',' order by t1.a)
+from csq_t1 t1
+where (select count(*) + 1 from csq_t1 t2 where t2.a = t1.d) = 1;
+ string_agg 
+------------
+ 100
+(1 row)
+
 --
 -- correlation in a func expr
 --