Reduce decompression for INSERT with unique constraints

On INSERT into compressed chunks with unique constraints we can
check for conflict without decompressing when no ON CONFLICT clause
is present and we only have one unique constraint. With ON CONFLICT
clause with DO NOTHING we can just skip the INSERT if we detect conflict
and return early. Only for ON CONFLICT DO UPDATE/UPSERT do we need
to decompress when there is a constraint conflict.
Doing the optimization in the presence of multiple constraints is
also possible but not part of this patch.

Author: svenklemm

src/nodes/chunk_dispatch/chunk_dispatch.h

@@ -76,6 +76,9 @@ typedef struct ChunkDispatchState
 	bool is_dropped_attr_exists;
 	int64 batches_decompressed;
 	int64 tuples_decompressed;
+
+	/* Should this INSERT be skipped due to ON CONFLICT DO NOTHING */
+	bool skip_insert;
 } ChunkDispatchState;
 
 extern TSDLLEXPORT bool ts_is_chunk_dispatch_state(PlanState *state);

src/nodes/chunk_dispatch/chunk_insert_state.h

@@ -65,5 +65,5 @@ extern ChunkInsertState *ts_chunk_insert_state_create(Oid chunk_relid,
 													  const ChunkDispatch *dispatch);
 extern void ts_chunk_insert_state_destroy(ChunkInsertState *state);
 
-OnConflictAction chunk_dispatch_get_on_conflict_action(const ChunkDispatch *dispatch);
+TSDLLEXPORT OnConflictAction chunk_dispatch_get_on_conflict_action(const ChunkDispatch *dispatch);
 void ts_set_compression_status(ChunkInsertState *state, const Chunk *chunk);

src/nodes/hypertable_modify.c

@@ -687,6 +687,14 @@ ExecModifyTable(CustomScanState *cs_node, PlanState *pstate)
 
 		context.planSlot = ExecProcNode(subplanstate);
 
+		if (cds && cds->rri && operation == CMD_INSERT && cds->skip_insert)
+		{
+			cds->skip_insert = false;
+			if (node->ps.instrument)
+				node->ps.instrument->ntuples2++;
+			return NULL;
+		}
+
 		/* No more tuples to process? */
 		if (TupIsNull(context.planSlot))
 			break;

tsl/src/compression/compression_dml.c

@@ -29,24 +29,25 @@
 #include <nodes/hypertable_modify.h>
 #include <ts_catalog/array_utils.h>
 
-static struct decompress_batches_stats
-decompress_batches_indexscan(Relation in_rel, Relation out_rel, Relation index_rel,
-							 Snapshot snapshot, ScanKeyData *index_scankeys, int num_index_scankeys,
-							 ScanKeyData *heap_scankeys, int num_heap_scankeys,
-							 ScanKeyData *mem_scankeys, int num_mem_scankeys,
-							 Bitmapset *null_columns, List *is_nulls);
+static struct decompress_batches_stats decompress_batches_indexscan(
+	Relation in_rel, Relation out_rel, Relation index_rel, Snapshot snapshot,
+	ScanKeyData *index_scankeys, int num_index_scankeys, ScanKeyData *heap_scankeys,
+	int num_heap_scankeys, ScanKeyData *mem_scankeys, int num_mem_scankeys,
+	OnConflictAction on_conflict, bool *skip_insert, Bitmapset *null_columns, List *is_nulls);
 static struct decompress_batches_stats
 decompress_batches_seqscan(Relation in_rel, Relation out_rel, Snapshot snapshot,
 						   ScanKeyData *scankeys, int num_scankeys, ScanKeyData *mem_scankeys,
-						   int num_mem_scankeys, Bitmapset *null_columns, List *is_nulls);
+						   int num_mem_scankeys, OnConflictAction on_conflict, bool *skip_insert,
+						   Bitmapset *null_columns, List *is_nulls);
 
-static bool batch_matches(RowDecompressor *decompressor, ScanKeyData *scankeys, int num_scankeys);
+static bool batch_matches(RowDecompressor *decompressor, ScanKeyData *scankeys, int num_scankeys,
+						  OnConflictAction on_conflict, bool *skip_insert);
 static void process_predicates(Chunk *ch, CompressionSettings *settings, List *predicates,
 							   ScanKeyData **mem_scankeys, int *num_mem_scankeys,
 							   List **heap_filters, List **index_filters, List **is_null);
 static Relation find_matching_index(Relation comp_chunk_rel, List **index_filters,
 									List **heap_filters);
-static Bitmapset *compressed_insert_key_columns(Relation relation);
+static Bitmapset *compressed_insert_key_columns(Relation relation, bool *covering);
 static BatchFilter *make_batchfilter(char *column_name, StrategyNumber strategy, Oid collation,
 									 RegProcedure opcode, Const *value, bool is_null_check,
 									 bool is_null, bool is_array_op);
@@ -84,7 +85,20 @@ decompress_batches_for_insert(const ChunkInsertState *cis, TupleTableSlot *slot)
 	CompressionSettings *settings = ts_compression_settings_get(cis->compressed_chunk_table_id);
 	Assert(settings);
 
-	Bitmapset *key_columns = compressed_insert_key_columns(out_rel);
+	bool covering;
+	Bitmapset *key_columns = compressed_insert_key_columns(out_rel, &covering);
+	OnConflictAction on_conflict = ONCONFLICT_UPDATE;
+	/*
+	 *  When no on conflict clause is specified and the index is covering, we can
+	 *  error out before decompressing anything.
+	 *  For ON CONFLICT DO NOTHING with covering index we can skip decompression
+	 *  and abort the insert when we find a matching tuple.
+	 *  For ON CONFLICT DO UPDATE we need to decompress the tuple on match.
+	 */
+	if (covering && cis->cds->dispatch)
+	{
+		on_conflict = chunk_dispatch_get_on_conflict_action(cis->cds->dispatch);
+	}
 	Bitmapset *index_columns = NULL;
 	Bitmapset *null_columns = NULL;
 	struct decompress_batches_stats stats;
@@ -113,6 +127,7 @@ decompress_batches_for_insert(const ChunkInsertState *cis, TupleTableSlot *slot)
 																  &index_columns,
 																  &num_index_scankeys);
 
+	bool skip_insert = false;
 	if (index_rel)
 	{
 		/*
@@ -147,6 +162,8 @@ decompress_batches_for_insert(const ChunkInsertState *cis, TupleTableSlot *slot)
 											 num_heap_scankeys,
 											 mem_scankeys,
 											 num_mem_scankeys,
+											 on_conflict,
+											 &skip_insert,
 											 NULL, /* no null column check for non-segmentby
 													  columns */
 											 NIL);
@@ -171,12 +188,19 @@ decompress_batches_for_insert(const ChunkInsertState *cis, TupleTableSlot *slot)
 										   num_heap_scankeys,
 										   mem_scankeys,
 										   num_mem_scankeys,
+										   on_conflict,
+										   &skip_insert,
 										   null_columns,
 										   NIL);
 		bms_free(key_columns);
 	}
 
 	Assert(cis->cds != NULL);
+	if (skip_insert)
+	{
+		cis->cds->skip_insert = true;
+	}
+
 	cis->cds->batches_decompressed += stats.batches_decompressed;
 	cis->cds->tuples_decompressed += stats.tuples_decompressed;
 
@@ -260,6 +284,8 @@ decompress_batches_for_update_delete(HypertableModifyState *ht_state, Chunk *chu
 											 num_scankeys,
 											 mem_scankeys,
 											 num_mem_scankeys,
+											 ONCONFLICT_UPDATE,
+											 NULL,
 											 null_columns,
 											 is_null);
 		/* close the selected index */
@@ -274,6 +300,8 @@ decompress_batches_for_update_delete(HypertableModifyState *ht_state, Chunk *chu
 										   num_scankeys,
 										   mem_scankeys,
 										   num_mem_scankeys,
+										   ONCONFLICT_UPDATE,
+										   NULL,
 										   null_columns,
 										   is_null);
 	}
@@ -318,6 +346,7 @@ decompress_batches_indexscan(Relation in_rel, Relation out_rel, Relation index_r
 							 Snapshot snapshot, ScanKeyData *index_scankeys, int num_index_scankeys,
 							 ScanKeyData *heap_scankeys, int num_heap_scankeys,
 							 ScanKeyData *mem_scankeys, int num_mem_scankeys,
+							 OnConflictAction on_conflict, bool *skip_insert,
 							 Bitmapset *null_columns, List *is_nulls)
 {
 	HeapTuple compressed_tuple;
@@ -404,12 +433,22 @@ decompress_batches_indexscan(Relation in_rel, Relation out_rel, Relation index_r
 						  decompressor.compressed_datums,
 						  decompressor.compressed_is_nulls);
 
-		if (num_mem_scankeys && !batch_matches(&decompressor, mem_scankeys, num_mem_scankeys))
+		if (num_mem_scankeys &&
+			!batch_matches(&decompressor, mem_scankeys, num_mem_scankeys, on_conflict, skip_insert))
 		{
 			row_decompressor_reset(&decompressor);
 			continue;
 		}
 
+		if (skip_insert && *skip_insert)
+		{
+			row_decompressor_close(&decompressor);
+			index_endscan(scan);
+			index_close(index_rel, AccessShareLock);
+			ExecDropSingleTupleTableSlot(slot);
+			return stats;
+		}
+
 		write_logical_replication_msg_decompression_start();
 		result = delete_compressed_tuple(&decompressor, snapshot, compressed_tuple);
 		/* skip reporting error if isolation level is < Repeatable Read
@@ -470,7 +509,8 @@ decompress_batches_indexscan(Relation in_rel, Relation out_rel, Relation index_r
 static struct decompress_batches_stats
 decompress_batches_seqscan(Relation in_rel, Relation out_rel, Snapshot snapshot,
 						   ScanKeyData *scankeys, int num_scankeys, ScanKeyData *mem_scankeys,
-						   int num_mem_scankeys, Bitmapset *null_columns, List *is_nulls)
+						   int num_mem_scankeys, OnConflictAction on_conflict, bool *skip_insert,
+						   Bitmapset *null_columns, List *is_nulls)
 {
 	RowDecompressor decompressor;
 	bool decompressor_initialized = false;
@@ -534,12 +574,21 @@ decompress_batches_seqscan(Relation in_rel, Relation out_rel, Snapshot snapshot,
 						  decompressor.compressed_datums,
 						  decompressor.compressed_is_nulls);
 
-		if (num_mem_scankeys && !batch_matches(&decompressor, mem_scankeys, num_mem_scankeys))
+		if (num_mem_scankeys &&
+			!batch_matches(&decompressor, mem_scankeys, num_mem_scankeys, on_conflict, skip_insert))
 		{
 			row_decompressor_reset(&decompressor);
 			continue;
 		}
 
+		if (skip_insert && *skip_insert)
+		{
+			row_decompressor_close(&decompressor);
+			ExecDropSingleTupleTableSlot(slot);
+			table_endscan(scan);
+			return stats;
+		}
+
 		write_logical_replication_msg_decompression_start();
 		result = delete_compressed_tuple(&decompressor, snapshot, compressed_tuple);
 		/* skip reporting error if isolation level is < Repeatable Read
@@ -582,7 +631,8 @@ decompress_batches_seqscan(Relation in_rel, Relation out_rel, Snapshot snapshot,
 }
 
 static bool
-batch_matches(RowDecompressor *decompressor, ScanKeyData *scankeys, int num_scankeys)
+batch_matches(RowDecompressor *decompressor, ScanKeyData *scankeys, int num_scankeys,
+			  OnConflictAction on_conflict, bool *skip_insert)
 {
 	int num_tuples = decompress_batch(decompressor);
 
@@ -599,6 +649,16 @@ batch_matches(RowDecompressor *decompressor, ScanKeyData *scankeys, int num_scan
 #endif
 		if (valid)
 		{
+			if (on_conflict == ONCONFLICT_NONE)
+			{
+				ereport(ERROR,
+						(errcode(ERRCODE_UNIQUE_VIOLATION),
+						 errmsg("duplicate key value violates unique constraint")));
+			}
+			if (on_conflict == ONCONFLICT_NOTHING)
+			{
+				*skip_insert = true;
+			}
 			return true;
 		}
 	}
@@ -743,9 +803,12 @@ decompress_chunk_walker(PlanState *ps, struct decompress_chunk_context *ctx)
  * In case of multiple unique indexes we have to return the shared columns.
  * For expression indexes we ignore the columns with expressions, for partial
  * indexes we ignore predicate.
+ *
+ * The covering flag is set to true if we have a single constraint that is covered
+ * by all the columns present in the Bitmapset.
  */
 static Bitmapset *
-compressed_insert_key_columns(Relation relation)
+compressed_insert_key_columns(Relation relation, bool *covering)
 {
 	Bitmapset *shared_attrs = NULL; /* indexed columns */
 	ListCell *l;
@@ -792,8 +855,24 @@ compressed_insert_key_columns(Relation relation)
 		}
 		index_close(indexDesc, AccessShareLock);
 
-		shared_attrs = shared_attrs ? bms_intersect(idx_attrs, shared_attrs) : idx_attrs;
+		if (!shared_attrs)
+		{
+			/* First iteration */
+			shared_attrs = idx_attrs;
+			/* We only optimize unique constraint checks for non-partial indexes. */
+			*covering = indexDesc->rd_indpred == NIL;
+		}
+		else
+		{
+			shared_attrs = bms_intersect(idx_attrs, shared_attrs);
+			*covering = true;
+		}
 
+		/* When multiple unique indexes are present, in theory there could be no shared
+		 * columns even though that is very unlikely as they will probably at least share
+		 * the partitioning columns. But since we are looking at chunk indexes here that
+		 * is not guaranteed.
+		 */
 		if (!shared_attrs)
 			return NULL;
 	}