diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c
index d366f6c56..9f9dde1cb 100644
--- a/src/backend/distributed/planner/insert_select_planner.c
+++ b/src/backend/distributed/planner/insert_select_planner.c
@@ -41,8 +41,6 @@
 static DistributedPlan * CreateDistributedInsertSelectPlan(Query *originalQuery,
 														   PlannerRestrictionContext *
 														   plannerRestrictionContext);
-static bool SafeToPushDownSubquery(PlannerRestrictionContext *plannerRestrictionContext,
-								   Query *originalQuery);
 static Task * RouterModifyTaskForShardInterval(Query *originalQuery,
 											   ShardInterval *shardInterval,
 											   RelationRestrictionContext *
@@ -217,7 +215,7 @@ CreateDistributedInsertSelectPlan(Query *originalQuery,
 	RelationRestrictionContext *relationRestrictionContext =
 		plannerRestrictionContext->relationRestrictionContext;
 	bool allReferenceTables = relationRestrictionContext->allReferenceTables;
-	bool safeToPushDownSubquery = false;
+	bool queryContainsDistributionKeyEquality = false;
 
 	distributedPlan->operation = originalQuery->commandType;
 
@@ -234,8 +232,8 @@ CreateDistributedInsertSelectPlan(Query *originalQuery,
 		return distributedPlan;
 	}
 
-	safeToPushDownSubquery = SafeToPushDownSubquery(plannerRestrictionContext,
-													originalQuery);
+	queryContainsDistributionKeyEquality =
+		QueryContainsDistributionKeyEquality(plannerRestrictionContext, originalQuery);
 
 	/*
 	 * Plan select query for each shard in the target table. Do so by replacing the
@@ -255,7 +253,7 @@ CreateDistributedInsertSelectPlan(Query *originalQuery,
 		modifyTask = RouterModifyTaskForShardInterval(originalQuery, targetShardInterval,
 													  relationRestrictionContext,
 													  taskIdIndex,
-													  safeToPushDownSubquery);
+													  queryContainsDistributionKeyEquality);
 
 		/* add the task if it could be created */
 		if (modifyTask != NULL)
@@ -392,34 +390,6 @@ DistributedInsertSelectSupported(Query *queryTree, RangeTblEntry *insertRte,
 }
 
 
-/*
- * SafeToPushDownSubquery returns true if either
- *    (i)  there exists join in the query and all relations joined on their
- *         partition keys
- *    (ii) there exists only union set operations and all relations has
- *         partition keys in the same ordinal position in the query
- */
-static bool
-SafeToPushDownSubquery(PlannerRestrictionContext *plannerRestrictionContext,
-					   Query *originalQuery)
-{
-	bool restrictionEquivalenceForPartitionKeys =
-		RestrictionEquivalenceForPartitionKeys(plannerRestrictionContext);
-
-	if (restrictionEquivalenceForPartitionKeys)
-	{
-		return true;
-	}
-
-	if (ContainsUnionSubquery(originalQuery))
-	{
-		return SafeToPushdownUnionSubquery(plannerRestrictionContext);
-	}
-
-	return false;
-}
-
-
 /*
  * RouterModifyTaskForShardInterval creates a modify task by
  * replacing the partitioning qual parameter added in distributed_planner()
diff --git a/src/backend/distributed/planner/relation_restriction_equivalence.c b/src/backend/distributed/planner/relation_restriction_equivalence.c
index f2ba53f14..5a6efe6e9 100644
--- a/src/backend/distributed/planner/relation_restriction_equivalence.c
+++ b/src/backend/distributed/planner/relation_restriction_equivalence.c
@@ -143,6 +143,34 @@ static bool JoinRestrictionListExistsInContext(JoinRestriction *joinRestrictionI
 											   joinRestrictionContext);
 
 
+/*
+ * QueryContainsDistributionKeyEquality returns true if either
+ *    (i)  there exists join in the query and all relations joined on their
+ *         partition keys
+ *    (ii) there exists only union set operations and all relations has
+ *         partition keys in the same ordinal position in the query
+ */
+bool
+QueryContainsDistributionKeyEquality(PlannerRestrictionContext *plannerRestrictionContext,
+									 Query *originalQuery)
+{
+	bool restrictionEquivalenceForPartitionKeys =
+		RestrictionEquivalenceForPartitionKeys(plannerRestrictionContext);
+
+	if (restrictionEquivalenceForPartitionKeys)
+	{
+		return true;
+	}
+
+	if (ContainsUnionSubquery(originalQuery))
+	{
+		return SafeToPushdownUnionSubquery(plannerRestrictionContext);
+	}
+
+	return false;
+}
+
+
 /*
  * SafeToPushdownUnionSubquery returns true if all the relations are returns
  * partition keys in the same ordinal position and there is no reference table
diff --git a/src/include/distributed/relation_restriction_equivalence.h b/src/include/distributed/relation_restriction_equivalence.h
index c80225963..1f9ef5245 100644
--- a/src/include/distributed/relation_restriction_equivalence.h
+++ b/src/include/distributed/relation_restriction_equivalence.h
@@ -15,14 +15,18 @@
 #include "distributed/distributed_planner.h"
 
 
+extern bool QueryContainsDistributionKeyEquality(PlannerRestrictionContext *
+												 plannerRestrictionContext,
+												 Query *originalQuery);
+extern bool SafeToPushdownUnionSubquery(PlannerRestrictionContext *
+										plannerRestrictionContext);
 extern bool ContainsUnionSubquery(Query *queryTree);
 extern bool RestrictionEquivalenceForPartitionKeys(PlannerRestrictionContext *
 												   plannerRestrictionContext);
 extern List * GenerateAllAttributeEquivalences(PlannerRestrictionContext *
 											   plannerRestrictionContext);
 extern uint32 ReferenceRelationCount(RelationRestrictionContext *restrictionContext);
-extern bool SafeToPushdownUnionSubquery(
-	PlannerRestrictionContext *plannerRestrictionContext);
+
 extern List * RelationIdList(Query *query);
 extern PlannerRestrictionContext * FilterPlannerRestrictionForQuery(
 	PlannerRestrictionContext *plannerRestrictionContext,