From b694db3e884b06b7f6e4508f41ee2a2288f73c20 Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <boekewurm+postgres@gmail.com>
Date: Fri, 8 Apr 2022 14:44:01 +0200
Subject: [PATCH v2 6/7] Implement specialized uncacheable attribute iteration

Uses an iterator to prevent doing duplicate work while iterating over
attributes.

Inspiration: https://www.postgresql.org/message-id/CAEze2WjE9ka8i%3Ds-Vv5oShro9xTrt5VQnQvFG9AaRwWpMm3-fg%40mail.gmail.com
---
 src/include/access/itup.h              | 179 +++++++++++++++++++++++++
 src/include/access/nbtree.h            |  12 +-
 src/include/access/nbtree_specialize.h |  34 +++++
 3 files changed, 223 insertions(+), 2 deletions(-)

diff --git a/src/include/access/itup.h b/src/include/access/itup.h
index 2c8877e991..cc29614107 100644
--- a/src/include/access/itup.h
+++ b/src/include/access/itup.h
@@ -59,6 +59,15 @@ typedef struct IndexAttributeBitMapData
 
 typedef IndexAttributeBitMapData * IndexAttributeBitMap;
 
+typedef struct IAttrIterStateData
+{
+	int			offset;
+	bool		slow;
+	bool		isNull;
+} IAttrIterStateData;
+
+typedef IAttrIterStateData * IAttrIterState;
+
 /*
  * t_info manipulation macros
  */
@@ -126,6 +135,42 @@ typedef IndexAttributeBitMapData * IndexAttributeBitMap;
 	) \
 )
 
+/* ----------------
+ *		index_attiterinit
+ *
+ *		This gets called many times, so we macro the cacheable and NULL
+ *		lookups, and call nocache_index_attiterinit() for the rest.
+ *
+ *		tup - the tuple being iterated on
+ *		attnum - the attribute number that we start the iteration with
+ *				 in the first index_attiternext call
+ *		tupdesc - the tuple description
+ *
+ * ----------------
+ */
+#define index_attiterinit(tup, attnum, tupleDesc, iter) \
+do { \
+	if ((attnum) == 1) \
+	{ \
+		*(iter) = ((IAttrIterStateData) { \
+			0 /* Offset of attribute 1 is always 0 */, \
+			false /* slow */, \
+			false /* isNull */ \
+		}); \
+	} \
+	else if (!IndexTupleHasNulls(tup) && \
+			 TupleDescAttr((tupleDesc), (attnum)-1)->attcacheoff >= 0) \
+	{ \
+		*(iter) = ((IAttrIterStateData) { \
+			TupleDescAttr((tupleDesc), (attnum)-1)->attcacheoff, /* offset */ \
+			false, /* slow */ \
+			false /* isNull */ \
+		}); \
+	} \
+	else \
+		nocache_index_attiterinit((tup), (attnum) - 1, (tupleDesc), (iter)); \
+} while (false);
+
 /*
  * MaxIndexTuplesPerPage is an upper bound on the number of tuples that can
  * fit on one index page.  An index tuple must have either data or a null
@@ -161,4 +206,138 @@ extern IndexTuple CopyIndexTuple(IndexTuple source);
 extern IndexTuple index_truncate_tuple(TupleDesc sourceDescriptor,
 									   IndexTuple source, int leavenatts);
 
+/*
+ * Initiate an index attribute iterator to attribute attnum,
+ * and return the corresponding datum.
+ *
+ * This is nearly the same as index_deform_tuple, except that this
+ * returns the internal state up to attnum, instead of populating the
+ * datum- and isnull-arrays
+ */
+static inline void
+nocache_index_attiterinit(IndexTuple tup, AttrNumber attnum, TupleDesc tupleDesc, IAttrIterState iter)
+{
+	bool		hasnulls = IndexTupleHasNulls(tup);
+	int			curatt;
+	char	   *tp;				/* ptr to tuple data */
+	int			off;			/* offset in tuple data */
+	bits8	   *bp;				/* ptr to null bitmap in tuple */
+	bool		slow = false;	/* can we use/set attcacheoff? */
+	bool		null = false;
+
+	/* Assert to protect callers */
+	Assert(PointerIsValid(iter));
+	Assert(tupleDesc->natts <= INDEX_MAX_KEYS);
+	Assert(attnum <= tupleDesc->natts);
+	Assert(attnum > 0);
+
+	/* XXX "knows" t_bits are just after fixed tuple header! */
+	bp = (bits8 *) ((char *) tup + sizeof(IndexTupleData));
+
+	tp = (char *) tup + IndexInfoFindDataOffset(tup->t_info);
+	off = 0;
+
+	for (curatt = 0; curatt < attnum; curatt++)
+	{
+		Form_pg_attribute thisatt = TupleDescAttr(tupleDesc, curatt);
+
+		if (hasnulls && att_isnull(curatt, bp))
+		{
+			null = true;
+			slow = true;		/* can't use attcacheoff anymore */
+			continue;
+		}
+
+		null = false;
+
+		if (!slow && thisatt->attcacheoff >= 0)
+			off = thisatt->attcacheoff;
+		else if (thisatt->attlen == -1)
+		{
+			off = att_align_pointer(off, thisatt->attalign, -1,
+									tp + off);
+			slow = true;
+		}
+		else
+		{
+			/* not varlena, so safe to use att_align_nominal */
+			off = att_align_nominal(off, thisatt->attalign);
+		}
+
+		off = att_addlength_pointer(off, thisatt->attlen, tp + off);
+
+		if (thisatt->attlen <= 0)
+			slow = true;		/* can't use attcacheoff anymore */
+	}
+
+	iter->isNull = null;
+	iter->offset = off;
+	iter->slow = slow;
+}
+
+/* ----------------
+ *		index_attiternext() - get the next attribute of an index tuple
+ *
+ *		This gets called many times, so we do the least amount of work
+ *		possible.
+ *
+ *		The code does not attempt to update attcacheoff; as it is unlikely
+ *		to reach a situation where the cached offset matters a lot.
+ *		If the cached offset do matter, the caller should make sure that
+ *		PopulateTupleDescCacheOffsets() was called on the tuple descriptor
+ *		to populate the attribute offset cache.
+ *
+ * ----------------
+ */
+static inline Datum
+index_attiternext(IndexTuple tup, AttrNumber attnum, TupleDesc tupleDesc, IAttrIterState iter)
+{
+	bool		hasnulls = IndexTupleHasNulls(tup);
+	char	   *tp;				/* ptr to tuple data */
+	bits8	   *bp;				/* ptr to null bitmap in tuple */
+	Datum		datum;
+	Form_pg_attribute thisatt = TupleDescAttr(tupleDesc, attnum - 1);
+
+	Assert(PointerIsValid(iter));
+	Assert(tupleDesc->natts <= INDEX_MAX_KEYS);
+	Assert(attnum <= tupleDesc->natts);
+	Assert(attnum > 0);
+
+	bp = (bits8 *) ((char *) tup + sizeof(IndexTupleData));
+
+	tp = (char *) tup + IndexInfoFindDataOffset(tup->t_info);
+
+	if (hasnulls && att_isnull(attnum - 1, bp))
+	{
+		iter->isNull = true;
+		iter->slow = true;
+		return (Datum) 0;
+	}
+
+	iter->isNull = false;
+
+	if (!iter->slow && thisatt->attcacheoff >= 0)
+		iter->offset = thisatt->attcacheoff;
+	else if (thisatt->attlen == -1)
+	{
+		iter->offset = att_align_pointer(iter->offset, thisatt->attalign, -1,
+										 tp + iter->offset);
+		iter->slow = true;
+	}
+	else
+	{
+		/* not varlena, so safe to use att_align_nominal */
+		iter->offset = att_align_nominal(iter->offset, thisatt->attalign);
+	}
+
+	datum = fetchatt(thisatt, tp + iter->offset);
+
+	iter->offset = att_addlength_pointer(iter->offset, thisatt->attlen, tp + iter->offset);
+
+	if (thisatt->attlen <= 0)
+		iter->slow = true;		/* can't use attcacheoff anymore */
+
+	return datum;
+}
+
 #endif							/* ITUP_H */
diff --git a/src/include/access/nbtree.h b/src/include/access/nbtree.h
index 1559399b0e..80f2575884 100644
--- a/src/include/access/nbtree.h
+++ b/src/include/access/nbtree.h
@@ -1122,6 +1122,7 @@ typedef struct BTOptions
  */
 #define NBTS_TYPE_SINGLE_COLUMN single
 #define NBTS_TYPE_CACHED cached
+#define NBTS_TYPE_UNCACHED uncached
 #define NBTS_TYPE_DEFAULT default
 
 
@@ -1152,12 +1153,19 @@ do { \
 
 #define NBT_SPECIALIZE_CALL(function, rel, ...) \
 ( \
-	IndexRelationGetNumberOfKeyAttributes(rel) == 1 ? (   \
+	IndexRelationGetNumberOfKeyAttributes(rel) == 1 ? ( \
 		NBTS_MAKE_NAME(function, NBTS_TYPE_SINGLE_COLUMN)(__VA_ARGS__) \
 	) \
 	: \
 	( \
-		NBTS_MAKE_NAME(function, NBTS_TYPE_CACHED)(__VA_ARGS__) \
+		TupleDescAttr(RelationGetDescr(rel), \
+					  IndexRelationGetNumberOfKeyAttributes(rel) - 1)->attcacheoff > 0 ? ( \
+			NBTS_MAKE_NAME(function, NBTS_TYPE_CACHED)(__VA_ARGS__) \
+		) \
+		: \
+		( \
+			NBTS_MAKE_NAME(function, NBTS_TYPE_UNCACHED)(__VA_ARGS__) \
+		) \
 	) \
 )
 
diff --git a/src/include/access/nbtree_specialize.h b/src/include/access/nbtree_specialize.h
index 642bc4c795..52739d390e 100644
--- a/src/include/access/nbtree_specialize.h
+++ b/src/include/access/nbtree_specialize.h
@@ -168,6 +168,40 @@
 #undef nbts_attiter_nextattdatum
 #undef nbts_attiter_curattisnull
 
+/*
+ * Multiple key columns, but attcacheoff -optimization doesn't apply.
+ */
+#define NBTS_SPECIALIZING_UNCACHED
+#define NBTS_TYPE NBTS_TYPE_UNCACHED
+
+#define nbts_attiterdeclare(itup) \
+	IAttrIterStateData	NBTS_MAKE_NAME(itup, iter)
+
+#define nbts_attiterinit(itup, initAttNum, tupDesc) \
+	index_attiterinit((itup), (initAttNum), (tupDesc), &(NBTS_MAKE_NAME(itup, iter)))
+
+#define nbts_foreachattr(initAttNum, endAttNum) \
+	for (int spec_i = (initAttNum); spec_i <= (endAttNum); spec_i++)
+
+#define nbts_attiter_attnum spec_i
+
+#define nbts_attiter_nextattdatum(itup, tupDesc) \
+	index_attiternext((itup), spec_i, (tupDesc), &(NBTS_MAKE_NAME(itup, iter)))
+
+#define nbts_attiter_curattisnull(itup) \
+	NBTS_MAKE_NAME(itup, iter).isNull
+
+#include NBT_SPECIALIZE_FILE
+
+#undef NBTS_TYPE
+#undef NBTS_SPECIALIZING_UNCACHED
+#undef nbts_attiterdeclare
+#undef nbts_attiterinit
+#undef nbts_foreachattr
+#undef nbts_attiter_attnum
+#undef nbts_attiter_nextattdatum
+#undef nbts_attiter_curattisnull
+
 /* reset call to SPECIALIZE_CALL for default behaviour */
 #undef nbts_call_norel
 #define nbts_call_norel(name, rel, ...) \
-- 
2.30.2

