Skip to content

Default to ZSTD compression when writing Parquet #981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 16 commits into from
Jan 11, 2025
Merged

Default to ZSTD compression when writing Parquet #981

Changes from 4 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
0c3fed9
fix: update default compression to ZSTD and improve documentation for…
kosiew Dec 24, 2024
b1db46c
fix: clarify compression level documentation for ZSTD in write_parque…
kosiew Dec 27, 2024
819de0d
fix: update default compression level for ZSTD to 4 in write_parquet …
kosiew Jan 7, 2025
56965f4
fix: improve docstring formatting for DataFrame parquet writing method
kosiew Jan 7, 2025
df7d65e
feat: implement Compression enum and update write_parquet method to u…
kosiew Jan 8, 2025
f62a7a8
add test
kosiew Jan 8, 2025
b5b3c47
fix: remove unused import and update default compression to ZSTD in r…
kosiew Jan 8, 2025
2362992
fix: update compression type strings to lowercase in DataFrame parque…
kosiew Jan 8, 2025
b86b142
test: update parquet compression tests to validate invalid and defaul…
kosiew Jan 8, 2025
41e1742
add comment on source of Compression
kosiew Jan 8, 2025
fe502e8
docs: enhance Compression enum documentation and add default level me…
kosiew Jan 8, 2025
67529b8
test: include gzip in default compression level tests for write_parquet
kosiew Jan 8, 2025
811f633
refactor: simplify Compression enum methods and improve type handling…
kosiew Jan 9, 2025
50a58b3
docs: update Compression enum methods to include return type descript…
kosiew Jan 9, 2025
55fc97e
move comment to within test
kosiew Jan 9, 2025
73519fe
Ruff format
timsaucer Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions python/datafusion/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -620,16 +620,25 @@ def write_csv(self, path: str | pathlib.Path, with_header: bool = False) -> None
def write_parquet(
self,
path: str | pathlib.Path,
compression: str = "uncompressed",
compression: str = "ZSTD",
compression_level: int | None = None,
) -> None:
"""Execute the :py:class:`DataFrame` and write the results to a Parquet file.

Args:
path: Path of the Parquet file to write.
compression: Compression type to use.
compression_level: Compression level to use.
"""
compression: Compression type to use. Default is "ZSTD".
compression_level: Compression level to use. For ZSTD, the
recommended range is 1 to 22, with the default being 4. Higher levels
provide better compression but slower speed.
"""
if compression == "ZSTD":
Copy link
Contributor

@ion-elgreco ion-elgreco Jan 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You might want to reuse this code I added in deltalake a while ago:

https://github.com/delta-io/delta-rs/blob/053601b72c2ee053f7ff0309c7d27e122c4e3852/python/deltalake/table.py#L79-L129

Copy link
Contributor Author

@kosiew kosiew Jan 8, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @ion-elgreco ,

I added the Compression Enum but omitted the check_valid_levels because these are already implemented in Rust DataFrame eg

datafusion-python/src/dataframe.rs

Lines 490 to 493 in 63b13da

"zstd" => Compression::ZSTD(
ZstdLevel::try_new(verify_compression_level(compression_level)? as i32)
.map_err(|e| PyValueError::new_err(format!("{e}")))?,
),

Compression levels are tested in:

datafusion-python/python/tests/test_dataframe.py

Lines 1093 to 1106 in 63b13da

@pytest.mark.parametrize(
"compression, compression_level",
[("gzip", 12), ("brotli", 15), ("zstd", 23), ("wrong", 12)],
)
def test_write_compressed_parquet_wrong_compression_level(
df, tmp_path, compression, compression_level
):
path = tmp_path
with pytest.raises(ValueError):
df.write_parquet(
str(path),
compression=compression,
compression_level=compression_level,

if compression_level is None:
# Default compression level for ZSTD is 4 like in delta-rs
# https://github.com/apache/datafusion-python/pull/981#discussion_r1899871918
compression_level = 4
elif not (1 <= compression_level <= 22):
raise ValueError("Compression level for ZSTD must be between 1 and 22")
self.df.write_parquet(str(path), compression, compression_level)

def write_json(self, path: str | pathlib.Path) -> None:
Expand Down
Loading