Add ExprMentionsPartitionColumn function to identify partition column references in queries

src/backend/distributed/planner/multi_logical_planner.c

@@ -136,6 +136,8 @@ static MultiNode * ApplyCartesianProductReferenceJoin(MultiNode *leftNode,
 static MultiNode * ApplyCartesianProduct(MultiNode *leftNode, MultiNode *rightNode,
 										 List *partitionColumnList, JoinType joinType,
 										 List *joinClauses);
+static bool
+ExprMentionsPartitionColumn(Node *node, Query *query);										 
 
 
 /*
@@ -224,14 +226,14 @@ TargetListOnPartitionColumn(Query *query, List *targetEntryList)
 		TargetEntry *targetEntry = (TargetEntry *) lfirst(targetEntryCell);
 		Expr *targetExpression = targetEntry->expr;
 
-		bool skipOuterVars = true;
-		bool isPartitionColumn = IsPartitionColumn(targetExpression, query,
-												   skipOuterVars);
+		bool isPartitionColumn = ExprMentionsPartitionColumn((Node *) targetExpression,
+                                query);
+
 		Var *column = NULL;
 		RangeTblEntry *rte = NULL;
 
 		FindReferencedTableColumn(targetExpression, NIL, query, &column, &rte,
-								  skipOuterVars);
+								  /*skipOuterVars=*/false);
 		Oid relationId = rte ? rte->relid : InvalidOid;
 
 		/*
@@ -292,6 +294,66 @@ TargetListOnPartitionColumn(Query *query, List *targetEntryList)
 }
 
 
+/*
+ * 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 >= 180000
+        /* 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);
+}
+
+
 /*
  * FindNodeMatchingCheckFunctionInRangeTableList finds a node for which the checker
  * function returns true.