diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c
index 195adc91a..32464032b 100644
--- a/src/backend/distributed/planner/distributed_planner.c
+++ b/src/backend/distributed/planner/distributed_planner.c
@@ -147,6 +147,10 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 
 			result = CreateDistributedPlannedStmt(planId, result, originalQuery, parse,
 												  boundParams, plannerRestrictionContext);
+
+			setPartitionedTablesInherited = true;
+			AdjustPartitioningForDistributedPlanning(parse,
+													 setPartitionedTablesInherited);
 		}
 	}
 	PG_CATCH();
@@ -156,13 +160,6 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	}
 	PG_END_TRY();
 
-	if (needsDistributedPlanning)
-	{
-		setPartitionedTablesInherited = true;
-
-		AdjustPartitioningForDistributedPlanning(parse, setPartitionedTablesInherited);
-	}
-
 	/* remove the context from the context list */
 	PopPlannerRestrictionContext();
 
@@ -1446,6 +1443,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;
 }