From 1e091fbe751a97ba296cc43fd66c057102b34b43 Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onderkalaci@gmail.com>
Date: Thu, 9 Mar 2023 14:38:40 +0300
Subject: [PATCH 2/2] WIP: Optimize for non-pkey / non-RI unique indexes

---
 src/backend/executor/execReplication.c     | 13 ++--
 src/backend/replication/logical/relation.c | 82 ++++++++++++++++++++--
 src/backend/replication/logical/worker.c   |  2 +-
 src/include/executor/executor.h            |  3 +-
 src/include/replication/logicalrelation.h  |  6 +-
 5 files changed, 89 insertions(+), 17 deletions(-)

diff --git a/src/backend/executor/execReplication.c b/src/backend/executor/execReplication.c
index d1be63d9de..6f9bfdac05 100644
--- a/src/backend/executor/execReplication.c
+++ b/src/backend/executor/execReplication.c
@@ -137,8 +137,8 @@ build_replindex_scan_key(ScanKey skey, Relation rel, Relation idxrel,
  * contents, and return true.  Return false otherwise.
  */
 bool
-RelationFindReplTupleByIndex(Relation rel, Oid idxoid,
-							 LockTupleMode lockmode,
+RelationFindReplTupleByIndex(Relation rel, LogicalRepRelation *remoterel,
+							 Oid idxoid, LockTupleMode lockmode,
 							 TupleTableSlot *searchslot,
 							 TupleTableSlot *outslot)
 {
@@ -150,12 +150,12 @@ RelationFindReplTupleByIndex(Relation rel, Oid idxoid,
 	Relation	idxrel;
 	bool		found;
 	TypeCacheEntry **eq = NULL;
-	bool		idxIsRelationIdentityOrPK;
+	bool		isIdxSafeToSkipDuplicates;
 
 	/* Open the index. */
 	idxrel = index_open(idxoid, RowExclusiveLock);
 
-	idxIsRelationIdentityOrPK = IdxIsRelationIdentityOrPK(rel, idxoid);
+	isIdxSafeToSkipDuplicates = IsIdxSafeToSkipDuplicates(rel, idxrel, remoterel);
 
 	InitDirtySnapshot(snap);
 
@@ -174,10 +174,9 @@ retry:
 	while (index_getnext_slot(scan, ForwardScanDirection, outslot))
 	{
 		/*
-		 * Avoid expensive equality check if the index is primary key or
-		 * replica identity index.
+		 * Avoid expensive equality check if the index is unique.
 		 */
-		if (!idxIsRelationIdentityOrPK)
+		if (!isIdxSafeToSkipDuplicates)
 		{
 			if (eq == NULL)
 			{
diff --git a/src/backend/replication/logical/relation.c b/src/backend/replication/logical/relation.c
index 1616bd5712..a7cc38bfab 100644
--- a/src/backend/replication/logical/relation.c
+++ b/src/backend/replication/logical/relation.c
@@ -27,6 +27,7 @@
 #include "replication/logicalrelation.h"
 #include "replication/worker_internal.h"
 #include "utils/inval.h"
+#include "utils/syscache.h"
 
 
 static MemoryContext LogicalRepRelMapContext = NULL;
@@ -771,6 +772,76 @@ IndexContainsAnyRemoteColumn(IndexInfo  *indexInfo,
 	return false;
 }
 
+/*
+ * WIP
+ */
+bool
+AllIndexColumnsProvided(IndexInfo  *indexInfo, LogicalRepRelation  *remoterel)
+{
+	int indexColumnMatchCount = 0;
+	for (int i = 0; i < indexInfo->ii_NumIndexAttrs; i++)
+	{
+		int			keycol = indexInfo->ii_IndexAttrNumbers[i];
+
+		if (!AttributeNumberIsValid(keycol))
+				return false;
+
+		if (bms_is_member( keycol- 1, remoterel->attkeys))
+			indexColumnMatchCount++;
+	}
+
+	if (indexColumnMatchCount != indexInfo->ii_NumIndexAttrs)
+		return false;
+
+	return true;
+}
+
+bool
+IndexHasPkeyCharacterstics(IndexInfo  *indexInfo, Relation localrel)
+{
+
+	if (!indexInfo->ii_Unique)
+			return false;
+
+	if (!indexInfo->ii_NullsNotDistinct)
+		return false;
+
+	/*
+	 * Check that all of the attributes in a primary key are marked as not
+	 * null.  (We don't really expect to see that; it'd mean the parser messed
+	 * up.  But it seems wise to check anyway.)
+	 */
+	for (int i = 0; i < indexInfo->ii_NumIndexKeyAttrs; i++)
+	{
+		AttrNumber	attnum = indexInfo->ii_IndexAttrNumbers[i];
+		HeapTuple	atttuple;
+		Form_pg_attribute attform;
+
+		if (attnum == 0)
+			return false;
+
+		/* System attributes are never null, so no need to check */
+		if (attnum < 0)
+			return false;
+
+		atttuple = SearchSysCache2(ATTNUM,
+								   ObjectIdGetDatum(RelationGetRelid(localrel)),
+								   Int16GetDatum(attnum));
+		if (!HeapTupleIsValid(atttuple))
+			return false;
+
+		attform = (Form_pg_attribute) GETSTRUCT(atttuple);
+
+		if (!attform->attnotnull)
+			return false;
+
+		ReleaseSysCache(atttuple);
+	}
+
+	return true;
+
+}
+
 /*
  * Returns the oid of an index that can be used by the apply worker to scan
  * the relation. The index must be btree, non-partial, and have at least
@@ -859,18 +930,17 @@ GetRelationIdentityOrPK(Relation rel)
 }
 
 /*
- * Given a relation and OID of an index, returns true if the index is relation's
- * replica identity index or relation's primary key's index.
+ * Given index relation, returns true if the index is unique.
  *
  * Returns false otherwise.
  */
 bool
-IdxIsRelationIdentityOrPK(Relation rel, Oid idxoid)
+IsIdxSafeToSkipDuplicates(Relation localrel, Relation idxrel, LogicalRepRelation *remoterel)
 {
-	Assert(OidIsValid(idxoid));
+	IndexInfo  *indexInfo = BuildIndexInfo(idxrel);
 
-	return RelationGetReplicaIndex(rel) == idxoid ||
-		RelationGetPrimaryKeyIndex(rel) == idxoid;
+	return	AllIndexColumnsProvided(indexInfo, remoterel) &&
+			IndexHasPkeyCharacterstics(indexInfo, localrel);
 }
 
 /*
diff --git a/src/backend/replication/logical/worker.c b/src/backend/replication/logical/worker.c
index 10f9711972..73faf7b39d 100644
--- a/src/backend/replication/logical/worker.c
+++ b/src/backend/replication/logical/worker.c
@@ -2846,7 +2846,7 @@ FindReplTupleInLocalRel(EState *estate, Relation localrel,
 		   (remoterel->replident == REPLICA_IDENTITY_FULL));
 
 	if (OidIsValid(localidxoid))
-		found = RelationFindReplTupleByIndex(localrel, localidxoid,
+		found = RelationFindReplTupleByIndex(localrel, remoterel, localidxoid,
 											 LockTupleExclusive,
 											 remoteslot, *localslot);
 	else
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 946abc0051..4a636d7ef3 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -634,7 +634,8 @@ extern void check_exclusion_constraint(Relation heap, Relation index,
 /*
  * prototypes from functions in execReplication.c
  */
-extern bool RelationFindReplTupleByIndex(Relation rel, Oid idxoid,
+#include "replication/logicalproto.h"
+extern bool RelationFindReplTupleByIndex(Relation rel, LogicalRepRelation *remoterel, Oid idxoid,
 										 LockTupleMode lockmode,
 										 TupleTableSlot *searchslot,
 										 TupleTableSlot *outslot);
diff --git a/src/include/replication/logicalrelation.h b/src/include/replication/logicalrelation.h
index 723fb6d707..5066f513f6 100644
--- a/src/include/replication/logicalrelation.h
+++ b/src/include/replication/logicalrelation.h
@@ -49,7 +49,9 @@ extern LogicalRepRelMapEntry *logicalrep_partition_open(LogicalRepRelMapEntry *r
 extern void logicalrep_rel_close(LogicalRepRelMapEntry *rel,
 								 LOCKMODE lockmode);
 extern bool IsIndexUsableForReplicaIdentityFull(IndexInfo *indexInfo);
+extern bool AllIndexColumnsProvided(IndexInfo  *indexInfo, LogicalRepRelation *remoterel);
 extern Oid GetRelationIdentityOrPK(Relation rel);
-extern bool IdxIsRelationIdentityOrPK(Relation rel, Oid idxoid);
-
+extern bool IsIdxSafeToSkipDuplicates(Relation localrel, Relation idxrel,
+									  LogicalRepRelation *remoterel);
+extern bool IndexHasPkeyCharacterstics(IndexInfo  *indexInfo, Relation localrel);
 #endif							/* LOGICALRELATION_H */
-- 
2.34.1

