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

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,12 @@ 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);
+static bool count_no_match_row_survives(PlannerInfo *root, OpExpr *opexp,
+										Expr *defaultExpr);
 
 typedef struct NonNullableVarsContext
 {
@@ -574,6 +581,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 +632,145 @@ 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.
+			 * The empty-input default is what the subquery would yield for an
+			 * outer row with no matching inner rows.
+			 */
+			Expr	   *defaultExpr = build_count_default_expr((Node *) origSubqueryTLE->expr);
+
+			if (count_no_match_row_survives(root, opexp, defaultExpr))
+			{
+				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);
+			}
+		}
+		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);
+}
+
+/*
+ * count_no_match_row_survives
+ *
+ * Could a no-match row satisfy "outerExpr OP (subquery)"? Plug defaultExpr in
+ * for the subquery and constant-fold: false if it folds to FALSE/NULL, else true.
+ */
+static bool
+count_no_match_row_survives(PlannerInfo *root, OpExpr *opexp, Expr *defaultExpr)
+{
+	OpExpr	   *testexpr = (OpExpr *) copyObject(opexp);
+	Node	   *folded;
+
+	list_nth_replace(testexpr->args, 1, copyObject(defaultExpr));
+	folded = eval_const_expressions(root, (Node *) testexpr);
+
+	if (IsA(folded, Const))
+	{
+		Const	   *c = (Const *) folded;
+
+		if (c->constisnull || !DatumGetBool(c->constvalue))
+			return false;
+	}
+
+	return true;
+}
+
 /* NOTIN subquery transformation -start */
 
 /* check if NOT IN conversion to antijoin is possible */

src/backend/optimizer/prep/prepjointree.c

@@ -691,6 +691,22 @@ 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).
+						 * opexp must run ABOVE the LEFT JOIN, not as its join
+						 * condition: as a join qual a matched row that fails it
+						 * would be treated as unmatched, null-extended, and let
+						 * back in by the COALESCE default. Wrap the join in a
+						 * FromExpr so opexp stays a post-join filter.
+						 */
+						*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,112 @@ 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;
+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;
+-- Scenario 2: optimizer=on, but query shape forces ORCA fallback to PostgreSQL planner
+-- (ordered aggregate disabled in ORCA by default)
+set optimizer=on;
+set optimizer_enable_orderedagg=off;
+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;
+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;
+-- Scenario 2: optimizer=on, but query shape forces ORCA fallback to PostgreSQL planner
+set optimizer=on;
+set optimizer_enable_orderedagg=off;
+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
 --