Properly detect key columns for unique expression indexes

tsl/src/compression/compression.c

@@ -2101,6 +2101,72 @@
 	return num_scankeys;
 }
 
+/*
+ * For insert into compressed chunks with unique index determine the
+ * columns which are safe to use for batch filtering.
+ *
+ * This is based on RelationGetIndexAttrBitmap from postgres with changes
+ * to also track unique expression indexes.
+ */
+static Bitmapset *
+compressed_insert_key_columns(Relation relation)
+{
+	Bitmapset *indexattrs = NULL; /* indexed columns */
+	ListCell *l;
+
+	/* Fast path if definitely no indexes */
+	if (!RelationGetForm(relation)->relhasindex)
+		return NULL;
+
+	List *indexoidlist = RelationGetIndexList(relation);
+
+	/* Fall out if no indexes (but relhasindex was set) */
+	if (indexoidlist == NIL)
+		return NULL;
+
+	/*
+	 * For each index, add referenced attributes to indexattrs.
+	 *
+	 * Note: we consider all indexes returned by RelationGetIndexList, even if
+	 * they are not indisready or indisvalid.  This is important because an
+	 * index for which CREATE INDEX CONCURRENTLY has just started must be
+	 * included in HOT-safety decisions (see README.HOT).  If a DROP INDEX
+	 * CONCURRENTLY is far enough along that we should ignore the index, it
+	 * won't be returned at all by RelationGetIndexList.
+	 */
+	foreach (l, indexoidlist)
+	{
+		Oid indexOid = lfirst_oid(l);
+		Relation indexDesc = index_open(indexOid, AccessShareLock);
+
+		if (!indexDesc->rd_index->indisunique)
+		{
+			index_close(indexDesc, AccessShareLock);
+			continue;
+		}
+
+		/* Collect simple attribute references.
+		 * For covering indexes we only need to collect the key attributes.
+		 * Unlike RelationGetIndexAttrBitmap we allow expression indexes
+		 * but we do not extract attributes from the expressions as that
+		 * would not be a safe filter as the expression can alter attributes
+		 * which would not make them sufficient for batch filtering.
+		 */
+		for (int i = 0; i < indexDesc->rd_index->indnkeyatts; i++)
+		{
+			int attrnum = indexDesc->rd_index->indkey.values[i];
+			if (attrnum != 0)
+			{
+				indexattrs =
+					bms_add_member(indexattrs, attrnum - FirstLowInvalidHeapAttributeNumber);
+			}
+		}
+		index_close(indexDesc, AccessShareLock);
+	}
+
+	return indexattrs;
+}
+
 void
 decompress_batches_for_insert(ChunkInsertState *cis, Chunk *chunk, TupleTableSlot *slot)
 {
@@ -2133,7 +2199,7 @@
 	Relation in_rel = relation_open(comp->table_id, RowExclusiveLock);
 
 	RowDecompressor decompressor = build_decompressor(in_rel, out_rel);
-	Bitmapset *key_columns = RelationGetIndexAttrBitmap(out_rel, INDEX_ATTR_BITMAP_KEY);
+	Bitmapset *key_columns = compressed_insert_key_columns(out_rel);
 	Bitmapset *null_columns = NULL;
 
 	int num_scankeys;

tsl/test/shared/expected/compress_unique_index.out

@@ -23,11 +23,33 @@ SELECT compress_chunk(show_chunks('compress_unique')) IS NOT NULL AS compress;
 INSERT INTO compress_unique VALUES ('2000-01-01','m1','c2','2000-01-01');
 ERROR:  duplicate key value violates unique constraint "_hyper_X_X_chunk_uniq_expr"
 \set ON_ERROR_STOP 1
+-- should only decompress 1 batch
+EXPLAIN (analyze,costs off,summary off,timing off) INSERT INTO compress_unique VALUES ('2000-01-01','m1','c2','2000-01-02');
+QUERY PLAN
+ Custom Scan (HypertableModify) (actual rows=0 loops=1)
+   Batches decompressed: 1
+   Tuples decompressed: 1
+   ->  Insert on compress_unique (actual rows=0 loops=1)
+         ->  Custom Scan (ChunkDispatch) (actual rows=1 loops=1)
+               ->  Result (actual rows=1 loops=1)
+(6 rows)
+
+-- should decompress no batches
+EXPLAIN (analyze,costs off,summary off,timing off) INSERT INTO compress_unique VALUES ('2000-01-01','m1','c3','2000-01-02');
+QUERY PLAN
+ Custom Scan (HypertableModify) (actual rows=0 loops=1)
+   ->  Insert on compress_unique (actual rows=0 loops=1)
+         ->  Custom Scan (ChunkDispatch) (actual rows=1 loops=1)
+               ->  Result (actual rows=1 loops=1)
+(4 rows)
+
 SELECT * FROM compress_unique ORDER BY compress_unique;
        offset_timestamp       | meter_id | meter_channel_id |          timestamp           
 ------------------------------+----------+------------------+------------------------------
  Sat Jan 01 00:00:00 2000 PST | m1       | c1               | Sat Jan 01 00:00:00 2000 PST
  Sat Jan 01 00:00:00 2000 PST | m1       | c2               | Sat Jan 01 00:00:00 2000 PST
-(2 rows)
+ Sat Jan 01 00:00:00 2000 PST | m1       | c2               | Sun Jan 02 00:00:00 2000 PST
+ Sat Jan 01 00:00:00 2000 PST | m1       | c3               | Sun Jan 02 00:00:00 2000 PST
+(4 rows)
 
 DROP TABLE compress_unique;

tsl/test/shared/sql/CMakeLists.txt

@@ -2,7 +2,6 @@ set(TEST_FILES_SHARED
     cagg_compression.sql
     classify_relation.sql
     compat.sql
-    compress_unique_index.sql
     constify_timestamptz_op_interval.sql
     constraint_aware_append.sql
     constraint_exclusion_prepared.sql
@@ -17,8 +16,8 @@ set(TEST_TEMPLATES_SHARED
     space_constraint.sql.in)
 
 if((${PG_VERSION_MAJOR} GREATER_EQUAL "14"))
-  list(APPEND TEST_FILES_SHARED compression_dml.sql decompress_tracking.sql
-       memoize.sql)
+  list(APPEND TEST_FILES_SHARED compress_unique_index.sql compression_dml.sql
+       decompress_tracking.sql memoize.sql)
 endif()
 
 # this test was changing the contents of tables in shared_setup.sql thus causing

tsl/test/shared/sql/compress_unique_index.sql

@@ -20,6 +20,12 @@ SELECT compress_chunk(show_chunks('compress_unique')) IS NOT NULL AS compress;
 INSERT INTO compress_unique VALUES ('2000-01-01','m1','c2','2000-01-01');
 \set ON_ERROR_STOP 1
 
+-- should only decompress 1 batch
+EXPLAIN (analyze,costs off,summary off,timing off) INSERT INTO compress_unique VALUES ('2000-01-01','m1','c2','2000-01-02');
+
+-- should decompress no batches
+EXPLAIN (analyze,costs off,summary off,timing off) INSERT INTO compress_unique VALUES ('2000-01-01','m1','c3','2000-01-02');
+
 SELECT * FROM compress_unique ORDER BY compress_unique;
 
 DROP TABLE compress_unique;