Use partial sort optimizatiob only for bounded sort

Author: knizhnik

src/backend/executor/nodeAgg.c

@@ -520,7 +520,9 @@ initialize_phase(AggState *aggstate, int newphase)
 												  sortnode->collations,
 												  sortnode->nullsFirst,
 												  work_mem,
-												  false);
+												  false,
+												  sortnode->prefix
+			);
 	}
 
 	aggstate->current_phase = newphase;
@@ -597,7 +599,7 @@ initialize_aggregate(AggState *aggstate, AggStatePerTrans pertrans,
 									 pertrans->sortOperators,
 									 pertrans->sortCollations,
 									 pertrans->sortNullsFirst,
-									 work_mem, false);
+									 work_mem, false, 0);
 	}
 
 	/*

src/backend/executor/nodeSort.c

@@ -89,7 +89,9 @@ ExecSort(SortState *node)
 											  plannode->collations,
 											  plannode->nullsFirst,
 											  work_mem,
-											  node->randomAccess);
+											  node->randomAccess,
+											  plannode->prefix
+			);
 		if (node->bounded)
 			tuplesort_set_bound(tuplesortstate, node->bound);
 		node->tuplesortstate = (void *) tuplesortstate;
@@ -98,15 +100,14 @@ ExecSort(SortState *node)
 		 * Scan the subplan and feed all the tuples to tuplesort.
 		 */
 
-		for (;;)
+		do
 		{
 			slot = ExecProcNode(outerNode);
 
 			if (TupIsNull(slot))
 				break;
 
-			tuplesort_puttupleslot(tuplesortstate, slot);
-		}
+		} while (tuplesort_puttupleslot(tuplesortstate, slot));
 
 		/*
 		 * Complete the sort.

src/backend/nodes/copyfuncs.c

@@ -827,6 +827,7 @@ _copySort(const Sort *from)
 	CopyPlanFields((const Plan *) from, (Plan *) newnode);
 
 	COPY_SCALAR_FIELD(numCols);
+	COPY_SCALAR_FIELD(prefix);
 	COPY_POINTER_FIELD(sortColIdx, from->numCols * sizeof(AttrNumber));
 	COPY_POINTER_FIELD(sortOperators, from->numCols * sizeof(Oid));
 	COPY_POINTER_FIELD(collations, from->numCols * sizeof(Oid));

src/backend/nodes/outfuncs.c

@@ -778,6 +778,7 @@ _outSort(StringInfo str, const Sort *node)
 	_outPlanInfo(str, (const Plan *) node);
 
 	WRITE_INT_FIELD(numCols);
+	WRITE_INT_FIELD(prefix);
 
 	appendStringInfoString(str, " :sortColIdx");
 	for (i = 0; i < node->numCols; i++)

src/backend/nodes/readfuncs.c

@@ -1951,6 +1951,7 @@ _readSort(void)
 	ReadCommonPlan(&local_node->plan);
 
 	READ_INT_FIELD(numCols);
+	READ_INT_FIELD(prefix);
 	READ_ATTRNUMBER_ARRAY(sortColIdx, local_node->numCols);
 	READ_OID_ARRAY(sortOperators, local_node->numCols);
 	READ_OID_ARRAY(collations, local_node->numCols);

src/backend/optimizer/plan/planner.c

@@ -1811,7 +1811,7 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 				/* Presorted path is a loser */
 				sorted_path = NULL;
 			}
-		} else if (partial_sorted_path != NULL) {
+		} else if (partial_sorted_path != NULL && limit_tuples > 0.0 && limit_tuples < path_rows) {
 			Path	sort_path;
 			Path	partial_sort_path;
 
@@ -1826,7 +1826,7 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 					  0.0, work_mem, root->limit_tuples);
 
 			partial_sort_path.total_cost -= partial_sort_path.startup_cost;
-			partial_sort_path.startup_cost /= partial_sorted_path->pathkeys->length+1;
+			partial_sort_path.startup_cost /= path_rows/limit_tuples;
 			partial_sort_path.total_cost += partial_sort_path.startup_cost;
 
 
@@ -2399,15 +2399,18 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 	{
 		if (!pathkeys_contained_in(root->sort_pathkeys, current_pathkeys))
 		{
+			double numRows = result_plan->plan_rows;
 			result_plan = (Plan *) make_sort_from_pathkeys(root,
 														   result_plan,
-														 root->sort_pathkeys,
+														   root->sort_pathkeys,
 														   limit_tuples);
-			if (partial_sorted_path && best_path == partial_sorted_path)
-			{
+			/* check if input data was partially sorted and there is limit clause */
+			if (partial_sorted_path && best_path == partial_sorted_path && limit_tuples > 0.0 && limit_tuples < numRows)
+			{								
 				result_plan->total_cost -= result_plan->startup_cost;
-				result_plan->startup_cost /= partial_sorted_path->pathkeys->length+1;
+				result_plan->startup_cost /= numRows/limit_tuples;
 				result_plan->total_cost += result_plan->startup_cost;				
+				((Sort*)result_plan)->prefix = partial_sorted_path->pathkeys->length;
 			}
 			current_pathkeys = root->sort_pathkeys;
 		}

src/backend/utils/adt/orderedsetaggs.c

@@ -276,7 +276,7 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 												   qstate->sortOperators,
 												   qstate->sortCollations,
 												   qstate->sortNullsFirsts,
-												   work_mem, false);
+												   work_mem, false, 0);
 	else
 		osastate->sortstate = tuplesort_begin_datum(qstate->sortColType,
 													qstate->sortOperator,

src/backend/utils/sort/tuplesort.c

@@ -219,6 +219,8 @@ struct Tuplesortstate
 								 * tuples to return? */
 	bool		boundUsed;		/* true if we made use of a bounded heap */
 	int			bound;			/* if bounded, the maximum number of tuples */
+	int         prefix;         /* number of prefix keys by which input data is already sorted */
+	int         matchedKeys;    /* number of matched keys */
 	int64		availMem;		/* remaining memory available, in bytes */
 	int64		allowedMem;		/* total memory allowed, in bytes */
 	int			maxTapes;		/* number of tapes (Knuth's T) */
@@ -458,7 +460,7 @@ struct Tuplesortstate
 
 
 static Tuplesortstate *tuplesort_begin_common(int workMem, bool randomAccess);
-static void puttuple_common(Tuplesortstate *state, SortTuple *tuple);
+static bool puttuple_common(Tuplesortstate *state, SortTuple *tuple);
 static bool consider_abort_common(Tuplesortstate *state);
 static void inittapes(Tuplesortstate *state);
 static void selectnewtape(Tuplesortstate *state);
@@ -575,7 +577,7 @@ tuplesort_begin_common(int workMem, bool randomAccess)
 	state->availMem = state->allowedMem;
 	state->sortcontext = sortcontext;
 	state->tapeset = NULL;
-
+	state->prefix = 0;
 	state->memtupcount = 0;
 
 	/*
@@ -613,7 +615,7 @@ tuplesort_begin_heap(TupleDesc tupDesc,
 					 int nkeys, AttrNumber *attNums,
 					 Oid *sortOperators, Oid *sortCollations,
 					 bool *nullsFirstFlags,
-					 int workMem, bool randomAccess)
+					 int workMem, bool randomAccess, int prefix)
 {
 	Tuplesortstate *state = tuplesort_begin_common(workMem, randomAccess);
 	MemoryContext oldcontext;
@@ -648,6 +650,7 @@ tuplesort_begin_heap(TupleDesc tupDesc,
 
 	/* Prepare SortSupport data for each column */
 	state->sortKeys = (SortSupport) palloc0(nkeys * sizeof(SortSupportData));
+	state->prefix = prefix;
 
 	for (i = 0; i < nkeys; i++)
 	{
@@ -1202,21 +1205,22 @@ grow_memtuples(Tuplesortstate *state)
  *
  * Note that the input data is always copied; the caller need not save it.
  */
-void
+bool
 tuplesort_puttupleslot(Tuplesortstate *state, TupleTableSlot *slot)
 {
 	MemoryContext oldcontext = MemoryContextSwitchTo(state->sortcontext);
 	SortTuple	stup;
-
+	bool needmore;
 	/*
 	 * Copy the given tuple into memory we control, and decrease availMem.
 	 * Then call the common code.
 	 */
 	COPYTUP(state, &stup, (void *) slot);
 
-	puttuple_common(state, &stup);
+	needmore = puttuple_common(state, &stup);
 
 	MemoryContextSwitchTo(oldcontext);
+	return needmore;
 }
 
 /*
@@ -1397,7 +1401,7 @@ tuplesort_putdatum(Tuplesortstate *state, Datum val, bool isNull)
 /*
  * Shared code for tuple and datum cases.
  */
-static void
+static bool
 puttuple_common(Tuplesortstate *state, SortTuple *tuple)
 {
 	switch (state->status)
@@ -1441,14 +1445,14 @@ puttuple_common(Tuplesortstate *state, SortTuple *tuple)
 						 pg_rusage_show(&state->ru_start));
 #endif
 				make_bounded_heap(state);
-				return;
+				break;
 			}
 
 			/*
 			 * Done if we still fit in available memory and have array slots.
 			 */
 			if (state->memtupcount < state->memtupsize && !LACKMEM(state))
-				return;
+				break;
 
 			/*
 			 * Nope; time to switch to tape-based operation.
@@ -1475,6 +1479,9 @@ puttuple_common(Tuplesortstate *state, SortTuple *tuple)
 			{
 				/* new tuple <= top of the heap, so we can discard it */
 				free_sort_tuple(state, tuple);
+				if (state->prefix > state->matchedKeys && state->memtupcount >= state->bound) { 
+					return false;
+				}
 				CHECK_FOR_INTERRUPTS();
 			}
 			else
@@ -1517,6 +1524,7 @@ puttuple_common(Tuplesortstate *state, SortTuple *tuple)
 			elog(ERROR, "invalid tuplesort state");
 			break;
 	}
+	return true;
 }
 
 static bool
@@ -3071,19 +3079,20 @@ comparetup_heap(const SortTuple *a, const SortTuple *b, Tuplesortstate *state)
 	HeapTupleData rtup;
 	TupleDesc	tupDesc;
 	int			nkey;
-	int32		compare;
+	int32		compare = 0;
 	AttrNumber	attno;
 	Datum		datum1,
 				datum2;
 	bool		isnull1,
 				isnull2;
 
+	state->matchedKeys = 0;
 
 	/* Compare the leading sort key */
 	compare = ApplySortComparator(a->datum1, a->isnull1,
 								  b->datum1, b->isnull1,
 								  sortKey);
-	if (compare != 0)
+	if (compare != 0) 
 		return compare;
 
 	/* Compare additional sort keys */
@@ -3119,10 +3128,11 @@ comparetup_heap(const SortTuple *a, const SortTuple *b, Tuplesortstate *state)
 									  datum2, isnull2,
 									  sortKey);
 		if (compare != 0)
-			return compare;
+			break;
 	}
+	state->matchedKeys = nkey;
 
-	return 0;
+	return compare;
 }
 
 static void

src/include/nodes/plannodes.h

@@ -679,6 +679,7 @@ typedef struct Sort
 {
 	Plan		plan;
 	int			numCols;		/* number of sort-key columns */
+	int         prefix;         /* sorted prefix: number of keys by which input stream is already sorted */
 	AttrNumber *sortColIdx;		/* their indexes in the target list */
 	Oid		   *sortOperators;	/* OIDs of operators to sort them by */
 	Oid		   *collations;		/* OIDs of collations */

src/include/utils/tuplesort.h

@@ -62,7 +62,7 @@ extern Tuplesortstate *tuplesort_begin_heap(TupleDesc tupDesc,
 					 int nkeys, AttrNumber *attNums,
 					 Oid *sortOperators, Oid *sortCollations,
 					 bool *nullsFirstFlags,
-					 int workMem, bool randomAccess);
+ 					 int workMem, bool randomAccess, int prefix);
 extern Tuplesortstate *tuplesort_begin_cluster(TupleDesc tupDesc,
 						Relation indexRel,
 						int workMem, bool randomAccess);
@@ -81,7 +81,7 @@ extern Tuplesortstate *tuplesort_begin_datum(Oid datumType,
 
 extern void tuplesort_set_bound(Tuplesortstate *state, int64 bound);
 
-extern void tuplesort_puttupleslot(Tuplesortstate *state,
+extern bool tuplesort_puttupleslot(Tuplesortstate *state,
 					   TupleTableSlot *slot);
 extern void tuplesort_putheaptuple(Tuplesortstate *state, HeapTuple tup);
 extern void tuplesort_putindextuplevalues(Tuplesortstate *state,