Add support for index-only scans in GiST.

This adds a new GiST opclass method, 'fetch', which is used to reconstruct
the original Datum from the value stored in the index. Also, the 'canreturn'
index AM interface function gains a new 'attno' argument. That makes it
possible to use index-only scans on a multi-column index where some of the
opclasses support index-only scans but some do not.

This patch adds support in the box and point opclasses. Other opclasses
can added later as follow-on patches (btree_gist would be particularly
interesting).

Anastasia Lubennikova, with additional fixes and modifications by me.

5 files changed, 190 insertions, 3 deletions

diff --git a/src/backend/access/gist/gist.c b/src/backend/access/gist/gist.c
index db2a452a4ab..96b7701633f 100644
--- a/src/backend/access/gist/gist.c
+++ b/src/backend/access/gist/gist.c
@@ -1404,6 +1404,14 @@ initGISTstate(Relation index)
 		else
 			giststate->distanceFn[i].fn_oid = InvalidOid;
 
+		/* opclasses are not required to provide a Fetch method */
+		if (OidIsValid(index_getprocid(index, i + 1, GIST_FETCH_PROC)))
+			fmgr_info_copy(&(giststate->fetchFn[i]),
+						 index_getprocinfo(index, i + 1, GIST_FETCH_PROC),
+						   scanCxt);
+		else
+			giststate->fetchFn[i].fn_oid = InvalidOid;
+
 		/*
 		 * If the index column has a specified collation, we should honor that
 		 * while doing comparisons.  However, we may have a collatable storage
diff --git a/src/backend/access/gist/gistget.c b/src/backend/access/gist/gistget.c
index 717cb85f773..e4c00c2c9f5 100644
--- a/src/backend/access/gist/gistget.c
+++ b/src/backend/access/gist/gistget.c
@@ -228,7 +228,9 @@ gistindex_keytest(IndexScanDesc scan,
  * tuples should be reported directly into the bitmap.  If they are NULL,
  * we're doing a plain or ordered indexscan.  For a plain indexscan, heap
  * tuple TIDs are returned into so->pageData[].  For an ordered indexscan,
- * heap tuple TIDs are pushed into individual search queue items.
+ * heap tuple TIDs are pushed into individual search queue items.  In an
+ * index-only scan, reconstructed index tuples are returned along with the
+ * TIDs.
  *
  * If we detect that the index page has split since we saw its downlink
  * in the parent, we push its new right sibling onto the queue so the
@@ -239,6 +241,8 @@ gistScanPage(IndexScanDesc scan, GISTSearchItem *pageItem, double *myDistances,
 			 TIDBitmap *tbm, int64 *ntids)
 {
 	GISTScanOpaque so = (GISTScanOpaque) scan->opaque;
+	GISTSTATE  *giststate = so->giststate;
+	Relation	r = scan->indexRelation;
 	Buffer		buffer;
 	Page		page;
 	GISTPageOpaque opaque;
@@ -288,6 +292,8 @@ gistScanPage(IndexScanDesc scan, GISTSearchItem *pageItem, double *myDistances,
 	}
 
 	so->nPageData = so->curPageData = 0;
+	if (so->pageDataCxt)
+		MemoryContextReset(so->pageDataCxt);
 
 	/*
 	 * check all tuples on page
@@ -326,10 +332,21 @@ gistScanPage(IndexScanDesc scan, GISTSearchItem *pageItem, double *myDistances,
 		else if (scan->numberOfOrderBys == 0 && GistPageIsLeaf(page))
 		{
 			/*
-			 * Non-ordered scan, so report heap tuples in so->pageData[]
+			 * Non-ordered scan, so report tuples in so->pageData[]
 			 */
 			so->pageData[so->nPageData].heapPtr = it->t_tid;
 			so->pageData[so->nPageData].recheck = recheck;
+
+			/*
+			 * In an index-only scan, also fetch the data from the tuple.
+			 */
+			if (scan->xs_want_itup)
+			{
+				oldcxt = MemoryContextSwitchTo(so->pageDataCxt);
+				so->pageData[so->nPageData].ftup =
+					gistFetchTuple(giststate, r, it);
+				MemoryContextSwitchTo(oldcxt);
+			}
 			so->nPageData++;
 		}
 		else
@@ -352,6 +369,12 @@ gistScanPage(IndexScanDesc scan, GISTSearchItem *pageItem, double *myDistances,
 				item->blkno = InvalidBlockNumber;
 				item->data.heap.heapPtr = it->t_tid;
 				item->data.heap.recheck = recheck;
+
+				/*
+				 * In an index-only scan, also fetch the data from the tuple.
+				 */
+				if (scan->xs_want_itup)
+					item->data.heap.ftup = gistFetchTuple(giststate, r, it);
 			}
 			else
 			{
@@ -412,6 +435,13 @@ getNextNearest(IndexScanDesc scan)
 	GISTScanOpaque so = (GISTScanOpaque) scan->opaque;
 	bool		res = false;
 
+	if (scan->xs_itup)
+	{
+		/* free previously returned tuple */
+		pfree(scan->xs_itup);
+		scan->xs_itup = NULL;
+	}
+
 	do
 	{
 		GISTSearchItem *item = getNextGISTSearchItem(so);
@@ -424,6 +454,10 @@ getNextNearest(IndexScanDesc scan)
 			/* found a heap item at currently minimal distance */
 			scan->xs_ctup.t_self = item->data.heap.heapPtr;
 			scan->xs_recheck = item->data.heap.recheck;
+
+			/* in an index-only scan, also return the reconstructed tuple. */
+			if (scan->xs_want_itup)
+				scan->xs_itup = item->data.heap.ftup;
 			res = true;
 		}
 		else
@@ -465,6 +499,8 @@ gistgettuple(PG_FUNCTION_ARGS)
 
 		so->firstCall = false;
 		so->curPageData = so->nPageData = 0;
+		if (so->pageDataCxt)
+			MemoryContextReset(so->pageDataCxt);
 
 		fakeItem.blkno = GIST_ROOT_BLKNO;
 		memset(&fakeItem.data.parentlsn, 0, sizeof(GistNSN));
@@ -483,10 +519,17 @@ gistgettuple(PG_FUNCTION_ARGS)
 		{
 			if (so->curPageData < so->nPageData)
 			{
+
 				/* continuing to return tuples from a leaf page */
 				scan->xs_ctup.t_self = so->pageData[so->curPageData].heapPtr;
 				scan->xs_recheck = so->pageData[so->curPageData].recheck;
+
+				/* in an index-only scan, also return the reconstructed tuple */
+				if (scan->xs_want_itup)
+					scan->xs_itup = so->pageData[so->curPageData].ftup;
+
 				so->curPageData++;
+
 				PG_RETURN_BOOL(true);
 			}
 
@@ -533,6 +576,8 @@ gistgetbitmap(PG_FUNCTION_ARGS)
 
 	/* Begin the scan by processing the root page */
 	so->curPageData = so->nPageData = 0;
+	if (so->pageDataCxt)
+		MemoryContextReset(so->pageDataCxt);
 
 	fakeItem.blkno = GIST_ROOT_BLKNO;
 	memset(&fakeItem.data.parentlsn, 0, sizeof(GistNSN));
@@ -558,3 +603,20 @@ gistgetbitmap(PG_FUNCTION_ARGS)
 
 	PG_RETURN_INT64(ntids);
 }
+
+/*
+ * Can we do index-only scans on the given index column?
+ *
+ * Opclasses that implement a fetch function support index-only scans.
+ */
+Datum
+gistcanreturn(PG_FUNCTION_ARGS)
+{
+	Relation	index = (Relation) PG_GETARG_POINTER(0);
+	int			attno = PG_GETARG_INT32(1);
+
+	if (OidIsValid(index_getprocid(index, attno, GIST_FETCH_PROC)))
+		PG_RETURN_BOOL(true);
+	else
+		PG_RETURN_BOOL(false);
+}
diff --git a/src/backend/access/gist/gistproc.c b/src/backend/access/gist/gistproc.c
index 9fab6c87c05..9d21e3fb947 100644
--- a/src/backend/access/gist/gistproc.c
+++ b/src/backend/access/gist/gistproc.c
@@ -152,6 +152,16 @@ gist_box_decompress(PG_FUNCTION_ARGS)
 }
 
 /*
+ * GiST Fetch method for boxes
+ * do not do anything --- we just return the stored box as is.
+ */
+Datum
+gist_box_fetch(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_POINTER(PG_GETARG_POINTER(0));
+}
+
+/*
  * The GiST Penalty method for boxes (also used for points)
  *
  * As in the R-tree paper, we use change in area as our penalty metric
@@ -1186,6 +1196,33 @@ gist_point_compress(PG_FUNCTION_ARGS)
 	PG_RETURN_POINTER(entry);
 }
 
+/*
+ * GiST Fetch method for point
+ *
+ * Get point coordinates from its bounding box coordinates and form new
+ * gistentry.
+ */
+Datum
+gist_point_fetch(PG_FUNCTION_ARGS)
+{
+	GISTENTRY  *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
+	BOX		   *in = DatumGetBoxP(entry->key);
+	Point	   *r;
+	GISTENTRY  *retval;
+
+	retval = palloc(sizeof(GISTENTRY));
+
+	r = (Point *) palloc(sizeof(Point));
+	r->x = in->high.x;
+	r->y = in->high.y;
+	gistentryinit(*retval, PointerGetDatum(r),
+				  entry->rel, entry->page,
+				  entry->offset, FALSE);
+
+	PG_RETURN_POINTER(retval);
+}
+
+
 #define point_point_distance(p1,p2) \
 	DatumGetFloat8(DirectFunctionCall2(point_distance, \
 									   PointPGetDatum(p1), PointPGetDatum(p2)))
diff --git a/src/backend/access/gist/gistscan.c b/src/backend/access/gist/gistscan.c
index 991858ff43f..3522d75a496 100644
--- a/src/backend/access/gist/gistscan.c
+++ b/src/backend/access/gist/gistscan.c
@@ -88,6 +88,13 @@ gistbeginscan(PG_FUNCTION_ARGS)
 
 	scan->opaque = so;
 
+	/*
+	 * All fields required for index-only scans are null until gistrescan.
+	 * However, we set up scan->xs_itupdesc whether we'll need it or not,
+	 * since that's cheap.
+	 */
+	scan->xs_itupdesc = RelationGetDescr(r);
+
 	MemoryContextSwitchTo(oldCxt);
 
 	PG_RETURN_POINTER(scan);
@@ -141,6 +148,17 @@ gistrescan(PG_FUNCTION_ARGS)
 		first_time = false;
 	}
 
+	/*
+	 * If we're doing an index-only scan, also create a memory context to hold
+	 * the returned tuples.
+	 */
+	if (scan->xs_want_itup && so->pageDataCxt == NULL)
+		so->pageDataCxt = AllocSetContextCreate(so->giststate->scanCxt,
+												"GiST page data context",
+												ALLOCSET_DEFAULT_MINSIZE,
+												ALLOCSET_DEFAULT_INITSIZE,
+												ALLOCSET_DEFAULT_MAXSIZE);
+
 	/* create new, empty RBTree for search queue */
 	oldCxt = MemoryContextSwitchTo(so->queueCxt);
 	so->queue = pairingheap_allocate(pairingheap_GISTSearchItem_cmp, scan);
diff --git a/src/backend/access/gist/gistutil.c b/src/backend/access/gist/gistutil.c
index 824c40eb203..1680251a18b 100644
--- a/src/backend/access/gist/gistutil.c
+++ b/src/backend/access/gist/gistutil.c
@@ -294,8 +294,9 @@ gistDeCompressAtt(GISTSTATE *giststate, Relation r, IndexTuple tuple, Page p,
 
 	for (i = 0; i < r->rd_att->natts; i++)
 	{
-		Datum		datum = index_getattr(tuple, i + 1, giststate->tupdesc, &isnull[i]);
+		Datum		datum;
 
+		datum = index_getattr(tuple, i + 1, giststate->tupdesc, &isnull[i]);
 		gistdentryinit(giststate, i, &attdata[i],
 					   datum, r, p, o,
 					   FALSE, isnull[i]);
@@ -598,6 +599,67 @@ gistFormTuple(GISTSTATE *giststate, Relation r,
 	return res;
 }
 
+/*
+ * initialize a GiST entry with fetched value in key field
+ */
+static Datum
+gistFetchAtt(GISTSTATE *giststate, int nkey, Datum k, Relation r)
+{
+	GISTENTRY	fentry;
+	GISTENTRY  *fep;
+
+	gistentryinit(fentry, k, r, NULL, (OffsetNumber) 0, false);
+
+	fep = (GISTENTRY *)
+		DatumGetPointer(FunctionCall1Coll(&giststate->fetchFn[nkey],
+										  giststate->supportCollation[nkey],
+										  PointerGetDatum(&fentry)));
+
+	/* fetchFn set 'key', return it to the caller */
+	return fep->key;
+}
+
+/*
+ * Fetch all keys in tuple.
+ * returns new IndexTuple that contains GISTENTRY with fetched data
+ */
+IndexTuple
+gistFetchTuple(GISTSTATE *giststate, Relation r, IndexTuple tuple)
+{
+	MemoryContext oldcxt = MemoryContextSwitchTo(giststate->tempCxt);
+	Datum		fetchatt[INDEX_MAX_KEYS];
+	bool		isnull[INDEX_MAX_KEYS];
+	int			i;
+
+	for (i = 0; i < r->rd_att->natts; i++)
+	{
+		Datum		datum;
+
+		datum = index_getattr(tuple, i + 1, giststate->tupdesc, &isnull[i]);
+
+		if (giststate->fetchFn[i].fn_oid != InvalidOid)
+		{
+			if (!isnull[i])
+				fetchatt[i] = gistFetchAtt(giststate, i, datum, r);
+			else
+				fetchatt[i] = (Datum) 0;
+		}
+		else
+		{
+			/*
+			 * Index-only scans not supported for this column. Since the
+			 * planner chose an index-only scan anyway, it is not interested
+			 * in this column, and we can replace it with a NULL.
+			 */
+			isnull[i] = true;
+			fetchatt[i] = (Datum) 0;
+		}
+	}
+	MemoryContextSwitchTo(oldcxt);
+
+	return index_form_tuple(giststate->tupdesc, fetchatt, isnull);
+}
+
 float
 gistpenalty(GISTSTATE *giststate, int attno,
 			GISTENTRY *orig, bool isNullOrig,