diff --git a/src/backend/distributed/planner/relation_restriction_equivalence.c b/src/backend/distributed/planner/relation_restriction_equivalence.c
index 4d6fa1faa..648b5957c 100644
--- a/src/backend/distributed/planner/relation_restriction_equivalence.c
+++ b/src/backend/distributed/planner/relation_restriction_equivalence.c
@@ -62,6 +62,8 @@ typedef struct AttributeEquivalenceClass
 typedef struct FindRteIdentityContext {
 	uint32 rteIdentity;
 	Query* query;
+	Query* temp;
+	bool found;
 }FindRteIdentityContext;
 
 /*
@@ -1775,7 +1777,7 @@ RelationRestrictionPartitionKeyIndex(Query* originalQuery, RelationRestriction *
 	FindRteIdentityContext* findRteIdentityContext = palloc0(sizeof(FindRteIdentityContext));
 	findRteIdentityContext->rteIdentity = GetRTEIdentity(relationRestriction->rte);
 	findRteIdentityWalker((Node*) originalQuery, findRteIdentityContext);
-	if (findRteIdentityContext->query == NULL) {
+	if (!findRteIdentityContext->found) {
 		return NULL;
 	}
 	List *relationTargetList = findRteIdentityContext->query->targetList;
@@ -1821,19 +1823,12 @@ findRteIdentityWalker(Node *node, FindRteIdentityContext* context)
 	if (IsA(node, Query))
 	{
 		Query *query = (Query *) node;
-		RangeTblEntry* rte = NULL;
-		foreach_ptr(rte, query->rtable) {
-			if (rte->rtekind == RTE_RELATION) {
-				if (GetRTEIdentity(rte) == context->rteIdentity) {
-					context->query = query;
-					return true;
-				}
-			}
-		}
+		Query* prev = context->temp;
+		context->temp = query;
 		query_tree_walker(query, findRteIdentityWalker, context,
 						  QTW_EXAMINE_RTES_BEFORE);
-
-		return false;
+		context->temp = prev;
+		return false;				  
 	}
 
 	if (!IsA(node, RangeTblEntry))
@@ -1841,6 +1836,14 @@ findRteIdentityWalker(Node *node, FindRteIdentityContext* context)
 		return expression_tree_walker(node, findRteIdentityWalker,
 									  context);
 	}
+	RangeTblEntry* rte = (RangeTblEntry*) node;
+	if (rte->rtekind == RTE_RELATION) {
+		if (GetRTEIdentity(rte) == context->rteIdentity) {
+			context->found = true;
+			context->query = context->temp;
+			return true;
+		}
+	}
 	return false;
 }
 
diff --git a/src/test/regress/expected/union_pushdown.out b/src/test/regress/expected/union_pushdown.out
index c1d5b0fa1..183c53e26 100644
--- a/src/test/regress/expected/union_pushdown.out
+++ b/src/test/regress/expected/union_pushdown.out
@@ -236,8 +236,12 @@ FROM
       UNION ALL SELECT user_id AS user_id FROM users_table_part JOIN events_table_part USING (user_id) WHERE users_table_part.value_1 = 8  GROUP BY user_id)) AS bar
       WHERE user_id < 2000 ) as level_1 ) as level_2 ) as level_3
 ORDER BY 1 DESC LIMIT 10;
-DEBUG:  generating subplan XXX_1 for subquery SELECT user_id FROM union_pushdown.users_table_part WHERE (value_2 OPERATOR(pg_catalog.=) 3)
-ERROR:  complex joins are only supported when all distributed tables are co-located and joined on their distribution columns
+DEBUG:  push down of limit count: 10
+ count
+---------------------------------------------------------------------
+    78
+(1 row)
+
 -- safe to pushdown
 SELECT DISTINCT user_id FROM (
     SELECT * FROM