diff --git a/src/backend/distributed/planner/multi_planner.c b/src/backend/distributed/planner/multi_planner.c
index 52681bcc3..e9b375862 100644
--- a/src/backend/distributed/planner/multi_planner.c
+++ b/src/backend/distributed/planner/multi_planner.c
@@ -127,6 +127,11 @@ multi_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 		{
 			result = CreateDistributedPlan(result, originalQuery, parse,
 										   boundParams, plannerRestrictionContext);
+
+			assignRTEIdentities = false;
+			setPartitionedTablesInherited = true;
+
+			AdjustParseTree(parse, assignRTEIdentities, setPartitionedTablesInherited);
 		}
 	}
 	PG_CATCH();
@@ -136,14 +141,6 @@ multi_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	}
 	PG_END_TRY();
 
-	if (needsDistributedPlanning)
-	{
-		assignRTEIdentities = false;
-		setPartitionedTablesInherited = true;
-
-		AdjustParseTree(parse, assignRTEIdentities, setPartitionedTablesInherited);
-	}
-
 	/* remove the context from the context list */
 	PopPlannerRestrictionContext();
 
@@ -971,6 +968,13 @@ CurrentPlannerRestrictionContext(void)
 	plannerRestrictionContext =
 		(PlannerRestrictionContext *) linitial(plannerRestrictionContextList);
 
+	if (plannerRestrictionContext == NULL)
+	{
+		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR),
+						errmsg("planner restriction context stack was empty"),
+						errdetail("Please report this to the Citus core team.")));
+	}
+
 	return plannerRestrictionContext;
 }