Add more compression types for to_json

[bhavitvyamalik]

This PR adds bz2, xz, and zip (WIP) for to_json. I also plan to add infer like how pandas does it

[bhavitvyamalik]

@lhoestq, I looked into how to compress with zipfile for which few methods exist, let me know which one looks good:

1. create the file in normal wb mode and then zip it separately
2. use ZipFile.write_str to write file into the archive. For this we'll need to change how we're writing files from _write method

How pandas handles it is that they have created a wrapper for standard library class ZipFile and allow the returned file-like handle to accept byte strings via write method instead of write_str (purpose was to change the name of function by creating that wrapper)

[lhoestq]

1. sounds not ideal since it creates an intermediary file.
   I like pandas' approach. Is it possible to implement 2. using the pandas class ? Or maybe we can have something similar ?

[bhavitvyamalik]

Definitely, @lhoestq! I've adapted that from original code and turns out it is faster than gz compression. Apart from that I've also added infer option to automatically infer compression type from path_or_buf given

[bhavitvyamalik]

One small thing, currently I'm assuming that user will provide compression extension in path_or_buf. Is it this also possible?
dataset.to_json("from_dataset.json", compression="zip")?
Should I put an assert to ensure the file name provided always has a compression extension?

[lhoestq]

Thanks !

One small thing, currently I'm assuming that user will provide compression extension in path_or_buf. Is it this also possible?
dataset.to_json("from_dataset.json", compression="zip")?
Should I put an assert to ensure the file name provided always has a compression extension?

I think it's fine as it is right now :) No need to check the extension of the filename passed to path_or_buf.

[lhoestq]

turns out it is faster than gz compression

I think the default compression level of gzip is 9 in python, which is very slow. Maybe we can switch to compression level 6 instead which is faster, like the gzip command on unix

[lhoestq]

I found that fsspec has something that may interest you: fsspec.open(..., compression=...). I don't remember if we've already mentioned it or not

It also has zip if I understand correctly ! see https://github.com/fsspec/filesystem_spec/blob/master/fsspec/compression.py#L70

Since fsspec is a dependency of datasets we can use all this :)

Let me know if you prefer using fsspec instead (I haven't tested this yet to write compressed files). IMO it sounds pretty easy to use and it would make the code base simpler

[bhavitvyamalik]

Just tried fsspec but I'm not able to write compressed zip files :/
gzip, xz, bz2 are all working fine and it's really simple (no need for FileWriteHandler now!)

[bhavitvyamalik]

somehow gzip test is failing due to few mismatches

[bhavitvyamalik]

update: instead of reading compressed files and comparing them directly, I uncompressed them using fsspec and then compared. The bug went away!

[lhoestq]

Thanks ! Just adding an error message in case someone passes compression for a buffer or file-like object: