From c507801c54acba596de14125f9d6451bb0c99624 Mon Sep 17 00:00:00 2001
From: Mehmet Yilmaz <mehmetyilmaz@microsoft.com>
Date: Tue, 24 Jun 2025 12:42:47 +0000
Subject: [PATCH] Add ExprReferencesPartitionColumn function to check partition
 column references in queries

---
 .../planner/query_pushdown_planning.c         | 102 ++++++++++++++++--
 1 file changed, 94 insertions(+), 8 deletions(-)

diff --git a/src/backend/distributed/planner/query_pushdown_planning.c b/src/backend/distributed/planner/query_pushdown_planning.c
index 6150f4987..e4c859d10 100644
--- a/src/backend/distributed/planner/query_pushdown_planning.c
+++ b/src/backend/distributed/planner/query_pushdown_planning.c
@@ -113,6 +113,8 @@ static Var * PartitionColumnForPushedDownSubquery(Query *query);
 static bool ContainsReferencesToRelids(Query *query, Relids relids, int *foundRelid);
 static bool ContainsReferencesToRelidsWalker(Node *node,
 											 RelidsReferenceWalkerContext *context);
+static bool
+ExprReferencesPartitionColumn(Node *node, Query *query);									 
 
 
 /*
@@ -1036,8 +1038,13 @@ DeferErrorIfCannotPushdownSubquery(Query *subqueryTree, bool outerMostQueryHasLi
 		errorDetail = "For Update/Share commands are currently unsupported";
 	}
 
-	/* grouping sets are not allowed in subqueries*/
-	if (subqueryTree->groupingSets)
+	/* real GROUPING SETS / ROLLUP / CUBE are still disallowed;
+	* but PG 18 also fills groupingSets with a single simple item
+	* when the GROUP BY list contains un-projected columns.  Allow that. */
+	if (subqueryTree->groupingSets &&
+		!(list_length(subqueryTree->groupingSets) == 1 &&
+		((GroupingSet *) linitial(subqueryTree->groupingSets))->kind
+			== GROUPING_SET_SIMPLE))
 	{
 		preconditionsSatisfied = false;
 		errorDetail = "could not run distributed query with GROUPING SETS, CUBE, "
@@ -1097,12 +1104,27 @@ DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral,
 	/* group clause list must include partition column */
 	if (subqueryTree->groupClause)
 	{
-		List *groupClauseList = subqueryTree->groupClause;
-		List *targetEntryList = subqueryTree->targetList;
-		List *groupTargetEntryList = GroupTargetEntryList(groupClauseList,
-														  targetEntryList);
-		bool groupOnPartitionColumn =
-			TargetListOnPartitionColumn(subqueryTree, groupTargetEntryList);
+		bool groupOnPartitionColumn = false;
+
+		/* Walk every GROUP BY expression. */
+		ListCell *lc;
+		foreach (lc, subqueryTree->groupClause)
+		{
+			Expr *expr = (Expr *)
+				get_sortgroupclause_expr(
+					(SortGroupClause *) lfirst(lc),
+					subqueryTree->targetList);
+
+			expr = (Expr *) strip_implicit_coercions((Node *) expr);
+
+			if (ExprReferencesPartitionColumn((Node *) expr, subqueryTree))
+			{
+				groupOnPartitionColumn = true;
+				break;
+			}
+		}
+
+
 		if (!groupOnPartitionColumn)
 		{
 			preconditionsSatisfied = false;
@@ -1171,6 +1193,70 @@ DeferErrorIfSubqueryRequiresMerge(Query *subqueryTree, bool lateral,
 }
 
 
+/*
+ * ExprReferencesPartitionColumn
+ *
+ * Recursively return true iff `node` ultimately refers to the partition
+ * (distribution) column of *any* distributed table in `query`.
+ *
+ * Handles:
+ *  • plain Vars that point straight at RTE_RELATION
+ *  • OUTER_VAR indirection (PG ≥ 18 resjunk columns)
+ *  • Vars that point at an RTE_GROUP (PG ≥ 18); we descend into
+ *    its groupexprs until we reach real relations.
+ */
+static bool
+ExprReferencesPartitionColumn(Node *node, Query *query)
+{
+    if (node == NULL)
+        return false;
+
+    /* 1. Plain or OUTER_VAR Var ---------------------------------------- */
+    if (IsA(node, Var))
+    {
+        Var *v = (Var *) node;
+
+        /* Follow the OUTER_VAR → targetlist indirection exactly once */
+        if (v->varno == OUTER_VAR)
+        {
+            TargetEntry *tle = get_tle_by_resno(query->targetList, v->varattno);
+            return tle && ExprReferencesPartitionColumn((Node *) tle->expr, query);
+        }
+
+        /* Sanity: varno must fit in rtable */
+        if (v->varno <= 0 || v->varno > list_length(query->rtable))
+            return false;
+
+        RangeTblEntry *rte = rt_fetch(v->varno, query->rtable);
+
+        /* 1a. Real relation: compare attnos ---------------------------- */
+        if (rte->rtekind == RTE_RELATION && HasDistributionKey(rte->relid))
+        {
+            Var *partcol = DistPartitionKey(rte->relid);
+            return partcol && partcol->varattno == v->varattno;
+        }
+
+        /* 1b. Synthetic GROUP entry: descend into its expressions ------- */
+        if (rte->rtekind == RTE_GROUP && rte->groupexprs)
+        {
+            ListCell *lc;
+            foreach (lc, rte->groupexprs)
+            {
+                if (ExprReferencesPartitionColumn((Node *) lfirst(lc), query))
+                    return true;
+            }
+        }
+
+        return false;      /* other RTE kinds are irrelevant here */
+    }
+
+    /* 2. Any other node type: walk its children ------------------------ */
+    return expression_tree_walker(node,
+           ExprReferencesPartitionColumn,
+           (void *) query);
+}
+
+
 /*
  * DeferErrorIfUnsupportedTableCombination checks if the given query tree contains any
  * unsupported range table combinations. For this, the function walks over all