From 3eebb28da2b1e2e8a130ee62b9e96981684eff95 Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <boekewurm+postgres@gmail.com>
Date: Wed, 31 Jan 2024 01:40:40 +0100
Subject: [PATCH v16 3/6] btree: Introduce attribute iterator specialization
 infrastructure

This allows the specialization framework to utilize different optimized
attribute iteration methods based on key shape, e.g. one that utilizes
attcacheoff; one that doens't need to calculate offset because of the
use of only one key attribute, or one that incrementally finds next
offsets for each attribute in the tuple.

Even though not all nbt*_spec functions have been updated, these functions
can now directly call (and potentially inline, and optimize for) the
specialized functions they call, instead of having to determine the right
specialization locally based on the (potentially locally unavailable) index
relation, making the specialization of those functions still worth
specializing/duplicating.
---
 src/backend/access/nbtree/nbtsearch_spec.c    | 17 +++--
 src/backend/access/nbtree/nbtsort_spec.c      | 24 +++----
 src/backend/access/nbtree/nbtutils_spec.c     | 62 ++++++++++++-------
 .../utils/sort/tuplesortvariants_spec.c       | 53 +++++++++-------
 src/include/access/nbtree_spec.h              | 49 +++++++++++++++
 5 files changed, 141 insertions(+), 64 deletions(-)

diff --git a/src/backend/access/nbtree/nbtsearch_spec.c b/src/backend/access/nbtree/nbtsearch_spec.c
index d65765ca23..cf5e559c28 100644
--- a/src/backend/access/nbtree/nbtsearch_spec.c
+++ b/src/backend/access/nbtree/nbtsearch_spec.c
@@ -581,6 +581,7 @@ _bt_compare(Relation rel,
 	int			ncmpkey;
 	int			ntupatts;
 	int32		result;
+	nbts_attiterdeclare(itup);
 
 	Assert(_bt_check_natts(rel, key->heapkeyspace, page, offnum));
 	Assert(key->keysz <= IndexRelationGetNumberOfKeyAttributes(rel));
@@ -611,24 +612,28 @@ _bt_compare(Relation rel,
 	ncmpkey = Min(ntupatts, key->keysz);
 	Assert(key->heapkeyspace || ncmpkey == key->keysz);
 	Assert(!BTreeTupleIsPosting(itup) || key->allequalimage);
+
 	scankey = key->scankeys;
-	for (int i = 1; i <= ncmpkey; i++)
+	nbts_attiterinit(itup, 1, itupdesc);
+
+	nbts_foreachattr(1, ncmpkey)
 	{
 		Datum		datum;
-		bool		isNull;
 
-		datum = index_getattr(itup, scankey->sk_attno, itupdesc, &isNull);
+		datum = nbts_attiter_nextattdatum(itup, itupdesc);
 
-		if (scankey->sk_flags & SK_ISNULL)	/* key is NULL */
+		/* key is NULL */
+		if (scankey->sk_flags & SK_ISNULL)
 		{
-			if (isNull)
+			if (nbts_attiter_curattisnull(itup))
 				result = 0;		/* NULL "=" NULL */
 			else if (scankey->sk_flags & SK_BT_NULLS_FIRST)
 				result = -1;	/* NULL "<" NOT_NULL */
 			else
 				result = 1;		/* NULL ">" NOT_NULL */
 		}
-		else if (isNull)		/* key is NOT_NULL and item is NULL */
+		/* key is NOT_NULL and item is NULL */
+		else if (nbts_attiter_curattisnull(itup))
 		{
 			if (scankey->sk_flags & SK_BT_NULLS_FIRST)
 				result = 1;		/* NOT_NULL ">" NULL */
diff --git a/src/backend/access/nbtree/nbtsort_spec.c b/src/backend/access/nbtree/nbtsort_spec.c
index 77829a23e0..8f5c255981 100644
--- a/src/backend/access/nbtree/nbtsort_spec.c
+++ b/src/backend/access/nbtree/nbtsort_spec.c
@@ -34,8 +34,7 @@ _bt_load(BTWriteState *wstate, BTSpool *btspool, BTSpool *btspool2)
 				itup2 = NULL;
 	bool		load1;
 	TupleDesc	tupdes = RelationGetDescr(wstate->index);
-	int			i,
-				keysz = IndexRelationGetNumberOfKeyAttributes(wstate->index);
+	int			keysz = IndexRelationGetNumberOfKeyAttributes(wstate->index);
 	SortSupport sortKeys;
 	int64		tuples_done = 0;
 	bool		deduplicate;
@@ -59,7 +58,7 @@ _bt_load(BTWriteState *wstate, BTSpool *btspool, BTSpool *btspool2)
 		/* Prepare SortSupport data for each column */
 		sortKeys = (SortSupport) palloc0(keysz * sizeof(SortSupportData));
 
-		for (i = 0; i < keysz; i++)
+		for (int i = 0; i < keysz; i++)
 		{
 			SortSupport sortKey = sortKeys + i;
 			ScanKey		scanKey = wstate->inskey->scankeys + i;
@@ -92,21 +91,24 @@ _bt_load(BTWriteState *wstate, BTSpool *btspool, BTSpool *btspool2)
 			else if (itup != NULL)
 			{
 				int32		compare = 0;
+				nbts_attiterdeclare(itup);
+				nbts_attiterdeclare(itup2);
 
-				for (i = 1; i <= keysz; i++)
+				nbts_attiterinit(itup, 1, tupdes);
+				nbts_attiterinit(itup2, 1, tupdes);
+
+				nbts_foreachattr(1, keysz)
 				{
 					SortSupport entry;
 					Datum		attrDatum1,
 								attrDatum2;
-					bool		isNull1,
-								isNull2;
 
-					entry = sortKeys + i - 1;
-					attrDatum1 = index_getattr(itup, i, tupdes, &isNull1);
-					attrDatum2 = index_getattr(itup2, i, tupdes, &isNull2);
+					entry = sortKeys + nbts_attiter_attnum - 1;
+					attrDatum1 = nbts_attiter_nextattdatum(itup, tupdes);
+					attrDatum2 = nbts_attiter_nextattdatum(itup2, tupdes);
 
-					compare = ApplySortComparator(attrDatum1, isNull1,
-												  attrDatum2, isNull2,
+					compare = ApplySortComparator(attrDatum1, nbts_attiter_curattisnull(itup),
+												  attrDatum2, nbts_attiter_curattisnull(itup2),
 												  entry);
 					if (compare > 0)
 					{
diff --git a/src/backend/access/nbtree/nbtutils_spec.c b/src/backend/access/nbtree/nbtutils_spec.c
index 38e27f1757..fad68250de 100644
--- a/src/backend/access/nbtree/nbtutils_spec.c
+++ b/src/backend/access/nbtree/nbtutils_spec.c
@@ -56,7 +56,7 @@ _bt_mkscankey(Relation rel, IndexTuple itup)
 	int			indnkeyatts;
 	int16	   *indoption;
 	int			tupnatts;
-	int			i;
+	nbts_attiterdeclare(itup);
 
 	itupdesc = RelationGetDescr(rel);
 	indnkeyatts = IndexRelationGetNumberOfKeyAttributes(rel);
@@ -87,7 +87,10 @@ _bt_mkscankey(Relation rel, IndexTuple itup)
 	key->scantid = key->heapkeyspace && itup ?
 		BTreeTupleGetHeapTID(itup) : NULL;
 	skey = key->scankeys;
-	for (i = 0; i < indnkeyatts; i++)
+
+	nbts_attiterinit(itup, 1, itupdesc);
+
+	nbts_foreachattr(1, indnkeyatts)
 	{
 		FmgrInfo   *procinfo;
 		Datum		arg;
@@ -98,27 +101,30 @@ _bt_mkscankey(Relation rel, IndexTuple itup)
 		 * We can use the cached (default) support procs since no cross-type
 		 * comparison can be needed.
 		 */
-		procinfo = index_getprocinfo(rel, i + 1, BTORDER_PROC);
+		procinfo = index_getprocinfo(rel, nbts_attiter_attnum, BTORDER_PROC);
 
 		/*
 		 * Key arguments built from truncated attributes (or when caller
 		 * provides no tuple) are defensively represented as NULL values. They
 		 * should never be used.
 		 */
-		if (i < tupnatts)
-			arg = index_getattr(itup, i + 1, itupdesc, &null);
+		if (nbts_attiter_attnum <= tupnatts)
+		{
+			arg = nbts_attiter_nextattdatum(itup, itupdesc);
+			null = nbts_attiter_curattisnull(itup);
+		}
 		else
 		{
 			arg = (Datum) 0;
 			null = true;
 		}
-		flags = (null ? SK_ISNULL : 0) | (indoption[i] << SK_BT_INDOPTION_SHIFT);
-		ScanKeyEntryInitializeWithInfo(&skey[i],
+		flags = (null ? SK_ISNULL : 0) | (indoption[nbts_attiter_attnum - 1] << SK_BT_INDOPTION_SHIFT);
+		ScanKeyEntryInitializeWithInfo(&skey[nbts_attiter_attnum - 1],
 									   flags,
-									   (AttrNumber) (i + 1),
+									   (AttrNumber) nbts_attiter_attnum,
 									   InvalidStrategy,
 									   InvalidOid,
-									   rel->rd_indcollation[i],
+									   rel->rd_indcollation[nbts_attiter_attnum - 1],
 									   procinfo,
 									   arg);
 		/* Record if any key attribute is NULL (or truncated) */
@@ -712,6 +718,8 @@ _bt_keep_natts(Relation rel, IndexTuple lastleft, IndexTuple firstright,
 	TupleDesc	itupdesc = RelationGetDescr(rel);
 	int			keepnatts;
 	ScanKey		scankey;
+	nbts_attiterdeclare(lastleft);
+	nbts_attiterdeclare(firstright);
 
 	/*
 	 * _bt_compare() treats truncated key attributes as having the value minus
@@ -723,20 +731,22 @@ _bt_keep_natts(Relation rel, IndexTuple lastleft, IndexTuple firstright,
 
 	scankey = itup_key->scankeys;
 	keepnatts = 1;
-	for (int attnum = 1; attnum <= nkeyatts; attnum++, scankey++)
+
+	nbts_attiterinit(lastleft, 1, itupdesc);
+	nbts_attiterinit(firstright, 1, itupdesc);
+
+	nbts_foreachattr(1, nkeyatts)
 	{
 		Datum		datum1,
 					datum2;
-		bool		isNull1,
-					isNull2;
 
-		datum1 = index_getattr(lastleft, attnum, itupdesc, &isNull1);
-		datum2 = index_getattr(firstright, attnum, itupdesc, &isNull2);
+		datum1 = nbts_attiter_nextattdatum(lastleft, itupdesc);
+		datum2 = nbts_attiter_nextattdatum(firstright, itupdesc);
 
-		if (isNull1 != isNull2)
+		if (nbts_attiter_curattisnull(lastleft) != nbts_attiter_curattisnull(firstright))
 			break;
 
-		if (!isNull1 &&
+		if (!nbts_attiter_curattisnull(lastleft) &&
 			DatumGetInt32(FunctionCall2Coll(&scankey->sk_func,
 											scankey->sk_collation,
 											datum1,
@@ -744,6 +754,7 @@ _bt_keep_natts(Relation rel, IndexTuple lastleft, IndexTuple firstright,
 			break;
 
 		keepnatts++;
+		scankey++;
 	}
 
 	/*
@@ -784,24 +795,27 @@ _bt_keep_natts_fast(Relation rel, IndexTuple lastleft, IndexTuple firstright)
 	TupleDesc	itupdesc = RelationGetDescr(rel);
 	int			keysz = IndexRelationGetNumberOfKeyAttributes(rel);
 	int			keepnatts;
+	nbts_attiterdeclare(lastleft);
+	nbts_attiterdeclare(firstright);
 
 	keepnatts = 1;
-	for (int attnum = 1; attnum <= keysz; attnum++)
+	nbts_attiterinit(lastleft, 1, itupdesc);
+	nbts_attiterinit(firstright, 1, itupdesc);
+
+	nbts_foreachattr(1, keysz)
 	{
 		Datum		datum1,
 					datum2;
-		bool		isNull1,
-					isNull2;
 		Form_pg_attribute att;
 
-		datum1 = index_getattr(lastleft, attnum, itupdesc, &isNull1);
-		datum2 = index_getattr(firstright, attnum, itupdesc, &isNull2);
-		att = TupleDescAttr(itupdesc, attnum - 1);
+		datum1 = nbts_attiter_nextattdatum(lastleft, itupdesc);
+		datum2 = nbts_attiter_nextattdatum(firstright, itupdesc);
+		att = TupleDescAttr(itupdesc, nbts_attiter_attnum - 1);
 
-		if (isNull1 != isNull2)
+		if (nbts_attiter_curattisnull(lastleft) != nbts_attiter_curattisnull(firstright))
 			break;
 
-		if (!isNull1 &&
+		if (!nbts_attiter_curattisnull(lastleft) &&
 			!datum_image_eq(datum1, datum2, att->attbyval, att->attlen))
 			break;
 
diff --git a/src/backend/utils/sort/tuplesortvariants_spec.c b/src/backend/utils/sort/tuplesortvariants_spec.c
index 705da09329..cf262eee2d 100644
--- a/src/backend/utils/sort/tuplesortvariants_spec.c
+++ b/src/backend/utils/sort/tuplesortvariants_spec.c
@@ -66,47 +66,54 @@ comparetup_index_btree_tiebreak(const SortTuple *a, const SortTuple *b,
 	bool		equal_hasnull = false;
 	int			nkey;
 	int32		compare;
-	Datum		datum1,
-				datum2;
-	bool		isnull1,
-				isnull2;
+	nbts_attiterdeclare(tuple1);
+	nbts_attiterdeclare(tuple2);
 
 	tuple1 = (IndexTuple) a->tuple;
 	tuple2 = (IndexTuple) b->tuple;
 	keysz = base->nKeys;
 	tupDes = RelationGetDescr(arg->index.indexRel);
 
-	if (sortKey->abbrev_converter)
+	if (!sortKey->abbrev_converter)
 	{
-		datum1 = index_getattr(tuple1, 1, tupDes, &isnull1);
-		datum2 = index_getattr(tuple2, 1, tupDes, &isnull2);
-
-		compare = ApplySortAbbrevFullComparator(datum1, isnull1,
-												datum2, isnull2,
-												sortKey);
-		if (compare != 0)
-			return compare;
+		nkey = 2;
+		sortKey++;
+	}
+	else
+	{
+		nkey = 1;
 	}
 
 	/* they are equal, so we only need to examine one null flag */
 	if (a->isnull1)
 		equal_hasnull = true;
 
-	sortKey++;
-	for (nkey = 2; nkey <= keysz; nkey++, sortKey++)
+	nbts_attiterinit(tuple1, nkey, tupDes);
+	nbts_attiterinit(tuple2, nkey, tupDes);
+
+	nbts_foreachattr(nkey, keysz)
 	{
-		datum1 = index_getattr(tuple1, nkey, tupDes, &isnull1);
-		datum2 = index_getattr(tuple2, nkey, tupDes, &isnull2);
+		Datum	datum1,
+				datum2;
+		datum1 = nbts_attiter_nextattdatum(tuple1, tupDes);
+		datum2 = nbts_attiter_nextattdatum(tuple2, tupDes);
+
+		if (nbts_attiter_attnum == 1)
+			compare = ApplySortAbbrevFullComparator(datum1, nbts_attiter_curattisnull(tuple1),
+													datum2, nbts_attiter_curattisnull(tuple2),
+													sortKey);
+		else
+			compare = ApplySortComparator(datum1, nbts_attiter_curattisnull(tuple1),
+										  datum2, nbts_attiter_curattisnull(tuple2),
+										  sortKey);
 
-		compare = ApplySortComparator(datum1, isnull1,
-									  datum2, isnull2,
-									  sortKey);
 		if (compare != 0)
-			return compare;		/* done when we find unequal attributes */
+			return compare;
 
-		/* they are equal, so we only need to examine one null flag */
-		if (isnull1)
+		if (nbts_attiter_curattisnull(tuple1))
 			equal_hasnull = true;
+
+		sortKey++;
 	}
 
 	/*
diff --git a/src/include/access/nbtree_spec.h b/src/include/access/nbtree_spec.h
index 731f2fb946..c16659e603 100644
--- a/src/include/access/nbtree_spec.h
+++ b/src/include/access/nbtree_spec.h
@@ -20,6 +20,33 @@
  * Specialized files define their functions like usual, with an additional
  * requirement that they #define their function name at the top, with
  * #define funcname NBTS_FUNCTION(funcname).
+ *
+ * Key attribute iteration is specialized through the use of the following
+ * macros:
+ *
+ * - nbts_attiterdeclare(itup)
+ *   Declare the variables required to iterate over the provided IndexTuple's
+ *   key attributes. Many tuples may have their attributes iterated over at the
+ *   same time.
+ * - nbts_attiterinit(itup, initAttNum, tupDesc)
+ *   Initialize the attribute iterator for the provided IndexTuple at
+ *   the provided AttributeNumber.
+ * - nbts_foreachattr(initAttNum, endAttNum)
+ *   Start a loop over the attributes, starting at initAttNum and ending at
+ *   endAttNum, inclusive. It also takes care of truncated attributes.
+ * - nbts_attiter_attnum
+ *   The current attribute number
+ * - nbts_attiter_nextattdatum(itup, tupDesc)
+ *   Updates the attribute iterator state to the next attribute. Returns the
+ *   datum of the next attribute, which might be null (see below)
+ * - nbts_attiter_curattisnull(itup)
+ *   Returns whether the result from the last nbts_attiter_nextattdatum is
+ *   null.
+ * - nbts_context(irel)
+ *   Constructs a context that is used to call specialized functions.
+ *   Note that this is unneeded in paths that are inaccessible to unspecialized
+ *   code paths (i.e. code included through nbtree_spec.h), because that
+ *   always calls the optimized functions directly.
  */
 
 #if defined(NBT_FILE)
@@ -57,9 +84,31 @@
  */
 #define NBTS_TYPE NBTS_TYPE_CACHED
 
+#define nbts_attiterdeclare(itup) \
+	bool	NBTS_MAKE_NAME(itup, isNull)
+
+#define nbts_attiterinit(itup, initAttNum, tupDesc) do {} while (false)
+
+#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_getattr((itup), spec_i, (tupDesc), &(NBTS_MAKE_NAME(itup, isNull)))
+
+#define nbts_attiter_curattisnull(itup) \
+	NBTS_MAKE_NAME(itup, isNull)
+
 #include NBT_SPECIALIZE_FILE
 
 #undef NBTS_TYPE
+#undef nbts_attiterdeclare
+#undef nbts_attiterinit
+#undef nbts_foreachattr
+#undef nbts_attiter_attnum
+#undef nbts_attiter_nextattdatum
+#undef nbts_attiter_curattisnull
 
 /*
  * We're done with templating, so restore or create the macros which can be
-- 
2.40.1

