Keep no-match rows when pulling up a correlated aggregate subquery (bugfix to main)

contrib/pax_storage/src/test/regress/expected/eagerfree.out

@@ -1398,21 +1398,20 @@ where i < (select count(*) from smallt where smallt.i = smallt2.i) order by 1,2,
 
 explain select smallt2.* from smallt2
 where i < (select count(*) from smallt where smallt.i = smallt2.i);
-                                      QUERY PLAN                                       
----------------------------------------------------------------------------------------
+                                                   QUERY PLAN                                                    
+-----------------------------------------------------------------------------------------------------------------
  Gather Motion 3:1  (slice1; segments: 3)  (cost=1.57..3.56 rows=11 width=15)
-   ->  Hash Join  (cost=1.57..2.82 rows=4 width=15)
+   ->  Hash Left Join  (cost=1.57..2.82 rows=4 width=15)
          Hash Cond: (smallt2.i = "Expr_SUBQUERY".csq_c0)
-         Join Filter: (smallt2.i < "Expr_SUBQUERY".csq_c1)
+         Filter: (smallt2.i < CASE WHEN "Expr_SUBQUERY".csq_c1 THEN "Expr_SUBQUERY".csq_c2 ELSE '0'::bigint END)
          ->  Seq Scan on smallt2  (cost=0.00..1.17 rows=17 width=15)
-         ->  Hash  (cost=1.55..1.55 rows=1 width=12)
-               ->  Subquery Scan on "Expr_SUBQUERY"  (cost=1.50..1.55 rows=1 width=12)
-                     ->  HashAggregate  (cost=1.50..1.54 rows=1 width=12)
+         ->  Hash  (cost=1.55..1.55 rows=1 width=13)
+               ->  Subquery Scan on "Expr_SUBQUERY"  (cost=1.50..1.55 rows=1 width=13)
+                     ->  HashAggregate  (cost=1.50..1.54 rows=1 width=13)
                            Group Key: smallt.i
-                           Filter: (smallt.i < count(*))
                            ->  Seq Scan on smallt  (cost=0.00..1.33 rows=33 width=4)
  Optimizer: Postgres query optimizer
-(12 rows)
+(11 rows)
 
 -- Sort in MergeJoin
 -- start_ignore

src/backend/cdb/cdbsubselect.c

@@ -29,6 +29,7 @@
 #include "parser/parse_relation.h"	/* addRangeTableEntryForSubquery() */
 #include "parser/parsetree.h"	/* rt_fetch() */
 #include "rewrite/rewriteManip.h"
+#include "utils/fmgroids.h"		/* F_COUNT_ANY, F_COUNT_ */
 #include "utils/lsyscache.h"	/* get_op_btree_interpretation() */
 #include "utils/syscache.h"
 #include "cdb/cdbsubselect.h"	/* me */
@@ -42,6 +43,11 @@ 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 Expr *build_match_flag_case_expr(Var *flagVar, Var *aggVar, Expr *defaultExpr);
+static Expr *build_empty_input_default_expr(Node *expr);
+static Node *replace_agg_with_empty_default_mutator(Node *node, void *context);
+static bool no_match_row_survives(PlannerInfo *root, OpExpr *opexp,
+								  Expr *defaultExpr);
 
 typedef struct NonNullableVarsContext
 {
@@ -560,6 +566,14 @@ safe_to_convert_EXPR(SubLink *sublink, ConvertSubqueryToJoinContext *ctx1)
 	if (list_length(subselect->targetList) != 1)
 		return false;
 
+	/**
+	 * Correlation in the targetlist cannot be handled: the pulled-up
+	 * expression (and the empty-input default derived from it) would carry
+	 * upper-level Vars out of the subquery.
+	 */
+	if (contain_vars_of_level_or_above((Node *) subselect->targetList, 1))
+		return false;
+
 
 	/**
 	 * Walk the quals of the subquery to do a more fine grained check as to whether this subquery
@@ -623,6 +637,51 @@ convert_EXPR_to_join(PlannerInfo *root, OpExpr *opexp)
 
 		subselect->jointree->quals = ctx1.innerQual;
 
+		/*
+		 * An INNER join drops outer rows that have no matching inner
+		 * rows.  Without the pull-up they are kept: the subquery
+		 * computes its expression over empty input (COUNT = 0, other
+		 * aggregates NULL) and the comparison may still pass.
+		 *
+		 * So plug the empty-input value into the comparison and run
+		 * eval_const_expressions() on it.  FALSE or NULL means no-match
+		 * rows cannot pass and the INNER join is correct; otherwise use
+		 * a LEFT join to keep them.
+		 */
+		Expr	   *defaultExpr;
+		TargetEntry *flagTLE = NULL;
+		bool		use_left_join;
+
+		defaultExpr = build_empty_input_default_expr((Node *) origSubqueryTLE->expr);
+		use_left_join = no_match_row_survives(root, opexp, defaultExpr);
+
+		if (use_left_join)
+		{
+			/*
+			 * After the LEFT join the expression column is NULL both for a
+			 * no-match row and for a matched group whose expression is
+			 * genuinely NULL.  To tell them apart, add a constant-TRUE
+			 * match-flag column to the subquery: it can be NULL only when
+			 * the LEFT join found no match and filled the subquery's
+			 * columns with NULLs.
+			 *
+			 * The flag goes BEFORE the expression column: with this
+			 * order the planner can drop the SubqueryScan node from the
+			 * plan.
+			 */
+			TargetEntry *aggTLE = (TargetEntry *) llast(subselect->targetList);
+
+			flagTLE = makeTargetEntry((Expr *) makeBoolConst(true, false),
+									  aggTLE->resno,
+									  pstrdup("csq_count_flag"),
+									  false);
+			aggTLE->resno++;
+			subselect->targetList = list_truncate(subselect->targetList,
+												  list_length(subselect->targetList) - 1);
+			subselect->targetList = lappend(subselect->targetList, flagTLE);
+			subselect->targetList = lappend(subselect->targetList, aggTLE);
+		}
+
 		/**
 		 * Construct a new range table entry for the new pulled up subquery.
 		 */
@@ -644,7 +703,8 @@ convert_EXPR_to_join(PlannerInfo *root, OpExpr *opexp)
 
 		join_expr->quals = joinQual;
 
-		TargetEntry *subselectAggTLE = (TargetEntry *) list_nth(subselect->targetList, list_length(subselect->targetList) - 1);
+		/* The pulled-up expression column is last in either layout. */
+		TargetEntry *subselectAggTLE = (TargetEntry *) llast(subselect->targetList);
 
 		/**
 		 *	modify the op expr to involve the column that has the computed aggregate that needs to compared.
@@ -656,14 +716,163 @@ convert_EXPR_to_join(PlannerInfo *root, OpExpr *opexp)
 											 exprCollation((Node *) subselectAggTLE->expr),
 											 0);
 
-		list_nth_replace(opexp->args, 1, aggVar);
+		if (use_left_join)
+		{
+			Var		   *flagVar;
+			RangeTblEntry *joinRTE;
+			int			joinRTIndex;
+
+			join_expr->jointype = JOIN_LEFT;
+
+			/*
+			 * Give the outer join a range table entry and mark the Vars the
+			 * comparison uses as nulled by it.  Since the removal of
+			 * outerjoin_delayed the planner keeps a clause above an outer
+			 * join only when the clause's Vars carry the join's relid in
+			 * varnullingrels; without this the comparison would be pushed
+			 * down to the subquery rel (or the join removed as useless) and
+			 * the no-match default would never apply.
+			 */
+			joinRTE = makeNode(RangeTblEntry);
+			joinRTE->rtekind = RTE_JOIN;
+			joinRTE->jointype = JOIN_LEFT;
+			joinRTE->joinmergedcols = 0;
+			joinRTE->eref = makeAlias("unnamed_join", NIL);
+			joinRTE->inFromCl = false;
+			root->parse->rtable = lappend(root->parse->rtable, joinRTE);
+			joinRTIndex = list_length(root->parse->rtable);
+			join_expr->rtindex = joinRTIndex;
+
+			flagVar = (Var *) makeVar(rteIndex, flagTLE->resno, BOOLOID, -1,
+									  InvalidOid, 0);
+			flagVar->varnullingrels = bms_make_singleton(joinRTIndex);
+			aggVar->varnullingrels = bms_make_singleton(joinRTIndex);
+			list_nth_replace(opexp->args, 1,
+							 build_match_flag_case_expr(flagVar, aggVar, defaultExpr));
+		}
+		else
+		{
+			list_nth_replace(opexp->args, 1, aggVar);
+		}
 
 		return join_expr;
 	}
 
 	return NULL;
 }
 
+/*
+ * Build "CASE WHEN flagVar THEN aggVar ELSE defaultExpr END".
+ *
+ * flagVar is the subquery's match-flag column: TRUE for a matched group,
+ * NULL for a null-extended no-match row.
+ */
+static Expr *
+build_match_flag_case_expr(Var *flagVar, Var *aggVar, Expr *defaultExpr)
+{
+	CaseWhen   *casewhen;
+	CaseExpr   *caseexpr;
+
+	Assert(flagVar != NULL);
+	Assert(aggVar != NULL);
+	Assert(defaultExpr != NULL);
+
+	casewhen = makeNode(CaseWhen);
+	casewhen->expr = (Expr *) flagVar;
+	casewhen->result = (Expr *) aggVar;
+	casewhen->location = -1;
+
+	caseexpr = makeNode(CaseExpr);
+	caseexpr->casetype = exprType((Node *) aggVar);
+	caseexpr->casecollid = exprCollation((Node *) aggVar);
+	caseexpr->arg = NULL;
+	caseexpr->args = list_make1(casewhen);
+	caseexpr->defresult = defaultExpr;
+	caseexpr->location = -1;
+
+	return (Expr *) caseexpr;
+}
+
+static Expr *
+build_empty_input_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))
+	{
+		bool		is_count;
+
+		aggref = (Aggref *) node;
+		is_count = (aggref->aggfnoid == F_COUNT_ANY ||
+					aggref->aggfnoid == F_COUNT_);
+		if (is_count)
+		{
+			default_type = INT8OID;
+			default_collation = InvalidOid;
+		}
+		else
+		{
+			default_type = aggref->aggtype;
+			default_collation = exprCollation((Node *) aggref);
+		}
+
+		/*
+		 * COUNT is 0 over empty input; every other aggregate is NULL.  The
+		 * choice must follow the aggregate, not its result type: sum(int4)
+		 * also returns int8 but its empty-input value is NULL.
+		 */
+		get_typlenbyval(default_type, &typlen, &typbyval);
+		return (Node *) makeConst(default_type, -1, default_collation, typlen,
+								  is_count ? Int64GetDatum(0) : (Datum) 0,
+								  !is_count, typbyval);
+	}
+
+	return expression_tree_mutator(node, replace_agg_with_empty_default_mutator,
+								   context);
+}
+
+/*
+ * 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
+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

@@ -858,11 +858,39 @@ pull_up_sublinks_qual_recurse(PlannerInfo *root, Node *node,
 
 			if (IsA(rarg, SubLink))
 			{
+				/*
+				 * The pulled-up join is spliced in at *jtlink1, and in the
+				 * LEFT-join case the comparison itself moves there too, so
+				 * every Var of this query level used by the clause must be
+				 * available at that attach point.  Otherwise (e.g. an outer
+				 * join's ON clause referencing the non-nullable side) leave
+				 * the sublink to be planned as a SubPlan.
+				 */
+				if (!bms_is_subset(pull_varnos(root, node), available_rels1))
+					return node;
+
 				j = convert_EXPR_to_join(root, opexp);
 				if (j)
 				{
 					/* 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 no-match default of the CASE built by
+						 * convert_EXPR_to_join. 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/eagerfree.out

@@ -1379,21 +1379,20 @@ where i < (select count(*) from smallt where smallt.i = smallt2.i) order by 1,2,
 
 explain select smallt2.* from smallt2
 where i < (select count(*) from smallt where smallt.i = smallt2.i);
-                                      QUERY PLAN                                       
----------------------------------------------------------------------------------------
- Gather Motion 3:1  (slice1; segments: 3)  (cost=5.10..8.08 rows=17 width=15)
-   ->  Hash Join  (cost=5.10..8.08 rows=6 width=15)
-         Hash Cond: smallt2.i = "Expr_SUBQUERY".csq_c0
-         Join Filter: smallt2.i < "Expr_SUBQUERY".csq_c1
-         ->  Seq Scan on smallt2  (cost=0.00..2.50 rows=17 width=15)
-         ->  Hash  (cost=4.97..4.97 rows=4 width=12)
-               ->  Subquery Scan on "Expr_SUBQUERY"  (cost=4.75..4.97 rows=4 width=12)
-                     ->  HashAggregate  (cost=4.75..4.88 rows=4 width=12)
-                           Filter: smallt.i < count(*)
+                                                   QUERY PLAN                                                    
+-----------------------------------------------------------------------------------------------------------------
+ Gather Motion 3:1  (slice1; segments: 3)  (cost=1.59..3.97 rows=17 width=15)
+   ->  Hash Left Join  (cost=1.59..2.86 rows=6 width=15)
+         Hash Cond: (smallt2.i = "Expr_SUBQUERY".csq_c0)
+         Filter: (smallt2.i < CASE WHEN "Expr_SUBQUERY".csq_c1 THEN "Expr_SUBQUERY".csq_c2 ELSE '0'::bigint END)
+         ->  Seq Scan on smallt2  (cost=0.00..1.17 rows=17 width=15)
+         ->  Hash  (cost=1.54..1.54 rows=3 width=13)
+               ->  Subquery Scan on "Expr_SUBQUERY"  (cost=1.50..1.54 rows=3 width=13)
+                     ->  HashAggregate  (cost=1.50..1.53 rows=3 width=13)
                            Group Key: smallt.i
-                           ->  Seq Scan on smallt  (cost=0.00..4.00 rows=34 width=4)
+                           ->  Seq Scan on smallt  (cost=0.00..1.33 rows=33 width=4)
  Optimizer: Postgres query optimizer
-(12 rows)
+(11 rows)
 
 -- Sort in MergeJoin
 -- start_ignore