From 1609509386ef18d5a22ebc79c9075e34573f0640 Mon Sep 17 00:00:00 2001
From: Matthias van de Meent <boekewurm+postgres@gmail.com>
Date: Thu, 7 Apr 2022 12:47:50 +0200
Subject: [PATCH v6 4/8] Optimize attribute iterator access for single-column
 btree keys

This removes the index_getattr_nocache call path, which has significant overhead.
---
 src/include/access/nbtree.h            |  9 +++-
 src/include/access/nbtree_specialize.h | 63 ++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/include/access/nbtree.h b/src/include/access/nbtree.h
index 489b623663..1559399b0e 100644
--- a/src/include/access/nbtree.h
+++ b/src/include/access/nbtree.h
@@ -1120,6 +1120,7 @@ typedef struct BTOptions
 /*
  * Macros used in the nbtree specialization code.
  */
+#define NBTS_TYPE_SINGLE_COLUMN single
 #define NBTS_TYPE_CACHED cached
 #define NBTS_TYPE_DEFAULT default
 
@@ -1151,7 +1152,13 @@ do { \
 
 #define NBT_SPECIALIZE_CALL(function, rel, ...) \
 ( \
-	NBTS_MAKE_NAME(function, NBTS_TYPE_CACHED)(__VA_ARGS__) \
+	IndexRelationGetNumberOfKeyAttributes(rel) == 1 ? (   \
+		NBTS_MAKE_NAME(function, NBTS_TYPE_SINGLE_COLUMN)(__VA_ARGS__) \
+	) \
+	: \
+	( \
+		NBTS_MAKE_NAME(function, NBTS_TYPE_CACHED)(__VA_ARGS__) \
+	) \
 )
 
 #else /* not defined NBTS_ENABLED */
diff --git a/src/include/access/nbtree_specialize.h b/src/include/access/nbtree_specialize.h
index 23fdda4f0e..9733a27bdd 100644
--- a/src/include/access/nbtree_specialize.h
+++ b/src/include/access/nbtree_specialize.h
@@ -79,6 +79,69 @@
 #define nbts_call_norel(name, rel, ...) \
 	(NBTS_FUNCTION(name)(__VA_ARGS__))
 
+/*
+ * 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 does not ever continue to a next column.
+ */
+
+#define NBTS_SPECIALIZING_SINGLE_COLUMN
+#define NBTS_TYPE NBTS_TYPE_SINGLE_COLUMN
+
+#define nbts_attiterdeclare(itup) \
+	bool	NBTS_MAKE_NAME(itup, isNull)
+
+#define nbts_attiterinit(itup, initAttNum, tupDesc)
+
+/*
+ * We void endAttNum to prevent unused variable warnings.
+ * The if- and for-loop are structured like this to make the compiler
+ * unroll the loop and detect only one single iteration. We need `break`
+ * in the following code block, so just a plain 'if' statement would
+ * not work.
+ */
+#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
+
+/*
+ * Simplified (optimized) variant of index_getattr specialized for extracting
+ * only the first attribute: cache offset is guaranteed to be 0, and as such
+ * no cache is required.
+ */
+#define nbts_attiter_nextattdatum(itup, tupDesc) \
+( \
+	AssertMacro(spec_i == 0), \
+	(IndexTupleHasNulls(itup) && att_isnull(0, (char *)(itup) + sizeof(IndexTupleData))) ? \
+	( \
+		(NBTS_MAKE_NAME(itup, isNull)) = true, \
+		(Datum)NULL \
+	) \
+	: \
+	( \
+		(NBTS_MAKE_NAME(itup, isNull) = false), \
+		(Datum) fetchatt(TupleDescAttr((tupDesc), 0), \
+		(char *) (itup) + IndexInfoFindDataOffset((itup)->t_info)) \
+	) \
+)
+
+#define nbts_attiter_curattisnull(tuple) \
+	NBTS_MAKE_NAME(tuple, isNull)
+
+#include NBT_SPECIALIZE_FILE
+
+#undef NBTS_TYPE
+#undef NBTS_SPECIALIZING_SINGLE_COLUMN
+#undef nbts_attiterdeclare
+#undef nbts_attiterinit
+#undef nbts_foreachattr
+#undef nbts_attiter_attnum
+#undef nbts_attiter_nextattdatum
+#undef nbts_attiter_curattisnull
+
 /*
  * Multiple key columns, optimized access for attcacheoff -cacheable offsets.
  */
-- 
2.30.2