From a5b5a2db52a05894c41be217dd18309a76767ab2 Mon Sep 17 00:00:00 2001
From: David Rowley <dgrowley@gmail.com>
Date: Fri, 24 Dec 2021 10:55:22 +1300
Subject: [PATCH v2 3/3] Improve heuristics for memory allocation sizes in
 tuplesort

---
 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         | 71 ++++++++++++++--------
 src/include/utils/tuplesort.h              |  2 +-
 6 files changed, 55 insertions(+), 34 deletions(-)

diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 5eb3fbf7f0..012cb09a07 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, 0, 0);
+												  NULL, false, 0, 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, 0, 0);
+									 work_mem, NULL, false, 0, 0.0);
 	}
 
 	/*
diff --git a/src/backend/executor/nodeIncrementalSort.c b/src/backend/executor/nodeIncrementalSort.c
index c259259403..a24c720f24 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, 0, 0);
+												false, 0, 0.0);
 		node->prefixsort_state = prefixsort_state;
 	}
 	else
@@ -616,7 +616,7 @@ ExecIncrementalSort(PlanState *pstate)
 												  plannode->sort.nullsFirst,
 												  work_mem,
 												  NULL,
-												  false, 0, 0);
+												  false, 0, 0.0);
 			node->fullsort_state = fullsort_state;
 		}
 		else
diff --git a/src/backend/executor/nodeSort.c b/src/backend/executor/nodeSort.c
index 36404d9465..c3a192028b 100644
--- a/src/backend/executor/nodeSort.c
+++ b/src/backend/executor/nodeSort.c
@@ -105,7 +105,7 @@ ExecSort(PlanState *pstate)
 												   NULL,
 												   node->randomAccess,
 												   plannode->plan.plan_width,
-												   (int64) plannode->plan.plan_rows);
+												   plannode->plan.plan_rows);
 		else
 			tuplesortstate = tuplesort_begin_heap(tupDesc,
 												  plannode->numCols,
@@ -117,7 +117,7 @@ ExecSort(PlanState *pstate)
 												  NULL,
 												  node->randomAccess,
 												  plannode->plan.plan_width,
-												  (int64) plannode->plan.plan_rows);
+												  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 594593d51c..f723d27abe 100644
--- a/src/backend/utils/adt/orderedsetaggs.c
+++ b/src/backend/utils/adt/orderedsetaggs.c
@@ -297,7 +297,7 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 												   NULL,
 												   qstate->rescan_needed,
 												   0,
-												   0);
+												   0.0);
 	else
 		osastate->sortstate = tuplesort_begin_datum(qstate->sortColType,
 													qstate->sortOperator,
@@ -307,7 +307,7 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 													NULL,
 													qstate->rescan_needed,
 													0,
-													0);
+													0.0);
 
 	osastate->number_of_rows = 0;
 	osastate->sort_done = false;
diff --git a/src/backend/utils/sort/tuplesort.c b/src/backend/utils/sort/tuplesort.c
index 896f9f81ba..fa82b5cf09 100644
--- a/src/backend/utils/sort/tuplesort.c
+++ b/src/backend/utils/sort/tuplesort.c
@@ -316,6 +316,7 @@ struct Tuplesortstate
 	SortTuple  *memtuples;		/* array of SortTuple structs */
 	int			memtupcount;	/* number of tuples currently present */
 	int			memtupsize;		/* allocated length of memtuples array */
+	int			initial_memtupsize;	/* initial size for memtuples array */
 	bool		growmemtuples;	/* memtuples' growth still underway? */
 
 	/*
@@ -474,8 +475,8 @@ struct Tuplesortstate
 	int			datumTypeLen;
 
 	int32		est_tupwidth;	/* Estimated avg width of tuples to sort */
-	int64		est_tuples;		/* Estimated number of tuples to sort */
-	int64		est_tuple_memory;	/* Estimated amount of memory required
+	double		est_tuples;		/* Estimated number of tuples to sort */
+	double		est_tuple_memory;	/* Estimated amount of memory required
 									 * to store all tuples at once */
 	/*
 	 * Resource snapshot for time of sort start.
@@ -613,7 +614,7 @@ static Tuplesortstate *tuplesort_begin_common(int workMem,
 											  SortCoordinate coordinate,
 											  bool randomAccess,
 											  int32 est_tupwidth,
-											  int64 est_tuples);
+											  double est_tuples);
 static void tuplesort_begin_batch(Tuplesortstate *state);
 static void puttuple_common(Tuplesortstate *state, SortTuple *tuple);
 static bool consider_abort_common(Tuplesortstate *state);
@@ -724,7 +725,7 @@ static void tuplesort_updatemax(Tuplesortstate *state);
 
 static Tuplesortstate *
 tuplesort_begin_common(int workMem, SortCoordinate coordinate,
-					   bool randomAccess, int32 est_tupwidth, int64 est_tuples)
+					   bool randomAccess, int32 est_tupwidth, double est_tuples)
 {
 	Tuplesortstate *state;
 	MemoryContext maincontext;
@@ -787,37 +788,47 @@ tuplesort_begin_common(int workMem, SortCoordinate coordinate,
 	state->est_tuples = est_tuples;
 
 	/* Estimate the amount of memory required to perform the sort */
-	if (state->est_tuples > 0 && state->est_tupwidth > 0)
-	{
-		state->est_tuple_memory = pg_nextpower2_64(state->est_tuples * (state->est_tupwidth + 48));
-	}
+	if (est_tuples > 0.0 && est_tupwidth > 0)
+		state->est_tuple_memory = est_tuples * (est_tupwidth + 48);
 	else
-	{
-		state->est_tuple_memory = 0;
-	}
+		state->est_tuple_memory = 0.0;
 
 	/*
 	 * Initial size of array must be more than ALLOCSET_SEPARATE_THRESHOLD;
 	 * see comments in grow_memtuples().
 	 */
-	state->memtupsize = INITIAL_MEMTUPSIZE;
+	state->initial_memtupsize = INITIAL_MEMTUPSIZE;
 
 	/*
 	 * See if we can come up with a better estimate for the size of the
 	 * memtuples array.
 	 */
-	if (state->est_tuple_memory > 0 && state->est_tuples < INT_MAX)
+	if (state->est_tuple_memory > 0.0 && est_tuples < INT_MAX)
 	{
-		int64	pow2tuples = pg_nextpower2_64(state->est_tuples);
+		int64	pow2tuples = pg_nextpower2_64((int64) est_tuples);
 
 		/*
-		 * Size the array to the next power of 2 after the est_tuples if the
-		 * size of all tuples and the array are both within allowedMem.
+		 * Attempt to set the initial_memtupsize at a size that will fit all tuples.
+		 * In cases where we expect to exceed allowedMem, we'll choose an
+		 * initial_memtupsize based on the percentage of the tuples we expect to
+		 * be able to sort in a single batch.
 		 */
 		if (state->est_tuple_memory + sizeof(SortTuple) * pow2tuples < state->allowedMem)
-			state->memtupsize = Max(pow2tuples, INITIAL_MEMTUPSIZE);
+			state->initial_memtupsize = Max(pow2tuples, INITIAL_MEMTUPSIZE);
+		else
+		{
+			state->initial_memtupsize = pg_prevpower2_64(est_tuples * (state->allowedMem / state->est_tuple_memory));
+
+			/*
+			 * Cap at pow2tuples for when est_tuple_memory is just smaller
+			 * than allowedMem.
+			 */
+			state->initial_memtupsize = Min(state->initial_memtupsize, pow2tuples);
+		}
 	}
+
 	state->memtuples = NULL;
+	state->memtupsize = state->initial_memtupsize;
 
 	/*
 	 * After all of the other non-parallel-related state, we setup all of the
@@ -869,6 +880,7 @@ tuplesort_begin_batch(Tuplesortstate *state)
 {
 	MemoryContext oldcontext;
 	int64		blockSize;
+	int64		tuparraymem = sizeof(SortTuple) * state->initial_memtupsize;
 
 	oldcontext = MemoryContextSwitchTo(state->maincontext);
 
@@ -877,12 +889,21 @@ tuplesort_begin_batch(Tuplesortstate *state)
 	 * would be the next power of 2 after est_tuple_memory.  However, just
 	 * in case we get very small estimates, clamp it to be at least
 	 * ALLOCSET_DEFAULT_INITSIZE, but no higher than the power of 2 below
-	 * our maximum memory cap.  Also, just to prevent insane memory size
-	 * requests, cap to 1GB.
+	 * our maximum memory cap.  Also, to prevent insane memory size requests,
+	 * cap to MaxAllocSize.
 	 */
-	blockSize = Max(state->est_tuple_memory, ALLOCSET_DEFAULT_INITSIZE);
-	blockSize = Min(blockSize, pg_prevpower2_64(state->allowedMem));
-	blockSize = Min(blockSize, ((int64) 1) * 1024 * 1024 * 1024);
+	if (state->est_tuple_memory > 0.0)
+	{
+		blockSize = pg_nextpower2_64((int64) Min(state->est_tuple_memory, MaxAllocSize));
+		blockSize = Max(blockSize, ALLOCSET_DEFAULT_INITSIZE);
+	}
+	else
+		blockSize = ALLOCSET_DEFAULT_INITSIZE;
+
+	/* don't exceed allowedMem */
+	blockSize = Min(blockSize, pg_prevpower2_64(state->allowedMem - tuparraymem));
+	/* allowedMem might be greater than MaxAllocSize */
+	blockSize = Min(blockSize, MaxAllocSize);
 
 	/*
 	 * Caller tuple (e.g. IndexTuple) memory context.
@@ -913,11 +934,11 @@ tuplesort_begin_batch(Tuplesortstate *state)
 	 */
 	state->growmemtuples = true;
 	state->slabAllocatorUsed = false;
-	if (state->memtuples != NULL && state->memtupsize != INITIAL_MEMTUPSIZE)
+	if (state->memtuples != NULL && state->memtupsize != state->initial_memtupsize)
 	{
 		pfree(state->memtuples);
 		state->memtuples = NULL;
-		state->memtupsize = INITIAL_MEMTUPSIZE;
+		state->memtupsize = state->initial_memtupsize;
 	}
 	if (state->memtuples == NULL)
 	{
@@ -963,7 +984,7 @@ tuplesort_begin_heap(TupleDesc tupDesc,
 					 Oid *sortOperators, Oid *sortCollations,
 					 bool *nullsFirstFlags,
 					 int workMem, SortCoordinate coordinate,
-					 bool randomAccess, int32 est_tupwidth, int64 est_tuples)
+					 bool randomAccess, int32 est_tupwidth, double est_tuples)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, coordinate,
 												   randomAccess, est_tupwidth,
diff --git a/src/include/utils/tuplesort.h b/src/include/utils/tuplesort.h
index 7dc5be0e9b..f3b146d745 100644
--- a/src/include/utils/tuplesort.h
+++ b/src/include/utils/tuplesort.h
@@ -202,7 +202,7 @@ extern Tuplesortstate *tuplesort_begin_heap(TupleDesc tupDesc,
 											bool *nullsFirstFlags,
 											int workMem, SortCoordinate coordinate,
 											bool randomAccess, int32 est_tupwidth,
-											int64 est_tuples);
+											double est_tuples);
 extern Tuplesortstate *tuplesort_begin_cluster(TupleDesc tupDesc,
 											   Relation indexRel, int workMem,
 											   SortCoordinate coordinate, bool randomAccess);
-- 
2.21.0.windows.1

