Extra memory copies in blosc, lz4, and zstd compress functions

[tomwhite]

To reproduce, put the following in repro.py:

import numpy as np
from numcodecs import blosc
from numcodecs.blosc import Blosc

if __name__ == "__main__":
    arr = np.random.random(25_000_000)  # 200MB
    blosc.compress(arr, b'lz4', 5, Blosc.SHUFFLE)

Then run the following (with numcodecs 0.15.1 and memray 1.16.0 installed):

memray run -f -o output.bin repro.py
memray flamegraph -f output.bin

The blosc.compress() function allocates 400MB, rather than the expected 200MB (this is in addition to the original 200MB arr). The reason is that on this line, where the destination buffer is resized by slicing, another memory allocation occurs, since slicing a bytes object creates a copy.

lz4 and zstd suffer from the same problem.

#656 is possibly related, but doesn't fix the issue.

[d-v-b]

thanks for finding this. do you have any ideas for how we could avoid the extra allocation when slicing?

[tomwhite]

do you have any ideas for how we could avoid the extra allocation when slicing?

Possibly by using a memoryview for dest?

[d-v-b]

#656 is possibly related, but doesn't fix the issue.

@tomwhite i believe #656 implements the switch to memoryviews, do you have any idea why it doesn't fix this issue? That might be valuable feedback for @jakirkham.

[jakirkham]

It's actually a different issue. However have gone ahead and added a fix in the same PR

For more context slicing bytes creates a copy as shown below

b1 = b"abcdef"
b2 = b1[:3]

There were places in the codecs where they truncated the result. However that copied it into a new buffer

To fix that we really don't even want to take a memoryview as it holds onto the old memory allocation, which is larger than we need. Further there are some code paths in Zarr & Numcodecs that can be sensitive to whether they get bytes or something else

So added logic to resize existing bytes objects instead of copying. This is effectively realloc, which can return the excess memory to the underlying memory pool

Noted how this works in this review: #656 (review)

[tomwhite]

@jakirkham thanks for the explanation and for the fix in #656! As I mentioned there I can see the memory saving with your code.

[tomwhite]

Fixed in #656