From ca293116fa6521e383147abd07d035fdc49f110e Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onder@citusdata.com>
Date: Thu, 19 Dec 2019 10:59:42 +0100
Subject: [PATCH] Reduce calls to FastPathRouterQuery()

Before this commit, we called it twice durning planning. Instead,
we save the information and pass it.
---
 .../distributed/planner/distributed_planner.c   | 17 ++++++++++++++++-
 .../distributed/planner/insert_select_planner.c |  3 +++
 .../distributed/planner/multi_router_planner.c  |  4 +++-
 src/include/distributed/distributed_planner.h   | 13 +++++++++++++
 4 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c
index 07641f272..2735d3562 100644
--- a/src/backend/distributed/planner/distributed_planner.c
+++ b/src/backend/distributed/planner/distributed_planner.c
@@ -128,6 +128,7 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	bool setPartitionedTablesInherited = false;
 	List *rangeTableList = ExtractRangeTableEntryList(parse);
 	int rteIdCounter = 1;
+	bool fastPathRouterQuery = false;
 
 	if (cursorOptions & CURSOR_OPT_FORCE_DISTRIBUTED)
 	{
@@ -151,6 +152,10 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 		else
 		{
 			needsDistributedPlanning = ListContainsDistributedTableRTE(rangeTableList);
+			if (needsDistributedPlanning)
+			{
+				fastPathRouterQuery = FastPathRouterQuery(parse);
+			}
 		}
 	}
 
@@ -215,8 +220,11 @@ distributed_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 		 * transformations made by postgres' planner.
 		 */
 
-		if (needsDistributedPlanning && FastPathRouterQuery(originalQuery))
+		if (needsDistributedPlanning && fastPathRouterQuery)
 		{
+			plannerRestrictionContext->fastPathRestrictionContext->fastPathRouterQuery =
+				true;
+
 			result = FastPathPlanner(originalQuery, parse, boundParams);
 		}
 		else
@@ -1870,6 +1878,9 @@ CreateAndPushPlannerRestrictionContext(void)
 	plannerRestrictionContext->joinRestrictionContext =
 		palloc0(sizeof(JoinRestrictionContext));
 
+	plannerRestrictionContext->fastPathRestrictionContext =
+		palloc0(sizeof(FastPathRestrictionContext));
+
 	plannerRestrictionContext->memoryContext = CurrentMemoryContext;
 
 	/* we'll apply logical AND as we add tables */
@@ -1929,6 +1940,10 @@ ResetPlannerRestrictionContext(PlannerRestrictionContext *plannerRestrictionCont
 	plannerRestrictionContext->joinRestrictionContext =
 		palloc0(sizeof(JoinRestrictionContext));
 
+	plannerRestrictionContext->fastPathRestrictionContext =
+		palloc0(sizeof(FastPathRestrictionContext));
+
+
 	/* we'll apply logical AND as we add tables */
 	plannerRestrictionContext->relationRestrictionContext->allReferenceTables = true;
 }
diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c
index b216a42d6..03ebd0e59 100644
--- a/src/backend/distributed/planner/insert_select_planner.c
+++ b/src/backend/distributed/planner/insert_select_planner.c
@@ -441,6 +441,9 @@ RouterModifyTaskForShardInterval(Query *originalQuery, ShardInterval *shardInter
 			plannerRestrictionContext->relationRestrictionContext);
 	copyOfPlannerRestrictionContext->joinRestrictionContext =
 		plannerRestrictionContext->joinRestrictionContext;
+	copyOfPlannerRestrictionContext->fastPathRestrictionContext =
+		plannerRestrictionContext->fastPathRestrictionContext;
+
 	relationRestrictionList =
 		copyOfPlannerRestrictionContext->relationRestrictionContext->
 		relationRestrictionList;
diff --git a/src/backend/distributed/planner/multi_router_planner.c b/src/backend/distributed/planner/multi_router_planner.c
index 1a2694a53..fe534c053 100644
--- a/src/backend/distributed/planner/multi_router_planner.c
+++ b/src/backend/distributed/planner/multi_router_planner.c
@@ -2014,6 +2014,8 @@ PlanRouterQuery(Query *originalQuery,
 	bool shardsPresent = false;
 	uint64 shardId = INVALID_SHARD_ID;
 	CmdType commandType = originalQuery->commandType;
+	bool fastPathRouterQuery =
+		plannerRestrictionContext->fastPathRestrictionContext->fastPathRouterQuery;
 
 	*placementList = NIL;
 
@@ -2022,7 +2024,7 @@ PlanRouterQuery(Query *originalQuery,
 	 * not been called. Thus, restriction information is not avaliable and we do the
 	 * shard pruning based on the distribution column in the quals of the query.
 	 */
-	if (FastPathRouterQuery(originalQuery))
+	if (fastPathRouterQuery)
 	{
 		List *shardIntervalList =
 			TargetShardIntervalForFastPathQuery(originalQuery, partitionValueConst,
diff --git a/src/include/distributed/distributed_planner.h b/src/include/distributed/distributed_planner.h
index d620e47e5..880781ec0 100644
--- a/src/include/distributed/distributed_planner.h
+++ b/src/include/distributed/distributed_planner.h
@@ -85,10 +85,23 @@ typedef struct JoinRestriction
 	RelOptInfo *outerrel;
 } JoinRestriction;
 
+typedef struct FastPathRestrictionContext
+{
+	bool fastPathRouterQuery;
+}FastPathRestrictionContext;
+
 typedef struct PlannerRestrictionContext
 {
 	RelationRestrictionContext *relationRestrictionContext;
 	JoinRestrictionContext *joinRestrictionContext;
+
+	/*
+	 * When the query is qualified for fast path, we don't have
+	 * the RelationRestrictionContext and JoinRestrictionContext
+	 * since those are dependent to calling standard_planner.
+	 * Instead, we keep this struct to pass some extra information.
+	 */
+	FastPathRestrictionContext *fastPathRestrictionContext;
 	bool hasSemiJoin;
 	MemoryContext memoryContext;
 } PlannerRestrictionContext;