Rich-T-kid commented on code in PR #12148:
URL: https://github.com/apache/arrow-rs/pull/12148#discussion_r4556744908
##########
arrow-ipc/src/compression.rs:
##########
@@ -52,6 +82,22 @@ impl IpcWriteContext {
std::mem::take(&mut self.scratch)
}
+ /// Set the lz4 frame block size used by the cached lz4 encoder.
+ ///
+ /// A cached encoder is dropped when the size changes so the next
compression call builds one
+ /// with the new block size. The options builder validates the size before
it reaches this
+ /// point.
Review Comment:
nit: making it only buildable from `from_bytes()` guarantees this so we can
omit it here.
```suggestion
/// Set the lz4 frame block size used by the cached lz4 encoder.
///
/// A cached encoder is dropped when the size changes so the next
compression call builds one
/// with the new block size.
```
##########
arrow-ipc/src/compression.rs:
##########
@@ -11,18 +11,48 @@ const LENGTH_NO_COMPRESSED_DATA: i64 = -1;
const LENGTH_OF_PREFIX_DATA: i64 = 8;
const DEFAULT_ZSTD_COMPRESSION_LEVEL: i32 = 3;
+/// LZ4 frame block sizes representable by the frame header.
+///
+/// `lz4_flex` also offers an 8 MiB block size, but encodes it by setting the
frame header's
+/// reserved bit, which the decoder rejects.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum Lz4BlockSize {
+ Max64KiB,
+ Max256KiB,
+ Max1MiB,
+ Max4MiB,
+}
Review Comment:
this is much cleaner, may be worth including the link to lz4 docs so future
readers understand why these are the only valid variants
##########
arrow-ipc/src/compression.rs:
##########
@@ -11,18 +11,48 @@ const LENGTH_NO_COMPRESSED_DATA: i64 = -1;
const LENGTH_OF_PREFIX_DATA: i64 = 8;
const DEFAULT_ZSTD_COMPRESSION_LEVEL: i32 = 3;
+/// LZ4 frame block sizes representable by the frame header.
+///
+/// `lz4_flex` also offers an 8 MiB block size, but encodes it by setting the
frame header's
+/// reserved bit, which the decoder rejects.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum Lz4BlockSize {
+ Max64KiB,
+ Max256KiB,
+ Max1MiB,
+ Max4MiB,
+}
+
+impl Lz4BlockSize {
+ /// Default block size balances compression throughput and retained
encoder buffers.
+ pub(crate) const DEFAULT: Self = Self::Max256KiB;
Review Comment:
restates what the code says
```suggestion
pub(crate) const DEFAULT: Self = Self::Max256KiB;
```
##########
arrow-ipc/src/writer.rs:
##########
@@ -54,6 +54,9 @@ pub struct IpcWriteOptions {
batch_compression_type: Option<crate::CompressionType>,
// Compression level
batch_compression_level: Option<i32>,
+ /// The lz4 frame block size; `None` selects 256 KiB. Set through
+ /// [`IpcWriteOptions::try_with_lz4_block_size`], which validates the
value.
+ batch_lz4_block_size: Option<Lz4BlockSize>,
Review Comment:
```suggestion
/// The lz4 frame block size; `None` selects 256 KiB.
batch_lz4_block_size: Option<Lz4BlockSize>,
```
##########
arrow-ipc/src/compression.rs:
##########
@@ -297,18 +345,21 @@ impl CompressionCodec {
#[cfg(feature = "lz4")]
struct Lz4Compressor {
encoder: lz4_flex::frame::FrameEncoder<Vec<u8>>,
+ block_size: Lz4BlockSize,
}
#[cfg(feature = "lz4")]
impl Lz4Compressor {
- fn new() -> Self {
+ fn new(block_size: Lz4BlockSize) -> Self {
// Pin the block size: with `BlockSize::Auto` the first buffer written
would fix it for
// every subsequent frame of the reused encoder. 256KB keeps large
buffers as fast as
- // with bigger blocks while the encoder retains only ~0.5MB of buffers.
+ // with bigger blocks while the encoder retains only ~0.5MB of
buffers; callers can pick
+ // another supported size via
`IpcWriteOptions::try_with_lz4_block_size`.
let frame_info =
-
lz4_flex::frame::FrameInfo::new().block_size(lz4_flex::frame::BlockSize::Max256KB);
+
lz4_flex::frame::FrameInfo::new().block_size(to_lz4_block_size(block_size));
Review Comment:
i think we need to update this entire comment block.
##########
arrow-ipc/src/writer.rs:
##########
@@ -3143,6 +3208,133 @@ mod tests {
}
}
+ /// Batch exercising nulls, a dictionary column, and enough bytes that lz4
emits several
Review Comment:
Can we retain a focused test that reuses one `IpcWriteContext` while
switching block sizes? Each current iteration creates a fresh StreamWriter, so
removing the cached-encoder reset still passes both new tests. This is the
behavior `set_lz4_block_size` exists to preserve.
##########
arrow-ipc/src/compression.rs:
##########
@@ -331,12 +382,24 @@ impl Lz4Compressor {
std::mem::swap(self.encoder.get_mut(), output);
if result.is_err() {
// The encoder may be left mid-frame, start from a fresh one next
time.
- *self = Self::new();
+ let block_size = self.block_size;
+ *self = Self::new(block_size);
}
result
}
}
+/// Maps a validated block size to the `lz4_flex` frame header code.
+#[cfg(feature = "lz4")]
+fn to_lz4_block_size(block_size: Lz4BlockSize) -> lz4_flex::frame::BlockSize {
Review Comment:
👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
