From c3eee14e92508cc614087dad155e2482cb4b71e0 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Fri, 21 Feb 2025 12:34:44 +0800
Subject: [PATCH v4 1/1] fix expanding virtual generated columns with grouping
 sets

---
 src/backend/optimizer/prep/prepjointree.c | 17 ++++++++++++++---
 src/test/regress/expected/join.out        |  8 ++++++++
 src/test/regress/sql/join.sql             |  4 ++++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c
index 8fdca35d087..aad51f3d0a0 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -70,6 +70,8 @@ typedef struct pullup_replace_vars_context
 	bool	   *outer_hasSubLinks;	/* -> outer query's hasSubLinks */
 	int			varno;			/* varno of subquery */
 	bool		wrap_non_vars;	/* do we need all non-Var outputs to be PHVs? */
+	bool		wrap_vars;		/* do we need all Var outputs to be PHVs?
+								 * this may needed for GroupingSet */
 	Node	  **rv_cache;		/* cache for results with PHVs */
 } pullup_replace_vars_context;
 
@@ -1295,6 +1297,7 @@ pull_up_simple_subquery(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte,
 	rvcontext.varno = varno;
 	/* this flag will be set below, if needed */
 	rvcontext.wrap_non_vars = false;
+	rvcontext.wrap_vars = false;
 	/* initialize cache array with indexes 0 .. length(tlist) */
 	rvcontext.rv_cache = palloc0((list_length(subquery->targetList) + 1) *
 								 sizeof(Node *));
@@ -1846,6 +1849,7 @@ pull_up_simple_values(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte)
 	rvcontext.outer_hasSubLinks = &parse->hasSubLinks;
 	rvcontext.varno = varno;
 	rvcontext.wrap_non_vars = false;
+	rvcontext.wrap_vars = false;
 	/* initialize cache array with indexes 0 .. length(tlist) */
 	rvcontext.rv_cache = palloc0((list_length(tlist) + 1) *
 								 sizeof(Node *));
@@ -2014,6 +2018,7 @@ pull_up_constant_function(PlannerInfo *root, Node *jtnode,
 	rvcontext.varno = ((RangeTblRef *) jtnode)->rtindex;
 	/* this flag will be set below, if needed */
 	rvcontext.wrap_non_vars = false;
+	rvcontext.wrap_vars = false;
 	/* initialize cache array with indexes 0 .. length(tlist) */
 	rvcontext.rv_cache = palloc0((list_length(rvcontext.targetlist) + 1) *
 								 sizeof(Node *));
@@ -2699,6 +2704,9 @@ pullup_replace_vars_callback(Var *var,
 				}
 			}
 
+			if (!wrap && rcon->wrap_vars && rcon->wrap_non_vars)
+				wrap = true;
+
 			if (wrap)
 			{
 				newnode = (Node *)
@@ -3002,7 +3010,7 @@ expand_virtual_generated_columns(PlannerInfo *root)
 
 					ChangeVarNodes(defexpr, 1, rt_index, 0);
 
-					te = makeTargetEntry((Expr *) defexpr, attnum, 0, false);
+					te = makeTargetEntry((Expr *) defexpr, attnum, NULL, false);
 					tlist = lappend(tlist, te);
 				}
 				else
@@ -3015,7 +3023,7 @@ expand_virtual_generated_columns(PlannerInfo *root)
 								  attr->atttypmod,
 								  attr->attcollation,
 								  0);
-					te = makeTargetEntry((Expr *) var, attnum, 0, false);
+					te = makeTargetEntry((Expr *) var, attnum, NULL, false);
 					tlist = lappend(tlist, te);
 				}
 			}
@@ -3036,12 +3044,15 @@ expand_virtual_generated_columns(PlannerInfo *root)
 
 			/*
 			 * If the query uses grouping sets, we need a PlaceHolderVar for
-			 * anything that's not a simple Var.  This ensures that
+			 * anything.  This ensures that
 			 * expressions retain their separate identity so that they will
 			 * match grouping set columns when appropriate.
 			 */
 			if (parse->groupingSets)
+			{
 				rvcontext.wrap_non_vars = true;
+				rvcontext.wrap_vars = true;
+			}
 
 			parse = (Query *) pullup_replace_vars((Node *) parse, &rvcontext);
 		}
diff --git a/src/test/regress/expected/join.out b/src/test/regress/expected/join.out
index e787792ece1..26e282a40dc 100644
--- a/src/test/regress/expected/join.out
+++ b/src/test/regress/expected/join.out
@@ -3580,6 +3580,14 @@ having c = 2;
    |   | 2 |   |  
 (1 row)
 
+select a,c from t
+group by grouping sets (a,c)
+having c = 2;
+ a | c 
+---+---
+   | 2
+(1 row)
+
 rollback;
 --
 -- test incorrect handling of placeholders that only appear in targetlists,
diff --git a/src/test/regress/sql/join.sql b/src/test/regress/sql/join.sql
index 828cbdfd9db..0abecd7e54e 100644
--- a/src/test/regress/sql/join.sql
+++ b/src/test/regress/sql/join.sql
@@ -1085,6 +1085,10 @@ select * from t
 group by grouping sets (a, b, c, d, e)
 having c = 2;
 
+select a,c from t
+group by grouping sets (a,c)
+having c = 2;
+
 rollback;
 
 --
-- 
2.34.1