commit 9ecde9d256d5f865f2cebe1c602bedbd04564da7
Author: Ronan Dunklau <ronan.dunklau@aiven.io>
Date:   Tue Jul 20 11:17:40 2021 +0200

    Teach sort and result nodes to perform a simplified version of a
    projection.
    
    Whenever we just need to remove some columns, we can use the resjunk
    infrastructure instead.

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index b099368809..72c9a20aa9 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -2323,7 +2323,9 @@ static void
 show_sort_keys(SortState *sortstate, List *ancestors, ExplainState *es)
 {
 	Sort	   *plan = (Sort *) sortstate->ss.ps.plan;
-	show_sort_group_keys((PlanState *) outerPlanState(sortstate), "Sort Key",
+	/* Sort target list does not necessarily contain the keys, use them from the
+	 * outer plan*/
+	show_sort_group_keys(outerPlanState(sortstate), "Sort Key",
 						 plan->numCols, 0, plan->sortColIdx,
 						 plan->sortOperators, plan->collations,
 						 plan->nullsFirst,
diff --git a/src/backend/executor/execJunk.c b/src/backend/executor/execJunk.c
index 9741897e83..92308b739b 100644
--- a/src/backend/executor/execJunk.c
+++ b/src/backend/executor/execJunk.c
@@ -15,6 +15,7 @@
 #include "postgres.h"
 
 #include "executor/executor.h"
+#include "optimizer/tlist.h"
 
 /*-------------------------------------------------------------------------
  *		XXX this stuff should be rewritten to take advantage
@@ -200,6 +201,64 @@ ExecInitJunkFilterConversion(List *targetList,
 	return junkfilter;
 }
 
+JunkFilter *
+ExecInitDiffJunkFilter(List *inputTargetList, List *outputTargetList, TupleTableSlot *slot)
+{
+	JunkFilter * junkfilter;
+	int			cleanLength;
+	AttrNumber *cleanMap;
+	ListCell   *t;
+	int			i;
+	TupleDesc  cleanTupType = ExecTypeFromTL(outputTargetList);
+
+	if (equal(inputTargetList, outputTargetList))
+		return NULL;
+
+	/*
+	 * Use the given slot, or make a new slot if we weren't given one.
+	 */
+	if (!slot)
+		slot = MakeSingleTupleTableSlot(cleanTupType, &TTSOpsVirtual);
+
+	cleanLength = cleanTupType->natts;
+	if (cleanLength > 0)
+	{
+		cleanMap = (AttrNumber *) palloc0(cleanLength * sizeof(AttrNumber));
+		i = 0;
+		foreach(t, outputTargetList)
+		{
+			int j;
+			TargetEntry *outTle = lfirst_node(TargetEntry, t);
+			TargetEntry * inTle;
+			switch (outTle->expr->type)
+			{
+				/* Const are kept as-is, so look it up underneath */
+				case T_Const:
+					inTle = tlist_member(outTle->expr, inputTargetList);
+					j = inTle->resno;
+					break;
+				case T_Var:
+					j = ((Var*) (outTle->expr))->varattno;
+					break;
+				default:
+					elog(ERROR, "Unsupported node type: %d", outTle->expr->type);
+			}
+			cleanMap[i] = j;
+			i++;
+		}
+	} else
+		cleanMap = NULL;
+	junkfilter = makeNode(JunkFilter);
+
+	junkfilter->jf_targetList = outputTargetList;
+	junkfilter->jf_cleanTupType = cleanTupType;
+	junkfilter->jf_cleanMap = cleanMap;
+	junkfilter->jf_resultSlot = slot;
+
+	return junkfilter;
+
+}
+
 /*
  * ExecFindJunkAttribute
  *
diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index ad11392b99..d2ca6e3acc 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -522,6 +522,22 @@ ExecGetResultSlotOps(PlanState *planstate, bool *isfixed)
 }
 
 
+bool
+ExecCanUseJunkFilter(PlanState *planstate)
+{
+	Plan * outerPlan= outerPlan(planstate->plan);
+	return outerPlan && tlist_is_subset(planstate->plan->targetlist, outerPlan->targetlist);
+}
+
+void
+ExecAssignJunkProjection(PlanState *planstate)
+{
+	/* The caller should have checked it first */
+	planstate->ps_junkFilter = ExecInitDiffJunkFilter(outerPlan(planstate->plan)->targetlist,
+											   planstate->plan->targetlist,
+											   planstate->ps_ResultTupleSlot);
+}
+
 /* ----------------
  *		ExecAssignProjectionInfo
  *
diff --git a/src/backend/executor/nodeResult.c b/src/backend/executor/nodeResult.c
index 0946af0a54..2a517637c2 100644
--- a/src/backend/executor/nodeResult.c
+++ b/src/backend/executor/nodeResult.c
@@ -132,8 +132,15 @@ ExecResult(PlanState *pstate)
 			node->rs_done = true;
 		}
 
-		/* form the result tuple using ExecProject(), and return it */
-		return ExecProject(node->ps.ps_ProjInfo);
+		/* form the result tuple using ExecProject() or ExecFilterJunk(), and return it
+		 * Only one of ps_junkFilter or ps_ProjInfo must be set at a time
+		 */
+		Assert((node->ps.ps_junkFilter == NULL && node->ps.ps_ProjInfo != NULL) ||
+			   (node->ps.ps_junkFilter != NULL && node->ps.ps_ProjInfo == NULL));
+		if (node->ps.ps_junkFilter != NULL)
+			return ExecFilterJunk(node->ps.ps_junkFilter, econtext->ecxt_outertuple);
+		else
+			return ExecProject(node->ps.ps_ProjInfo);
 	}
 
 	return NULL;
@@ -218,7 +225,14 @@ ExecInitResult(Result *node, EState *estate, int eflags)
 	 * Initialize result slot, type and projection.
 	 */
 	ExecInitResultTupleSlotTL(&resstate->ps, &TTSOpsVirtual);
-	ExecAssignProjectionInfo(&resstate->ps, NULL);
+	/*
+	 * If a real projection is needed, assign a projection info.
+	 * If we only need a subset of the column, assign a resjunkfilter
+	 */
+	if (ExecCanUseJunkFilter(&resstate->ps))
+		ExecAssignJunkProjection(&resstate->ps);
+	else
+		ExecAssignProjectionInfo(&resstate->ps, NULL);
 
 	/*
 	 * initialize child expressions
diff --git a/src/backend/executor/nodeSort.c b/src/backend/executor/nodeSort.c
index b99027e0d7..9c98928f70 100644
--- a/src/backend/executor/nodeSort.c
+++ b/src/backend/executor/nodeSort.c
@@ -149,10 +149,12 @@ ExecSort(PlanState *pstate)
 	 * tuples.  Note that we only rely on slot tuple remaining valid until the
 	 * next fetch from the tuplesort.
 	 */
-	slot = node->ss.ps.ps_ResultTupleSlot;
+	slot = node->outsortTupleSlot;
 	(void) tuplesort_gettupleslot(tuplesortstate,
 								  ScanDirectionIsForward(dir),
 								  false, slot, NULL);
+	if (!TTS_EMPTY(slot) && node->ss.ps.ps_junkFilter)
+		slot = ExecFilterJunk(node->ss.ps.ps_junkFilter, slot);
 	return slot;
 }
 
@@ -221,6 +223,18 @@ ExecInitSort(Sort *node, EState *estate, int eflags)
 	ExecInitResultTupleSlotTL(&sortstate->ss.ps, &TTSOpsMinimalTuple);
 	sortstate->ss.ps.ps_ProjInfo = NULL;
 
+	/* Initialize the junkfilter if any is needed.
+	 */
+	sortstate->ss.ps.ps_junkFilter = ExecInitDiffJunkFilter(outerPlan(sortstate->ss.ps.plan)->targetlist,
+											   sortstate->ss.ps.plan->targetlist,
+											   sortstate->ss.ps.ps_ResultTupleSlot);
+	if (sortstate->ss.ps.ps_junkFilter)
+		sortstate->outsortTupleSlot = ExecInitExtraTupleSlot(estate, sortstate->ss.ss_ScanTupleSlot->tts_tupleDescriptor, &TTSOpsMinimalTuple);
+	else
+		sortstate->outsortTupleSlot = sortstate->ss.ps.ps_ResultTupleSlot;
+
+
+
 	SO1_printf("ExecInitSort: %s\n",
 			   "sort node initialized");
 
@@ -243,6 +257,7 @@ ExecEndSort(SortState *node)
 	ExecClearTuple(node->ss.ss_ScanTupleSlot);
 	/* must drop pointer to sort result tuple */
 	ExecClearTuple(node->ss.ps.ps_ResultTupleSlot);
+	ExecClearTuple(node->outsortTupleSlot);
 
 	/*
 	 * Release tuplesort resources
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 90f430834f..e8a65d332c 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -1963,7 +1963,7 @@ create_projection_plan(PlannerInfo *root, ProjectionPath *best_path, int flags)
 			apply_pathtarget_labeling_to_tlist(tlist,
 											   best_path->path.pathtarget);
 	}
-	else if (is_projection_capable_path(best_path->subpath))
+	else if (is_projection_capable_path(best_path->subpath, NULL))
 	{
 		/*
 		 * Our caller requires that we return the exact tlist, but no separate
@@ -7028,17 +7028,20 @@ make_modifytable(PlannerInfo *root, Plan *subplan,
 /*
  * is_projection_capable_path
  *		Check whether a given Path node is able to do projection.
+ *		If the target argument is NULL, the node must be able to perform any
+ *		projection, otherwise we check that the given projection is possible.
  */
 bool
-is_projection_capable_path(Path *path)
+is_projection_capable_path(Path *path, PathTarget *target)
 {
 	/* Most plan types can project, so just list the ones that can't */
 	switch (path->pathtype)
 	{
+		case T_Sort:
+			return target && pathtarget_is_subset(target, path->pathtarget);
 		case T_Hash:
 		case T_Material:
 		case T_Memoize:
-		case T_Sort:
 		case T_IncrementalSort:
 		case T_Unique:
 		case T_SetOp:
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index bf0ccf155c..de6c8232df 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -4519,8 +4519,15 @@ create_ordered_paths(PlannerInfo *root,
 														limit_tuples);
 				/* Add projection step if needed */
 				if (sorted_path->pathtarget != target)
+				{
+					/* In our case, we should raise the target cost by the
+					 * previous cost, since all previous columns are also
+					 * required before the sort.
+					 */
+					target->cost.per_tuple += sorted_path->pathtarget->cost.per_tuple;
 					sorted_path = apply_projection_to_path(root, ordered_rel,
 														   sorted_path, target);
+				}
 
 				add_path(ordered_rel, sorted_path);
 			}
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index b145c5f45f..46d130cabe 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -770,6 +770,8 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 
 		case T_Material:
 		case T_Sort:
+			set_upper_references(root, plan, rtoffset);
+			break;
 		case T_IncrementalSort:
 		case T_Unique:
 		case T_SetOp:
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 96440c01ef..2a78a974ab 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -2673,7 +2673,7 @@ create_projection_path(PlannerInfo *root,
 	 * conclusion; see comments therein.
 	 */
 	oldtarget = subpath->pathtarget;
-	if (is_projection_capable_path(subpath) ||
+	if (is_projection_capable_path(subpath, target) ||
 		equal(oldtarget->exprs, target->exprs))
 	{
 		/* No separate Result node needed */
@@ -2744,7 +2744,7 @@ apply_projection_to_path(PlannerInfo *root,
 	 * If given path can't project, we might need a Result node, so make a
 	 * separate ProjectionPath.
 	 */
-	if (!is_projection_capable_path(path))
+	if (!is_projection_capable_path(path, target))
 		return (Path *) create_projection_path(root, rel, path, target);
 
 	/*
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 3dc03c913e..7c447e8544 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -159,6 +159,9 @@ extern JunkFilter *ExecInitJunkFilter(List *targetList,
 extern JunkFilter *ExecInitJunkFilterConversion(List *targetList,
 												TupleDesc cleanTupType,
 												TupleTableSlot *slot);
+extern JunkFilter *ExecInitDiffJunkFilter(List *inputTargetList,
+										  List *outputTargetList,
+										  TupleTableSlot *slot);
 extern AttrNumber ExecFindJunkAttribute(JunkFilter *junkfilter,
 										const char *attrName);
 extern AttrNumber ExecFindJunkAttributeInTlist(List *targetlist,
@@ -551,6 +554,8 @@ extern const TupleTableSlotOps *ExecGetResultSlotOps(PlanState *planstate,
 													 bool *isfixed);
 extern void ExecAssignProjectionInfo(PlanState *planstate,
 									 TupleDesc inputDesc);
+extern bool ExecCanUseJunkFilter(PlanState *planstate);
+extern void ExecAssignJunkProjection(PlanState *planstate);
 extern void ExecConditionalAssignProjectionInfo(PlanState *planstate,
 												TupleDesc inputDesc, Index varno);
 extern void ExecFreeExprContext(PlanState *planstate);
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 105180764e..3674c8f4f0 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -1004,6 +1004,7 @@ typedef struct PlanState
 	TupleTableSlot *ps_ResultTupleSlot; /* slot for my result tuples */
 	ExprContext *ps_ExprContext;	/* node's expression-evaluation context */
 	ProjectionInfo *ps_ProjInfo;	/* info for doing tuple projection */
+	JunkFilter *ps_junkFilter;
 
 	bool		async_capable;	/* true if node is async-capable */
 
@@ -2151,6 +2152,7 @@ typedef struct SortState
 	int64		bound_Done;		/* value of bound we did the sort with */
 	void	   *tuplesortstate; /* private state of tuplesort.c */
 	bool		am_worker;		/* are we a worker? */
+	TupleTableSlot *outsortTupleSlot;
 	SharedSortInfo *shared_info;	/* one entry per worker */
 } SortState;
 
diff --git a/src/include/optimizer/planmain.h b/src/include/optimizer/planmain.h
index bf1adfc52a..ac5d769e4b 100644
--- a/src/include/optimizer/planmain.h
+++ b/src/include/optimizer/planmain.h
@@ -46,7 +46,7 @@ extern ForeignScan *make_foreignscan(List *qptlist, List *qpqual,
 extern Plan *change_plan_targetlist(Plan *subplan, List *tlist,
 									bool tlist_parallel_safe);
 extern Plan *materialize_finished_plan(Plan *subplan);
-extern bool is_projection_capable_path(Path *path);
+extern bool is_projection_capable_path(Path *path, PathTarget *target);
 extern bool is_projection_capable_plan(Plan *plan);
 
 /* External use of these functions is deprecated: */
