Implement lazy-loading of data for CompImageHDU - #14353

Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

• Do the proposed changes actually accomplish desired goals?
• Do the proposed changes follow the Astropy coding guidelines?
• Are tests added/updated as required? If so, do they follow the Astropy testing guidelines?
• Are docs added/updated as required? If so, do they follow the Astropy documentation guidelines?
• Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see "When to rebase and squash commits".
• Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the bot.
• Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
• Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
• Is a milestone set? Milestone must be set but we cannot check for it on Actions; do not let the green checkmark fool you.
• At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) before merge.

I think I'll update this PR shortly changing it to be private API, so no need for anyone to review at the moment.

I've simplified the implementation to now use the same code regardless of whether we are accessing a single pixel, section, or the whole array, to avoid having three different decompress_* functions. I don't seem to notice any performance impact from doing this.

@saimn - assuming the CI passes, this is basically ready from my side, so feel free to review whenever you have time.

Oh, already phrase as numcodecs! This is not released yet, correct?

I have verified that I can get the codec successfully working with some SDO AIA data:

codec = tiled.codecs.Rice1(blocksize=blocksize, bytepix=bytepix, tilesize=tilesize)

f.seek(data_start)
for i in range(n_rows):
    block0 = f.read(table["size"][i])
    out = codec.decode(np.frombuffer(block0, "uint8"))
    assert (hdul[1].data[i] == out).all()

The length of each input block is different; is there a way to pass a concatenated list of N of these byte block and get N uncompressed blocks back?

f.seek(data_start)
block10 = f.read(table["size"][0:10].sum())
out = codec.decode(np.frombuffer(block10, "uint8"))

I get out with a size of 4096 (only one output block) or, sometimes, segfault.

'decompression warning: unused bytes at end of compressed buffer'

@martindurant just a side note, could you keep track of/share examples that trigger segfaults?

I seem to get the segfault if I try to decompress again following the decompression warning above.

@martindurant - thanks, that's useful to know! Unfortunately decode is not able to take a concatenated byte string of multiple tiles, just a single tile.

Do you happen to know if cfitsio's Rice is the same as https://sourceforge.net/p/lpcrice/code/HEAD/tree/trunk/rice/ ?(that being extracted from HDF5, and also implemented in imagecodecs)

if cfitsio's Rice is the same as

I have no idea, but I would probably guess it's subtly different!

This is not released yet, correct?

correct.

Could you open issues for the segfaults and any other things you find? (one tracking issue would be fine). I am happy to run them down next week if it means you can keep working on this.

Is it a valid expectation that numcodecs should be able to take more bytes than the codec is expecting to decode? I don't think that would be something we have tested or something the cfitsio layer would be ok with so that's probably the root of the segfault.

Sorry, not really sure what to put in a new issue except to rehash the conversation in this thread already!

Please see gist: https://gist.github.com/martindurant/16bee4256595d3b6814be139ab1bd54e
It demonstrated chunk-wise loading of a compressed FITS image (SDO/AIA) using zarr and kerchunk. Some notes:

• the codec currently in astropy does not accept the buf= argument
• the codec currently in astropy assumes the input is already an array, when it's likely a byte buffer
• this particular file has 4096x1 tiles and about 2x compression, which is pretty bad on both counts
• by using kerchunk/zarr, we would get parallel decompression if dask were invoked. With kerchunk, this parallelism could cross files too, but only selecting required selection from each
• even with effectively repeatedly calling cfitsio decompression in a loop with small tiles, the load speed is essentially the same. This would be better if the loop could be in C, just decoding each new block in a single bytestream as it comes, and if we could write the output directly into the final array.

[ ] %timeit g.data[:]
222 ms ± 426 µs
[ ] %timeit fits.getdata(fn, 1)
212 ms ± 3.09 ms

• parallelised by dask (local threads)

[ ] arr = da.from_array(g.data, chunks=(128, 4096))
... %timeit arr.compute()
138 ms ± 565 µs

FWIW, imagecodecs includes a rcomp codec based on cfitsio source code. Maybe that is compatible?

from imagecodecs.numcodecs import Rcomp
x = np.random.randint(0, 100, size=(256, 256), dtype="int16")
rice = Rcomp(shape=x.shape, dtype=x.dtype, nblock=32)
out = rice.encode(x) + b"extra"

I took a block from the one RICE-compressed file I have and gave it to Rcomp. It did not crash, but the output values are completely different.

Thanks for trying. How can I access a compressed chunk?

In this gist: https://gist.github.com/martindurant/16bee4256595d3b6814be139ab1bd54e
I show how to get the offset/sizess from the FITS extension table, which are relative to the end of the same table. The test_codec() function loops over them to test the Rice codec.

Thanks. I tried the Rcomp codec with the chunks from https://github.com/astropy/astropy/blob/main/astropy/io/fits/tests/data/compressed_image.fits and it worked. Are there other variations of the rice compression used in FITS?

Aha, I got the dtype endinaness wrong. :|

So, can is there a way to pass N blocks to Rcomp, and have it fill in N tiles/rows of output in one go? Just passing more bytes yields only the first output block (but no exception/crash).

Interestingly, I can pass an out= to Rcomp.decode, but only if it is a memoryview as opposed to a numpy array (getting AttributeError: 'str' object has no attribute 'kind').

Sorry everyone for keeping this thread going!

@martindurant - if you wanted to be able to pass compressed bytes for multiple tiles to a codec in one go wouldn't you need to also pass information about how long each section of the compressed bytes is? (that is, I don't know if this could be inferred from concatenated compressed bytes) If so, does the numcodecs API allow that?

is there a way to pass N blocks to Rcomp, and have it fill in N tiles/rows of output in one go?

I don't think so. Isn't that the job of Zarr to take decoded chunks and arrange them in the output array? I guess you want to avoid the detrimental overhead of Zarr handling many small chunks?

Interestingly, I can pass an out= to Rcomp.decode, but only if it is a memoryview as opposed to a numpy array (getting AttributeError: 'str' object has no attribute 'kind').

That's a bug. Try dtype=None or dtype=numpy.dtype(dtype) when passing a numpy array as output.

To some extent the real solution is to not have FITS files with tiny chunks, and I know @Cadair has been working on figuring out a better tile shape/size for his data, so maybe we should use an example with more sensible chunks to do any kind of profiling?

Yeah the AIA data made bad choices with respect to chunk sizes 😦

Yeah the AIA data made bad choices with respect to chunk sizes

Kerchunk's use case is exactly to index over archival data, so we don't get to dream of optimal chunk sizes. There are too many chunks in the AIA data to store them all, so we need to store fewer references, each to multiple compressed blocks.

I'll note that most compressors allow you to concatenate blocks

gzip.decompress(gzip.compress(b"hello") + gzip.compress(b"hello")) == b'hellohello'

wouldn't you need to also pass information about how long each section of the compressed bytes is

No, the codec should know when it gets to the end of each block (and it apparently does, since the astropy version gives an extra bytes warning and the imagecodecs one just decompresses the first block). All we need is to know how many bytes got used for a tile, as it is consumed.

Try dtype=None or dtype=numpy.dtype(dtype) when passing a numpy array as output.

The output array is used by zarr in zarr.core.Array._process_chunk; so although I can make it work when calling the codec directly, it won't work with zarr unless I subclass.

so we need to store fewer references, each to multiple compressed blocks

Ah ok, this explains what you are trying to achieve a little more. Can you elaborate on how this would work between kerchunk / zarr and the codecs?

I am really keen to see this support for RICE compressed FITS land in kerchunk, and we can put some of the Astropy funding towards working on it. What I am struggling to wrap my head around is what the todo list looks like for getting it working. Here are some things I know probably need doing, or questions I have, but I am sure you will tell me there are many others:

• A way to have kerchunk support per-tile codecs (needed for Quantization if nothing else, which isn't used by AIA but is by my DKIST data).
  • Figure out if it's possible to not have some or all of the FITS compression codecs take the output tile shape.
• Should we move the FITS specific codecs currently implemented in kerchunk into Astropy?
• Make the codecs take an out= buffer / array. (What types does this need to support?)
• Make the codecs handle decompressing more than one tile at once, if possible.
• There's a lot of subtlety / complexity in decompress_hdu_section which we should refactor so kerchunk can make use of it. It would be insane for us to both have a full implementation of decoding tile compressed files.

Can you elaborate on how this would work between kerchunk / zarr and the codecs?

For the AIA example, let's suppose we wish to view the first two (4096x1) tiles as a (4096x2) chunk, so kerchunk stores the start of the first tile and the end of the second as one reference. Process:

• zarr will read a dataset with chunksize set to (4096x2)
• when accessing some data overlapping the first zarr chunk, it will ask the storage for bytes, and be given the two tiles' worth
• this buffer is passed to codec.decode. Only if requesting the whole of the region and with no further decoding steps, it is passed out=; we don't in general know the output size of the chunk unless we make this a parameter saved with the codec. out is a numpy array (potentially cupy in the future) with shape of chunksize and final dtype
• we decode a tile and set the first 4096pix of the output, then decode the second tile from the remaining bytes

Note that I am seeing the imagecodecs' Rcomp as significantly faster that astropy's RICE, and if safely stops decompression after one tile without segfault. Maybe imagecodecs is a natural place for this to live, not kerchunk or astropy?

A way to have kerchunk support per-tile codecs

You mean that each tile has some different parametrisation? In numcodecs, these are normally saved along with the output chunk as a manner of binary header. We can't do that, so we would need to store the parameters some other way, to be decided.

let's suppose we wish to view the first two (4096x1) tiles as a (4096x2) chunk

I think there's quite a lot of things which we would need to do to get the simpler one tile == one chunk use case working first for any valid RICE file. I think working through that would help me get a better grip on the requirements for the codecs and other components.

the imagecodecs' Rcomp as significantly faster that astropy's RICE

Interesting, will have to have a look at that. We tried to not modify the C code from cfitsio to make it easy to take any upstream changes, perhaps that's the root of the problem. The segfault is something we want to fix.

Maybe imagecodecs is a natural place for this to live, not kerchunk or astropy?

I would be surprised if there is appetite for having astropy depend on imagecodecs to make FITS reading work, but I am not personally against it.

@martindurant - a quick question about the following comment:

For the AIA example, let's suppose we wish to view the first two (4096x1) tiles as a (4096x2) chunk, so kerchunk stores the start of the first tile and the end of the second as one reference

There is no guarantee that the first two tiles would be contiguous in the heap - the FITS standard doesn't require offsets to be monotonically increasing in the underlying binary table representing the compressed image. Does this cause issues with this approach?

Yes, it would critically kill the idea. However, AIA does seem to always be contiguous. I would be hard pressed to imagine why it ever would not be - isn't compression done from the first pixel to the last?

@martindurant - in practice I think this will probably be a case of the majority of files actually being fine, but one could imagine for instance having some kind of parallel compression where chunks are written to a heap as they are completed and therefore where the chunks end up not appearing sequentially in the heap. I'm not saying such files exist in the wild, but at the very least any kerchunk implementation that groups chunks should check that they are in fact sequential.

I have often fallen foul of over-specifying my code to the first example dataset :). Totally agree that kerchunk should only make multi-tile chunks after checking they are contiguous. Also hoping that future data will have more natural tile sizes to make that unnecessary. With the arrival of parquet storage for kerchunk references, I wonder whether we may be ok anyway (fsspec will combine adjacent reads within a file at access time rather that send off myriad 1500byte calls).

I have collated all the issues and comments that I can find and think of into this issue: #14575.

If we can continue the discussion there that would be great.