Partial decompression in TOAST iterators

Nikita Glukhov · Nikita Glukhov · commit 30743fec146c · 2021-03-29T20:40:14.000+03:00
diff --git a/src/backend/access/common/detoast.c b/src/backend/access/common/detoast.c
@@ -365,6 +365,8 @@ create_detoast_iterator(struct varlena *attr)
 			/* initialize state for pglz_decompress_iterate() */
 			iter->ctrl = 0;
 			iter->ctrlc = INVALID_CTRLC;
+			iter->len = 0;
+			iter->off = 0;
 		}
 		else
 		{
@@ -390,7 +392,7 @@ create_detoast_iterator(struct varlena *attr)
 		return create_detoast_iterator(attr);
 
 	}
-	else if (VARATT_IS_COMPRESSED(attr))
+	else if (1 && VARATT_IS_COMPRESSED(attr))
 	{
 		ToastBuffer *buf;
 
diff --git a/src/backend/access/common/toast_compression.c b/src/backend/access/common/toast_compression.c
@@ -247,6 +247,82 @@ lz4_decompress_datum_slice(const struct varlena *value, int32 slicelength)
 #endif
 }
 
+void
+toast_decompress_iterate(ToastBuffer *source, ToastBuffer *dest,
+						 DetoastIterator iter, const char *destend)
+{
+	const char *sp;
+	const char *srcend;
+	char	   *dp;
+	int32		dlen;
+	int32		slen;
+	bool		last_source_chunk;
+
+	/*
+	 * In the while loop, sp may be incremented such that it points beyond
+	 * srcend. To guard against reading beyond the end of the current chunk,
+	 * we set srcend such that we exit the loop when we are within four bytes
+	 * of the end of the current chunk. When source->limit reaches
+	 * source->capacity, we are decompressing the last chunk, so we can (and
+	 * need to) read every byte.
+	 */
+	last_source_chunk = source->limit == source->capacity;
+	srcend = last_source_chunk ? source->limit : source->limit - 4;
+	sp = source->position;
+	dp = dest->limit;
+	if (destend > dest->capacity)
+		destend = dest->capacity;
+
+	slen = srcend - source->position;
+
+	/*
+	 * Decompress the data using the appropriate decompression routine.
+	 */
+	switch (iter->compression_method)
+	{
+		case TOAST_PGLZ_COMPRESSION_ID:
+			dlen = pglz_decompress_state(sp, &slen, dp, destend - dp,
+										 last_source_chunk && destend == dest->capacity,
+										 last_source_chunk,
+										 &iter->decompression_state);
+			break;
+		case TOAST_LZ4_COMPRESSION_ID:
+			if (source->limit < source->capacity)
+				dlen = 0;	/* LZ4 needs need full data to decompress */
+			else
+			{
+				/* decompress the data */
+#ifndef USE_LZ4
+				NO_LZ4_SUPPORT();
+				dlen = 0;
+#else
+				dlen = LZ4_decompress_safe(source->buf + VARHDRSZ_COMPRESSED,
+										   VARDATA(dest->buf),
+										   VARSIZE(source->buf) - VARHDRSZ_COMPRESSED,
+										   VARDATA_COMPRESSED_GET_EXTSIZE(source->buf));
+
+				if (dlen < 0)
+					ereport(ERROR,
+							(errcode(ERRCODE_DATA_CORRUPTED),
+							 errmsg_internal("compressed lz4 data is corrupt")));
+#endif
+				slen = 0;
+				break;
+			}
+		default:
+			elog(ERROR, "invalid compression method id %d", iter->compression_method);
+			return;		/* keep compiler quiet */
+	}
+
+	if (dlen < 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATA_CORRUPTED),
+				 errmsg_internal("compressed data is corrupt")));
+
+	source->position += slen;
+	dest->limit += dlen;
+}
+
 /*
  * Extract compression ID from a varlena.
  *
diff --git a/src/backend/access/common/toast_internals.c b/src/backend/access/common/toast_internals.c
@@ -919,9 +919,33 @@ pglz_decompress_iterate(ToastBuffer *source, ToastBuffer *dest,
 		(source->limit == source->capacity ? source->limit : (source->limit - 4));
 	sp = (const unsigned char *) source->position;
 	dp = (unsigned char *) dest->limit;
-	if (destend < (unsigned char *) dest->capacity)
+	if (destend > (unsigned char *) dest->capacity)
 		destend = (unsigned char *) dest->capacity;
 
+	if (iter->len)
+	{
+		int32		len = iter->len;
+		int32		off = iter->off;
+		int32		copylen = Min(len, destend - dp);
+		int32		remlen = len - copylen;
+
+		while (copylen--)
+		{
+			*dp = dp[-off];
+			dp++;
+		}
+
+		iter->len = remlen;
+
+		if (dp >= destend)
+		{
+			dest->limit = (char *) dp;
+			return;
+		}
+
+		Assert(remlen == 0);
+	}
+
 	while (sp < srcend && dp < destend)
 	{
 		/*
@@ -942,7 +966,6 @@ pglz_decompress_iterate(ToastBuffer *source, ToastBuffer *dest,
 			ctrlc = 0;
 		}
 
-
 		for (; ctrlc < INVALID_CTRLC && sp < srcend && dp < destend; ctrlc++)
 		{
 
@@ -959,6 +982,7 @@ pglz_decompress_iterate(ToastBuffer *source, ToastBuffer *dest,
 				 */
 				int32		len;
 				int32		off;
+				int32		copylen;
 
 				len = (sp[0] & 0x0f) + 3;
 				off = ((sp[0] & 0xf0) << 4) | sp[1];
@@ -972,17 +996,21 @@ pglz_decompress_iterate(ToastBuffer *source, ToastBuffer *dest,
 				 * areas could overlap; to prevent possible uncertainty, we
 				 * copy only non-overlapping regions.
 				 */
-				len = Min(len, destend - dp);
-				while (off < len)
+				copylen = Min(len, destend - dp);
+				iter->len = len - copylen;
+
+				while (off < copylen)
 				{
 					/* see comments in common/pg_lzcompress.c */
 					memcpy(dp, dp - off, off);
-					len -= off;
+					copylen -= off;
 					dp += off;
 					off += off;
 				}
-				memcpy(dp, dp - off, len);
-				dp += len;
+				memcpy(dp, dp - off, copylen);
+				dp += copylen;
+
+				iter->off = off;
 			}
 			else
 			{
diff --git a/src/common/pg_lzcompress.c b/src/common/pg_lzcompress.c
@@ -700,7 +700,8 @@ typedef struct pglz_state
  */
 int32
 pglz_decompress_state(const char *source, int32 slen, char *dest,
-					  int32 rawsize, bool check_complete, void **pstate)
+					  int32 dlen, bool check_complete, bool last_cource_chunk,
+					  void **pstate)
 {
 	pglz_state *state = pstate ? *pstate : NULL;
 	const unsigned char *sp;
@@ -800,7 +801,7 @@ pglz_decompress_state(const char *source, int32 slen, char *dest,
 				 * must check this, else we risk an infinite loop below in the
 				 * face of corrupt data.)
 				 */
-				if (unlikely(sp > srcend || off == 0))
+				if (unlikely((sp > srcend && last_cource_chunk) || off == 0))
 					return -1;
 
 				/*
@@ -947,5 +948,5 @@ int32
 pglz_decompress(const char *source, int32 slen, char *dest, int32 rawsize,
 				bool check_complete)
 {
-	return pglz_decompress_state(source, slen, dest, rawsize, check_complete, NULL);
+	return pglz_decompress_state(source, slen, dest, rawsize, check_complete, true, NULL);
 }
diff --git a/src/include/access/detoast.h b/src/include/access/detoast.h
@@ -123,6 +123,8 @@ typedef struct DetoastIteratorData
 	unsigned char		ctrl;
 	int					ctrlc;
 	int					nrefs;
+	int32				len;
+	int32				off;
 	bool				compressed;		/* toast value is compressed? */
 	bool				done;
 }			DetoastIteratorData;
@@ -174,16 +176,16 @@ detoast_iterate(DetoastIterator detoast_iter, const char *destend)
 	if (!detoast_iter->compressed)
 		destend = NULL;
 
-	if (destend)
+	if (1 && destend)
 	{
 		const char *srcend = (const char *)
 			(fetch_iter->buf->limit == fetch_iter->buf->capacity ?
 			fetch_iter->buf->limit : fetch_iter->buf->limit - 4);
 
-		if (fetch_iter->buf->position >= srcend)
+		if (fetch_iter->buf->position >= srcend && !fetch_iter->done)
 			fetch_datum_iterate(fetch_iter);
 	}
-	else
+	else if (!fetch_iter->done)
 		fetch_datum_iterate(fetch_iter);
 
 	if (detoast_iter->compressed)
diff --git a/src/include/common/pg_lzcompress.h b/src/include/common/pg_lzcompress.h
@@ -87,8 +87,9 @@ extern int32 pglz_compress(const char *source, int32 slen, char *dest,
 						   const PGLZ_Strategy *strategy);
 extern int32 pglz_decompress(const char *source, int32 slen, char *dest,
 							 int32 rawsize, bool check_complete);
-extern int32 pglz_decompress_state(const char *source, int32 slen, char *dest,
-								   int32 rawsize, bool check_complete,
+extern int32 pglz_decompress_state(const char *source, int32 slen,
+								   char *dest, int32 dlen,
+								   bool check_complete, bool last_source_chunk,
 								   void **state);
 extern int32 pglz_maximum_compressed_size(int32 rawsize,
 										  int32 total_compressed_size);

Original file line number	Diff line number	Diff line change
`@@ -365,6 +365,8 @@ create_detoast_iterator(struct varlena *attr)`
`365`	`365`	`/* initialize state for pglz_decompress_iterate() */`
`366`	`366`	`iter->ctrl = 0;`
`367`	`367`	`iter->ctrlc = INVALID_CTRLC;`
	`368`	`+ iter->len = 0;`
	`369`	`+ iter->off = 0;`
`368`	`370`	`}`
`369`	`371`	`else`
`370`	`372`	`{`
`@@ -390,7 +392,7 @@ create_detoast_iterator(struct varlena *attr)`
`390`	`392`	`return create_detoast_iterator(attr);`
`391`	`393`
`392`	`394`	`}`
`393`		`- else if (VARATT_IS_COMPRESSED(attr))`
	`395`	`+ else if (1 && VARATT_IS_COMPRESSED(attr))`
`394`	`396`	`{`
`395`	`397`	`ToastBuffer *buf;`
`396`	`398`