From 7d2c7f08d9c5ad02101c0fb47cd8e859dd89083b Mon Sep 17 00:00:00 2001
From: David Rowley <drowley@postgresql.org>
Date: Sun, 5 May 2024 12:54:46 +1200
Subject: [PATCH] Fix query pullup issue with WindowClause runCondition

94985c210 added code to detect when WindowFuncs were monotonic and
allowed additional quals to be "pushed down" into the subquery to be
used as WindowClause runConditions in order to short-circuit execution
in nodeWindowAgg.c.

The Node representation of runConditions wasn't well selected and
because we do qual pushdown before planning the subquery, the planning
of the subquery could perform subquery pull-up of nested subqueries.
For WindowFuncs with args, the arguments could be changed after pushing
the qual down to the subquery.

This was made more difficult by the fact that the code duplicated the
WindowFunc inside an OpExpr to include in the WindowClauses runCondition
field.  This could result in duplication of subqueries and a pull-up of
such a subquery could result in another initplan parameter being issued
for the 2nd version of the subplan.  This could result in errors such as:

ERROR:  WindowFunc not found in subplan target lists

To fix this, we change the node representation of these run conditions
and instead of storing an OpExpr containing the WindowFunc in a list
inside WindowClause, we now store a new node type named
WindowFuncRunCondition within a new field in the WindowFunc.  These get
transformed into OpExprs later in planning once subquery pull-up has been
performed.

This problem did exist in v15 and v16, but that was fixed by 9d36b883b
and e5d20bbd.

Cat version bump due to new node type and modifying WindowFunc struct.

Bug: #18305
Reported-by: Zuming Jiang
Discussion: https://postgr.es/m/18305-33c49b4c830b37b3%40postgresql.org
---
 src/backend/nodes/nodeFuncs.c             | 26 ++++++++---
 src/backend/optimizer/path/allpaths.c     | 36 +++++----------
 src/backend/optimizer/plan/createplan.c   |  2 +-
 src/backend/optimizer/plan/planner.c      | 54 ++++++++++++++++++-----
 src/backend/optimizer/prep/prepjointree.c |  8 ----
 src/backend/optimizer/util/clauses.c      |  1 +
 src/backend/optimizer/util/pathnode.c     |  3 ++
 src/backend/parser/parse_clause.c         |  1 -
 src/backend/parser/parse_expr.c           |  1 +
 src/backend/parser/parse_func.c           |  1 +
 src/include/catalog/catversion.h          |  2 +-
 src/include/nodes/parsenodes.h            |  2 -
 src/include/nodes/pathnodes.h             |  1 +
 src/include/nodes/primnodes.h             | 30 +++++++++++++
 src/include/optimizer/pathnode.h          |  1 +
 src/test/regress/expected/window.out      | 18 ++++++++
 src/test/regress/sql/window.sql           |  8 ++++
 src/tools/pgindent/typedefs.list          |  1 +
 18 files changed, 143 insertions(+), 53 deletions(-)

diff --git a/src/backend/nodes/nodeFuncs.c b/src/backend/nodes/nodeFuncs.c
index e1df1894b6..89ee4b61f2 100644
--- a/src/backend/nodes/nodeFuncs.c
+++ b/src/backend/nodes/nodeFuncs.c
@@ -2163,6 +2163,16 @@ expression_tree_walker_impl(Node *node,
 					return true;
 				if (WALK(expr->aggfilter))
 					return true;
+				if (WALK(expr->runCondition))
+					return true;
+			}
+			break;
+		case T_WindowFuncRunCondition:
+			{
+				WindowFuncRunCondition *expr = (WindowFuncRunCondition *) node;
+
+				if (WALK(expr->arg))
+					return true;
 			}
 			break;
 		case T_SubscriptingRef:
@@ -2400,8 +2410,6 @@ expression_tree_walker_impl(Node *node,
 					return true;
 				if (WALK(wc->endOffset))
 					return true;
-				if (WALK(wc->runCondition))
-					return true;
 			}
 			break;
 		case T_CTECycleClause:
@@ -2752,8 +2760,6 @@ query_tree_walker_impl(Query *query,
 				return true;
 			if (WALK(wc->endOffset))
 				return true;
-			if (WALK(wc->runCondition))
-				return true;
 		}
 	}
 
@@ -3053,6 +3059,16 @@ expression_tree_mutator_impl(Node *node,
 				return (Node *) newnode;
 			}
 			break;
+		case T_WindowFuncRunCondition:
+			{
+				WindowFuncRunCondition *wfuncrc = (WindowFuncRunCondition *) node;
+				WindowFuncRunCondition *newnode;
+
+				FLATCOPY(newnode, wfuncrc, WindowFuncRunCondition);
+				MUTATE(newnode->arg, wfuncrc->arg, Expr *);
+				return (Node *) newnode;
+			}
+			break;
 		case T_SubscriptingRef:
 			{
 				SubscriptingRef *sbsref = (SubscriptingRef *) node;
@@ -3466,7 +3482,6 @@ expression_tree_mutator_impl(Node *node,
 				MUTATE(newnode->orderClause, wc->orderClause, List *);
 				MUTATE(newnode->startOffset, wc->startOffset, Node *);
 				MUTATE(newnode->endOffset, wc->endOffset, Node *);
-				MUTATE(newnode->runCondition, wc->runCondition, List *);
 				return (Node *) newnode;
 			}
 			break;
@@ -3799,7 +3814,6 @@ query_tree_mutator_impl(Query *query,
 			FLATCOPY(newnode, wc, WindowClause);
 			MUTATE(newnode->startOffset, wc->startOffset, Node *);
 			MUTATE(newnode->endOffset, wc->endOffset, Node *);
-			MUTATE(newnode->runCondition, wc->runCondition, List *);
 
 			resultlist = lappend(resultlist, (Node *) newnode);
 		}
diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index cc51ae1757..4895cee994 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -2205,7 +2205,7 @@ set_dummy_rel_pathlist(RelOptInfo *rel)
  * the run condition will handle all of the required filtering.
  *
  * Returns true if 'opexpr' was found to be useful and was added to the
- * WindowClauses runCondition.  We also set *keep_original accordingly and add
+ * WindowFunc's runCondition.  We also set *keep_original accordingly and add
  * 'attno' to *run_cond_attrs offset by FirstLowInvalidHeapAttributeNumber.
  * If the 'opexpr' cannot be used then we set *keep_original to true and
  * return false.
@@ -2358,7 +2358,7 @@ find_window_run_conditions(Query *subquery, RangeTblEntry *rte, Index rti,
 			*keep_original = true;
 			runopexpr = opexpr;
 
-			/* determine the operator to use for the runCondition qual */
+			/* determine the operator to use for the WindowFuncRunCondition */
 			runoperator = get_opfamily_member(opinfo->opfamily_id,
 											  opinfo->oplefttype,
 											  opinfo->oprighttype,
@@ -2369,27 +2369,15 @@ find_window_run_conditions(Query *subquery, RangeTblEntry *rte, Index rti,
 
 	if (runopexpr != NULL)
 	{
-		Expr	   *newexpr;
+		WindowFuncRunCondition *wfuncrc;
 
-		/*
-		 * Build the qual required for the run condition keeping the
-		 * WindowFunc on the same side as it was originally.
-		 */
-		if (wfunc_left)
-			newexpr = make_opclause(runoperator,
-									runopexpr->opresulttype,
-									runopexpr->opretset, (Expr *) wfunc,
-									otherexpr, runopexpr->opcollid,
-									runopexpr->inputcollid);
-		else
-			newexpr = make_opclause(runoperator,
-									runopexpr->opresulttype,
-									runopexpr->opretset,
-									otherexpr, (Expr *) wfunc,
-									runopexpr->opcollid,
-									runopexpr->inputcollid);
+		wfuncrc = makeNode(WindowFuncRunCondition);
+		wfuncrc->opno = runoperator;
+		wfuncrc->inputcollid = runopexpr->inputcollid;
+		wfuncrc->wfunc_left = wfunc_left;
+		wfuncrc->arg = copyObject(otherexpr);
 
-		wclause->runCondition = lappend(wclause->runCondition, newexpr);
+		wfunc->runCondition = lappend(wfunc->runCondition, wfuncrc);
 
 		/* record that this attno was used in a run condition */
 		*run_cond_attrs = bms_add_member(*run_cond_attrs,
@@ -2403,9 +2391,9 @@ find_window_run_conditions(Query *subquery, RangeTblEntry *rte, Index rti,
 
 /*
  * check_and_push_window_quals
- *		Check if 'clause' is a qual that can be pushed into a WindowFunc's
- *		WindowClause as a 'runCondition' qual.  These, when present, allow
- *		some unnecessary work to be skipped during execution.
+ *		Check if 'clause' is a qual that can be pushed into a WindowFunc
+ *		as a 'runCondition' qual.  These, when present, allow some unnecessary
+ *		work to be skipped during execution.
  *
  * 'run_cond_attrs' will be populated with all targetlist resnos of subquery
  * targets (offset by FirstLowInvalidHeapAttributeNumber) that we pushed
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 3b77886567..6b64c4a362 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -2699,7 +2699,7 @@ create_windowagg_plan(PlannerInfo *root, WindowAggPath *best_path)
 						  wc->inRangeColl,
 						  wc->inRangeAsc,
 						  wc->inRangeNullsFirst,
-						  wc->runCondition,
+						  best_path->runCondition,
 						  best_path->qual,
 						  best_path->topwindow,
 						  subplan);
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 5320da51a0..032818423f 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -870,9 +870,6 @@ subquery_planner(PlannerGlobal *glob, Query *parse, PlannerInfo *parent_root,
 												EXPRKIND_LIMIT);
 		wc->endOffset = preprocess_expression(root, wc->endOffset,
 											  EXPRKIND_LIMIT);
-		wc->runCondition = (List *) preprocess_expression(root,
-														  (Node *) wc->runCondition,
-														  EXPRKIND_TARGET);
 	}
 
 	parse->limitOffset = preprocess_expression(root, parse->limitOffset,
@@ -4527,9 +4524,11 @@ create_one_window_path(PlannerInfo *root,
 	{
 		WindowClause *wc = lfirst_node(WindowClause, l);
 		List	   *window_pathkeys;
+		List	   *runcondition = NIL;
 		int			presorted_keys;
 		bool		is_sorted;
 		bool		topwindow;
+		ListCell   *lc2;
 
 		window_pathkeys = make_pathkeys_for_window(root,
 												   wc,
@@ -4577,7 +4576,6 @@ create_one_window_path(PlannerInfo *root,
 			 * we do need to account for the increase in tlist width.
 			 */
 			int64		tuple_width = window_target->width;
-			ListCell   *lc2;
 
 			window_target = copy_pathtarget(window_target);
 			foreach(lc2, wflists->windowFuncs[wc->winref])
@@ -4599,17 +4597,53 @@ create_one_window_path(PlannerInfo *root,
 		topwindow = foreach_current_index(l) == list_length(activeWindows) - 1;
 
 		/*
-		 * Accumulate all of the runConditions from each intermediate
-		 * WindowClause.  The top-level WindowAgg must pass these as a qual so
-		 * that it filters out unwanted tuples correctly.
+		 * Collect the WindowFuncRunConditions from each WindowFunc and
+		 * convert them into OpExprs
 		 */
-		if (!topwindow)
-			topqual = list_concat(topqual, wc->runCondition);
+		foreach(lc2, wflists->windowFuncs[wc->winref])
+		{
+			ListCell   *lc3;
+			WindowFunc *wfunc = lfirst_node(WindowFunc, lc2);
+
+			foreach(lc3, wfunc->runCondition)
+			{
+				WindowFuncRunCondition *wfuncrc =
+					lfirst_node(WindowFuncRunCondition, lc3);
+				Expr	   *opexpr;
+				Expr	   *leftop;
+				Expr	   *rightop;
+
+				if (wfuncrc->wfunc_left)
+				{
+					leftop = (Expr *) copyObject(wfunc);
+					rightop = copyObject(wfuncrc->arg);
+				}
+				else
+				{
+					leftop = copyObject(wfuncrc->arg);
+					rightop = (Expr *) copyObject(wfunc);
+				}
+
+				opexpr = make_opclause(wfuncrc->opno,
+									   BOOLOID,
+									   false,
+									   leftop,
+									   rightop,
+									   InvalidOid,
+									   wfuncrc->inputcollid);
+
+				runcondition = lappend(runcondition, opexpr);
+
+				if (!topwindow)
+					topqual = lappend(topqual, opexpr);
+			}
+		}
 
 		path = (Path *)
 			create_windowagg_path(root, window_rel, path, window_target,
 								  wflists->windowFuncs[wc->winref],
-								  wc, topwindow ? topqual : NIL, topwindow);
+								  runcondition, wc,
+								  topwindow ? topqual : NIL, topwindow);
 	}
 
 	add_path(window_rel, path);
diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c
index 41da670f15..5482ab85a7 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -2175,14 +2175,6 @@ perform_pullup_replace_vars(PlannerInfo *root,
 	parse->returningList = (List *)
 		pullup_replace_vars((Node *) parse->returningList, rvcontext);
 
-	foreach(lc, parse->windowClause)
-	{
-		WindowClause *wc = lfirst_node(WindowClause, lc);
-
-		if (wc->runCondition != NIL)
-			wc->runCondition = (List *)
-				pullup_replace_vars((Node *) wc->runCondition, rvcontext);
-	}
 	if (parse->onConflict)
 	{
 		parse->onConflict->onConflictSet = (List *)
diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 59487cbd79..b4e085e9d4 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -2566,6 +2566,7 @@ eval_const_expressions_mutator(Node *node,
 				newexpr->inputcollid = expr->inputcollid;
 				newexpr->args = args;
 				newexpr->aggfilter = aggfilter;
+				newexpr->runCondition = expr->runCondition;
 				newexpr->winref = expr->winref;
 				newexpr->winstar = expr->winstar;
 				newexpr->winagg = expr->winagg;
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 3cf1dac087..3491c3af1c 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -3471,6 +3471,7 @@ create_minmaxagg_path(PlannerInfo *root,
  * 'subpath' is the path representing the source of data
  * 'target' is the PathTarget to be computed
  * 'windowFuncs' is a list of WindowFunc structs
+ * 'runCondition' is a list of OpExprs to short-circuit WindowAgg execution
  * 'winclause' is a WindowClause that is common to all the WindowFuncs
  * 'qual' WindowClause.runconditions from lower-level WindowAggPaths.
  *		Must always be NIL when topwindow == false
@@ -3486,6 +3487,7 @@ create_windowagg_path(PlannerInfo *root,
 					  Path *subpath,
 					  PathTarget *target,
 					  List *windowFuncs,
+					  List *runCondition,
 					  WindowClause *winclause,
 					  List *qual,
 					  bool topwindow)
@@ -3510,6 +3512,7 @@ create_windowagg_path(PlannerInfo *root,
 	pathnode->subpath = subpath;
 	pathnode->winclause = winclause;
 	pathnode->qual = qual;
+	pathnode->runCondition = runCondition;
 	pathnode->topwindow = topwindow;
 
 	/*
diff --git a/src/backend/parser/parse_clause.c b/src/backend/parser/parse_clause.c
index 4fc5fc87e0..8118036495 100644
--- a/src/backend/parser/parse_clause.c
+++ b/src/backend/parser/parse_clause.c
@@ -2956,7 +2956,6 @@ transformWindowDefinitions(ParseState *pstate,
 											 rangeopfamily, rangeopcintype,
 											 &wc->endInRangeFunc,
 											 windef->endOffset);
-		wc->runCondition = NIL;
 		wc->winref = winref;
 
 		result = lappend(result, wc);
diff --git a/src/backend/parser/parse_expr.c b/src/backend/parser/parse_expr.c
index 1c1c86aa3e..aba3546ed1 100644
--- a/src/backend/parser/parse_expr.c
+++ b/src/backend/parser/parse_expr.c
@@ -3826,6 +3826,7 @@ transformJsonAggConstructor(ParseState *pstate, JsonAggConstructor *agg_ctor,
 		/* wincollid and inputcollid will be set by parse_collate.c */
 		wfunc->args = args;
 		wfunc->aggfilter = aggfilter;
+		wfunc->runCondition = NIL;
 		/* winref will be set by transformWindowFuncCall */
 		wfunc->winstar = false;
 		wfunc->winagg = true;
diff --git a/src/backend/parser/parse_func.c b/src/backend/parser/parse_func.c
index 0cbc950c95..9b23344a3b 100644
--- a/src/backend/parser/parse_func.c
+++ b/src/backend/parser/parse_func.c
@@ -834,6 +834,7 @@ ParseFuncOrColumn(ParseState *pstate, List *funcname, List *fargs,
 		wfunc->winstar = agg_star;
 		wfunc->winagg = (fdresult == FUNCDETAIL_AGGREGATE);
 		wfunc->aggfilter = agg_filter;
+		wfunc->runCondition = NIL;
 		wfunc->location = location;
 
 		/*
diff --git a/src/include/catalog/catversion.h b/src/include/catalog/catversion.h
index 8618ea5907..5a1dd1cb8f 100644
--- a/src/include/catalog/catversion.h
+++ b/src/include/catalog/catversion.h
@@ -57,6 +57,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	202404291
+#define CATALOG_VERSION_NO	202405051
 
 #endif
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index af80a5d38e..3ca06fc3af 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -1549,8 +1549,6 @@ typedef struct WindowClause
 	int			frameOptions;	/* frame_clause options, see WindowDef */
 	Node	   *startOffset;	/* expression for starting bound, if any */
 	Node	   *endOffset;		/* expression for ending bound, if any */
-	/* qual to help short-circuit execution */
-	List	   *runCondition pg_node_attr(query_jumble_ignore);
 	/* in_range function for startOffset */
 	Oid			startInRangeFunc pg_node_attr(query_jumble_ignore);
 	/* in_range function for endOffset */
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 91a6ce90d8..c7a415b23d 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -2308,6 +2308,7 @@ typedef struct WindowAggPath
 	Path	   *subpath;		/* path representing input source */
 	WindowClause *winclause;	/* WindowClause we'll be using */
 	List	   *qual;			/* lower-level WindowAgg runconditions */
+	List	   *runCondition;	/* OpExpr List to short-circuit execution */
 	bool		topwindow;		/* false for all apart from the WindowAgg
 								 * that's closest to the root of the plan */
 } WindowAggPath;
diff --git a/src/include/nodes/primnodes.h b/src/include/nodes/primnodes.h
index 43f5766d73..4830efc573 100644
--- a/src/include/nodes/primnodes.h
+++ b/src/include/nodes/primnodes.h
@@ -575,6 +575,8 @@ typedef struct WindowFunc
 	List	   *args;
 	/* FILTER expression, if any */
 	Expr	   *aggfilter;
+	/* List of WindowFuncRunConditions to help short-circuit execution */
+	List	   *runCondition pg_node_attr(query_jumble_ignore);
 	/* index of associated WindowClause */
 	Index		winref;
 	/* true if argument list was really '*' */
@@ -585,6 +587,34 @@ typedef struct WindowFunc
 	ParseLoc	location;
 } WindowFunc;
 
+/*
+ * WindowFuncRunCondition
+ *
+ * Represents intermediate OpExprs which will be used by WindowAgg to
+ * short-circuit execution.
+ */
+typedef struct WindowFuncRunCondition
+{
+	Expr		xpr;
+
+	/* PG_OPERATOR OID of the operator */
+	Oid			opno;
+	/* OID of collation that operator should use */
+	Oid			inputcollid pg_node_attr(query_jumble_ignore);
+
+	/*
+	 * true of WindowFunc belongs on the left of the resulting OpExpr or false
+	 * if the WindowFunc is on the right.
+	 */
+	bool		wfunc_left;
+
+	/*
+	 * The Expr being compared to the WindowFunc to use in the OpExpr in the
+	 * WindowAgg's runCondition
+	 */
+	Expr	   *arg;
+} WindowFuncRunCondition;
+
 /*
  * MergeSupportFunc
  *
diff --git a/src/include/optimizer/pathnode.h b/src/include/optimizer/pathnode.h
index c5c4756b0f..112e7c23d4 100644
--- a/src/include/optimizer/pathnode.h
+++ b/src/include/optimizer/pathnode.h
@@ -250,6 +250,7 @@ extern WindowAggPath *create_windowagg_path(PlannerInfo *root,
 											Path *subpath,
 											PathTarget *target,
 											List *windowFuncs,
+											List *runCondition,
 											WindowClause *winclause,
 											List *qual,
 											bool topwindow);
diff --git a/src/test/regress/expected/window.out b/src/test/regress/expected/window.out
index e46710cf31..ae4e8851f8 100644
--- a/src/test/regress/expected/window.out
+++ b/src/test/regress/expected/window.out
@@ -4207,6 +4207,24 @@ WHERE s.c = 1;
                            ->  Seq Scan on empsalary e2
 (14 rows)
 
+-- Ensure the run condition optimization is used in cases where the WindowFunc
+-- has a Var from another query level
+EXPLAIN (COSTS OFF)
+SELECT 1 FROM
+  (SELECT ntile(s1.x) OVER () AS c
+   FROM (SELECT (SELECT 1) AS x) AS s1) s
+WHERE s.c = 1;
+                           QUERY PLAN                            
+-----------------------------------------------------------------
+ Subquery Scan on s
+   Filter: (s.c = 1)
+   ->  WindowAgg
+         Run Condition: (ntile((InitPlan 1).col1) OVER (?) <= 1)
+         InitPlan 1
+           ->  Result
+         ->  Result
+(7 rows)
+
 -- Tests to ensure we don't push down the run condition when it's not valid to
 -- do so.
 -- Ensure we don't push down when the frame options show that the window
diff --git a/src/test/regress/sql/window.sql b/src/test/regress/sql/window.sql
index 678aa1a21c..6de5493b05 100644
--- a/src/test/regress/sql/window.sql
+++ b/src/test/regress/sql/window.sql
@@ -1377,6 +1377,14 @@ SELECT 1 FROM
    WHERE e1.empno = e2.empno) s
 WHERE s.c = 1;
 
+-- Ensure the run condition optimization is used in cases where the WindowFunc
+-- has a Var from another query level
+EXPLAIN (COSTS OFF)
+SELECT 1 FROM
+  (SELECT ntile(s1.x) OVER () AS c
+   FROM (SELECT (SELECT 1) AS x) AS s1) s
+WHERE s.c = 1;
+
 -- Tests to ensure we don't push down the run condition when it's not valid to
 -- do so.
 
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index e10ff28ee5..eee989bba1 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3120,6 +3120,7 @@ WindowDef
 WindowFunc
 WindowFuncExprState
 WindowFuncLists
+WindowFuncRunCondition
 WindowObject
 WindowObjectData
 WindowStatePerAgg