From 3883b43328abceb1b10bc5d9463982ed6a381e5f Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <boekewurm+postgres@gmail.com>
Date: Wed, 31 Jan 2024 02:06:41 +0100
Subject: [PATCH v15 5/6] btree: Optimize nbts_attiter for indexes with a
 single key attribute

This removes the index_getattr_nocache call path, which can have significant
overhead with branching, and replaces access to att->attcacheoff=0 with
constant 0.
---
 src/include/access/nbtree.h      | 46 ++++++++++++++++++++++++++++++--
 src/include/access/nbtree_spec.h | 43 +++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/src/include/access/nbtree.h b/src/include/access/nbtree.h
index e1e80a8830..9579fed492 100644
--- a/src/include/access/nbtree.h
+++ b/src/include/access/nbtree.h
@@ -1119,6 +1119,7 @@ typedef struct BTOptions
 #define PROGRESS_BTREE_PHASE_LEAF_LOAD					5
 
 #define NBTS_TYPE_CACHED CACHED
+#define NBTS_TYPE_SINGLE_KEYATT SINGLE_KEYATT
 
 /* Generate the specialized function headers */
 #define NBT_FILE "access/nbtree_specfuncs.h"
@@ -1130,20 +1131,61 @@ typedef struct BTOptions
 
 typedef enum NBTS_CTX {
 	NBTS_CTX_CACHED = 1, /* Equivalent to unspecialized code */
+	NBTS_CTX_SINGLE_KEYATT, /* Single key column index; cannot handle more */
 } NBTS_CTX;
 
 static inline NBTS_CTX _nbt_spec_context(Relation irel)
 {
+	Assert(PointerIsValid(irel));
+
+	if (IndexRelationGetNumberOfKeyAttributes(irel) == 1)
+		return NBTS_CTX_SINGLE_KEYATT;
+
 	return NBTS_CTX_CACHED;
 }
 
+static inline Datum _bt_getfirstatt(IndexTuple tuple, TupleDesc tupleDesc,
+									bool *isNull)
+{
+	Datum	result;
+	if (IndexTupleHasNulls(tuple))
+	{
+		if (att_isnull(0, (bits8 *)(tuple) + sizeof(IndexTupleData)))
+		{
+			*isNull = true;
+			result = (Datum) 0;
+		}
+		else
+		{
+			*isNull = false;
+			result = fetchatt(TupleDescAttr(tupleDesc, 0),
+							  ((char *) tuple)
+							  + MAXALIGN(sizeof(IndexTupleData)
+							  + sizeof(IndexAttributeBitMapData)));
+		}
+	}
+	else
+	{
+		*isNull = false;
+		result = fetchatt(TupleDescAttr(tupleDesc, 0),
+						  ((char *) tuple)
+						  + MAXALIGN(sizeof(IndexTupleData)));
+	}
+
+	return result;
+}
+
 #define NBTS_CTX_NAME __nbts_ctx
 
 /* contextual specialization macros */
 #define NBTS_MAKE_CTX(rel) const NBTS_CTX NBTS_CTX_NAME PG_USED_FOR_ASSERTS_ONLY = _nbt_spec_context(rel)
 #define NBTS_SPECIALIZE_NAME(name) ( \
-	AssertMacro((NBTS_CTX_NAME) == NBTS_CTX_CACHED), \
-	NBTS_MAKE_NAME(name, NBTS_TYPE_CACHED) \
+	(NBTS_CTX_NAME == NBTS_CTX_SINGLE_KEYATT) ? ( \
+		NBTS_MAKE_NAME(name, NBTS_TYPE_SINGLE_KEYATT) \
+	) : ( \
+		AssertMacro((NBTS_CTX_NAME) == NBTS_CTX_CACHED), \
+		NBTS_MAKE_NAME(name, NBTS_TYPE_CACHED) \
+	) \
 )
 
 extern bool _btinsert_dispatch(Relation rel, Datum *values, bool *isnull,
diff --git a/src/include/access/nbtree_spec.h b/src/include/access/nbtree_spec.h
index c16659e603..4919f95e79 100644
--- a/src/include/access/nbtree_spec.h
+++ b/src/include/access/nbtree_spec.h
@@ -110,6 +110,49 @@
 #undef nbts_attiter_nextattdatum
 #undef nbts_attiter_curattisnull
 
+/*
+ * Specialization 2: SINGLE_KEYATT
+ *
+ * Optimized access for indexes with a single key column.
+ *
+ * Note that this path may never be used for indexes with multiple key
+ * columns, because it only considers the first column (plus any TIDs for tie
+ * breaking).
+ */
+#define NBTS_TYPE NBTS_TYPE_SINGLE_KEYATT
+
+#define nbts_attiterdeclare(itup) \
+	bool	NBTS_MAKE_NAME(itup, isNull)
+
+#define nbts_attiterinit(itup, initAttNum, tupDesc)
+
+#define nbts_foreachattr(initAttNum, endAttNum) \
+	Assert((endAttNum) == 1); ((void) (endAttNum)); \
+	if ((initAttNum) == 1) for (int spec_i = 0; spec_i < 1; spec_i++)
+
+#define nbts_attiter_attnum 1
+
+#define nbts_attiter_nextattdatum(itup, tupDesc) \
+( \
+	AssertMacro(spec_i == 0), \
+	_bt_getfirstatt(itup, tupDesc, &NBTS_MAKE_NAME(itup, isNull)) \
+)
+
+#define nbts_attiter_curattisnull(itup) \
+	NBTS_MAKE_NAME(itup, isNull)
+
+#include NBT_SPECIALIZE_FILE
+
+#undef NBTS_TYPE
+
+/* un-define the optimization macros */
+#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
  * used in non-template sources.
-- 
2.42.1