0004-Use-SK_SEARCHARRAY-for-batched-fast-path-FK-probes.patch

From a803c6dce04b122822e05768bf2290ac9275824b Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Wed, 25 Feb 2026 21:25:25 +0900
Subject: [PATCH 4/4] Use SK_SEARCHARRAY for batched fast-path FK probes

For single-column foreign keys, replace the per-row index probe loop
in ri_FastPathBatchFlush() with a single SK_SEARCHARRAY scan key.
The btree AM sorts and deduplicates the array internally, then walks
the matching leaf pages in one ordered traversal instead of descending
from the root once per row.

ri_FastPathBatchFlush() now dispatches to ri_FastPathFlushArray() for
single-column FKs and ri_FastPathFlushLoop() for multi-column FKs,
which retains the per-row probe loop.

ri_FastPathFlushArray() builds an ArrayType from the buffered FK
values (casting to the PK-side type if needed), constructs a scan key
with the SK_SEARCHARRAY flag, and iterates the matches.  Each matched
PK tuple is locked and rechecked as before.  A matched[] bitmap tracks
which batch items were satisfied; unmatched items are reported as
violations.

With a batch size of 64 and int/int FK, this gives a 3.3x speedup
over unpatched master (vs 2.2x with per-row probing alone).
---
 src/backend/utils/adt/ri_triggers.c       | 288 +++++++++++++++++-----
 src/test/regress/expected/foreign_key.out |  17 ++
 src/test/regress/sql/foreign_key.sql      |  17 ++
 3 files changed, 262 insertions(+), 60 deletions(-)

diff --git a/src/backend/utils/adt/ri_triggers.c b/src/backend/utils/adt/ri_triggers.c
index d27d82a1e9f..b6a2ad4dc96 100644
--- a/src/backend/utils/adt/ri_triggers.c
+++ b/src/backend/utils/adt/ri_triggers.c
@@ -224,6 +224,10 @@ typedef struct RI_FastPathEntry
 
 	/* For ri_FastPathEndBatch() */
 	const RI_ConstraintInfo *riinfo;
+
+	/* For ri_FastPathFlushArray() */
+	Datum		search_vals[RI_FASTPATH_BATCH_SIZE];
+	bool		matched[RI_FASTPATH_BATCH_SIZE];
 } RI_FastPathEntry;
 
 /*
@@ -315,6 +319,10 @@ static void ri_FastPathEndBatch(void *arg);
 static void ri_FastPathTeardown(void);
 static void ri_FastPathBatchAdd(const RI_ConstraintInfo *riinfo,
 								Relation fk_rel, TupleTableSlot *newslot);
+static void ri_FastPathFlushArray(RI_FastPathEntry *fpentry, TupleTableSlot *fk_slot,
+								  const RI_ConstraintInfo *riinfo, Relation fk_rel);
+static void ri_FastPathFlushLoop(RI_FastPathEntry *fpentry, TupleTableSlot *fk_slot,
+								 const RI_ConstraintInfo *riinfo, Relation fk_rel);
 static void ri_FastPathBatchFlush(RI_FastPathEntry *fpentry,
 								  Relation fk_rel);
 
@@ -4036,102 +4044,262 @@ ri_FastPathGetEntry(const RI_ConstraintInfo *riinfo)
 	return entry;
 }
 
+/*
+ * ri_FastPathFlushLoop
+ *		Multi-column fallback: probe the index once per buffered row.
+ *
+ * Used for composite foreign keys where SK_SEARCHARRAY does not
+ * apply.
+ */
 static void
-ri_FastPathBatchFlush(RI_FastPathEntry *fpentry, Relation fk_rel)
+ri_FastPathFlushLoop(RI_FastPathEntry *fpentry, TupleTableSlot *fk_slot,
+					 const RI_ConstraintInfo *riinfo, Relation fk_rel)
 {
-	const RI_ConstraintInfo *riinfo = fpentry->riinfo;
 	Relation	pk_rel = fpentry->pk_rel;
 	Relation	idx_rel = fpentry->idx_rel;
 	IndexScanDesc scandesc = fpentry->scandesc;
 	TupleTableSlot *slot = fpentry->slot;
 	Snapshot	snapshot = fpentry->snapshot;
-	TupleTableSlot *fk_slot;
 	Datum		pk_vals[INDEX_MAX_KEYS];
 	char		pk_nulls[INDEX_MAX_KEYS];
 	ScanKeyData skey[INDEX_MAX_KEYS];
-	Oid			saved_userid;
-	int			saved_sec_context;
-	MemoryContext oldcxt;
 
-	if (fpentry->batch_count == 0)
-		return;
+	for (int i = 0; i < fpentry->batch_count; i++)
+	{
+		bool	found = false;
 
-	if (riinfo->fpmeta == NULL)
-		ri_populate_fastpath_metadata((RI_ConstraintInfo *) riinfo,
-									  fk_rel, idx_rel);
-	Assert(riinfo->fpmeta);
+		ExecStoreHeapTuple(fpentry->batch[i], fk_slot, false);
 
-	CommandCounterIncrement();
-	snapshot->curcid = GetCurrentCommandId(false);
+		ri_ExtractValues(fk_rel, fk_slot, riinfo, false, pk_vals, pk_nulls);
+		build_index_scankeys(riinfo, idx_rel, pk_vals, pk_nulls, skey);
+		found = ri_FastPathProbeOne(pk_rel, idx_rel, scandesc,
+									fpentry->xact_scan, slot,
+									snapshot, fpentry->xact_snap,
+									riinfo, skey, riinfo->nkeys,
+									true);
 
-	GetUserIdAndSecContext(&saved_userid, &saved_sec_context);
-	SetUserIdAndSecContext(RelationGetForm(pk_rel)->relowner,
-						   saved_sec_context |
-						   SECURITY_LOCAL_USERID_CHANGE |
-						   SECURITY_NOFORCE_RLS);
+		if (!found)
+			ri_ReportViolation(riinfo, pk_rel, fk_rel,
+							   fk_slot, NULL,
+							   RI_PLAN_CHECK_LOOKUPPK, false, false);
+	}
+}
 
-	fk_slot = MakeSingleTupleTableSlot(RelationGetDescr(fk_rel),
-									   &TTSOpsHeapTuple);
+/*
+ * ri_FastPathFlushArray
+ *		Single-column fast path using SK_SEARCHARRAY.
+ *
+ * Builds an array of FK values and does one index scan with
+ * SK_SEARCHARRAY.  The index AM sorts and deduplicates the array
+ * internally, then walks matching leaf pages in order.  Each
+ * matched PK tuple is locked and rechecked as before; a matched[]
+ * bitmap tracks which batch items were satisfied.
+ */
+static void
+ri_FastPathFlushArray(RI_FastPathEntry *fpentry, TupleTableSlot *fk_slot,
+					  const RI_ConstraintInfo *riinfo, Relation fk_rel)
+{
+	FastPathMeta *fpmeta = riinfo->fpmeta;
+	Relation	pk_rel = fpentry->pk_rel;
+	Relation	idx_rel = fpentry->idx_rel;
+	IndexScanDesc scandesc = fpentry->scandesc;
+	TupleTableSlot *slot = fpentry->slot;
+	Snapshot	snapshot = fpentry->snapshot;
+	Datum	   *search_vals = fpentry->search_vals;
+	bool	   *matched = fpentry->matched;
+	int			nvals = fpentry->batch_count;
+	Datum		pk_vals[INDEX_MAX_KEYS];
+	char		pk_nulls[INDEX_MAX_KEYS];
+	ScanKeyData skey[1];
+	RI_CompareHashEntry *entry;
+	Oid			elem_type;
+	int16		elem_len;
+	bool		elem_byval;
+	char		elem_align;
+	ArrayType  *arr;
+	MemoryContext oldcxt;
 
-	oldcxt = MemoryContextSwitchTo(TopTransactionContext);
-	for (int i = 0; i < fpentry->batch_count; i++)
-	{
-		HeapTuple fktuple = fpentry->batch[i];
-		bool	found = false;
+	Assert(fpmeta);
 
-		ExecStoreHeapTuple(fktuple, fk_slot, false);
+	memset(matched, 0, nvals * sizeof(bool));
 
+	/*
+	 * Extract and cast FK values.  We need the PK-side type for
+	 * the array element type since the scan key compares against
+	 * the index which stores PK-typed values.
+	 */
+	entry = fpmeta->compare_entries[0];
+	for (int i = 0; i < nvals; i++)
+	{
+		ExecStoreHeapTuple(fpentry->batch[i], fk_slot, false);
 		ri_ExtractValues(fk_rel, fk_slot, riinfo, false, pk_vals, pk_nulls);
-		build_index_scankeys(riinfo, idx_rel, pk_vals, pk_nulls, skey);
 
-		index_rescan(scandesc, skey, riinfo->nkeys, NULL, 0);
+		/* Cast if needed (e.g. int8 FK -> numeric PK) */
+		if (OidIsValid(entry->cast_func_finfo.fn_oid))
+			search_vals[i] = FunctionCall3(&entry->cast_func_finfo,
+										   pk_vals[0],
+										   Int32GetDatum(-1),
+										   BoolGetDatum(false));
+		else
+			search_vals[i] = pk_vals[0];
+	}
 
-		if (index_getnext_slot(scandesc, ForwardScanDirection, slot))
-		{
-			bool	concurrently_updated;
+	/*
+	 * Array element type must match the operator's right-hand input
+	 * type, which is what the index comparison expects on the search
+	 * side.  ri_populate_fastpath_metadata() stores exactly this via
+	 * get_op_opfamily_properties(), which returns the operator's
+	 * right-hand type as the subtype for cross-type operators (e.g.
+	 * int8 for int48eq) and the common type for same-type operators.
+	 */
+	elem_type = fpmeta->subtypes[0];
+	Assert(OidIsValid(elem_type));
+	get_typlenbyvalalign(elem_type, &elem_len, &elem_byval, &elem_align);
 
-			if (ri_LockPKTuple(pk_rel, slot, snapshot,
-							   &concurrently_updated))
-			{
-				if (concurrently_updated)
-					found = recheck_matched_pk_tuple(idx_rel, skey, slot);
-				else
-					found = true;
-			}
-		}
+	arr = construct_array(search_vals, nvals,
+						  elem_type, elem_len, elem_byval, elem_align);
 
-		if (found && IsolationUsesXactSnapshot())
-		{
-			IndexScanDesc xact_scan;
-			TupleTableSlot *xact_slot;
-			Snapshot	xact_snap = GetTransactionSnapshot();
+	/*
+	 * Build scan key with SK_SEARCHARRAY.  The btree code will
+	 * internally sort and deduplicate, then walk leaf pages in order.
+	 */
+	ScanKeyEntryInitialize(&skey[0],
+						   SK_SEARCHARRAY,
+						   1,				/* attno */
+						   fpmeta->strats[0],
+						   fpmeta->subtypes[0],
+						   idx_rel->rd_indcollation[0],
+						   fpmeta->regops[0],
+						   PointerGetDatum(arr));
 
-			xact_slot = table_slot_create(pk_rel, NULL);
-			xact_scan = index_beginscan(pk_rel, idx_rel,
-										xact_snap, NULL,
-										riinfo->nkeys, 0);
-			index_rescan(xact_scan, skey, riinfo->nkeys, NULL, 0);
+	oldcxt = MemoryContextSwitchTo(TopTransactionContext);
 
-			if (!index_getnext_slot(xact_scan, ForwardScanDirection,
-									xact_slot))
-				found = false;
+	index_rescan(scandesc, skey, 1, NULL, 0);
 
-			index_endscan(xact_scan);
-			ExecDropSingleTupleTableSlot(xact_slot);
+	/*
+	 * Walk all matches.  The btree returns them in index order.
+	 * For each match, find which batch item(s) it satisfies.
+	 */
+	while (index_getnext_slot(scandesc, ForwardScanDirection, slot))
+	{
+		Datum		found_val;
+		bool		found_null;
+		bool		concurrently_updated;
+		ScanKeyData recheck_skey[1];
+		bool		recheck_skey_valid = false;
+
+		if (!ri_LockPKTuple(pk_rel, slot, snapshot, &concurrently_updated))
+			continue;
+
+		/* Extract the PK value from the matched and locked tuple */
+		found_val = slot_getattr(slot, riinfo->pk_attnums[0], &found_null);
+		Assert(!found_null);
+
+		if (concurrently_updated)
+		{
+			/*
+			 * Build a single-key scankey for recheck.  We need the
+			 * actual PK value that was found, not the FK search value.
+			 */
+			ScanKeyEntryInitialize(&recheck_skey[0], 0, 1,
+								   fpmeta->strats[0],
+								   fpmeta->subtypes[0],
+								   idx_rel->rd_indcollation[0],
+								   fpmeta->regops[0],
+								   found_val);
+			recheck_skey_valid = true;
+			if (!recheck_matched_pk_tuple(idx_rel, recheck_skey, slot))
+				continue;
+		}
+
+		/* RR/SERIALIZABLE crosscheck */
+		if (IsolationUsesXactSnapshot())
+		{
+			IndexScanDesc xact_scan = fpentry->xact_scan;
+
+			if (!recheck_skey_valid)
+				ScanKeyEntryInitialize(&recheck_skey[0], 0, 1,
+									   fpmeta->strats[0],
+									   fpmeta->subtypes[0],
+									   idx_rel->rd_indcollation[0],
+									   fpmeta->regops[0],
+									   found_val);
+
+			index_rescan(xact_scan, recheck_skey, 1, NULL, 0);
+			if (!index_getnext_slot(xact_scan, ForwardScanDirection, slot))
+				continue;
 		}
 
 		/*
-		 * Report immediately.  ri_ReportViolation calls ereport(ERROR)
-		 * which doesn't return, so remaining batch items and cleanup
-		 * are handled by the error path (ResourceOwner + XactCallback).
+		 * Linear scan to mark all batch items matching this PK value.
+		 * O(batch_size) per match, O(batch_size^2) worst case -- fine
+		 * for the current batch size of 64.
 		 */
-		if (!found)
+		for (int i = 0; i < nvals; i++)
+		{
+			if (!matched[i] &&
+				DatumGetBool(FunctionCall2Coll(&entry->eq_opr_finfo,
+											   idx_rel->rd_indcollation[0],
+											   found_val,
+											   search_vals[i])))
+				matched[i] = true;
+		}
+	}
+
+	/* Report first unmatched row */
+	for (int i = 0; i < nvals; i++)
+	{
+		if (!matched[i])
+		{
+			ExecStoreHeapTuple(fpentry->batch[i], fk_slot, false);
 			ri_ReportViolation(riinfo, pk_rel, fk_rel,
 							   fk_slot, NULL,
 							   RI_PLAN_CHECK_LOOKUPPK, false, false);
+		}
 	}
 
 	MemoryContextSwitchTo(oldcxt);
+
+	pfree(arr);
+}
+
+
+static void
+ri_FastPathBatchFlush(RI_FastPathEntry *fpentry, Relation fk_rel)
+{
+	const RI_ConstraintInfo *riinfo = fpentry->riinfo;
+	Relation	pk_rel = fpentry->pk_rel;
+	Relation	idx_rel = fpentry->idx_rel;
+	Snapshot	snapshot = fpentry->snapshot;
+	TupleTableSlot *fk_slot;
+	Oid			saved_userid;
+	int			saved_sec_context;
+
+	if (fpentry->batch_count == 0)
+		return;
+
+	if (riinfo->fpmeta == NULL)
+		ri_populate_fastpath_metadata((RI_ConstraintInfo *) riinfo,
+									  fk_rel, idx_rel);
+	Assert(riinfo->fpmeta);
+
+	CommandCounterIncrement();
+	snapshot->curcid = GetCurrentCommandId(false);
+
+	GetUserIdAndSecContext(&saved_userid, &saved_sec_context);
+	SetUserIdAndSecContext(RelationGetForm(pk_rel)->relowner,
+						   saved_sec_context |
+						   SECURITY_LOCAL_USERID_CHANGE |
+						   SECURITY_NOFORCE_RLS);
+
+	fk_slot = MakeSingleTupleTableSlot(RelationGetDescr(fk_rel),
+									   &TTSOpsHeapTuple);
+
+	if (riinfo->nkeys == 1)
+		ri_FastPathFlushArray(fpentry, fk_slot, riinfo, fk_rel);
+	else
+		ri_FastPathFlushLoop(fpentry, fk_slot, riinfo, fk_rel);
+
 	SetUserIdAndSecContext(saved_userid, saved_sec_context);
 
 	/* Free materialized tuples and reset */
diff --git a/src/test/regress/expected/foreign_key.out b/src/test/regress/expected/foreign_key.out
index 16bb6370a97..0a24acdb138 100644
--- a/src/test/regress/expected/foreign_key.out
+++ b/src/test/regress/expected/foreign_key.out
@@ -3593,3 +3593,20 @@ COMMIT;
 ERROR:  insert or update on table "fp_fk_commit" violates foreign key constraint "fp_fk_commit_a_fkey"
 DETAIL:  Key (a)=(999) is not present in table "fp_pk_commit".
 DROP TABLE fp_fk_commit, fp_pk_commit;
+-- Cross-type FK with bulk insert: int8 FK referencing int4 PK,
+-- values cast during array construction
+CREATE TABLE fp_pk_cross (a int4 PRIMARY KEY);
+INSERT INTO fp_pk_cross SELECT generate_series(1, 200);
+CREATE TABLE fp_fk_cross (a int8 REFERENCES fp_pk_cross);
+INSERT INTO fp_fk_cross SELECT generate_series(1, 200);
+INSERT INTO fp_fk_cross VALUES (999);
+ERROR:  insert or update on table "fp_fk_cross" violates foreign key constraint "fp_fk_cross_a_fkey"
+DETAIL:  Key (a)=(999) is not present in table "fp_pk_cross".
+DROP TABLE fp_fk_cross, fp_pk_cross;
+-- Duplicate FK values: when using the batched SAOP path, every
+-- row must be recognized as satisfied, not just the first match
+CREATE TABLE fp_pk_dup (a int PRIMARY KEY);
+INSERT INTO fp_pk_dup VALUES (1);
+CREATE TABLE fp_fk_dup (a int REFERENCES fp_pk_dup);
+INSERT INTO fp_fk_dup SELECT 1 FROM generate_series(1, 100);
+DROP TABLE fp_fk_dup, fp_pk_dup;
diff --git a/src/test/regress/sql/foreign_key.sql b/src/test/regress/sql/foreign_key.sql
index bc24272df20..ce85a21fc79 100644
--- a/src/test/regress/sql/foreign_key.sql
+++ b/src/test/regress/sql/foreign_key.sql
@@ -2577,3 +2577,20 @@ INSERT INTO fp_fk_commit VALUES (1);
 INSERT INTO fp_fk_commit VALUES (999);
 COMMIT;
 DROP TABLE fp_fk_commit, fp_pk_commit;
+
+-- Cross-type FK with bulk insert: int8 FK referencing int4 PK,
+-- values cast during array construction
+CREATE TABLE fp_pk_cross (a int4 PRIMARY KEY);
+INSERT INTO fp_pk_cross SELECT generate_series(1, 200);
+CREATE TABLE fp_fk_cross (a int8 REFERENCES fp_pk_cross);
+INSERT INTO fp_fk_cross SELECT generate_series(1, 200);
+INSERT INTO fp_fk_cross VALUES (999);
+DROP TABLE fp_fk_cross, fp_pk_cross;
+
+-- Duplicate FK values: when using the batched SAOP path, every
+-- row must be recognized as satisfied, not just the first match
+CREATE TABLE fp_pk_dup (a int PRIMARY KEY);
+INSERT INTO fp_pk_dup VALUES (1);
+CREATE TABLE fp_fk_dup (a int REFERENCES fp_pk_dup);
+INSERT INTO fp_fk_dup SELECT 1 FROM generate_series(1, 100);
+DROP TABLE fp_fk_dup, fp_pk_dup;
-- 
2.47.3