This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 8de5640254 feat(ipc): Supports compression level configuration (#10133)
8de5640254 is described below
commit 8de56402540c27469e15286e474a2aaa965d23f9
Author: Zhen Wang <[email protected]>
AuthorDate: Sun Jun 28 05:47:59 2026 +0800
feat(ipc): Supports compression level configuration (#10133)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- Closes #10132.
# Rationale for this change
Make compression level of ipc writer configurable
# What changes are included in this PR?
- Add a compression level option to `IpcWriteOptions`
- Add `try_with_compression_level` to configure IPC batch compression
levels
- Pass the configured compression level through the IPC write path when
constructing the compression codec
# Are these changes tested?
I locally modified the test_write_file_with_zstd_compression test case
to benchmark write time and output file size across different
compression levels.
<img width="1171" height="765" alt="image"
src="https://github.com/user-attachments/assets/1a08197c-c732-4094-8e34-62ca37f7f6d3";
/>
print logs:
```
ZSTD compression level -999: wrote 4001530 bytes in 1.658125ms
ZSTD compression level 1: wrote 3657530 bytes in 38.953042ms
ZSTD compression level 22: wrote 3638762 bytes in 869.19925ms
```
# Are there any user-facing changes?
~~Yes. This PR adds a new `batch_compression_level` option to
`IpcWriteOptions`.~~ No public API changes
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-ipc/src/compression.rs | 39 ++++++++++++++++-----
arrow-ipc/src/writer.rs | 83 ++++++++++++++++++++++++++++++++++++++++++--
2 files changed, 110 insertions(+), 12 deletions(-)
diff --git a/arrow-ipc/src/compression.rs b/arrow-ipc/src/compression.rs
index ff6e83dfdd..45558adb44 100644
--- a/arrow-ipc/src/compression.rs
+++ b/arrow-ipc/src/compression.rs
@@ -21,6 +21,7 @@ use arrow_schema::ArrowError;
const LENGTH_NO_COMPRESSED_DATA: i64 = -1;
const LENGTH_OF_PREFIX_DATA: i64 = 8;
+const DEFAULT_ZSTD_COMPRESSION_LEVEL: i32 = 3;
/// Additional context that may be needed for compression.
///
@@ -35,10 +36,9 @@ pub struct CompressionContext {
impl CompressionContext {
#[cfg(feature = "zstd")]
- fn zstd_compressor(&mut self) -> &mut zstd::bulk::Compressor<'static> {
+ fn zstd_compressor(&mut self, level: i32) -> &mut
zstd::bulk::Compressor<'static> {
self.compressor.get_or_insert_with(|| {
- zstd::bulk::Compressor::new(zstd::DEFAULT_COMPRESSION_LEVEL)
- .expect("can use default compression level")
+ zstd::bulk::Compressor::new(level).expect("can use default
compression level")
})
}
}
@@ -111,7 +111,7 @@ impl std::fmt::Debug for DecompressionContext {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionCodec {
Lz4Frame,
- Zstd,
+ Zstd(i32),
}
impl TryFrom<CompressionType> for CompressionCodec {
@@ -119,7 +119,7 @@ impl TryFrom<CompressionType> for CompressionCodec {
fn try_from(compression_type: CompressionType) -> Result<Self, ArrowError>
{
match compression_type {
- CompressionType::ZSTD => Ok(CompressionCodec::Zstd),
+ CompressionType::ZSTD =>
Ok(CompressionCodec::Zstd(DEFAULT_ZSTD_COMPRESSION_LEVEL)),
CompressionType::LZ4_FRAME => Ok(CompressionCodec::Lz4Frame),
other_type => Err(ArrowError::NotYetImplemented(format!(
"compression type {other_type:?} not supported "
@@ -129,6 +129,25 @@ impl TryFrom<CompressionType> for CompressionCodec {
}
impl CompressionCodec {
+ /// Creates a [`CompressionCodec`] with an explicit compression level.
+ ///
+ /// The level is used for [`CompressionType::ZSTD`].
+ /// [`CompressionType::LZ4_FRAME`] does not yet support compression levels
+ /// and ignores this value. Returns an error for unsupported compression
+ /// types.
+ pub(crate) fn try_new_with_compression_level(
+ compression_type: CompressionType,
+ compression_level: i32,
+ ) -> Result<Self, ArrowError> {
+ match compression_type {
+ CompressionType::ZSTD =>
Ok(CompressionCodec::Zstd(compression_level)),
+ CompressionType::LZ4_FRAME => Ok(CompressionCodec::Lz4Frame),
+ other_type => Err(ArrowError::NotYetImplemented(format!(
+ "compression type {other_type:?} not supported "
+ ))),
+ }
+ }
+
/// Compresses the data in `input` to `output` and appends the
/// data using the specified compression mechanism.
///
@@ -213,7 +232,7 @@ impl CompressionCodec {
) -> Result<(), ArrowError> {
match self {
CompressionCodec::Lz4Frame => compress_lz4(input, output),
- CompressionCodec::Zstd => compress_zstd(input, output, context),
+ CompressionCodec::Zstd(level) => compress_zstd(input, output,
context, *level),
}
}
@@ -227,7 +246,7 @@ impl CompressionCodec {
) -> Result<Vec<u8>, ArrowError> {
let ret = match self {
CompressionCodec::Lz4Frame => decompress_lz4(input,
decompressed_size)?,
- CompressionCodec::Zstd => decompress_zstd(input,
decompressed_size, context)?,
+ CompressionCodec::Zstd(_) => decompress_zstd(input,
decompressed_size, context)?,
};
if ret.len() != decompressed_size {
return Err(ArrowError::IpcError(format!(
@@ -279,8 +298,9 @@ fn compress_zstd(
input: &[u8],
output: &mut Vec<u8>,
context: &mut CompressionContext,
+ level: i32,
) -> Result<(), ArrowError> {
- let result = context.zstd_compressor().compress(input)?;
+ let result = context.zstd_compressor(level).compress(input)?;
output.extend_from_slice(&result);
Ok(())
}
@@ -291,6 +311,7 @@ fn compress_zstd(
_input: &[u8],
_output: &mut Vec<u8>,
_context: &mut CompressionContext,
+ _level: i32,
) -> Result<(), ArrowError> {
Err(ArrowError::InvalidArgumentError(
"zstd IPC compression requires the zstd feature".to_string(),
@@ -363,7 +384,7 @@ mod tests {
#[cfg(feature = "zstd")]
fn test_zstd_compression() {
let input_bytes = b"hello zstd";
- let codec = super::CompressionCodec::Zstd;
+ let codec =
super::CompressionCodec::Zstd(super::DEFAULT_ZSTD_COMPRESSION_LEVEL);
let mut output_bytes: Vec<u8> = Vec::new();
codec
.compress(input_bytes, &mut output_bytes, &mut Default::default())
diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs
index 4142858ce8..9ab86807d4 100644
--- a/arrow-ipc/src/writer.rs
+++ b/arrow-ipc/src/writer.rs
@@ -66,6 +66,8 @@ pub struct IpcWriteOptions {
/// Compression, if desired. Will result in a runtime error
/// if the corresponding feature is not enabled
batch_compression_type: Option<crate::CompressionType>,
+ // Compression level
+ batch_compression_level: Option<i32>,
/// How to handle updating dictionaries in IPC messages
dictionary_handling: DictionaryHandling,
}
@@ -166,6 +168,63 @@ impl IpcWriteOptions {
}
Ok(self)
}
+
+ /// Configures the compression level used when writing compressed IPC
batches.
+ ///
+ /// Compression levels require metadata V5 or newer and are currently only
+ /// supported for ZSTD compression.
+ pub fn try_with_compression_level(
+ mut self,
+ batch_compression_level: Option<i32>,
+ ) -> Result<Self, ArrowError> {
+ self.batch_compression_level = batch_compression_level;
+
+ if self.batch_compression_level.is_some()
+ && self.metadata_version < crate::MetadataVersion::V5
+ {
+ return Err(ArrowError::InvalidArgumentError(
+ "Compression only supported in metadata v5 and
above".to_string(),
+ ));
+ }
+
+ match (self.batch_compression_type, self.batch_compression_level) {
+ (Some(crate::CompressionType::ZSTD), Some(level)) => {
+ return self.check_zstd_level(level);
+ }
+ (Some(crate::CompressionType::LZ4_FRAME), Some(_)) => {
+ return Err(ArrowError::InvalidArgumentError(
+ "LZ4 Frame compression does not support configurable
compression levels"
+ .to_string(),
+ ));
+ }
+ _ => {}
+ }
+
+ Ok(self)
+ }
+
+ #[cfg(not(feature = "zstd"))]
+ fn check_zstd_level(self, _level: i32) -> Result<Self, ArrowError> {
+ Err(ArrowError::InvalidArgumentError(
+ "zstd IPC compression requires the zstd feature".to_string(),
+ ))
+ }
+
+ #[cfg(feature = "zstd")]
+ fn check_zstd_level(self, level: i32) -> Result<Self, ArrowError> {
+ let range = zstd::compression_level_range();
+ if !range.contains(&(level as zstd::zstd_safe::CompressionLevel)) {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "ZSTD compression level must be between {} and {}, got {}",
+ range.start(),
+ range.end(),
+ level,
+ )));
+ }
+
+ Ok(self)
+ }
+
/// Try to create IpcWriteOptions, checking for incompatible settings
pub fn try_new(
alignment: usize,
@@ -192,6 +251,7 @@ impl IpcWriteOptions {
write_legacy_ipc_format,
metadata_version,
batch_compression_type: None,
+ batch_compression_level: None,
dictionary_handling: DictionaryHandling::default(),
}),
crate::MetadataVersion::V5 => {
@@ -205,6 +265,7 @@ impl IpcWriteOptions {
write_legacy_ipc_format,
metadata_version,
batch_compression_type: None,
+ batch_compression_level: None,
dictionary_handling: DictionaryHandling::default(),
})
}
@@ -229,6 +290,7 @@ impl Default for IpcWriteOptions {
write_legacy_ipc_format: false,
metadata_version: crate::MetadataVersion::V5,
batch_compression_type: None,
+ batch_compression_level: None,
dictionary_handling: DictionaryHandling::default(),
}
}
@@ -703,8 +765,15 @@ impl IpcDataGenerator {
c.finish()
});
- let compression_codec: Option<CompressionCodec> =
- batch_compression_type.map(TryInto::try_into).transpose()?;
+ let batch_compression_level = write_options.batch_compression_level;
+ let compression_codec: Option<CompressionCodec> =
batch_compression_type
+ .map(|compression_type| match batch_compression_level {
+ Some(level) => {
+
CompressionCodec::try_new_with_compression_level(compression_type, level)
+ }
+ None => compression_type.try_into(),
+ })
+ .transpose()?;
let alignment = write_options.alignment;
let mut variadic_buffer_counts = vec![];
@@ -785,8 +854,14 @@ impl IpcDataGenerator {
c.finish()
});
+ let batch_compression_level = write_options.batch_compression_level;
let compression_codec: Option<CompressionCodec> =
batch_compression_type
- .map(|batch_compression_type| batch_compression_type.try_into())
+ .map(|batch_compression_type| match batch_compression_level {
+ Some(level) => {
+
CompressionCodec::try_new_with_compression_level(batch_compression_type, level)
+ }
+ None => batch_compression_type.try_into(),
+ })
.transpose()?;
let alignment = write_options.alignment;
@@ -2495,6 +2570,8 @@ mod tests {
let write_option = IpcWriteOptions::try_new(8, false,
crate::MetadataVersion::V5)
.unwrap()
.try_with_compression(Some(crate::CompressionType::ZSTD))
+ .unwrap()
+ .try_with_compression_level(Some(1))
.unwrap();
let mut writer =
