From e8d891adb0f4bdcf091edc5e02155303503ba603 Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <boekewurm+postgres@gmail.com>
Date: Wed, 11 Jan 2023 20:04:56 +0100
Subject: [PATCH v10 4/6] Optimize nbts_attiter for nkeyatts==1 btrees

This removes the index_getattr_nocache call path, which has significant overhead, and uses constant 0 offset.
---
 src/backend/access/nbtree/README        |  1 +
 src/backend/access/nbtree/nbtree_spec.c |  3 ++
 src/include/access/nbtree.h             | 35 ++++++++++++++++
 src/include/access/nbtree_spec.h        | 56 ++++++++++++++++++++++++-
 4 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/src/backend/access/nbtree/README b/src/backend/access/nbtree/README
index 4b11ea9ad7..6864902637 100644
--- a/src/backend/access/nbtree/README
+++ b/src/backend/access/nbtree/README
@@ -1104,6 +1104,7 @@ in the index AM to call the specialized functions, increasing the
 performance of those hot paths.
 
 Optimized code paths exist for the following cases, in order of preference:
+ - indexes with only a single key attribute
  - multi-column indexes that could benefit from the attcacheoff optimization
    NB: This is also the default path, and is comparatively slow for uncachable
    attribute offsets.
diff --git a/src/backend/access/nbtree/nbtree_spec.c b/src/backend/access/nbtree/nbtree_spec.c
index 6b766581ab..21635397ed 100644
--- a/src/backend/access/nbtree/nbtree_spec.c
+++ b/src/backend/access/nbtree/nbtree_spec.c
@@ -33,6 +33,9 @@ _bt_specialize(Relation rel)
 		case NBTS_CTX_CACHED:
 			_bt_specialize_cached(rel);
 			break;
+		case NBTS_CTX_SINGLE_KEYATT:
+			_bt_specialize_single_keyatt(rel);
+			break;
 		case NBTS_CTX_DEFAULT:
 			break;
 	}
diff --git a/src/include/access/nbtree.h b/src/include/access/nbtree.h
index f3f0961052..4628c41e9a 100644
--- a/src/include/access/nbtree.h
+++ b/src/include/access/nbtree.h
@@ -1123,6 +1123,7 @@ typedef struct BTOptions
 #define PROGRESS_BTREE_PHASE_LEAF_LOAD					5
 
 typedef enum NBTS_CTX {
+	NBTS_CTX_SINGLE_KEYATT, 
 	NBTS_CTX_CACHED, 
 	NBTS_CTX_DEFAULT, /* fallback */
 } NBTS_CTX;
@@ -1132,9 +1133,43 @@ static inline NBTS_CTX _nbt_spec_context(Relation irel)
 	if (!PointerIsValid(irel))
 		return NBTS_CTX_DEFAULT;
 
+	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 NBT_SPECIALIZE_FILE "access/nbtree_specfuncs.h"
 #include "nbtree_spec.h"
diff --git a/src/include/access/nbtree_spec.h b/src/include/access/nbtree_spec.h
index fa38b09c6e..8e476c300d 100644
--- a/src/include/access/nbtree_spec.h
+++ b/src/include/access/nbtree_spec.h
@@ -44,6 +44,7 @@
 /*
  * Macros used in the nbtree specialization code.
  */
+#define NBTS_TYPE_SINGLE_KEYATT single_keyatt
 #define NBTS_TYPE_CACHED cached
 #define NBTS_TYPE_DEFAULT default
 #define NBTS_CTX_NAME __nbts_ctx
@@ -51,8 +52,10 @@
 /* contextual specializations */
 #define NBTS_MAKE_CTX(rel) const NBTS_CTX NBTS_CTX_NAME = _nbt_spec_context(rel)
 #define NBTS_SPECIALIZE_NAME(name) ( \
-	(NBTS_CTX_NAME) == NBTS_CTX_CACHED ? (NBTS_MAKE_NAME(name, NBTS_TYPE_CACHED)) : ( \
-		NBTS_MAKE_NAME(name, NBTS_TYPE_DEFAULT) \
+	(NBTS_CTX_NAME) == NBTS_CTX_SINGLE_KEYATT ? (NBTS_MAKE_NAME(name, NBTS_TYPE_SINGLE_KEYATT)) : ( \
+		(NBTS_CTX_NAME) == NBTS_CTX_CACHED ? (NBTS_MAKE_NAME(name, NBTS_TYPE_CACHED)) : ( \
+			NBTS_MAKE_NAME(name, NBTS_TYPE_DEFAULT) \
+		) \
 	) \
 )
 
@@ -164,6 +167,55 @@ do { \
 #undef nbts_attiter_nextattdatum
 #undef nbts_attiter_curattisnull
 
+/*
+ * Specialization 3: SINGLE_KEYATT
+ *
+ * Optimized access for indexes with a single key column.
+ *
+ * Note that this path cannot be used for indexes with multiple key
++ *  columns, because it never considers the next column.
+ */
+
+/* the default context (and later contexts) do need to specialize, so here's that */
+#undef nbts_prep_ctx
+#define nbts_prep_ctx(rel)
+
+#define NBTS_SPECIALIZING_SINGLE_KEYATT
+#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
+#undef NBTS_SPECIALIZING_SINGLE_KEYATT
+
+/* 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
+
 /*
  * All next uses of nbts_prep_ctx are in non-templated code, so here we make
  * sure we actually create the context.
-- 
2.39.0

