v2-0001-Use-a-generation-context-for-tuple-storage-to-spe.patch

From 6fc22095e3567118f906c90fc39522b47e04b689 Mon Sep 17 00:00:00 2001
From: David Rowley <dgrowley@gmail.com>
Date: Sun, 1 Aug 2021 18:35:57 +1200
Subject: [PATCH v2 1/3] Use a generation context for tuple storage to speedup
 sorting

---
 src/backend/executor/nodeAgg.c             |  4 +-
 src/backend/executor/nodeIncrementalSort.c |  4 +-
 src/backend/executor/nodeSort.c            |  4 +-
 src/backend/utils/adt/orderedsetaggs.c     |  4 +-
 src/backend/utils/sort/tuplesort.c         | 48 ++++++++++++++++------
 src/include/utils/tuplesort.h              |  3 +-
 6 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index f5a187cae3..7da7a9ecb4 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -530,7 +530,7 @@ initialize_phase(AggState *aggstate, int newphase)
 												  sortnode->collations,
 												  sortnode->nullsFirst,
 												  work_mem,
-												  NULL, false);
+												  NULL, false, 0, 0);
 	}
 
 	aggstate->current_phase = newphase;
@@ -617,7 +617,7 @@ initialize_aggregate(AggState *aggstate, AggStatePerTrans pertrans,
 									 pertrans->sortOperators,
 									 pertrans->sortCollations,
 									 pertrans->sortNullsFirst,
-									 work_mem, NULL, false);
+									 work_mem, NULL, false, 0, 0);
 	}
 
 	/*
diff --git a/src/backend/executor/nodeIncrementalSort.c b/src/backend/executor/nodeIncrementalSort.c
index 934426a667..c259259403 100644
--- a/src/backend/executor/nodeIncrementalSort.c
+++ b/src/backend/executor/nodeIncrementalSort.c
@@ -315,7 +315,7 @@ switchToPresortedPrefixMode(PlanState *pstate)
 												&(plannode->sort.nullsFirst[nPresortedCols]),
 												work_mem,
 												NULL,
-												false);
+												false, 0, 0);
 		node->prefixsort_state = prefixsort_state;
 	}
 	else
@@ -616,7 +616,7 @@ ExecIncrementalSort(PlanState *pstate)
 												  plannode->sort.nullsFirst,
 												  work_mem,
 												  NULL,
-												  false);
+												  false, 0, 0);
 			node->fullsort_state = fullsort_state;
 		}
 		else
diff --git a/src/backend/executor/nodeSort.c b/src/backend/executor/nodeSort.c
index c68d9e41c2..cdba5e4814 100644
--- a/src/backend/executor/nodeSort.c
+++ b/src/backend/executor/nodeSort.c
@@ -113,7 +113,9 @@ ExecSort(PlanState *pstate)
 												  plannode->nullsFirst,
 												  work_mem,
 												  NULL,
-												  node->randomAccess);
+												  node->randomAccess,
+												  plannode->plan.plan_width,
+												  (int64) plannode->plan.plan_rows);
 		if (node->bounded)
 			tuplesort_set_bound(tuplesortstate, node->bound);
 		node->tuplesortstate = (void *) tuplesortstate;
diff --git a/src/backend/utils/adt/orderedsetaggs.c b/src/backend/utils/adt/orderedsetaggs.c
index 89f1c3b3a0..6033721590 100644
--- a/src/backend/utils/adt/orderedsetaggs.c
+++ b/src/backend/utils/adt/orderedsetaggs.c
@@ -295,7 +295,9 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 												   qstate->sortNullsFirsts,
 												   work_mem,
 												   NULL,
-												   qstate->rescan_needed);
+												   qstate->rescan_needed,
+												   0,
+												   0);
 	else
 		osastate->sortstate = tuplesort_begin_datum(qstate->sortColType,
 													qstate->sortOperator,
diff --git a/src/backend/utils/sort/tuplesort.c b/src/backend/utils/sort/tuplesort.c
index 90e26745df..9156b63d9e 100644
--- a/src/backend/utils/sort/tuplesort.c
+++ b/src/backend/utils/sort/tuplesort.c
@@ -473,6 +473,9 @@ struct Tuplesortstate
 	/* we need typelen in order to know how to copy the Datums. */
 	int			datumTypeLen;
 
+	int32		est_tupwidth;	/* Estimated avg width of tuples to sort */
+	int64		est_tuples;		/* Estimated number of tuples to sort */
+
 	/*
 	 * Resource snapshot for time of sort start.
 	 */
@@ -607,7 +610,9 @@ struct Sharedsort
 
 static Tuplesortstate *tuplesort_begin_common(int workMem,
 											  SortCoordinate coordinate,
-											  bool randomAccess);
+											  bool randomAccess,
+											  int32 est_tupwidth,
+											  int64 est_tuples);
 static void tuplesort_begin_batch(Tuplesortstate *state);
 static void puttuple_common(Tuplesortstate *state, SortTuple *tuple);
 static bool consider_abort_common(Tuplesortstate *state);
@@ -718,7 +723,7 @@ static void tuplesort_updatemax(Tuplesortstate *state);
 
 static Tuplesortstate *
 tuplesort_begin_common(int workMem, SortCoordinate coordinate,
-					   bool randomAccess)
+					   bool randomAccess, int32 est_tupwidth, int64 est_tuples)
 {
 	Tuplesortstate *state;
 	MemoryContext maincontext;
@@ -783,6 +788,10 @@ tuplesort_begin_common(int workMem, SortCoordinate coordinate,
 	state->memtupsize = INITIAL_MEMTUPSIZE;
 	state->memtuples = NULL;
 
+	/* Fill in any size estimates we've received */
+	state->est_tupwidth = est_tupwidth;
+	state->est_tuples = est_tuples;
+
 	/*
 	 * After all of the other non-parallel-related state, we setup all of the
 	 * state needed for each batch.
@@ -832,8 +841,21 @@ static void
 tuplesort_begin_batch(Tuplesortstate *state)
 {
 	MemoryContext oldcontext;
+	int64		mem_est;
 
 	oldcontext = MemoryContextSwitchTo(state->maincontext);
+	
+	/* Estimate the amount of memory required to perform the sort */
+	if (state->est_tuples > 0 && state->est_tupwidth > 0)
+	{
+		mem_est = pg_nextpower2_64(state->est_tuples * (state->est_tupwidth + 48));
+		mem_est = Max(mem_est, ALLOCSET_DEFAULT_INITSIZE);
+		mem_est = Min(mem_est, pg_prevpower2_64(state->allowedMem));
+	}
+	else
+	{
+		mem_est = ALLOCSET_DEFAULT_MAXSIZE;
+	}
 
 	/*
 	 * Caller tuple (e.g. IndexTuple) memory context.
@@ -844,9 +866,9 @@ tuplesort_begin_batch(Tuplesortstate *state)
 	 * in the parent context, not this context, because there is no need to
 	 * free memtuples early.
 	 */
-	state->tuplecontext = AllocSetContextCreate(state->sortcontext,
-												"Caller tuples",
-												ALLOCSET_DEFAULT_SIZES);
+	state->tuplecontext = GenerationContextCreate(state->sortcontext,
+												  "Caller tuples",
+												  (Size) mem_est);
 
 	state->status = TSS_INITIAL;
 	state->bounded = false;
@@ -897,10 +919,12 @@ tuplesort_begin_heap(TupleDesc tupDesc,
 					 int nkeys, AttrNumber *attNums,
 					 Oid *sortOperators, Oid *sortCollations,
 					 bool *nullsFirstFlags,
-					 int workMem, SortCoordinate coordinate, bool randomAccess)
+					 int workMem, SortCoordinate coordinate,
+					 bool randomAccess, int32 est_tupwidth, int64 est_tuples)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
-												   randomAccess);
+												   randomAccess, est_tupwidth,
+												   est_tuples);
 	MemoryContext oldcontext;
 	int			i;
 
@@ -973,7 +997,7 @@ tuplesort_begin_cluster(TupleDesc tupDesc,
 						SortCoordinate coordinate, bool randomAccess)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
-												   randomAccess);
+												   randomAccess, 0, 0);
 	BTScanInsert indexScanKey;
 	MemoryContext oldcontext;
 	int			i;
@@ -1070,7 +1094,7 @@ tuplesort_begin_index_btree(Relation heapRel,
 							bool randomAccess)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
-												   randomAccess);
+												   randomAccess, 0, 0);
 	BTScanInsert indexScanKey;
 	MemoryContext oldcontext;
 	int			i;
@@ -1150,7 +1174,7 @@ tuplesort_begin_index_hash(Relation heapRel,
 						   bool randomAccess)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
-												   randomAccess);
+												   randomAccess, 0, 0);
 	MemoryContext oldcontext;
 
 	oldcontext = MemoryContextSwitchTo(state->maincontext);
@@ -1193,7 +1217,7 @@ tuplesort_begin_index_gist(Relation heapRel,
 						   bool randomAccess)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
-												   randomAccess);
+												   randomAccess, 0, 0);
 	MemoryContext oldcontext;
 	int			i;
 
@@ -1248,7 +1272,7 @@ tuplesort_begin_datum(Oid datumType, Oid sortOperator, Oid sortCollation,
 					  SortCoordinate coordinate, bool randomAccess)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
-												   randomAccess);
+												   randomAccess, 0, 0);
 	MemoryContext oldcontext;
 	int16		typlen;
 	bool		typbyval;
diff --git a/src/include/utils/tuplesort.h b/src/include/utils/tuplesort.h
index f94949370b..3a722796e7 100644
--- a/src/include/utils/tuplesort.h
+++ b/src/include/utils/tuplesort.h
@@ -201,7 +201,8 @@ extern Tuplesortstate *tuplesort_begin_heap(TupleDesc tupDesc,
 											Oid *sortOperators, Oid *sortCollations,
 											bool *nullsFirstFlags,
 											int workMem, SortCoordinate coordinate,
-											bool randomAccess);
+											bool randomAccess, int32 est_tupwidth,
+											int64 est_tuples);
 extern Tuplesortstate *tuplesort_begin_cluster(TupleDesc tupDesc,
 											   Relation indexRel, int workMem,
 											   SortCoordinate coordinate, bool randomAccess);
-- 
2.21.0.windows.1