Further fallout from the MergeAppend patch.

Fix things so that top-N sorting can be used in child Sort nodes of a
MergeAppend node, when there is a LIMIT and no intervening joins or
grouping.  Actually doing this on the executor side isn't too bad,
but it's a bit messier to get the planner to cost it properly.
Per gripe from Robert Haas.

In passing, fix an oversight in the original top-N-sorting patch:
query_planner should not assume that a LIMIT can be used to make an
explicit sort cheaper when there will be grouping or aggregation in
between.  Possibly this should be back-patched, but I'm not sure the
mistake is serious enough to be a real problem in practice.

4 files changed, 59 insertions, 4 deletions

diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 7a84bd9123..41ad512a29 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -714,7 +714,7 @@ create_merge_append_plan(PlannerInfo *root, MergeAppendPath *best_path)
 		if (!pathkeys_contained_in(pathkeys, subpath->pathkeys))
 			subplan = (Plan *) make_sort(root, subplan, numsortkeys,
 										 sortColIdx, sortOperators, nullsFirst,
-										 -1.0);
+										 best_path->limit_tuples);
 
 		subplans = lappend(subplans, subplan);
 	}
diff --git a/src/backend/optimizer/plan/planmain.c b/src/backend/optimizer/plan/planmain.c
index cab6e9e25a..9e6b0b724c 100644
--- a/src/backend/optimizer/plan/planmain.c
+++ b/src/backend/optimizer/plan/planmain.c
@@ -101,8 +101,9 @@ query_planner(PlannerInfo *root, List *tlist,
 	ListCell   *lc;
 	double		total_pages;
 
-	/* Make tuple_fraction accessible to lower-level routines */
+	/* Make tuple_fraction, limit_tuples accessible to lower-level routines */
 	root->tuple_fraction = tuple_fraction;
+	root->limit_tuples = limit_tuples;
 
 	*num_groups = 1;			/* default result */
 
@@ -315,6 +316,9 @@ query_planner(PlannerInfo *root, List *tlist,
 			!pathkeys_contained_in(root->distinct_pathkeys, root->group_pathkeys) ||
 		 !pathkeys_contained_in(root->window_pathkeys, root->group_pathkeys))
 			tuple_fraction = 0.0;
+
+		/* In any case, limit_tuples shouldn't be specified here */
+		Assert(limit_tuples < 0);
 	}
 	else if (parse->hasAggs || root->hasHavingQual)
 	{
@@ -323,6 +327,9 @@ query_planner(PlannerInfo *root, List *tlist,
 		 * it will deliver a single result row (so leave *num_groups 1).
 		 */
 		tuple_fraction = 0.0;
+
+		/* limit_tuples shouldn't be specified here */
+		Assert(limit_tuples < 0);
 	}
 	else if (parse->distinctClause)
 	{
@@ -347,6 +354,9 @@ query_planner(PlannerInfo *root, List *tlist,
 		 */
 		if (tuple_fraction >= 1.0)
 			tuple_fraction /= *num_groups;
+
+		/* limit_tuples shouldn't be specified here */
+		Assert(limit_tuples < 0);
 	}
 	else
 	{
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 620888cbb8..6324bce240 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -968,6 +968,7 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 	{
 		/* No set operations, do regular planning */
 		List	   *sub_tlist;
+		double		sub_limit_tuples;
 		AttrNumber *groupColIdx = NULL;
 		bool		need_tlist_eval = true;
 		QualCost	tlist_cost;
@@ -1120,12 +1121,27 @@ grouping_planner(PlannerInfo *root, double tuple_fraction)
 			root->query_pathkeys = NIL;
 
 		/*
+		 * Figure out whether there's a hard limit on the number of rows that
+		 * query_planner's result subplan needs to return.  Even if we know a
+		 * hard limit overall, it doesn't apply if the query has any
+		 * grouping/aggregation operations.
+		 */
+		if (parse->groupClause ||
+			parse->distinctClause ||
+			parse->hasAggs ||
+			parse->hasWindowFuncs ||
+			root->hasHavingQual)
+			sub_limit_tuples = -1.0;
+		else
+			sub_limit_tuples = limit_tuples;
+
+		/*
 		 * Generate the best unsorted and presorted paths for this Query (but
 		 * note there may not be any presorted path).  query_planner will also
 		 * estimate the number of groups in the query, and canonicalize all
 		 * the pathkeys.
 		 */
-		query_planner(root, sub_tlist, tuple_fraction, limit_tuples,
+		query_planner(root, sub_tlist, tuple_fraction, sub_limit_tuples,
 					  &cheapest_path, &sorted_path, &dNumGroups);
 
 		/*
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 20d8074bce..231d221b21 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -694,6 +694,35 @@ create_merge_append_path(PlannerInfo *root,
 	pathnode->path.pathkeys = pathkeys;
 	pathnode->subpaths = subpaths;
 
+	/*
+	 * Apply query-wide LIMIT if known and path is for sole base relation.
+	 * Finding out the latter at this low level is a bit klugy.
+	 */
+	pathnode->limit_tuples = root->limit_tuples;
+	if (pathnode->limit_tuples >= 0)
+	{
+		Index		rti;
+
+		for (rti = 1; rti < root->simple_rel_array_size; rti++)
+		{
+			RelOptInfo *brel = root->simple_rel_array[rti];
+
+			if (brel == NULL)
+				continue;
+
+			/* ignore RTEs that are "other rels" */
+			if (brel->reloptkind != RELOPT_BASEREL)
+				continue;
+
+			if (brel != rel)
+			{
+				/* Oops, it's a join query */
+				pathnode->limit_tuples = -1.0;
+				break;
+			}
+		}
+	}
+
 	/* Add up all the costs of the input paths */
 	input_startup_cost = 0;
 	input_total_cost = 0;
@@ -720,7 +749,7 @@ create_merge_append_path(PlannerInfo *root,
 					  subpath->parent->width,
 					  0.0,
 					  work_mem,
-					  -1.0);
+					  pathnode->limit_tuples);
 			input_startup_cost += sort_path.startup_cost;
 			input_total_cost += sort_path.total_cost;
 		}