Implement zstd Compression Support for JSONL and Parquet Files

[justHungryMan]

Support for zstd compression in both JSONL and Parquet file formats.

Parquet Files:

• The implementation applies compression directly within the internal write function (pq.ParquetWriter) using the compression option.
• To maintain clarity and avoid redundancy in DiskWriter implementation, we continue to pass the compression option at the top level but set compression=None in the actual DiskWriter call. This ensures that DiskWriter does not attempt to compress already compressed data.
• Despite specifying the compression option in ParquetWriter, it is feasible to pass super().init(compression=None) to control behavior from a higher abstraction level. This design choice allows flexibility depending on how users might want to integrate or extend our functionality.

Filename Convention:

• For clarity and standardization, files processed with zstd compression are named with the suffix {file_name}.zstd.parquet (or .zstd.jsonl), aligning with typical conventions and aiding in clear differentiation of compressed files.
• If the file extension is set to .parquet.zstd, the pq.ParquetFile function fails to properly recognize and read the file. To ensure seamless integration and functionality, the recommended file naming convention is {file_name}.zstd.parquet, which has been tested and confirmed to work effectively with our current setup.

[justHungryMan]

Or we can just use .parquet (not .{compression}.parquet) without adding here.
Then, we do not pass compression option at src/datatrove/pipeline/writers/parquet.py line 26.

change to

super().__init__(
    output_folder,
    output_filename,
    compression=None,
    adapter=adapter,
    mode="wb",
    expand_metadata=expand_metadata,
    max_file_size=max_file_size,
)

[guipenedo]

Isn't the usual convention to append .zst (instead of .zstd) at the end of the filename rather than in the middle?
Edit: ok I see this is a parquet specific thing

[guipenedo]

I think we should go with your idea of handling everything on the ParquetWriter side, by changing https://github.com/huggingface/datatrove/pull/230/files#diff-6b0424e98052b42ca1d50f9fe6008cfb0b4191bbcd33ef7529112db08b5a2b4dR26 to None

[justHungryMan]

@guipenedo

Should we include the compression type in the filename within the ParquetWriter code, such as {filename}.snappy.parquet?

[justHungryMan]

add validation check for supported protocols

[justHungryMan]

Set superclass compression parameter to None to leverage pyarrow's native compression functions.

[justHungryMan]

HI @guipenedo ,

I used the branch requested in the pull request to process the jsonl.zst file found at dclm-baseline-1.0, and I can confirm that it works correctly. Please review this pull request.

[guipenedo]

LGTM, sorry for the delay and thanks a lot!