Support FETCH FIRST WITH TIES

WITH TIES is an option to the FETCH FIRST N ROWS clause (the SQL
standard's spelling of LIMIT), where you additionally get rows that
compare equal to the last of those N rows by the columns in the
mandatory ORDER BY clause.

There was a proposal by Andrew Gierth to implement this functionality in
a more powerful way that would yield more features, but the other patch
had not been finished at this time, so we decided to use this one for
now in the spirit of incremental development.

Author: Surafel Temesgen <surafel3000@gmail.com>
Reviewed-by: Álvaro Herrera <alvherre@alvh.no-ip.org>
Reviewed-by: Tomas Vondra <tomas.vondra@2ndquadrant.com>
Discussion: https://postgr.es/m/CALAY4q9ky7rD_A4vf=FVQvCGngm3LOes-ky0J6euMrg=_Se+ag@mail.gmail.com
Discussion: https://postgr.es/m/87o8wvz253.fsf@news-spur.riddles.org.uk

1 files changed, 144 insertions, 16 deletions

diff --git a/src/backend/executor/nodeLimit.c b/src/backend/executor/nodeLimit.c
index d792e97cfab..12cfc5177bb 100644
--- a/src/backend/executor/nodeLimit.c
+++ b/src/backend/executor/nodeLimit.c
@@ -41,6 +41,7 @@ static TupleTableSlot *			/* return: a tuple or NULL */
 ExecLimit(PlanState *pstate)
 {
 	LimitState *node = castNode(LimitState, pstate);
+	ExprContext *econtext = node->ps.ps_ExprContext;
 	ScanDirection direction;
 	TupleTableSlot *slot;
 	PlanState  *outerPlan;
@@ -102,6 +103,16 @@ ExecLimit(PlanState *pstate)
 					node->lstate = LIMIT_EMPTY;
 					return NULL;
 				}
+
+				/*
+				 * Tuple at limit is needed for comparation in subsequent
+				 * execution to detect ties.
+				 */
+				if (node->limitOption == LIMIT_OPTION_WITH_TIES &&
+					node->position - node->offset == node->count - 1)
+				{
+					ExecCopySlot(node->last_slot, slot);
+				}
 				node->subSlot = slot;
 				if (++node->position > node->offset)
 					break;
@@ -125,10 +136,13 @@ ExecLimit(PlanState *pstate)
 			if (ScanDirectionIsForward(direction))
 			{
 				/*
-				 * Forwards scan, so check for stepping off end of window. If
-				 * we are at the end of the window, return NULL without
-				 * advancing the subplan or the position variable; but change
-				 * the state machine state to record having done so.
+				 * Forwards scan, so check for stepping off end of window.  At
+				 * the end of the window, the behavior depends on whether WITH
+				 * TIES was specified: in that case, we need to change the
+				 * state machine to LIMIT_WINDOWTIES.  If not (nothing was
+				 * specified, or ONLY was) return NULL without advancing the
+				 * subplan or the position variable but change the state
+				 * machine to record having done so
 				 *
 				 * Once at the end, ideally, we can shut down parallel
 				 * resources but that would destroy the parallel context which
@@ -139,12 +153,72 @@ ExecLimit(PlanState *pstate)
 				if (!node->noCount &&
 					node->position - node->offset >= node->count)
 				{
-					node->lstate = LIMIT_WINDOWEND;
+					if (node->limitOption == LIMIT_OPTION_COUNT)
+					{
+						node->lstate = LIMIT_WINDOWEND;
+						return NULL;
+					}
+					else
+					{
+						node->lstate = LIMIT_WINDOWEND_TIES;
+						/* fall-through */
+					}
+				}
+				else
+				{
+					/*
+					 * Get next tuple from subplan, if any.
+					 */
+					slot = ExecProcNode(outerPlan);
+					if (TupIsNull(slot))
+					{
+						node->lstate = LIMIT_SUBPLANEOF;
+						return NULL;
+					}
+
+					/*
+					 * Tuple at limit is needed for comparation in subsequent
+					 * execution to detect ties.
+					 */
+					if (node->limitOption == LIMIT_OPTION_WITH_TIES &&
+						node->position - node->offset == node->count - 1)
+					{
+						ExecCopySlot(node->last_slot, slot);
+					}
+					node->subSlot = slot;
+					node->position++;
+					break;
+				}
+			}
+			else
+			{
+				/*
+				 * Backwards scan, so check for stepping off start of window.
+				 * As above, change only state-machine status if so.
+				 */
+				if (node->position <= node->offset + 1)
+				{
+					node->lstate = LIMIT_WINDOWSTART;
 					return NULL;
 				}
 
 				/*
-				 * Get next tuple from subplan, if any.
+				 * Get previous tuple from subplan; there should be one!
+				 */
+				slot = ExecProcNode(outerPlan);
+				if (TupIsNull(slot))
+					elog(ERROR, "LIMIT subplan failed to run backwards");
+				node->subSlot = slot;
+				node->position--;
+				break;
+			}
+
+		case LIMIT_WINDOWEND_TIES:
+			if (ScanDirectionIsForward(direction))
+			{
+				/*
+				 * Advance the subplan until we find the first row with
+				 * different ORDER BY pathkeys.
 				 */
 				slot = ExecProcNode(outerPlan);
 				if (TupIsNull(slot))
@@ -152,14 +226,29 @@ ExecLimit(PlanState *pstate)
 					node->lstate = LIMIT_SUBPLANEOF;
 					return NULL;
 				}
-				node->subSlot = slot;
-				node->position++;
+
+				/*
+				 * Test if the new tuple and the last tuple match. If so we
+				 * return the tuple.
+				 */
+				econtext->ecxt_innertuple = slot;
+				econtext->ecxt_outertuple = node->last_slot;
+				if (ExecQualAndReset(node->eqfunction, econtext))
+				{
+					node->subSlot = slot;
+					node->position++;
+				}
+				else
+				{
+					node->lstate = LIMIT_WINDOWEND;
+					return NULL;
+				}
 			}
 			else
 			{
 				/*
 				 * Backwards scan, so check for stepping off start of window.
-				 * As above, change only state-machine status if so.
+				 * Change only state-machine status if so.
 				 */
 				if (node->position <= node->offset + 1)
 				{
@@ -168,13 +257,15 @@ ExecLimit(PlanState *pstate)
 				}
 
 				/*
-				 * Get previous tuple from subplan; there should be one!
+				 * Get previous tuple from subplan; there should be one! And
+				 * change state-machine status.
 				 */
 				slot = ExecProcNode(outerPlan);
 				if (TupIsNull(slot))
 					elog(ERROR, "LIMIT subplan failed to run backwards");
 				node->subSlot = slot;
 				node->position--;
+				node->lstate = LIMIT_INWINDOW;
 			}
 			break;
 
@@ -199,12 +290,28 @@ ExecLimit(PlanState *pstate)
 				return NULL;
 
 			/*
-			 * Backing up from window end: simply re-return the last tuple
-			 * fetched from the subplan.
+			 * We already past one position to detect ties so re-fetch
+			 * previous tuple; there should be one!  Note previous tuple must
+			 * be in window.
 			 */
-			slot = node->subSlot;
-			node->lstate = LIMIT_INWINDOW;
-			/* position does not change 'cause we didn't advance it before */
+			if (node->limitOption == LIMIT_OPTION_WITH_TIES)
+			{
+				slot = ExecProcNode(outerPlan);
+				if (TupIsNull(slot))
+					elog(ERROR, "LIMIT subplan failed to run backwards");
+				node->subSlot = slot;
+				node->lstate = LIMIT_INWINDOW;
+			}
+			else
+			{
+				/*
+				 * Backing up from window end: simply re-return the last tuple
+				 * fetched from the subplan.
+				 */
+				slot = node->subSlot;
+				node->lstate = LIMIT_INWINDOW;
+				/* position does not change 'cause we didn't advance it before */
+			}
 			break;
 
 		case LIMIT_WINDOWSTART:
@@ -319,7 +426,7 @@ recompute_limits(LimitState *node)
 static int64
 compute_tuples_needed(LimitState *node)
 {
-	if (node->noCount)
+	if ((node->noCount) || (node->limitOption == LIMIT_OPTION_WITH_TIES))
 		return -1;
 	/* Note: if this overflows, we'll return a negative value, which is OK */
 	return node->count + node->offset;
@@ -372,6 +479,7 @@ ExecInitLimit(Limit *node, EState *estate, int eflags)
 										   (PlanState *) limitstate);
 	limitstate->limitCount = ExecInitExpr((Expr *) node->limitCount,
 										  (PlanState *) limitstate);
+	limitstate->limitOption = node->limitOption;
 
 	/*
 	 * Initialize result type.
@@ -388,6 +496,26 @@ ExecInitLimit(Limit *node, EState *estate, int eflags)
 	 */
 	limitstate->ps.ps_ProjInfo = NULL;
 
+	/*
+	 * Initialize the equality evaluation, to detect ties.
+	 */
+	if (node->limitOption == LIMIT_OPTION_WITH_TIES)
+	{
+		TupleDesc	desc;
+		const TupleTableSlotOps *ops;
+
+		desc = ExecGetResultType(outerPlanState(limitstate));
+		ops = ExecGetResultSlotOps(outerPlanState(limitstate), NULL);
+
+		limitstate->last_slot = ExecInitExtraTupleSlot(estate, desc, ops);
+		limitstate->eqfunction = execTuplesMatchPrepare(desc,
+														node->uniqNumCols,
+														node->uniqColIdx,
+														node->uniqOperators,
+														node->uniqCollations,
+														&limitstate->ps);
+	}
+
 	return limitstate;
 }