Rogdham / pyzstd

Python bindings to Zstandard (zstd) compression library, the API style is similar to Python's bz2/lzma/zlib modules.
https://pyzstd.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
19 stars 3 forks source link

Seek to frame numbers #19

Open bebehei opened 2 months ago

bebehei commented 2 months ago

I'm currently testing the zstd seekable frame format. I'm missing to seek to a specific frame instead to its byte position.

from pyzstd import SeekableZstdFile, ZstdFile

FILENAME="framed.zst"

with SeekableZstdFile(FILENAME, 'w') as filewriter:
    filewriter.write(b"first_frame")
    filewriter.flush(mode=ZstdFile.FLUSH_FRAME)
    filewriter.write(b"second_frame")
    filewriter.flush(mode=ZstdFile.FLUSH_FRAME)
    filewriter.write(b"third_frame")
    filewriter.close()

with SeekableZstdFile(FILENAME, 'r') as filereader:
    # Feature Request: Seek to frame numbers
    filereader.seek_to_frame(-1)
    # Prints out `third_frame`
    print(filereader.read())

I've got a big file (about 600GB uncompressed), which I read and segment logically into zstd frames via SeekableZstdFile.flush as written in the example above.

After all data is compressed, I flush again the data and then I also add an index of every segement with the metadata as an additional frame. This index ist just a JSON serialized text.

To find that metadata JSON, the program must do some weird binary search now. Since it's already exposed in its own frame, simply jumping to the last frame would do it as O(1).


zstd exposes the frame data in its .h files too: https://github.com/facebook/zstd/blob/20707e3718ee14250fb8a44b3bf023ea36bd88df/contrib/seekable_format/zstd_seekable.h#L199-L220

Is it possible to expose this in the python module, too?

Rogdham commented 2 months ago

Hello, unfortunately pyzstd does not have a public API for this as of now.

In the meanwhile, could you try the following?

with SeekableZstdFile(FILENAME, 'r') as filereader:
    # Seek to last frame
    filereader.seek(filereader._buffer.raw._seek_table.get_frame_sizes(-1)[1])
    # Prints out `third_frame`
    print(filereader.read())
bebehei commented 1 month ago

Thanks @Rogdham, this works flawlessly.

Hadn't been able to come up with this myself by introspection. Therefore would definitely vote for a public API of this. Others will probably have this problem, too.