Add max_batches to compact_chunk

Adds a max_batches parameter to the compaction policy
and compact_chunk function. When set, compaction stops
between merge groups after decompressing that many
batches, bounding work per chunk. Zero means unlimited.

Author: Poroma-Banerjee

.unreleased/pr_10266

@@ -0,0 +1 @@
+Implements: #10266 Add max_batches to compact_chunk

sql/maintenance_utils.sql

@@ -90,7 +90,8 @@ CREATE OR REPLACE FUNCTION _timescaledb_functions.recompress_chunk_segmentwise(
 ) RETURNS REGCLASS AS '@MODULE_PATHNAME@', 'ts_recompress_chunk_segmentwise' LANGUAGE C STRICT VOLATILE;
 
 CREATE OR REPLACE FUNCTION _timescaledb_functions.compact_chunk(
-    uncompressed_chunk REGCLASS
+    uncompressed_chunk REGCLASS,
+    max_batches INTEGER DEFAULT 0
 ) RETURNS REGCLASS AS '@MODULE_PATHNAME@', 'ts_compact_chunk' LANGUAGE C STRICT VOLATILE;
 
 -- find the index on the compressed chunk that can be used to recompress efficiently

sql/policy_api.sql

@@ -52,6 +52,7 @@ CREATE OR REPLACE FUNCTION @extschema@.add_compaction_policy(
     initial_start TIMESTAMPTZ = NULL,
     timezone TEXT = NULL,
     max_chunks INTEGER = NULL,
+    max_batches INTEGER = NULL,
     inactive_for INTERVAL = NULL
 ) RETURNS INTEGER
 AS '@MODULE_PATHNAME@', 'ts_policy_compaction_add'

sql/policy_internal.sql

@@ -27,6 +27,7 @@ AS $$
 DECLARE
   htid               INTEGER;
   max_chunks         INTEGER := 0;
+  max_batches        INTEGER := 0;
   numchunks          INTEGER := 0;
   processed          INTEGER := 0;
   chunk_rec          RECORD;
@@ -53,6 +54,7 @@ BEGIN
 
   verbose_log := COALESCE(jsonb_object_field_text(config, 'verbose_log')::BOOLEAN, FALSE);
   max_chunks  := COALESCE(jsonb_object_field_text(config, 'max_chunks')::INTEGER, 0);
+  max_batches := COALESCE(jsonb_object_field_text(config, 'max_batches')::INTEGER, 0);
   -- when set, skip chunks written within this window; NULL disables the gate
   inactive_for := jsonb_object_field_text(config, 'inactive_for')::INTERVAL;
 
@@ -74,7 +76,7 @@ BEGIN
         WHERE s.last_update >= now() - inactive_for))
   LOOP
     BEGIN
-      PERFORM _timescaledb_functions.compact_chunk(chunk_rec.relid);
+      PERFORM _timescaledb_functions.compact_chunk(chunk_rec.relid, max_batches);
       numchunks := numchunks + 1;
     EXCEPTION WHEN OTHERS THEN
       GET STACKED DIAGNOSTICS

sql/updates/reverse-dev.sql

@@ -3,7 +3,7 @@ DROP FUNCTION IF EXISTS _timescaledb_functions.hypertable_status_text(int);
 DROP FUNCTION IF EXISTS _timescaledb_functions.hypertable_status_text(regclass);
 DROP FUNCTION IF EXISTS _timescaledb_functions.decompress_batch(record);
 DROP FUNCTION IF EXISTS _timescaledb_functions.estimate_uncompressed_size(regclass, double precision);
-DROP FUNCTION IF EXISTS _timescaledb_functions.compact_chunk(REGCLASS);
+DROP FUNCTION IF EXISTS _timescaledb_functions.compact_chunk(REGCLASS, INTEGER);
 DROP PROCEDURE IF EXISTS _timescaledb_functions.policy_compression_execute(INTEGER, INTEGER, ANYELEMENT, INTEGER, BOOLEAN, BOOLEAN, BOOLEAN, BOOLEAN, BOOLEAN);
 
 --
@@ -243,7 +243,7 @@ ALTER TABLE _timescaledb_catalog.compression_chunk_size
 
 -- Remove compaction policy jobs since the policy does not exist in the older version.
 DELETE FROM _timescaledb_config.bgw_job WHERE proc_schema = '_timescaledb_functions' AND proc_name = 'policy_compaction';
-DROP FUNCTION IF EXISTS @extschema@.add_compaction_policy(REGCLASS, BOOL, INTERVAL, TIMESTAMPTZ, TEXT, INTEGER, INTERVAL);
+DROP FUNCTION IF EXISTS @extschema@.add_compaction_policy(REGCLASS, BOOL, INTERVAL, TIMESTAMPTZ, TEXT, INTEGER, INTEGER, INTERVAL);
 DROP FUNCTION IF EXISTS @extschema@.remove_compaction_policy(REGCLASS, BOOL);
 DROP PROCEDURE IF EXISTS _timescaledb_functions.policy_compaction(INTEGER, JSONB);
 DROP FUNCTION IF EXISTS _timescaledb_functions.policy_compaction_check(JSONB);

tsl/src/bgw_policy/compaction_api.c

@@ -70,6 +70,16 @@ policy_compaction_check(PG_FUNCTION_ARGS)
 						POL_COMPACTION_CONF_KEY_MAX_CHUNKS)));
 	}
 
+	int32 max_batches =
+		ts_jsonb_get_int32_field(config, POL_COMPACTION_CONF_KEY_MAX_BATCHES, &found);
+	if (found && max_batches < 0)
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("%s must be greater than or equal to 0",
+						POL_COMPACTION_CONF_KEY_MAX_BATCHES)));
+	}
+
 	/* Reading the field parses it, raising on a malformed interval */
 	Interval *inactive_for =
 		ts_jsonb_get_interval_field(config, POL_COMPACTION_CONF_KEY_INACTIVE_FOR);
@@ -106,7 +116,8 @@ policy_compaction_add(PG_FUNCTION_ARGS)
 	bool if_not_exists = PG_GETARG_BOOL(1);
 	bool user_defined_schedule_interval = !PG_ARGISNULL(2);
 	int32 max_chunks = PG_ARGISNULL(5) ? 0 : PG_GETARG_INT32(5);
-	Interval *inactive_for = PG_ARGISNULL(6) ? NULL : PG_GETARG_INTERVAL_P(6);
+	int32 max_batches = PG_ARGISNULL(6) ? 0 : PG_GETARG_INT32(6);
+	Interval *inactive_for = PG_ARGISNULL(7) ? NULL : PG_GETARG_INTERVAL_P(7);
 	Cache *hcache;
 	Hypertable *ht;
 	int32 hypertable_id;
@@ -137,6 +148,13 @@ policy_compaction_add(PG_FUNCTION_ARGS)
 				 errmsg("max_chunks must be greater than or equal to 0")));
 	}
 
+	if (max_batches < 0)
+	{
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("max_batches must be greater than or equal to 0")));
+	}
+
 	ht = ts_hypertable_cache_get_cache_and_entry(ht_oid, CACHE_FLAG_NONE, &hcache);
 	Assert(ht != NULL);
 	hypertable_id = ht->fd.id;
@@ -205,6 +223,10 @@ policy_compaction_add(PG_FUNCTION_ARGS)
 	{
 		ts_jsonb_add_int32(&parse_state, POL_COMPACTION_CONF_KEY_MAX_CHUNKS, max_chunks);
 	}
+	if (max_batches > 0)
+	{
+		ts_jsonb_add_int32(&parse_state, POL_COMPACTION_CONF_KEY_MAX_BATCHES, max_batches);
+	}
 	if (inactive_for != NULL)
 	{
 		ts_jsonb_add_interval(&parse_state, POL_COMPACTION_CONF_KEY_INACTIVE_FOR, inactive_for);

tsl/src/bgw_policy/policies_v2.h

@@ -42,6 +42,7 @@
 #define POLICY_COMPACTION_PROC_NAME "policy_compaction"
 #define POLICY_COMPACTION_CHECK_NAME "policy_compaction_check"
 #define POL_COMPACTION_CONF_KEY_MAX_CHUNKS "max_chunks"
+#define POL_COMPACTION_CONF_KEY_MAX_BATCHES "max_batches"
 #define POL_COMPACTION_CONF_KEY_INACTIVE_FOR "inactive_for"
 
 #define SHOW_POLICY_KEY_POLICY_NAME "policy_name"

tsl/src/compression/recompress.c

@@ -155,15 +155,17 @@ static bool batches_overlap_firstlast(RecompressContext *recompress_ctx, Datum *
 static void decompress_batch_to_tuplesort(TupleTableSlot *slot, TupleDesc tupdesc,
 										  RowDecompressor *decompressor,
 										  Tuplesortstate *recompress_tuplesortstate,
-										  Relation compressed_chunk_rel, Snapshot snapshot);
+										  Relation compressed_chunk_rel, Snapshot snapshot,
+										  int *processed_batches);
 static bool compact_chunk_find_overlapping_batches(Relation compressed_chunk_rel,
 												   IndexScanDesc index_scan,
 												   RecompressContext *recompress_ctx,
 												   CompactChunkScanState *state);
 static bool compact_chunk_recompress_overlapping_batches(
 	Relation compressed_chunk_rel, IndexScanDesc index_scan, Snapshot snapshot,
 	RecompressContext *recompress_ctx, CompactChunkScanState *state, RowCompressor *compressor,
-	RowDecompressor *decompressor, Tuplesortstate *recompress_tuplesortstate, BulkWriter *writer);
+	RowDecompressor *decompressor, Tuplesortstate *recompress_tuplesortstate, BulkWriter *writer,
+	int max_batches);
 static void try_updating_chunk_status(Chunk *uncompressed_chunk, Relation uncompressed_chunk_rel);
 
 /*
@@ -272,7 +274,10 @@ tsl_compact_chunk(PG_FUNCTION_ARGS)
 						ts_chunk_get_table_name(chunk))));
 	}
 
-	uncompressed_relid = compact_chunk_impl(chunk);
+	int max_batches = PG_GETARG_INT32(1);
+	Assert(max_batches >= 0);
+
+	uncompressed_relid = compact_chunk_impl(chunk, max_batches);
 
 	PG_RETURN_OID(uncompressed_relid);
 }
@@ -1061,7 +1066,8 @@ static void
 decompress_batch_to_tuplesort(TupleTableSlot *slot, TupleDesc tupdesc,
 							  RowDecompressor *decompressor,
 							  Tuplesortstate *recompress_tuplesortstate,
-							  Relation compressed_chunk_rel, Snapshot snapshot)
+							  Relation compressed_chunk_rel, Snapshot snapshot,
+							  int *processed_batches)
 {
 	bool should_free;
 	HeapTuple compressed_tuple = ExecFetchSlotHeapTuple(slot, false, &should_free);
@@ -1092,6 +1098,11 @@ decompress_batch_to_tuplesort(TupleTableSlot *slot, TupleDesc tupdesc,
 	{
 		heap_freetuple(compressed_tuple);
 	}
+
+	if (processed_batches)
+	{
+		(*processed_batches)++;
+	}
 }
 
 /*
@@ -1169,14 +1180,17 @@ static bool
 compact_chunk_recompress_overlapping_batches(
 	Relation compressed_chunk_rel, IndexScanDesc index_scan, Snapshot snapshot,
 	RecompressContext *recompress_ctx, CompactChunkScanState *state, RowCompressor *compressor,
-	RowDecompressor *decompressor, Tuplesortstate *recompress_tuplesortstate, BulkWriter *writer)
+	RowDecompressor *decompressor, Tuplesortstate *recompress_tuplesortstate, BulkWriter *writer,
+	int max_batches)
 {
 	TupleTableSlot *previous_compressed_slot = table_slot_create(compressed_chunk_rel, NULL);
 	TupleTableSlot *compressed_slot = table_slot_create(compressed_chunk_rel, NULL);
 
 	TupleDesc compressed_rel_tupdesc = RelationGetDescr(compressed_chunk_rel);
 	bool overlapping = false;
 	bool found_overlaps = false;
+	/* Counts decompressed batches */
+	int processed_batches = 0;
 
 	/*
 	 * The find pass identified the first overlapping pair. Fetch both batches by
@@ -1200,7 +1214,8 @@ compact_chunk_recompress_overlapping_batches(
 									  decompressor,
 									  recompress_tuplesortstate,
 									  compressed_chunk_rel,
-									  snapshot);
+									  snapshot,
+									  &processed_batches);
 
 		found = table_index_fetch_tuple(index_scan->xs_heapfetch,
 										&state->previous_tid,
@@ -1214,7 +1229,8 @@ compact_chunk_recompress_overlapping_batches(
 									  decompressor,
 									  recompress_tuplesortstate,
 									  compressed_chunk_rel,
-									  snapshot);
+									  snapshot,
+									  &processed_batches);
 
 		overlapping = true;
 		found_overlaps = true;
@@ -1248,6 +1264,12 @@ compact_chunk_recompress_overlapping_batches(
 								   compressor,
 								   writer);
 				overlapping = false;
+
+				/* Check only after a full flush so we never leave partial work */
+				if (max_batches > 0 && processed_batches >= max_batches)
+				{
+					break;
+				}
 			}
 
 			ItemPointerCopy(&index_scan->xs_heaptid, &state->previous_tid);
@@ -1285,7 +1307,8 @@ compact_chunk_recompress_overlapping_batches(
 											  decompressor,
 											  recompress_tuplesortstate,
 											  compressed_chunk_rel,
-											  snapshot);
+											  snapshot,
+											  &processed_batches);
 
 				overlapping = true;
 				found_overlaps = true;
@@ -1296,7 +1319,8 @@ compact_chunk_recompress_overlapping_batches(
 										  decompressor,
 										  recompress_tuplesortstate,
 										  compressed_chunk_rel,
-										  snapshot);
+										  snapshot,
+										  &processed_batches);
 
 			CommandCounterIncrement();
 		}
@@ -1307,6 +1331,12 @@ compact_chunk_recompress_overlapping_batches(
 			recompress_segment(recompress_tuplesortstate, compressed_chunk_rel, compressor, writer);
 			overlapping = false;
 			CommandCounterIncrement();
+
+			/* Check only after a full flush so we never leave partial work */
+			if (max_batches > 0 && processed_batches >= max_batches)
+			{
+				break;
+			}
 		}
 
 		ItemPointerCopy(&index_scan->xs_heaptid, &state->previous_tid);
@@ -1325,7 +1355,7 @@ compact_chunk_recompress_overlapping_batches(
 }
 
 Oid
-compact_chunk_impl(Chunk *uncompressed_chunk)
+compact_chunk_impl(Chunk *uncompressed_chunk, int max_batches)
 {
 	Oid uncompressed_chunk_id = uncompressed_chunk->fd.relid;
 
@@ -1496,7 +1526,8 @@ compact_chunk_impl(Chunk *uncompressed_chunk)
 													 &compressor,
 													 &decompressor,
 													 recompress_tuplesortstate,
-													 &writer);
+													 &writer,
+													 max_batches);
 		row_compressor_close(&compressor);
 		row_decompressor_close(&decompressor);
 		tuplesort_end(recompress_tuplesortstate);

tsl/src/compression/recompress.h

@@ -43,7 +43,7 @@ extern Datum tsl_recompress_chunk_segmentwise(PG_FUNCTION_ARGS);
 extern Datum tsl_compact_chunk(PG_FUNCTION_ARGS);
 
 void recompress_chunk_segmentwise_impl(Chunk *chunk, bool fullrecompress);
-Oid compact_chunk_impl(Chunk *chunk);
+Oid compact_chunk_impl(Chunk *chunk, int max_batches);
 bool recompress_chunk_in_memory_impl(Chunk *uncompressed_chunk);
 void rebuild_sparse_index_impl(Chunk *uncompressed_chunk, bool force);