Repository navigation
Extra memory copies in blosc, lz4, and zstd compress functions #717
Description
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.binThe 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.
Reacted by Josh Moore
Activity
Opened zarr-developers/numcodecs#717 to track numcodecs issue
d-v-b commentedon Mar 25, 2025
thanks for finding this. do you have any ideas for how we could avoid the extra allocation when slicing?
tomwhite commentedon Mar 25, 2025
Possibly by using a
memoryviewfordest?d-v-b commentedon Mar 26, 2025
@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 commentedon Mar 27, 2025
It's actually a different issue. However have gone ahead and added a fix in the same PR
For more context slicing
bytescreates a copy as shown belowThere 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
memoryviewas 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 getbytesor something elseSo added logic to resize existing
bytesobjects instead of copying. This is effectivelyrealloc, which can return the excess memory to the underlying memory poolNoted how this works in this review: #656 (review)
tomwhite commentedon Mar 27, 2025
@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 commentedon Mar 31, 2025
Fixed in #656