13
\$\begingroup\$

I'm relatively new to Python 3, especially IO programming, writing a Blender add-on to import model data.

The model data is available in a custom compression, and I originally wrote code in C# to decompress it in memory, porting it to Python 3.

However, being a little unsure about the "optimal" usage of IO classes and functions in Python, I got a little speed problem. The code runs 10 times slower in Python compared to the C# equivalent, and I don't find any more optimization potential, because of my limited Python knowledge.

A test yielded the following speed results at decompressing the same file (around 50 megabytes of data):

  • C#: ~4-5 seconds
  • Python: ~43 seconds

I wonder if anyone can point out where I'm failing hard in Python and need to learn a lot...

Here is the C# code. It uses an extended System.IO.BinaryReader and System.IO.BinaryWriter, operating the same; just with handling endianness more flexible.

public static int Decompress(Stream input, MemoryStream output)
{
    using (BinaryDataReader reader = new BinaryDataReader(input, true))
    using (BinaryDataWriter writer = new BinaryDataWriter(output, true))
    {
        reader.ByteOrder = ByteOrder.BigEndian;
        uint decompressedSize = reader.ReadUInt32();
        // Decompress the data.
        int decompressedBytes = 0;
        while (decompressedBytes < decompressedSize)
        {
            // Read the configuration byte of a decompression setting group, and go through each bit of it.
            byte groupConfig = reader.ReadByte();
            for (int i = 7; i >= 0; i--)
            {
                // Check if bit of the current chunk is set.
                if ((groupConfig & (1 << i)) == (1 << i))
                {
                    // Bit is set, copy 1 raw byte to the output.
                    writer.Write(reader.ReadByte());
                    decompressedBytes++;
                }
                else if (decompressedBytes < decompressedSize) // This does not make sense for last byte.
                {
                    // Bit is not set and data copying configuration follows, either 2 or 3 bytes long.
                    ushort dataBackSeekOffset = reader.ReadUInt16();
                    int dataSize;
                    // If the nibble of the first back seek offset byte is 0, the config is 3 bytes long.
                    byte nibble = (byte)(dataBackSeekOffset >> 12/*1 byte (8 bits) + 1 nibble (4 bits)*/);
                    if (nibble == 0)
                    {
                        // Nibble is 0, the number of bytes to read is in third byte, which is (size + 0x12).
                        dataSize = reader.ReadByte() + 0x12;
                    }
                    else
                    {
                        // Nibble is not 0, and determines (size + 0x02) of bytes to read.
                        dataSize = nibble + 0x02;
                        // Remaining bits are the real back seek offset.
                        dataBackSeekOffset &= 0x0FFF;
                    }
                    // Since bytes can be reread right after they were written, write and read bytes one by one.
                    for (int j = 0; j < dataSize; j++)
                    {
                        // Read one byte from the current back seek position.
                        writer.Position -= dataBackSeekOffset + 1;
                        byte readByte = (byte)writer.BaseStream.ReadByte();
                        // Write the byte to the end of the memory stream.
                        writer.Seek(0, SeekOrigin.End);
                        writer.Write(readByte);
                        decompressedBytes++;
                    }
                }
            }
        }
        return decompressedBytes;
    }
}

I translated it very closely into Python, using struct to read non-byte data from the input stream, otherwise just using io.BytesIO for the in-memory decompression:

def decompress(compressed):
    decompressed_size = struct.unpack(">I", compressed.read(4))[0]
    # Use an in-memory stream and open a reader/writer on it to decompress in.
    decompressed = io.BytesIO()
    # Decompress the data.
    decompressed_bytes = 0
    while decompressed_bytes < decompressed_size:
        # Read the configuration byte of a decompression setting group, and go through each bit of it.
        group_config = compressed.read(1)[0]
        for i in range(7, -1, -1):
            # Check if the bit of the current chunk is set.
            if group_config & (1 << i) == 1 << i:
                # Bit is set, copy 1 raw byte to the output.
                decompressed.write(compressed.read(1))
                decompressed_bytes += 1
            elif decompressed_bytes < decompressed_size: # This does not make sense for the last byte.
                # Bit is not set and data copying configuration follows, either 2 or 3 bytes long.
                data_back_seek_offset = struct.unpack(">H", compressed.read(2))[0]
                # If the nibble of the first back seek offset byte is 0, the config is 3 bytes long.
                nibble = data_back_seek_offset >> 12 # 1 byte (8 bits) + 1 nibble (4 bits)
                if nibble:
                    # Nibble is not 0, and determines (size + 0x02) of bytes to read.
                    data_size = nibble + 0x02
                    # Remaining bits are the real back seek offset
                    data_back_seek_offset &= 0x0FFF
                else:
                    # Nibble is 0, the number of bytes to read is in third byte, which is (size + 0x12).
                    data_size = compressed.read(1)[0] + 0x12
                # Since bytes can be re-read right after they were written, write and read bytes one by one.
                for j in range(0, data_size):
                    # Read one byte from the current back seek position.
                    decompressed.seek(-data_back_seek_offset - 1, io.SEEK_CUR)
                    read_byte = decompressed.read(1)
                    # Write the byte to the end of the memory stream.
                    decompressed.seek(0, io.SEEK_END)
                    decompressed.write(read_byte)
                    decompressed_bytes += 1
    # Seek back to the start of the in-memory stream and return it.
    decompressed.seek(0)
    return decompressed

I hope this question is not too broad and doesn't require too much work... some general optimization tips when using BytesIO (or using a class suiting this case better?) would be great already!

A test file for decompression can be found here (~42 MB). Python needs ~22 seconds to decompress it, while C# only requires ~2 seconds. When decompressed successfully, it results in a 4000x4000 bitmap looking like this:

enter image description here

\$\endgroup\$
  • \$\begingroup\$ Can you provide some compressed data? It's hard to test this code with no data to test it against. \$\endgroup\$ – Gareth Rees May 23 '16 at 10:53
  • 1
    \$\begingroup\$ @GarethRees I'll do and compress some random junk when I'm home in a few hours, stand by! (I'll leave another edit / comment when I'm done). \$\endgroup\$ – Ray Koopa May 23 '16 at 11:33
  • \$\begingroup\$ @GarethRees I added a 42 MB test file, a compress bitmap image, as described in the last paragraph. Thanks for any help in advance! \$\endgroup\$ – Ray Koopa May 23 '16 at 18:19
13
\$\begingroup\$

This kind of bit-slinging code is one of Python's weak spots, I'm afraid. But it is possible to make substantial improvements.

  1. First of all, let's establish a baseline. This is in Python 3:

    >>> benchmark = lambda:decompress(open('image.compressed', 'rb'))
    >>> from timeit import timeit
    >>> timeit(benchmark, number=1)
    80.23794965818524
    
  2. One of the reasons why C# runs so much faster is that it has a JIT compiler. So we can try Python's JIT compiler, PyPy. Unfortunately, this does not yet support Python 3, so we have to backport the code by adding a couple of calls to ord.

    This reduces the runtime by 80%:

    $ pypy cr129054.py
    14.9790380001
    

    But I suspect that PyPy does not work in the context of a Blender extension, so how much speedup can we get in plain Python 3?

  3. Replace io.BytesIO with bytearray. The seek and read calls become indexing operations, and the write method calls become extend or append.

    def decompress(compressed):
        decompressed_size = struct.unpack(">I", compressed.read(4))[0]
        decompressed = bytearray()
        # Decompress the data.
        while len(decompressed) < decompressed_size:
            # Read the configuration byte of a decompression setting
            # group, and go through each bit of it.
            group_config = compressed.read(1)[0]
            for i in range(7, -1, -1):
                # Check if the bit of the current chunk is set.
                if group_config & (1 << i) == 1 << i:
                    # Bit is set, copy 1 raw byte to the output.
                    decompressed.extend(compressed.read(1))
                elif len(decompressed) < decompressed_size:
                    # Bit is not set and data copying configuration
                    # follows, either 2 or 3 bytes long.
                    offset = struct.unpack(">H", compressed.read(2))[0]
                    # If the nibble of the first back seek offset byte is
                    # 0, the config is 3 bytes long.
                    nibble = offset >> 12 # 1 byte (8 bits) + 1 nibble (4 bits)
                    if nibble:
                        # Nibble is not 0, and determines (size + 0x02) of
                        # bytes to read.
                        data_size = nibble + 0x02
                        # Remaining bits are the real back seek offset
                        offset &= 0x0FFF
                    else:
                        # Nibble is 0, the number of bytes to read is in
                        # third byte, which is (size + 0x12).
                        data_size = compressed.read(1)[0] + 0x12
                    for j in range(0, data_size):
                        decompressed.append(decompressed[-offset])
        return decompressed
    

    I haven't investigated exactly why this is faster, but it brings the runtime down by about 30%:

    >>> timeit(benchmark, number=1)
    55.275238760281354
    
  4. Instead of reading and writing one byte at a time from the decompressed output:

    for j in range(0, data_size):
        decompressed.append(decompressed[-offset])
    

    read as much as possible to minimize the number of operations:

    offset += 1
    if data_size == offset:
        chunk = decompressed[-offset:]
    elif data_size < offset:
        chunk = decompressed[-offset:data_size-offset]
    else:
        copies, remainder = divmod(data_size, offset)
        chunk = decompressed[-offset:] * copies
        if remainder:
            chunk += decompressed[-offset:-offset + remainder]
    decompressed.extend(chunk)
    

    (I didn't check this very carefully, so it's possible that there's an off-by-one error. But this idea is right even if the code is wrong.)

    This cuts another 15% off the runtime:

    >>> timeit(benchmark, number=1)
    43.47242012480274
    
  5. Avoid having to look up modules, function, and methods by caching them in local variables:

    _unpack = struct.unpack
    _divmod = divmod
    _read = compressed.read
    _extend = decompressed.extend
    

    This saves a little bit:

    >>> timeit(benchmark, number=1)
    39.8517839522101
    
  6. Precompute the bit values. Instead of:

    for i in range(7, -1, -1):
        # Check if the bit of the current chunk is set.
        if group_config & (1 << i) == 1 << i:
            # Bit is set, copy 1 raw byte to the output.
            decompressed.extend(compressed.read(1))
    

    write:

    for i in (128, 64, 32, 16, 8, 4, 2, 1):
        # Check if the bit of the current chunk is set.
        if group_config & i:
            # Bit is set, copy 1 raw byte to the output.
            _extend(_read(1))
    

    This reduces the runtime to about a third of the original:

    >>> timeit(benchmark, number=1)
    27.295489253941923
    

    That's as far as I got in plain Python. I think that if this was a bottleneck in my application, I would consider switching to the C API.

\$\endgroup\$
  • \$\begingroup\$ This is really amazing! Sometimes the simple things like a bytearray are so much better than the "complicated" BytesIO ones... since this is part of a Blender add-on, I sadly cannot switch to PyPy (s. blender.stackexchange.com/questions/3158/…), but that's totally okay. The speed is really enjoyable now (not all files are that big, so most files (just being like 5 MB) are decompressed almost at once!). I'll check the buffering chunk code soon if it could be off by one. Thanks so much! \$\endgroup\$ – Ray Koopa May 23 '16 at 21:03
  • \$\begingroup\$ I could optimize it very slightly more, by using the below mentioned custom struct.unpack alternative, checking if the offset is bigger than 4095 to see if I have to bitshift the nibble out of it at all, and also cache the decompressed.append method. I actually had to fix the offset computation by - 1 in the last lines. \$\endgroup\$ – Ray Koopa May 24 '16 at 17:38
2
\$\begingroup\$

Edit: Guess I was too slow. I'll remove duplicate remarks later.

Firstly, just a 10x slowdown sounds kinda good. We're talking about Python with an I/O and bit-twiddling heavy task, likely you're also using CPython so this comes at no surprise. I'm assuming you might not be able to use other implementations, e.g. PyPy, which could be faster than the bytecode interpreter at this task.

In general you should take a profiler and look where the most time is spent. If you can, do all this in C (or Cython, or whathaveyou) and keep the performance oriented stuff out of Python entirely.

Since you're already decompressing into memory you could possibly mmap the file and get rid of the read calls and buffering - that said I don't know if that would necessarily improve performance.

Another suggestion would be not to use struct, but use your own inlined decoding functions instead of going through the dynamic logic of that function.

The byte-by-byte decoding loop is also horrible for performance. It's also completely unnecessary if you restructure the logic a bit.

Consider the values for how much to copy and the offset you can easily optimise the repeated reads away and just repeat a chunk you've read initially a couple of times to form the output block.

With that in mind I get the following, which is a bit faster at a glance and returns the same result. I'm pretty confident it's correct in general too. You might want to change the names (maybe include "big-endian" somewhere, or even "int32" etc.) and clean up the logic a bit.

def read_int(file):
    buffer = file.read(4)
    return (buffer[0] << 24) + (buffer[1] << 16) + (buffer[2] << 8) + buffer[3]


def read_short(file):
    buffer = file.read(2)
    return (buffer[0] << 8) + buffer[1]


def decompress(compressed):
    decompressed_size = read_int(compressed)
    # Use an in-memory stream and open a reader/writer on it to decompress in.
    decompressed = io.BytesIO()
    # Decompress the data.
    decompressed_bytes = 0
    while decompressed_bytes < decompressed_size:
        # Read the configuration byte of a decompression setting group, and go through each bit of it.
        group_config = compressed.read(1)[0]
        for i in range(7, -1, -1):
            # Check if the bit of the current chunk is set.
            if group_config & 1 << i:
                # Bit is set, copy 1 raw byte to the output.
                decompressed.write(compressed.read(1))
                decompressed_bytes += 1
            elif decompressed_bytes < decompressed_size: # This does not make sense for the last byte.
                # Bit is not set and data copying configuration follows, either 2 or 3 bytes long.
                data_back_seek_offset = read_short(compressed)
                # If the nibble of the first back seek offset byte is 0, the config is 3 bytes long.
                nibble = data_back_seek_offset >> 12 # 1 byte (8 bits) + 1 nibble (4 bits)
                if nibble:
                    # Nibble is not 0, and determines (size + 0x02) of bytes to read.
                    data_size = nibble + 0x02
                    # Remaining bits are the real back seek offset
                    data_back_seek_offset &= 0x0FFF
                else:
                    # Nibble is 0, the number of bytes to read is in third byte, which is (size + 0x12).
                    data_size = compressed.read(1)[0] + 0x12
                seek_to = -data_back_seek_offset - 1
                decompressed.seek(seek_to, io.SEEK_CUR)
                chunk_length = min(data_size, data_back_seek_offset + 1)
                repeat = decompressed.read(chunk_length)
                decompressed.seek(0, io.SEEK_END)
                repeat_chunk = data_size // chunk_length
                for _ in range(0, repeat_chunk):
                    decompressed.write(repeat)
                rest = data_size % chunk_length
                if rest > 0:
                    decompressed.write(repeat[:rest])
                decompressed_bytes += data_size
    # Seek back to the start of the in-memory stream and return it.
    decompressed.seek(0)
    return decompressed
\$\endgroup\$
  • \$\begingroup\$ Thanks for the additional ideas, especially about optimizing struct.unpack away. I think they can be easily combined with the above solution! =) I had a read on mmap too, but found it too complicated for now ;3 \$\endgroup\$ – Ray Koopa May 23 '16 at 22:04

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.