diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c
index 4a7e5e31a..1d8361360 100644
--- a/src/backend/distributed/planner/multi_router_planner.c
+++ b/src/backend/distributed/planner/multi_router_planner.c
@@ -1134,6 +1134,23 @@ ErrorIfOnConflictNotSupported(Query *queryTree)
 	Node *arbiterWhere = queryTree->onConflict->arbiterWhere;
 	Node *onConflictWhere = queryTree->onConflict->onConflictWhere;
 
+
+	bool setTargetEntryPartitionColumn = false;
+
+	if (partitionColumn) {
+		RangeTblEntry* resultRTE = ExtractResultRelationRTE(queryTree);
+		/*
+			* FirstLowInvalidHeapAttributeNumber is added as an offset to rte->updatedCols.
+			* So we substract that to get the column no for an updated column that matches
+			* resultRTE->updatedcols.
+			*/
+		int updatedColNoWithOffset = partitionColumn->varattno - FirstLowInvalidHeapAttributeNumber;
+		if (bms_is_member(updatedColNoWithOffset, resultRTE->updatedCols)) {
+			setTargetEntryPartitionColumn = true;
+		}
+	}
+
+
 	/*
 	 * onConflictSet is expanded via expand_targetlist() on the standard planner.
 	 * This ends up adding all the columns to the onConflictSet even if the user
@@ -1146,17 +1163,6 @@ ErrorIfOnConflictNotSupported(Query *queryTree)
 	foreach(setTargetCell, onConflictSet)
 	{
 		TargetEntry *setTargetEntry = (TargetEntry *) lfirst(setTargetCell);
-		bool setTargetEntryPartitionColumn = false;
-
-		/* reference tables do not have partition column */
-		if (partitionColumn == NULL)
-		{
-			setTargetEntryPartitionColumn = false;
-		}
-		else if (setTargetEntry->resno == partitionColumn->varattno)
-		{
-			setTargetEntryPartitionColumn = true;
-		}
 
 		if (setTargetEntryPartitionColumn)
 		{