From 3bf4f3aee0f6fd9d0f7b29602113e97f35ef6a6f Mon Sep 17 00:00:00 2001
From: Bharath Rupireddy <bharath.rupireddy@enterprisedb.com>
Date: Wed, 9 Dec 2020 18:31:55 +0530
Subject: [PATCH v10] Tuple Cost Adjustment for Parallel Inserts in CTAS

---
 src/backend/commands/createas.c       |  9 ++++
 src/backend/commands/explain.c        | 10 +++++
 src/backend/optimizer/path/costsize.c | 19 ++++++++-
 src/backend/optimizer/plan/planner.c  | 61 +++++++++++++++++++++++++++
 src/include/commands/createas.h       | 16 +++++++
 src/include/nodes/parsenodes.h        |  1 +
 6 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/src/backend/commands/createas.c b/src/backend/commands/createas.c
index 20f59cc2b8..eee8d19259 100644
--- a/src/backend/commands/createas.c
+++ b/src/backend/commands/createas.c
@@ -316,10 +316,19 @@ ExecCreateTableAs(ParseState *pstate, CreateTableAsStmt *stmt,
 		query = linitial_node(Query, rewritten);
 		Assert(query->commandType == CMD_SELECT);
 
+		/*
+		 * Indication to the planner that the SELECT is from CTAS so that it
+		 * can adjust the parallel tuple cost if possible.
+		 */
+		if (IsParallelInsertInCTASAllowed(into, NULL))
+			query->CTASParallelInsInfo |= CTAS_PARALLEL_INS_SELECT;
+
 		/* plan the query */
 		plan = pg_plan_query(query, pstate->p_sourcetext,
 							 CURSOR_OPT_PARALLEL_OK, params);
 
+		query->CTASParallelInsInfo &= CTAS_PARALLEL_INS_UNDEF;
+
 		/*
 		 * Use a snapshot with an updated command ID to ensure this query sees
 		 * results of any previously executed queries.  (This could only
diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 03ac29cd64..8bd231a0c3 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -387,9 +387,19 @@ ExplainOneQuery(Query *query, int cursorOptions,
 			bufusage_start = pgBufferUsage;
 		INSTR_TIME_SET_CURRENT(planstart);
 
+		/*
+		 * Indication to the planner that the SELECT is from CTAS so that it
+		 * can adjust the parallel tuple cost if possible.
+		 */
+		if (IsParallelInsertInCTASAllowed(into, NULL))
+			query->CTASParallelInsInfo |= CTAS_PARALLEL_INS_SELECT;
+
 		/* plan the query */
 		plan = pg_plan_query(query, queryString, cursorOptions, params);
 
+		if (into)
+			query->CTASParallelInsInfo &= CTAS_PARALLEL_INS_UNDEF;
+
 		INSTR_TIME_SET_CURRENT(planduration);
 		INSTR_TIME_SUBTRACT(planduration, planstart);
 
diff --git a/src/backend/optimizer/path/costsize.c b/src/backend/optimizer/path/costsize.c
index 22d6935824..3a316f25f1 100644
--- a/src/backend/optimizer/path/costsize.c
+++ b/src/backend/optimizer/path/costsize.c
@@ -76,6 +76,7 @@
 #include "access/amapi.h"
 #include "access/htup_details.h"
 #include "access/tsmapi.h"
+#include "commands/createas.h"
 #include "executor/executor.h"
 #include "executor/nodeAgg.h"
 #include "executor/nodeHash.h"
@@ -378,6 +379,7 @@ cost_gather(GatherPath *path, PlannerInfo *root,
 {
 	Cost		startup_cost = 0;
 	Cost		run_cost = 0;
+	bool		ignore_tuple_cost = false;
 
 	/* Mark the path with the correct row estimate */
 	if (rows)
@@ -393,7 +395,22 @@ cost_gather(GatherPath *path, PlannerInfo *root,
 
 	/* Parallel setup and communication cost. */
 	startup_cost += parallel_setup_cost;
-	run_cost += parallel_tuple_cost * path->path.rows;
+
+	/*
+	 * Do not consider tuple cost in case of parallel inserts by workers. We
+	 * would have set ignore flag in apply_scanjoin_target_to_paths before
+	 * generating Gather path for the upper level SELECT part of the CTAS.
+	 */
+	if ((root->parse->CTASParallelInsInfo & CTAS_PARALLEL_INS_SELECT) &&
+		(root->parse->CTASParallelInsInfo & CTAS_PARALLEL_INS_IGN_TUP_COST))
+	{
+		ignore_tuple_cost = true;
+		/* Reset the ignore flag. */
+		root->parse->CTASParallelInsInfo &= ~CTAS_PARALLEL_INS_IGN_TUP_COST;
+	}
+
+	if (!ignore_tuple_cost)
+		run_cost += parallel_tuple_cost * path->path.rows;
 
 	path->path.startup_cost = startup_cost;
 	path->path.total_cost = (startup_cost + run_cost);
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 1a94b58f8b..1041593237 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -28,6 +28,7 @@
 #include "catalog/pg_inherits.h"
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
+#include "commands/createas.h"
 #include "executor/executor.h"
 #include "executor/nodeAgg.h"
 #include "foreign/fdwapi.h"
@@ -7338,6 +7339,45 @@ can_partial_agg(PlannerInfo *root)
 	return true;
 }
 
+/*
+ * ignore_parallel_tuple_cost
+ *
+ * Gather node will not receive any tuples from the workers in case each worker
+ * inserts them in parallel. So, we set a flag to ignore parallel tuple cost by
+ * the Gather path in cost_gather if the SELECT is for CTAS and we are
+ * generating an upper level Gather path.
+*/
+static bool
+ignore_parallel_tuple_cost(PlannerInfo *root)
+{
+	if (root->query_level == 1 &&
+		(root->parse->CTASParallelInsInfo & CTAS_PARALLEL_INS_SELECT))
+	{
+		/*
+		 * In each of following cases, a parent path will be generated for the
+		 * upper Gather path(in grouping_planner), in which case we can not
+		 * let parallel inserts happen. So we do not set ignore tuple cost
+		 * flag.
+		 */
+		if (root->parse->rowMarks ||
+			limit_needed(root->parse) ||
+			root->parse->sortClause ||
+			root->parse->distinctClause ||
+			root->parse->hasWindowFuncs ||
+			root->parse->groupClause ||
+			root->parse->groupingSets ||
+			root->parse->hasAggs ||
+			root->hasHavingQual)
+			return false;
+
+		root->parse->CTASParallelInsInfo |= CTAS_PARALLEL_INS_IGN_TUP_COST;
+
+		return true;
+	}
+
+	return false;
+}
+
 /*
  * apply_scanjoin_target_to_paths
  *
@@ -7557,8 +7597,29 @@ apply_scanjoin_target_to_paths(PlannerInfo *root,
 	 * one of the generated paths may turn out to be the cheapest one.
 	 */
 	if (rel->consider_parallel && !IS_OTHER_REL(rel))
+	{
+		/*
+		 * Set a flag to ignore parallel tuple cost by the Gather path in
+		 * cost_gather if the SELECT is for CTAS and we are generating an upper
+		 * level Gather path.
+		 */
+		bool	ignore = ignore_parallel_tuple_cost(root);
+
 		generate_useful_gather_paths(root, rel, false);
 
+		/*
+		 * Reset the ignore flag, in case we set it but
+		 * generate_useful_gather_paths returned without reaching cost_gather.
+		 */
+		if (ignore &&
+			(root->parse->CTASParallelInsInfo &
+			 CTAS_PARALLEL_INS_IGN_TUP_COST))
+		{
+			root->parse->CTASParallelInsInfo &=
+											~CTAS_PARALLEL_INS_IGN_TUP_COST;
+		}
+	}
+
 	/*
 	 * Reassess which paths are the cheapest, now that we've potentially added
 	 * new Gather (or Gather Merge) and/or Append (or MergeAppend) paths to
diff --git a/src/include/commands/createas.h b/src/include/commands/createas.h
index ab3aab58c5..6e722f0ac0 100644
--- a/src/include/commands/createas.h
+++ b/src/include/commands/createas.h
@@ -39,6 +39,22 @@ typedef struct
 	Oid			object_id;
 } DR_intorel;
 
+/*
+ * Information sent to the planner from CTAS to account for the cost
+ * calculations in cost_gather. We need to do this because, no tuples will be
+ * received by the Gather node if the workers insert the tuples in parallel.
+ */
+typedef enum CTASParallelInsertOpt
+{
+	CTAS_PARALLEL_INS_UNDEF = 0, /* undefined */
+	CTAS_PARALLEL_INS_SELECT = 1 << 0, /* set to this before planning */
+	/*
+	 * Set to this while planning for upper Gather path to ignore parallel
+	 * tuple cost in cost_gather.
+	 */
+	CTAS_PARALLEL_INS_IGN_TUP_COST = 1 << 1
+} CTASParallelInsertOpt;
+
 #define IS_CTAS(intoclause) (intoclause && IsA(intoclause, IntoClause))
 #define IS_PARALLEL_CTAS_DEST(dest) (dest && dest->mydest == DestIntoRel && \
 									 IS_CTAS(((DR_intorel *) dest)->into) && \
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index ec14fc2036..b140a42551 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -180,6 +180,7 @@ typedef struct Query
 	 */
 	int			stmt_location;	/* start location, or -1 if unknown */
 	int			stmt_len;		/* length in bytes; 0 means "rest of string" */
+	uint8 		CTASParallelInsInfo; /* parallel insert in CTAS info */
 } Query;
 
 
-- 
2.25.1

