[FEA] Add zstd compression support for to_csv - #21518
Conversation
github-actions
Bot
added
libcudf
37b9318 to
e5a1fc2
Compare
There was a problem hiding this comment.
Pull request overview
This pull request adds GPU-accelerated ZSTD compression support to cuDF's CSV writer, providing API compatibility with pandas' compression="zstd" parameter. The implementation leverages nvCOMP for GPU-accelerated ZSTD compression and uses concatenated ZSTD frames to enable progressive chunk-based compression.
Changes:
- Added ZSTD compression support to CSV writer at C++, pylibcudf, and Python API levels
- Implemented compression using nvCOMP's batched compression API with fallback to uncompressed on failure
- Added compression parameter to write_csv with validation to only allow ZSTD
- Added Python and C++ tests for ZSTD compression (both chunked and non-chunked)
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| cpp/include/cudf/io/csv.hpp | Added compression member variable, getter, setter, and builder method to csv_writer_options |
| cpp/src/io/csv/writer_impl.cu | Implemented ZSTD compression logic including compress_chunk(), write_data_with_compression() helper, and integration into write_chunked() and write_chunked_begin() |
| cpp/tests/io/csv_test.cpp | Added two C++ tests for ZSTD compression (basic and chunked) |
| python/pylibcudf/pylibcudf/libcudf/io/csv.pxd | Added compression() method declaration to csv_writer_options_builder |
| python/pylibcudf/pylibcudf/io/csv.pyx | Implemented compression() method in CsvWriterOptionsBuilder |
| python/pylibcudf/pylibcudf/io/csv.pyi | Added type hint for compression() method |
| python/pylibcudf/pylibcudf/io/csv.pxd | Added compression() method declaration |
| python/cudf/cudf/utils/ioutils.py | Updated documentation to describe ZSTD compression support |
| python/cudf/cudf/io/csv.py | Added compression parameter validation and mapping to pylibcudf enum |
| python/cudf/cudf/tests/input_output/test_csv.py | Added Python tests for ZSTD compression including unsupported compression types, basic compression, and chunked compression |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
2aa9404 to
0a58680
Compare
vuule
added
feature request
|
Pasting the review by Claude here (still manually validating if they make sense). Please exercise caution: Click to expandFindingsCritical
Suggestions
Nits
|
|
Thank you for the PR, @a-hirota! I ran some benchmarks locally and it performs surprisingly well (given that we compress one block at a time). Would be nice to decouple chunk size from ZSDT block size, but this can be a separate PR. Please let us know if you need help addressing the comments in the review that @mhaseeb123 has posted :) |
|
@vuule Should we take this over and merge it? |
d52e473 to
3d09384
Compare
|
Taking this over so it can land — thanks @a-hirota for the original work, authorship is preserved on the commit. I rebased onto current
One correctness hole not in the original review: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughCSV writer options now support validated NONE or ZSTD compression. The setting flows through libcudf, pylibcudf, and cuDF APIs, with compressed device output, updated documentation, and C++ and Python round-trip coverage. ChangesCSV ZSTD writer compression
Estimated code review effort: 3 (Moderate) | ~30 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tests/io/csv_test.cpp`:
- Around line 2974-3037: Expand the ZSTD CSV coverage in
cpp/tests/io/csv_test.cpp at lines 2974-3037 with cases for empty inputs,
null-containing and sliced columns, boundary and multi-block row sizes, and
non-ASCII UTF-8 strings; reuse existing round-trip helpers and preserve
header/no-header coverage. Also update
python/cudf/cudf/tests/input_output/test_csv.py at lines 2073-2107 to add empty
and all-null DataFrame cases alongside the existing single-element and
mixed-null tests, covering the required Python input categories.
In `@python/cudf/cudf/tests/input_output/test_csv.py`:
- Around line 2103-2107: Add the missing "zstd" entry to the compression map
used by the high-level read_csv implementation in csv.py, mapping it to
plc.io.types.CompressionType.ZSTD. Preserve the existing mappings and ensure
read_csv accepts compression="zstd" without raising KeyError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: de4841dd-a92c-4a33-a974-5488e543a4c9
📒 Files selected for processing (10)
cpp/include/cudf/io/csv.hppcpp/src/io/csv/writer_impl.cucpp/tests/io/csv_test.cpppython/cudf/cudf/io/csv.pypython/cudf/cudf/tests/input_output/test_csv.pypython/cudf/cudf/utils/ioutils.pypython/pylibcudf/pylibcudf/io/csv.pxdpython/pylibcudf/pylibcudf/io/csv.pyipython/pylibcudf/pylibcudf/io/csv.pyxpython/pylibcudf/pylibcudf/libcudf/io/csv.pxd
vuule
changed the title
to_csv
3d09384 to
5816f91
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
python/cudf/cudf/tests/input_output/test_csv.py (1)
2073-2097: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd empty, all-null, and no-header cases.
Extend the matrix with an empty DataFrame, an explicitly typed all-null column, and
header=False; adjust the read-back assertion for the no-header case. These branches are currently untested. As per coding guidelines, Python test files must provide empty, all-null, single-element, and mixed-type coverage.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf/cudf/tests/input_output/test_csv.py` around lines 2073 - 2097, Extend test_to_csv_zstd_compression’s parameter matrix with an empty DataFrame, an explicitly typed all-null column, and a header=False case. Add a header parameter to the test and pass it to both CSV writes, adjusting the read-back comparison so no-header output is validated without expecting a header.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf/cudf/tests/input_output/test_csv.py`:
- Around line 2098-2101: Update the ZSTD decompression setup in the affected CSV
test to call stream_reader with read_across_frames=True, ensuring concatenated
frames are fully read. Store the reader, read and decode its contents, then
explicitly close the reader without closing or otherwise touching the file
handle still used by cuDF.
---
Nitpick comments:
In `@python/cudf/cudf/tests/input_output/test_csv.py`:
- Around line 2073-2097: Extend test_to_csv_zstd_compression’s parameter matrix
with an empty DataFrame, an explicitly typed all-null column, and a header=False
case. Add a header parameter to the test and pass it to both CSV writes,
adjusting the read-back comparison so no-header output is validated without
expecting a header.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 798f3d14-8380-4171-a41c-00fcbaed3b54
📒 Files selected for processing (10)
cpp/include/cudf/io/csv.hppcpp/src/io/csv/writer_impl.cucpp/tests/io/csv_test.cpppython/cudf/cudf/io/csv.pypython/cudf/cudf/tests/input_output/test_csv.pypython/cudf/cudf/utils/ioutils.pypython/pylibcudf/pylibcudf/io/csv.pxdpython/pylibcudf/pylibcudf/io/csv.pyipython/pylibcudf/pylibcudf/io/csv.pyxpython/pylibcudf/pylibcudf/libcudf/io/csv.pxd
🚧 Files skipped from review as they are similar to previous changes (9)
- python/pylibcudf/pylibcudf/libcudf/io/csv.pxd
- python/pylibcudf/pylibcudf/io/csv.pyi
- python/pylibcudf/pylibcudf/io/csv.pxd
- python/pylibcudf/pylibcudf/io/csv.pyx
- python/cudf/cudf/utils/ioutils.py
- python/cudf/cudf/io/csv.py
- cpp/src/io/csv/writer_impl.cu
- cpp/include/cudf/io/csv.hpp
- cpp/tests/io/csv_test.cpp
| with open(fname, "rb") as f: | ||
| # the writer emits concatenated frames, so the streaming API is required | ||
| decompressed = zstd.ZstdDecompressor().stream_reader(f).read() | ||
| assert decompressed.decode("utf-8") == expected |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
python - <<'PY'
import io
import zstandard as zstd
frames = (
zstd.ZstdCompressor().compress(b"frame-1")
+ zstd.ZstdCompressor().compress(b"frame-2")
)
with zstd.ZstdDecompressor().stream_reader(
io.BytesIO(frames), read_across_frames=True
) as reader:
assert reader.read() == b"frame-1frame-2"
PYRepository: rapidsai/cudf
Length of output: 273
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate test_csv.py =="
git ls-files | grep -F 'python/cudf/cudf/tests/input_output/test_csv.py' || true
echo "== relevant lines =="
sed -n '2040,2125p' python/cudf/cudf/tests/input_output/test_csv.py | cat -n | sed 's/^/2040+/'
echo "== zstandard usages in test_csv.py =="
rg -n "zstd|read_across_frames|stream_reader|chunksize|csv_writer" python/cudf/cudf/tests/input_output/test_csv.py | sed -n '1,200p'
echo "== zstandard requirements/docs references =="
rg -n "python-zstandard|zstandard|read_across_frames" -S pyproject.toml setup.py setup.cfg requirements* constraints* python -g '*txt' -g '*.toml' -g '*.cfg' -g '*.py' 2>/dev/null | head -n 200Repository: rapidsai/cudf
Length of output: 6354
🌐 Web query:
python-zstandard ZstdDecompressor stream_reader read_across_frames default
💡 Result:
In the python-zstandard library, the ZstdDecompressor.stream_reader method includes a read_across_frames argument that controls how the reader handles input containing multiple Zstandard frames [1][2]. The default value for read_across_frames is False [1][2]. When this is set to False, the reader stops decompressing when it encounters the end of a single Zstandard frame [1][3][4]. If set to True, the reader can transparently continue decompressing data that spans multiple Zstandard frames [1][3][4]. The library documentation and maintainers have noted that the default behavior may change to True in a future release [1][2][5]. Consequently, it is strongly recommended that users explicitly specify read_across_frames=True or read_across_frames=False in their code to ensure consistent, future-proof behavior [2][5][6].
Citations:
- 1: https://python-zstandard.readthedocs.io/en/latest/decompressor.html
- 2: https://python-zstandard.readthedocs.io/en/latest/news.html
- 3: https://pypi.org/project/zstandard/0.14.0/
- 4: Files with multiple frames have frames skipped? indygreg/python-zstandard#59
- 5: https://github.com/indygreg/python-zstandard/blob/main/docs/news.rst
- 6: https://repo.mercurial-scm.org/hg/file/675775c33ab6/contrib/python-zstandard/NEWS.rst
🌐 Web query:
python-zstandard stream_reader context manager close __enter__ __exit__
💡 Result:
In the python-zstandard library, the object returned by ZstdDecompressor.stream_reader (an instance of ZstdDecompressionReader) supports the context manager protocol using the enter and exit methods [1][2]. Context Manager Behavior: - enter: Marks the reader as having been entered (it cannot be entered multiple times) and returns the instance itself [2]. - exit: Explicitly calls the close method on the stream [1][3]. This releases underlying resources, such as the buffer and the decompressor reference, and marks the stream as closed [2]. Once closed, any subsequent I/O operations will fail [1]. Important Considerations: - Automatic Closure: Using the stream as a context manager is a convenient way to ensure that resources are cleaned up promptly when processing is finished [1][4]. - Underlying Stream: The stream_reader method accepts a closefd argument (defaulting to True) [1][3]. If closefd is True, calling close (including via the context manager's exit) will also close the underlying source stream [1]. If you need to keep the source stream open after the reader is finished, set closefd=False [1][3]. - Not Mandatory: While context manager usage is supported and recommended for resource management, it is not strictly required [5]. You can manually call close on the reader instance if you prefer [1][2].
Citations:
- 1: https://python-zstandard.readthedocs.io/en/latest/decompressor.html
- 2: https://repo.mercurial-scm.org/hg/file/675775c33ab6/contrib/python-zstandard/c-ext/decompressionreader.c
- 3: Closing a ZstdCompressionWriter or ZstdDecompressionWriter should not close underlying file indygreg/python-zstandard#76
- 4: https://pypi.org/project/zstandard/0.10.1/
- 5: Is it possible to relax
withusage requirements? indygreg/python-zstandard#57
Read across concatenated ZSTD frames.
stream_reader() defaults to read_across_frames=False, so the chunksize=10 case reads only the first ZSTD frame and cannot match the full uncompressed CSV. Pass read_across_frames=True and close the reader explicitly to avoid touching the underlying file handle that cuDF is still using.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf/cudf/tests/input_output/test_csv.py` around lines 2098 - 2101,
Update the ZSTD decompression setup in the affected CSV test to call
stream_reader with read_across_frames=True, ensuring concatenated frames are
fully read. Store the reader, read and decode its contents, then explicitly
close the reader without closing or otherwise touching the file handle still
used by cuDF.
Source: MCP tools
5816f91 to
d589b0e
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf/cudf/io/csv.py`:
- Around line 406-411: Extend the compression guard in the CSV writer path to
reject StringIO and other text-mode sinks, not only cases returning output as a
string without path_or_buf. Ensure compressed binary output never reaches
SinkInfo with a text stream, while preserving supported binary sinks and
uncompressed StringIO behavior. Add a regression test covering compressed output
directed to an in-memory string buffer.
In `@python/cudf/cudf/tests/input_output/test_csv.py`:
- Around line 2077-2089: Extend the CSV test input matrix around the existing
DataFrame cases to include an empty frame such as {} and an all-null column case
such as {"a": [None, None]}. Keep the existing single-element, mixed-type, and
partially-null cases unchanged so the matrix covers both empty and all-null
writer paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 64f3d8bd-50b9-471f-a73e-00e9a2b78c45
📒 Files selected for processing (10)
cpp/include/cudf/io/csv.hppcpp/src/io/csv/writer_impl.cucpp/tests/io/csv_test.cpppython/cudf/cudf/io/csv.pypython/cudf/cudf/tests/input_output/test_csv.pypython/cudf/cudf/utils/ioutils.pypython/pylibcudf/pylibcudf/io/csv.pxdpython/pylibcudf/pylibcudf/io/csv.pyipython/pylibcudf/pylibcudf/io/csv.pyxpython/pylibcudf/pylibcudf/libcudf/io/csv.pxd
🚧 Files skipped from review as they are similar to previous changes (7)
- python/pylibcudf/pylibcudf/io/csv.pxd
- python/pylibcudf/pylibcudf/libcudf/io/csv.pxd
- python/cudf/cudf/utils/ioutils.py
- python/pylibcudf/pylibcudf/io/csv.pyx
- cpp/tests/io/csv_test.cpp
- cpp/include/cudf/io/csv.hpp
- cpp/src/io/csv/writer_impl.cu
Adds GPU-accelerated ZSTD compression to `write_csv` and `to_csv`, matching the pandas `compression="zstd"` API. ZSTD is the only supported codec because it allows concatenated frames: each chunk can be compressed independently while the resulting file remains a single stream that standard tools (`zstd -d`) can decompress. - Add `compression` to `csv_writer_options` and its builder, rejecting codecs other than NONE/ZSTD - Compress via the existing `io::detail::compress` API so the writer inherits the shared host/device dispatch, and fail loudly rather than silently emitting uncompressed bytes into a `.zst` file - Plumb `compression` through the pylibcudf bindings and `to_csv`, rejecting the combination with a string return value since compressed output is binary Co-authored-by: Hirota Akio <33370421+a-hirota@users.noreply.github.com>
d589b0e to
6aeb4c1
Compare
|
/ok to test 6aeb4c1 |
The repo test guidelines call for empty, all-null, single-element and mixed-type inputs. A frame with no columns writes a header consisting of just the line terminator, which is still emitted as its own ZSTD frame, so it exercises a distinct writer path.
The writer compressed each row chunk into a single ZSTD frame, so the amount of work handed to the codec was tied to `rows_per_chunk` and a large chunk became one serial compression task. Concatenated ZSTD frames decompress to the concatenation of their payloads, so the byte stream can be split at arbitrary offsets without regard to row boundaries. Split each chunk into fixed-size blocks and compress them in one batched call, then pack the resulting frames into a contiguous buffer for a single sink write. Add `compression_block_size` to the writer options so the block size can be tuned independently of `rows_per_chunk`, defaulting to 1 MB and capped at the codec's maximum input size. Capping at the codec limit also makes the previous "chunk is too large" failure unreachable.
Mirror the new `compression_block_size` writer option on `CsvWriterOptionsBuilder` so the block size can be tuned from Python.
|
@vuule how do you want to deal with reviews here since you've taken over the PR? |
yeah, just struggling to find time to continue working on this. Should come back to this within a week. |
No rush, just checking in before I review (or ping for reviews). |
Not quite ready for review :) |
The trailing newline that separates a chunk from the next one was compressed and written on its own, so a chunked write cost two compress calls, two packing passes and two sink writes per chunk. Pass it to `write_compressed_to_sink` as a tail span instead, where it becomes one more block of the same batched call. Appending it still does not require copying the chunk. Since the blocks are no longer uniformly sized, lay the compression buffer out from the maximum compressed size of each individual block, which also drops the padding the uniform layout reserved for partial blocks.
Description
Adds GPU-accelerated ZSTD compression to the CSV writer, matching the pandas
compression="zstd"API.Why ZSTD only:
zstd -d) can decompressImplementation:
compressiontocsv_writer_optionsand its builder, rejecting codecs other thanNONE/ZSTDat option-construction timecudf::io::detail::compressAPI, so the writer inherits the shared host/device dispatch rather than re-implementing nvCOMP setup.zstfilecompressionthrough the pylibcudf bindings andcudf.to_csv, rejecting the combination with a string return value since compressed output is binaryUsage:
Notes
Taken over from @a-hirota (original authorship preserved). Rebased onto current
mainand the outstanding review items are addressed; see the comment below for details.Known follow-ups, deliberately out of scope here so the feature can land and be measured first:
rows_per_chunkfrom the ZSTD frame size (raised by @vuule)Checklist