Add support for synthetic GROUP RTE in FindReferencedTableColumn for PostgreSQL 18

src/backend/distributed/planner/insert_select_planner.c

@@ -108,6 +108,9 @@ static void ProcessEntryPair(TargetEntry *insertEntry, TargetEntry *selectEntry,
 							 Form_pg_attribute attr, int targetEntryIndex,
 							 List **projectedEntries, List **nonProjectedEntries);
 
+#if PG_VERSION_NUM >= PG_VERSION_18
+static bool ExprMentionsPartitionColumn(Node *node, Query *query);
+#endif
 
 /* depth of current insert/select planner. */
 static int insertSelectPlannerLevel = 0;
@@ -1404,6 +1407,7 @@ InsertPartitionColumnMatchesSelect(Query *query, RangeTblEntry *insertRte,
 		/* we can set the select relation id */
 		*selectPartitionColumnTableId = subqueryPartitionColumnRelationId;
 
+
 		break;
 	}
 
@@ -1422,6 +1426,79 @@ InsertPartitionColumnMatchesSelect(Query *query, RangeTblEntry *insertRte,
 }
 
 
+#if PG_VERSION_NUM >= PG_VERSION_18
+
+/*
+ * ExprMentionsPartitionColumn
+ *
+ * Return true iff `node` ultimately resolves to the partition column of
+ * *any* distributed table referenced by `query`.
+ *
+ *  • Understands OUTER_VAR indirection (PG ≥ 17)
+ *  • Understands synthetic RTE_GROUP entries (PG ≥ 18)
+ *  • Falls back to original logic for PG 15/16 automatically
+ */
+static bool
+ExprMentionsPartitionColumn(Node *node, Query *query)
+{
+	if (node == NULL)
+	{
+		return false;
+	}
+
+	if (IsA(node, Var))
+	{
+		Var *v = (Var *) node;
+
+		/* Follow OUTER_VAR → target-list indirection, if present */
+		if (v->varno == OUTER_VAR)
+		{
+			TargetEntry *tle = get_tle_by_resno(query->targetList, v->varattno);
+			return tle && ExprMentionsPartitionColumn((Node *) tle->expr, query);
+		}
+
+		/* Sanity-check varno */
+		if (v->varno <= 0 || v->varno > list_length(query->rtable))
+		{
+			return false;
+		}
+
+		RangeTblEntry *rte = rt_fetch(v->varno, query->rtable);
+
+#if PG_VERSION_NUM >= PG_VERSION_18
+
+		/* Synthetic GROUP RTE – examine its expressions instead */
+		if (rte->rtekind == RTE_GROUP && rte->groupexprs)
+		{
+			ListCell *lc;
+			foreach(lc, rte->groupexprs)
+			if (ExprMentionsPartitionColumn((Node *) lfirst(lc), query))
+			{
+				return true;
+			}
+			return false;
+		}
+#endif
+
+		/* Real table? — compare against its dist key */
+		if (rte->rtekind == RTE_RELATION && HasDistributionKey(rte->relid))
+		{
+			Var *partcol = DistPartitionKey(rte->relid);
+			return partcol && partcol->varattno == v->varattno;
+		}
+		return false;
+	}
+
+	/* Recurse through any other node type */
+	return expression_tree_walker(node,
+								  ExprMentionsPartitionColumn,
+								  (void *) query);
+}
+
+
+#endif
+
+
 /*
  * CreateNonPushableInsertSelectPlan creates a query plan for a SELECT into a
  * distributed table. The query plan can also be executed on a worker in MX.

src/backend/distributed/planner/multi_logical_optimizer.c

@@ -4557,6 +4557,31 @@ FindReferencedTableColumn(Expr *columnExpression, List *parentQueryList, Query *
 		FindReferencedTableColumn(joinColumn, parentQueryList, query, column,
 								  rteContainingReferencedColumn, skipOuterVars);
 	}
+#if PG_VERSION_NUM >= 180000
+	else if (rangeTableEntry->rtekind == RTE_GROUP)
+	{
+		/*
+		 * PG 18: synthetic GROUP RTE.  Each groupexprs item corresponds to the
+		 * columns produced by the grouping step, in the *same ordinal order* as
+		 * the Vars that reference them.
+		 */
+		List *groupexprs = rangeTableEntry->groupexprs;
+		AttrNumber groupIndex = candidateColumn->varattno - 1;
+
+		/* Safety check */
+		if (groupIndex < 0 || groupIndex >= list_length(groupexprs))
+		{
+			return;  /* malformed Var */
+		}
+		Expr *groupExpr = (Expr *) list_nth(groupexprs, groupIndex);
+
+		/* Recurse on the underlying expression (stay in the same query) */
+		FindReferencedTableColumn(groupExpr, parentQueryList, query,
+								  column, rteContainingReferencedColumn,
+								  skipOuterVars);
+	}
+#endif   /* PG_VERSION_NUM >= 180000 */
+
 	else if (rangeTableEntry->rtekind == RTE_CTE)
 	{
 		/*