recycle_mem_part_prune.patch

diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c
index 23a74bc3d9..d06b718c64 100644
--- a/src/backend/executor/execPartition.c
+++ b/src/backend/executor/execPartition.c
@@ -1438,18 +1438,6 @@ ExecSetupPartitionPruneState(PlanState *planstate, List *partitionpruneinfo)
 	prunestate->extparams = NULL;
 	prunestate->execparams = NULL;
 
-	/*
-	 * Create a sub memory context which we'll use when making calls to the
-	 * query planner's function to determine which partitions will match.  The
-	 * planner is not too careful about freeing memory, so we'll ensure we
-	 * call the function in this context to avoid any memory leaking in the
-	 * executor's memory context.
-	 */
-	prunestate->prune_context =
-		AllocSetContextCreate(CurrentMemoryContext,
-							  "Partition Prune",
-							  ALLOCSET_DEFAULT_SIZES);
-
 	i = 0;
 	foreach(lc, partitionpruneinfo)
 	{
@@ -1493,6 +1481,8 @@ ExecSetupPartitionPruneState(PlanState *planstate, List *partitionpruneinfo)
 		context->boundinfo = partition_bounds_copy(partdesc->boundinfo, partkey);
 		context->planstate = planstate;
 		context->safeparams = NULL; /* empty for now */
+		context->exprstates = palloc0(sizeof(ExprState *) *
+									  list_length(pinfo->pruning_steps));
 
 		pprune->pruning_steps = pinfo->pruning_steps;
 		pprune->extparams = bms_copy(pinfo->extparams);
@@ -1546,7 +1536,6 @@ Bitmapset *
 ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubnodes)
 {
 	PartitionPruningData *pprune;
-	MemoryContext oldcontext;
 	Bitmapset  *result = NULL;
 
 	/*
@@ -1557,22 +1546,9 @@ ExecFindInitialMatchingSubPlans(PartitionPruneState *prunestate, int nsubnodes)
 
 	pprune = prunestate->partprunedata;
 
-	/*
-	 * Switch to a temp context to avoid leaking memory in the executor's
-	 * memory context.
-	 */
-	oldcontext = MemoryContextSwitchTo(prunestate->prune_context);
-
 	/* Determine which subnodes match the external params */
 	find_subplans_for_params_recurse(prunestate, pprune, false, &result);
 
-	MemoryContextSwitchTo(oldcontext);
-
-	/* Move to the correct memory context */
-	result = bms_copy(result);
-
-	MemoryContextReset(prunestate->prune_context);
-
 	/*
 	 * Record that partition pruning has been performed for external params.
 	 * These are not required again afterwards, and nullifying them helps
@@ -1670,26 +1646,12 @@ Bitmapset *
 ExecFindMatchingSubPlans(PartitionPruneState *prunestate)
 {
 	PartitionPruningData *pprune;
-	MemoryContext oldcontext;
 	Bitmapset  *result = NULL;
 
 	pprune = prunestate->partprunedata;
 
-	/*
-	 * Switch to a temp context to avoid leaking memory in the executor's
-	 * memory context.
-	 */
-	oldcontext = MemoryContextSwitchTo(prunestate->prune_context);
-
 	find_subplans_for_params_recurse(prunestate, pprune, true, &result);
 
-	MemoryContextSwitchTo(oldcontext);
-
-	/* Move to the correct memory context */
-	result = bms_copy(result);
-
-	MemoryContextReset(prunestate->prune_context);
-
 	return result;
 }
 
diff --git a/src/backend/executor/nodeAppend.c b/src/backend/executor/nodeAppend.c
index d062cfddac..144b8651a3 100644
--- a/src/backend/executor/nodeAppend.c
+++ b/src/backend/executor/nodeAppend.c
@@ -321,6 +321,8 @@ ExecEndAppend(AppendState *node)
 	 */
 	for (i = 0; i < nplans; i++)
 		ExecEndNode(appendplans[i]);
+
+	MemoryContextStatsDetail(TopMemoryContext, INT_MAX);
 }
 
 void
diff --git a/src/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index 7666c6c412..d1a90b05fb 100644
--- a/src/backend/partitioning/partprune.c
+++ b/src/backend/partitioning/partprune.c
@@ -169,7 +169,7 @@ static PruneStepResult *perform_pruning_combine_step(PartitionPruneContext *cont
 static bool match_boolean_partition_clause(Oid partopfamily, Expr *clause,
 							   Expr *partkey, Expr **outconst);
 static bool partkey_datum_from_expr(PartitionPruneContext *context,
-						Expr *expr, Datum *value);
+						int step_id, Expr *expr, Datum *value);
 
 /*
  * make_partition_pruneinfo
@@ -444,6 +444,7 @@ prune_append_rel_partitions(RelOptInfo *rel)
 	/* Not valid when being called from the planner */
 	context.planstate = NULL;
 	context.safeparams = NULL;
+	context.exprstates = NULL;
 
 	/* Actual pruning happens here. */
 	partindexes = get_matching_partitions(&context, pruning_steps);
@@ -463,6 +464,10 @@ prune_append_rel_partitions(RelOptInfo *rel)
  *
  * Returns a Bitmapset of the RelOptInfo->part_rels indexes of the surviving
  * partitions.
+ *
+ * Caution: This is used inside the executor so we must be careful here to
+ * ensure we free all memory allocated within this function and any functions
+ * called from here.
  */
 Bitmapset *
 get_matching_partitions(PartitionPruneContext *context, List *pruning_steps)
@@ -552,6 +557,14 @@ get_matching_partitions(PartitionPruneContext *context, List *pruning_steps)
 		result = bms_add_member(result, context->boundinfo->default_index);
 	}
 
+	/* Ensure we free all allocations */
+	for (i = 0; i < num_steps; i++)
+	{
+		bms_free(results[i]->bound_offsets);
+		pfree(results[i]);
+	}
+	pfree(results);
+
 	return result;
 }
 
@@ -2785,7 +2798,8 @@ perform_pruning_base_step(PartitionPruneContext *context,
 			Datum		datum;
 
 			expr = lfirst(lc1);
-			if (partkey_datum_from_expr(context, expr, &datum))
+			if (partkey_datum_from_expr(context, opstep->step.step_id, expr,
+										&datum))
 			{
 				Oid			cmpfn;
 
@@ -3023,7 +3037,7 @@ match_boolean_partition_clause(Oid partopfamily, Expr *clause, Expr *partkey,
  *		evaluation was possible, otherwise false.
  */
 static bool
-partkey_datum_from_expr(PartitionPruneContext *context,
+partkey_datum_from_expr(PartitionPruneContext *context, int step_id,
 						Expr *expr, Datum *value)
 {
 	switch (nodeTag(expr))
@@ -3041,12 +3055,15 @@ partkey_datum_from_expr(PartitionPruneContext *context,
 			if (context->planstate &&
 				bms_is_member(((Param *) expr)->paramid, context->safeparams))
 			{
-				ExprState  *exprstate;
 				bool		isNull;
 
-				exprstate = ExecInitExpr(expr, context->planstate);
+				Assert(context->exprstates);
+
+				if (!context->exprstates[step_id])
+					context->exprstates[step_id] =
+									ExecInitExpr(expr, context->planstate);
 
-				*value = ExecEvalExprSwitchContext(exprstate,
+				*value = ExecEvalExprSwitchContext(context->exprstates[step_id],
 												   context->planstate->ps_ExprContext,
 												   &isNull);
 				if (isNull)
diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h
index e81bdc4a0a..4c7653e66f 100644
--- a/src/include/executor/execPartition.h
+++ b/src/include/executor/execPartition.h
@@ -159,8 +159,6 @@ typedef struct PartitionPruningData
  *						partitioned relation. First element contains the
  *						details for the target partitioned table.
  * num_partprunedata	Number of items in 'partprunedata' array.
- * prune_context		A memory context which can be used to call the query
- *						planner's partition prune functions.
  * extparams			All PARAM_EXTERN paramids which were found to match a
  *						partition key in each of the contained
  *						PartitionPruningData structs.
@@ -173,7 +171,6 @@ typedef struct PartitionPruneState
 {
 	PartitionPruningData *partprunedata;
 	int			num_partprunedata;
-	MemoryContext prune_context;
 	Bitmapset  *extparams;
 	Bitmapset  *execparams;
 	Bitmapset  *allparams;
diff --git a/src/include/partitioning/partprune.h b/src/include/partitioning/partprune.h
index a5568abce6..9f5fd198e7 100644
--- a/src/include/partitioning/partprune.h
+++ b/src/include/partitioning/partprune.h
@@ -50,6 +50,12 @@ typedef struct PartitionPruneContext
 	 * are not safe to use until the executor is running.
 	 */
 	Bitmapset  *safeparams;
+
+	/*
+	 * Array of ExprStates, one per pruning step.  Must be allocated if
+	 * planstate is non-NULL, otherwise can be NULL
+	 */
+	ExprState **exprstates;
 } PartitionPruneContext;