bpo-41486: Faster bz2/lzma/zlib via new output buffering (GH-21740)

Faster bz2/lzma/zlib via new output buffering. Also adds .readall() function to _compression.DecompressReader class to take best advantage of this in the consume-all-output at once scenario. Often a 5-20% speedup in common scenarios due to less data copying. Contributed by Ma Lin.

Author: Ma Lin

Doc/whatsnew/3.10.rst

@@ -1243,6 +1243,12 @@ Optimizations
  for more details. (Contributed by Victor Stinner and Pablo Galindo in
  :issue:`38980`.)
 
+* Use a new output buffer management code for :mod:`bz2` / :mod:`lzma` /
+ :mod:`zlib` modules, and add ``.readall()`` function to
+ ``_compression.DecompressReader`` class. bz2 decompression 1.09x ~ 1.17x
+ faster, lzma decompression 1.20x ~ 1.32x faster, ``GzipFile.read(-1)`` 1.11x
+ ~ 1.18x faster. (Contributed by Ma Lin, reviewed by Gregory P. Smith, in :issue:`41486`)
+
 * Function parameters and their annotations are no longer computed at runtime,
  but rather at compilation time. They are stored as a tuple of strings at the
  bytecode level. It is now around 2 times faster to create a function with

Include/internal/pycore_blocks_output_buffer.h

@@ -0,0 +1,317 @@
+/*
+ _BlocksOutputBuffer is used to maintain an output buffer
+ that has unpredictable size. Suitable for compression/decompression
+ API (bz2/lzma/zlib) that has stream->next_out and stream->avail_out:
+
+ stream->next_out: point to the next output position.
+ stream->avail_out: the number of available bytes left in the buffer.
+
+ It maintains a list of bytes object, so there is no overhead of resizing
+ the buffer.
+
+ Usage:
+
+ 1, Initialize the struct instance like this:
+ _BlocksOutputBuffer buffer = {.list = NULL};
+ Set .list to NULL for _BlocksOutputBuffer_OnError()
+
+ 2, Initialize the buffer use one of these functions:
+ _BlocksOutputBuffer_InitAndGrow()
+ _BlocksOutputBuffer_InitWithSize()
+
+ 3, If (avail_out == 0), grow the buffer:
+ _BlocksOutputBuffer_Grow()
+
+ 4, Get the current outputted data size:
+ _BlocksOutputBuffer_GetDataSize()
+
+ 5, Finish the buffer, and return a bytes object:
+ _BlocksOutputBuffer_Finish()
+
+ 6, Clean up the buffer when an error occurred:
+ _BlocksOutputBuffer_OnError()
+*/
+
+#ifndef Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
+#define Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include "Python.h"
+
+typedef struct {
+ // List of bytes objects
+ PyObject *list;
+ // Number of whole allocated size
+ Py_ssize_t allocated;
+ // Max length of the buffer, negative number means unlimited length.
+ Py_ssize_t max_length;
+} _BlocksOutputBuffer;
+
+static const char unable_allocate_msg[] = "Unable to allocate output buffer.";
+
+/* In 32-bit build, the max block size should <= INT32_MAX. */
+#define OUTPUT_BUFFER_MAX_BLOCK_SIZE (256*1024*1024)
+
+/* Block size sequence */
+#define KB (1024)
+#define MB (1024*1024)
+const Py_ssize_t BUFFER_BLOCK_SIZE[] =
+ { 32*KB, 64*KB, 256*KB, 1*MB, 4*MB, 8*MB, 16*MB, 16*MB,
+ 32*MB, 32*MB, 32*MB, 32*MB, 64*MB, 64*MB, 128*MB, 128*MB,
+ OUTPUT_BUFFER_MAX_BLOCK_SIZE };
+#undef KB
+#undef MB
+
+/* According to the block sizes defined by BUFFER_BLOCK_SIZE, the whole
+ allocated size growth step is:
+ 1 32 KB +32 KB
+ 2 96 KB +64 KB
+ 3 352 KB +256 KB
+ 4 1.34 MB +1 MB
+ 5 5.34 MB +4 MB
+ 6 13.34 MB +8 MB
+ 7 29.34 MB +16 MB
+ 8 45.34 MB +16 MB
+ 9 77.34 MB +32 MB
+ 10 109.34 MB +32 MB
+ 11 141.34 MB +32 MB
+ 12 173.34 MB +32 MB
+ 13 237.34 MB +64 MB
+ 14 301.34 MB +64 MB
+ 15 429.34 MB +128 MB
+ 16 557.34 MB +128 MB
+ 17 813.34 MB +256 MB
+ 18 1069.34 MB +256 MB
+ 19 1325.34 MB +256 MB
+ 20 1581.34 MB +256 MB
+ 21 1837.34 MB +256 MB
+ 22 2093.34 MB +256 MB
+ ...
+*/
+
+/* Initialize the buffer, and grow the buffer.
+
+ max_length: Max length of the buffer, -1 for unlimited length.
+
+ On success, return allocated size (>=0)
+ On failure, return -1
+*/
+static inline Py_ssize_t
+_BlocksOutputBuffer_InitAndGrow(_BlocksOutputBuffer *buffer,
+ const Py_ssize_t max_length,
+ void **next_out)
+{
+ PyObject *b;
+ Py_ssize_t block_size;
+
+ // ensure .list was set to NULL
+ assert(buffer->list == NULL);
+
+ // get block size
+ if (0 <= max_length && max_length < BUFFER_BLOCK_SIZE[0]) {
+ block_size = max_length;
+ } else {
+ block_size = BUFFER_BLOCK_SIZE[0];
+ }
+
+ // the first block
+ b = PyBytes_FromStringAndSize(NULL, block_size);
+ if (b == NULL) {
+ return -1;
+ }
+
+ // create the list
+ buffer->list = PyList_New(1);
+ if (buffer->list == NULL) {
+ Py_DECREF(b);
+ return -1;
+ }
+ PyList_SET_ITEM(buffer->list, 0, b);
+
+ // set variables
+ buffer->allocated = block_size;
+ buffer->max_length = max_length;
+
+ *next_out = PyBytes_AS_STRING(b);
+ return block_size;
+}
+
+/* Initialize the buffer, with an initial size.
+
+ Check block size limit in the outer wrapper function. For example, some libs
+ accept UINT32_MAX as the maximum block size, then init_size should <= it.
+
+ On success, return allocated size (>=0)
+ On failure, return -1
+*/
+static inline Py_ssize_t
+_BlocksOutputBuffer_InitWithSize(_BlocksOutputBuffer *buffer,
+ const Py_ssize_t init_size,
+ void **next_out)
+{
+ PyObject *b;
+
+ // ensure .list was set to NULL
+ assert(buffer->list == NULL);
+
+ // the first block
+ b = PyBytes_FromStringAndSize(NULL, init_size);
+ if (b == NULL) {
+ PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+ return -1;
+ }
+
+ // create the list
+ buffer->list = PyList_New(1);
+ if (buffer->list == NULL) {
+ Py_DECREF(b);
+ return -1;
+ }
+ PyList_SET_ITEM(buffer->list, 0, b);
+
+ // set variables
+ buffer->allocated = init_size;
+ buffer->max_length = -1;
+
+ *next_out = PyBytes_AS_STRING(b);
+ return init_size;
+}
+
+/* Grow the buffer. The avail_out must be 0, please check it before calling.
+
+ On success, return allocated size (>=0)
+ On failure, return -1
+*/
+static inline Py_ssize_t
+_BlocksOutputBuffer_Grow(_BlocksOutputBuffer *buffer,
+ void **next_out,
+ const Py_ssize_t avail_out)
+{
+ PyObject *b;
+ const Py_ssize_t list_len = Py_SIZE(buffer->list);
+ Py_ssize_t block_size;
+
+ // ensure no gaps in the data
+ if (avail_out != 0) {
+ PyErr_SetString(PyExc_SystemError,
+ "avail_out is non-zero in _BlocksOutputBuffer_Grow().");
+ return -1;
+ }
+
+ // get block size
+ if (list_len < (Py_ssize_t) Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE)) {
+ block_size = BUFFER_BLOCK_SIZE[list_len];
+ } else {
+ block_size = BUFFER_BLOCK_SIZE[Py_ARRAY_LENGTH(BUFFER_BLOCK_SIZE) - 1];
+ }
+
+ // check max_length
+ if (buffer->max_length >= 0) {
+ // if (rest == 0), should not grow the buffer.
+ Py_ssize_t rest = buffer->max_length - buffer->allocated;
+ assert(rest > 0);
+
+ // block_size of the last block
+ if (block_size > rest) {
+ block_size = rest;
+ }
+ }
+
+ // check buffer->allocated overflow
+ if (block_size > PY_SSIZE_T_MAX - buffer->allocated) {
+ PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+ return -1;
+ }
+
+ // create the block
+ b = PyBytes_FromStringAndSize(NULL, block_size);
+ if (b == NULL) {
+ PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+ return -1;
+ }
+ if (PyList_Append(buffer->list, b) < 0) {
+ Py_DECREF(b);
+ return -1;
+ }
+ Py_DECREF(b);
+
+ // set variables
+ buffer->allocated += block_size;
+
+ *next_out = PyBytes_AS_STRING(b);
+ return block_size;
+}
+
+/* Return the current outputted data size. */
+static inline Py_ssize_t
+_BlocksOutputBuffer_GetDataSize(_BlocksOutputBuffer *buffer,
+ const Py_ssize_t avail_out)
+{
+ return buffer->allocated - avail_out;
+}
+
+/* Finish the buffer.
+
+ Return a bytes object on success
+ Return NULL on failure
+*/
+static inline PyObject *
+_BlocksOutputBuffer_Finish(_BlocksOutputBuffer *buffer,
+ const Py_ssize_t avail_out)
+{
+ PyObject *result, *block;
+ const Py_ssize_t list_len = Py_SIZE(buffer->list);
+
+ // fast path for single block
+ if ((list_len == 1 && avail_out == 0) ||
+ (list_len == 2 && Py_SIZE(PyList_GET_ITEM(buffer->list, 1)) == avail_out))
+ {
+ block = PyList_GET_ITEM(buffer->list, 0);
+ Py_INCREF(block);
+
+ Py_CLEAR(buffer->list);
+ return block;
+ }
+
+ // final bytes object
+ result = PyBytes_FromStringAndSize(NULL, buffer->allocated - avail_out);
+ if (result == NULL) {
+ PyErr_SetString(PyExc_MemoryError, unable_allocate_msg);
+ return NULL;
+ }
+
+ // memory copy
+ if (list_len > 0) {
+ char *posi = PyBytes_AS_STRING(result);
+
+ // blocks except the last one
+ Py_ssize_t i = 0;
+ for (; i < list_len-1; i++) {
+ block = PyList_GET_ITEM(buffer->list, i);
+ memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block));
+ posi += Py_SIZE(block);
+ }
+ // the last block
+ block = PyList_GET_ITEM(buffer->list, i);
+ memcpy(posi, PyBytes_AS_STRING(block), Py_SIZE(block) - avail_out);
+ } else {
+ assert(Py_SIZE(result) == 0);
+ }
+
+ Py_CLEAR(buffer->list);
+ return result;
+}
+
+/* Clean up the buffer when an error occurred. */
+static inline void
+_BlocksOutputBuffer_OnError(_BlocksOutputBuffer *buffer)
+{
+ Py_CLEAR(buffer->list);
+}
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* Py_INTERNAL_BLOCKS_OUTPUT_BUFFER_H */

Lib/_compression.py

@@ -1,7 +1,7 @@
 """Internal classes used by the gzip, lzma and bz2 modules"""
 
 import io
-
+import sys
 
 BUFFER_SIZE = io.DEFAULT_BUFFER_SIZE # Compressed data read chunk size
 
@@ -110,6 +110,16 @@ def read(self, size=-1):
  self._pos += len(data)
  return data
 
+ def readall(self):
+ chunks = []
+ # sys.maxsize means the max length of output buffer is unlimited,
+ # so that the whole input buffer can be decompressed within one
+ # .decompress() call.
+ while data := self.read(sys.maxsize):
+ chunks.append(data)
+
+ return b"".join(chunks)
+
  # Rewind the file to the beginning of the data stream.
  def _rewind(self):
  self._fp.seek(0)

Misc/NEWS.d/next/Library/2020-10-16-15-34-30.bpo-41486.Mu9Iit.rst

@@ -0,0 +1,4 @@
+Use a new output buffer management code for :mod:`bz2` / :mod:`lzma` /
+:mod:`zlib` modules, and add ``.readall()`` function to
+``_compression.DecompressReader`` class. These bring some performance
+improvements. Patch by Ma Lin.