Refactor parts of the delta-delta compression

to split up the UUID compression PR.
In this refactor, I make changes such that we can serialize
compressed but data into a buffer, rather than let finalize()
allocate the space and then we copy the data. This saves a few
allocations and memcpy calls.

To enable the above changes we need a precise way to know
the size of the buffer we will need without actully finalizing
the data. A new function was added to do this as well as
extensive unit tests to verify the size calculations.

Disable-check: force-changelog-file

Author: dbeck

tsl/src/compression/algorithms/deltadelta.c

@@ -301,69 +301,99 @@ delta_delta_compressor_alloc(void)
 	return compressor;
 }
 
-static DeltaDeltaCompressed *
-delta_delta_from_parts(uint64 last_value, uint64 last_delta, Simple8bRleSerialized *deltas,
-					   Simple8bRleSerialized *nulls)
+static void *
+delta_delta_set_header_and_advance(uint64 last_value, uint64 last_delta, bool has_nulls,
+								   size_t compressed_size, void *dest)
 {
-	uint32 nulls_size = 0;
-	Size compressed_size;
-	char *compressed_data;
-	DeltaDeltaCompressed *compressed;
-
-	if (nulls != NULL)
-		nulls_size = simple8brle_serialized_total_size(nulls);
-
-	compressed_size =
-		sizeof(DeltaDeltaCompressed) + simple8brle_serialized_total_size(deltas) + nulls_size;
-
-	if (!AllocSizeIsValid(compressed_size))
-		ereport(ERROR,
-				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
-				 errmsg("compressed size exceeds the maximum allowed (%d)", (int) MaxAllocSize)));
-
-	compressed_data = palloc(compressed_size);
-	compressed = (DeltaDeltaCompressed *) compressed_data;
+	DeltaDeltaCompressed *compressed = (DeltaDeltaCompressed *) dest;
 	SET_VARSIZE(&compressed->vl_len_, compressed_size);
-
 	compressed->compression_algorithm = COMPRESSION_ALGORITHM_DELTADELTA;
 	compressed->last_value = last_value;
 	compressed->last_delta = last_delta;
-	compressed->has_nulls = nulls_size != 0 ? 1 : 0;
+	compressed->has_nulls = has_nulls ? 1 : 0;
+	compressed->padding[0] = 0;
+	compressed->padding[1] = 0;
+	return (char *) compressed + sizeof(*compressed);
+}
 
-	compressed_data += sizeof(*compressed);
-	compressed_data =
-		bytes_serialize_simple8b_and_advance(compressed_data,
-											 simple8brle_serialized_total_size(deltas),
-											 deltas);
+size_t
+delta_delta_compressor_compressed_size(DeltaDeltaCompressor *compressor, size_t *nulls_size_out)
+{
+	size_t compressed_size = sizeof(DeltaDeltaCompressed);
+	size_t nulls_size_actual;
 
-	if (compressed->has_nulls == 1 && nulls != NULL)
+	/* If there are no elements, the compressed size is 0 even if there are nulls */
+	if ((compressor->delta_delta.num_elements +
+		 compressor->delta_delta.num_uncompressed_elements) == 0)
 	{
-		CheckCompressedData(nulls->num_elements > deltas->num_elements);
-		bytes_serialize_simple8b_and_advance(compressed_data, nulls_size, nulls);
+		if (nulls_size_out != NULL)
+			*nulls_size_out = 0;
+		return 0;
 	}
 
-	return compressed;
+	compressed_size += simple8brle_compressor_compressed_const_size(&compressor->delta_delta);
+
+	if (compressor->has_nulls)
+	{
+		nulls_size_actual = simple8brle_compressor_compressed_const_size(&compressor->nulls);
+		compressed_size += nulls_size_actual;
+		if (nulls_size_out != NULL)
+			*nulls_size_out = nulls_size_actual;
+	}
+	else if (nulls_size_out != NULL)
+		*nulls_size_out = 0;
+
+	return compressed_size;
 }
 
 void *
 delta_delta_compressor_finish(DeltaDeltaCompressor *compressor)
 {
-	Simple8bRleSerialized *deltas = simple8brle_compressor_finish(&compressor->delta_delta);
-	Simple8bRleSerialized *nulls = simple8brle_compressor_finish(&compressor->nulls);
-	DeltaDeltaCompressed *compressed;
+	size_t total_size = delta_delta_compressor_compressed_size(compressor, NULL);
+	char *compressed = NULL;
 
-	if (deltas == NULL)
+	if (total_size == 0)
 		return NULL;
 
-	compressed = delta_delta_from_parts(compressor->prev_val,
-										compressor->prev_delta,
-										deltas,
-										compressor->has_nulls ? nulls : NULL);
-
-	Assert(compressed->compression_algorithm == COMPRESSION_ALGORITHM_DELTADELTA);
+	compressed = palloc(total_size);
+	delta_delta_compressor_finish_into(compressor, compressed);
 	return compressed;
 }
 
+void *
+delta_delta_compressor_finish_into(DeltaDeltaCompressor *compressor, void *dest)
+{
+	size_t data_size;
+	size_t nulls_size;
+	size_t compressed_size;
+	char *result = (char *) dest;
+
+	/* The compressed size includes the header and the nulls */
+	compressed_size = delta_delta_compressor_compressed_size(compressor, &nulls_size);
+	if (compressed_size == 0)
+		return dest;
+
+	/* Check if the data size is valid */
+	data_size = compressed_size - sizeof(DeltaDeltaCompressed) - nulls_size;
+	Assert(compressed_size > (sizeof(DeltaDeltaCompressed) + nulls_size));
+
+	result = delta_delta_set_header_and_advance(compressor->prev_val,
+												compressor->prev_delta,
+												compressor->has_nulls,
+												compressed_size,
+												dest);
+
+	result = simple8brle_compressor_finish_into(&compressor->delta_delta, result, data_size);
+
+	if (compressor->has_nulls)
+	{
+		Assert(nulls_size > 0);
+		result = simple8brle_compressor_finish_into(&compressor->nulls, result, nulls_size);
+	}
+
+	return result;
+}
+
 Datum
 tsl_deltadelta_compressor_finish(PG_FUNCTION_ARGS)
 {
@@ -423,7 +453,11 @@ int64_decompression_iterator_init_forward(DeltaDeltaDecompressionIterator *iter,
 										  Oid element_type)
 {
 	StringInfoData si = { .data = compressed, .len = VARSIZE(compressed) };
+	CheckCompressedData(VARSIZE(compressed) >= sizeof(DeltaDeltaCompressed));
+	CheckCompressedData(((DeltaDeltaCompressed *) compressed)->compression_algorithm ==
+						COMPRESSION_ALGORITHM_DELTADELTA);
 	DeltaDeltaCompressed *header = consumeCompressedData(&si, sizeof(DeltaDeltaCompressed));
+
 	Simple8bRleSerialized *deltas = bytes_deserialize_simple8b_and_advance(&si);
 
 	const bool has_nulls = header->has_nulls == 1;
@@ -456,6 +490,9 @@ int64_decompression_iterator_init_reverse(DeltaDeltaDecompressionIterator *iter,
 										  Oid element_type)
 {
 	StringInfoData si = { .data = compressed, .len = VARSIZE(compressed) };
+	CheckCompressedData(VARSIZE(compressed) >= sizeof(DeltaDeltaCompressed));
+	CheckCompressedData(((DeltaDeltaCompressed *) compressed)->compression_algorithm ==
+						COMPRESSION_ALGORITHM_DELTADELTA);
 	DeltaDeltaCompressed *header = consumeCompressedData(&si, sizeof(DeltaDeltaCompressed));
 	Simple8bRleSerialized *deltas = bytes_deserialize_simple8b_and_advance(&si);
 
@@ -723,20 +760,54 @@ deltadelta_compressed_recv(StringInfo buffer)
 	uint8 has_nulls;
 	uint64 last_value;
 	uint64 last_delta;
-	Simple8bRleSerialized *delta_deltas;
+	Simple8bRleSerialized *delta_delta_values;
 	Simple8bRleSerialized *nulls = NULL;
 	DeltaDeltaCompressed *compressed;
+	void *buf_ptr = NULL;
+	int delta_size = 0;
+	int nulls_size = 0;
+	size_t compressed_size = 0;
+
+	size_t allocated_size = 2 * sizeof(Simple8bRleSerialized) + sizeof(DeltaDeltaCompressed) +
+							(buffer->len - buffer->cursor);
+	compressed = palloc(allocated_size);
 
 	has_nulls = pq_getmsgbyte(buffer);
 	CheckCompressedData(has_nulls == 0 || has_nulls == 1);
 
 	last_value = pq_getmsgint64(buffer);
 	last_delta = pq_getmsgint64(buffer);
-	delta_deltas = simple8brle_serialized_recv(buffer);
+
+	/* Leave space for the header, but we don't yet know the size of the compressed data */
+	buf_ptr = (char *) compressed + sizeof(DeltaDeltaCompressed);
+
+	/* Calculate the size of the delta delta values based on the number of bytes read */
+	delta_size = buffer->cursor;
+	buf_ptr = simple8brle_serialized_recv_into(buffer, buf_ptr, &delta_delta_values);
+	delta_size = buffer->cursor - delta_size;
+
 	if (has_nulls)
-		nulls = simple8brle_serialized_recv(buffer);
+	{
+		/* Calculate the size of the nulls based on the number of bytes read */
+		nulls_size = buffer->cursor;
+		buf_ptr = simple8brle_serialized_recv_into(buffer, buf_ptr, &nulls);
+		nulls_size = buffer->cursor - nulls_size;
+		CheckCompressedData(delta_delta_values->num_elements < nulls->num_elements);
+	}
+
+	compressed_size = sizeof(DeltaDeltaCompressed) + delta_size + nulls_size;
+
+	if (!AllocSizeIsValid(compressed_size))
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("compressed size exceeds the maximum allowed (%d)", (int) MaxAllocSize)));
 
-	compressed = delta_delta_from_parts(last_value, last_delta, delta_deltas, nulls);
+	/* Set header but don't change the buffer pointer */
+	delta_delta_set_header_and_advance(last_value,
+									   last_delta,
+									   has_nulls,
+									   compressed_size,
+									   compressed);
 
 	PG_RETURN_POINTER(compressed);
 }

tsl/src/compression/algorithms/deltadelta.h

@@ -33,6 +33,9 @@ extern DeltaDeltaCompressor *delta_delta_compressor_alloc(void);
 extern void delta_delta_compressor_append_null(DeltaDeltaCompressor *compressor);
 extern void delta_delta_compressor_append_value(DeltaDeltaCompressor *compressor, int64 next_val);
 extern void *delta_delta_compressor_finish(DeltaDeltaCompressor *compressor);
+extern size_t delta_delta_compressor_compressed_size(DeltaDeltaCompressor *compressor,
+													 size_t *nulls_size /* out */);
+extern void *delta_delta_compressor_finish_into(DeltaDeltaCompressor *compressor, void *dest);
 
 extern DecompressionIterator *
 delta_delta_decompression_iterator_from_datum_forward(Datum deltadelta_compressed,

tsl/src/compression/algorithms/simple8b_rle.h

@@ -138,6 +138,8 @@ typedef struct Simple8bRleDecompressResult
 static inline void simple8brle_compressor_init(Simple8bRleCompressor *compressor);
 static inline Simple8bRleSerialized *
 simple8brle_compressor_finish(Simple8bRleCompressor *compressor);
+static inline char *simple8brle_compressor_finish_into(Simple8bRleCompressor *compressor,
+													   char *dest, size_t expected_size);
 static inline void simple8brle_compressor_append(Simple8bRleCompressor *compressor, uint64 val);
 static inline bool simple8brle_compressor_is_empty(Simple8bRleCompressor *compressor);
 
@@ -245,6 +247,33 @@ simple8brle_serialized_recv(StringInfo buffer)
 	return data;
 }
 
+static inline void *
+simple8brle_serialized_recv_into(StringInfo buffer, void *dest, Simple8bRleSerialized **data_out)
+{
+	uint32 i;
+	uint32 num_elements = pq_getmsgint32(buffer);
+	CheckCompressedData(num_elements <= GLOBAL_MAX_ROWS_PER_COMPRESSION);
+	uint32 num_blocks = pq_getmsgint32(buffer);
+	CheckCompressedData(num_blocks <= GLOBAL_MAX_ROWS_PER_COMPRESSION);
+	uint32 num_selector_slots = simple8brle_num_selector_slots_for_num_blocks(num_blocks);
+
+	Size compressed_size =
+		sizeof(Simple8bRleSerialized) + (num_blocks + num_selector_slots) * sizeof(uint64);
+	if (!AllocSizeIsValid(compressed_size))
+		ereport(ERROR,
+				(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+				 errmsg("compressed size exceeds the maximum allowed (%d)", (int) MaxAllocSize)));
+
+	*data_out = (Simple8bRleSerialized *) dest;
+	(*data_out)->num_elements = num_elements;
+	(*data_out)->num_blocks = num_blocks;
+
+	for (i = 0; i < num_blocks + num_selector_slots; i++)
+		(*data_out)->slots[i] = pq_getmsgint64(buffer);
+
+	return (char *) *data_out + compressed_size;
+}
+
 static void
 simple8brle_serialized_send(StringInfo buffer, const Simple8bRleSerialized *data)
 {
@@ -503,6 +532,63 @@ simple8brle_compressor_finish(Simple8bRleCompressor *compressor)
 	return compressed;
 }
 
+static char *
+simple8brle_compressor_finish_into(Simple8bRleCompressor *compressor, char *dest,
+								   size_t expected_size)
+{
+	size_t size_left;
+	size_t selector_size;
+	size_t compressed_size;
+	char *end_ptr;
+	Simple8bRleSerialized *compressed;
+	uint64 bits;
+
+	simple8brle_compressor_flush(compressor);
+	if (compressor->num_elements == 0)
+		return dest;
+
+	Ensure(dest != NULL, "dest is NULL");
+
+	Assert(compressor->last_block_set);
+	simple8brle_compressor_push_block(compressor, compressor->last_block);
+
+	compressed_size = simple8brle_compressor_compressed_size(compressor);
+	Ensure(expected_size == compressed_size,
+		   "expected_size: %zu, compressed_size: %zu",
+		   expected_size,
+		   compressed_size);
+
+	compressed = (Simple8bRleSerialized *) dest;
+	Assert(bit_array_num_buckets(&compressor->selectors) > 0);
+	Assert(compressor->compressed_data.num_elements > 0);
+	Assert(compressor->compressed_data.num_elements ==
+		   simple8brle_compressor_num_selectors(compressor));
+
+	*compressed = (Simple8bRleSerialized){
+		.num_elements = compressor->num_elements,
+		.num_blocks = compressor->compressed_data.num_elements,
+	};
+
+	size_left = compressed_size - sizeof(*compressed);
+	Assert(size_left >= bit_array_data_bytes_used(&compressor->selectors));
+	selector_size = bit_array_output(&compressor->selectors, compressed->slots, size_left, &bits);
+
+	size_left -= selector_size;
+	Assert(size_left ==
+		   (compressor->compressed_data.num_elements * sizeof(*compressor->compressed_data.data)));
+	Assert(compressor->selectors.buckets.num_elements ==
+		   simple8brle_num_selector_slots_for_num_blocks(compressor->compressed_data.num_elements));
+
+	memcpy(compressed->slots + compressor->selectors.buckets.num_elements,
+		   compressor->compressed_data.data,
+		   size_left);
+
+	end_ptr = (char *) (compressed->slots + compressor->selectors.buckets.num_elements) + size_left;
+	Assert(end_ptr == dest + expected_size);
+
+	return end_ptr;
+}
+
 static void
 simple8brle_compressor_flush(Simple8bRleCompressor *compressor)
 {