Skip to content

Avoid over-reserving Vecs in compression codecs #10483

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

Open
MassivePizza wants to merge 5 commits into
base: main
Choose a base branch
from
+81 −58
Open

Avoid over-reserving Vecs in compression codecs #10483

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
cb73239
avoid over-reserving Vecs in ZSTDCodec
MassivePizza Jul 29, 2026
139281d
simplify LZ4Codec writing and reading
MassivePizza Jul 29, 2026
65223f0
fix clippy
MassivePizza Jul 29, 2026
0ee7ac3
reserve later when compressing zstd
MassivePizza Jul 29, 2026
26d5ad0
use uncompress_size in codecs
MassivePizza Jul 29, 2026
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
139 changes: 81 additions & 58 deletions parquet/src/compression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,8 +281,11 @@ mod gzip_codec {
&mut self,
input_buf: &[u8],
output_buf: &mut Vec<u8>,
_uncompress_size: Option<usize>,
uncompress_size: Option<usize>,
) -> Result<usize> {
if let Some(len) = uncompress_size {
output_buf.reserve(len);
}
let mut decoder = read::MultiGzDecoder::new(input_buf);
decoder.read_to_end(output_buf).map_err(|e| e.into())
}
Expand Down Expand Up @@ -389,8 +392,10 @@ mod brotli_codec {
output_buf: &mut Vec<u8>,
uncompress_size: Option<usize>,
) -> Result<usize> {
let buffer_size = uncompress_size.unwrap_or(BROTLI_DEFAULT_BUFFER_SIZE);
brotli::Decompressor::new(input_buf, buffer_size)
if let Some(len) = uncompress_size {
output_buf.reserve(len);
}
brotli::Decompressor::new(input_buf, BROTLI_DEFAULT_BUFFER_SIZE)
.read_to_end(output_buf)
.map_err(|e| e.into())
}
Expand Down Expand Up @@ -446,8 +451,6 @@ mod lz4_codec {
use crate::compression::Codec;
use crate::errors::{ParquetError, Result};

const LZ4_BUFFER_SIZE: usize = 4096;

/// Codec for LZ4 compression algorithm.
pub struct LZ4Codec {}

Expand All @@ -463,33 +466,19 @@ mod lz4_codec {
&mut self,
input_buf: &[u8],
output_buf: &mut Vec<u8>,
_uncompress_size: Option<usize>,
uncompress_size: Option<usize>,
) -> Result<usize> {
let mut decoder = lz4_flex::frame::FrameDecoder::new(input_buf);
let mut buffer: [u8; LZ4_BUFFER_SIZE] = [0; LZ4_BUFFER_SIZE];
let mut total_len = 0;
loop {
let len = decoder.read(&mut buffer)?;
if len == 0 {
break;
}
total_len += len;
output_buf.write_all(&buffer[0..len])?;
if let Some(len) = uncompress_size {
output_buf.reserve(len);
}
let mut decoder = lz4_flex::frame::FrameDecoder::new(input_buf);
let total_len = decoder.read_to_end(output_buf)?;
Ok(total_len)
}

fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
let mut encoder = lz4_flex::frame::FrameEncoder::new(output_buf);
let mut from = 0;
loop {
let to = std::cmp::min(from + LZ4_BUFFER_SIZE, input_buf.len());
encoder.write_all(&input_buf[from..to])?;
from += LZ4_BUFFER_SIZE;
if from >= input_buf.len() {
break;
}
}
encoder.write_all(input_buf)?;
match encoder.finish() {
Ok(_) => Ok(()),
Err(e) => Err(ParquetError::External(Box::new(e))),
Expand All @@ -503,29 +492,66 @@ pub use lz4_codec::*;

#[cfg(any(feature = "zstd", test))]
mod zstd_codec {
use zstd::zstd_safe;

use crate::compression::{Codec, ZstdLevel};
use crate::errors::Result;
use std::io::Cursor;
use std::io::Read;

/// Codec for Zstandard compression algorithm.
///
/// Uses `zstd::bulk` API with reusable compressor/decompressor contexts
/// to avoid the overhead of reinitializing contexts for each operation.
pub struct ZSTDCodec {
compressor: zstd::bulk::Compressor<'static>,
decompressor: zstd::bulk::Decompressor<'static>,
cctx: zstd_safe::CCtx<'static>,
dctx: zstd_safe::DCtx<'static>,
}

impl ZSTDCodec {
/// Creates new Zstandard compression codec.
pub(crate) fn new(level: ZstdLevel) -> Self {
Self {
compressor: zstd::bulk::Compressor::new(level.compression_level())
.expect("valid zstd compression level"),
decompressor: zstd::bulk::Decompressor::new()
.expect("can create zstd decompressor"),
let mut cctx = zstd_safe::CCtx::create();
cctx.set_parameter(zstd_safe::CParameter::CompressionLevel(
level.compression_level(),
))
.expect("valid zstd compression level");

let dctx = zstd_safe::DCtx::create();

Self { cctx, dctx }
}
}

/// Avoids zstd crate abstractions to minimize redundant copies;
/// [zstd::stream::Encoder] uses a buffered writer which is pointless for Vec.
fn compress_to_vec(
cctx: &mut zstd_safe::CCtx<'static>,
input_buf: &[u8],
output_buf: &mut Vec<u8>,
) -> std::result::Result<(), zstd_safe::ErrorCode> {
cctx.reset(zstd_safe::ResetDirective::SessionOnly)?;
cctx.set_pledged_src_size(Some(input_buf.len() as u64))?;

let mut input = zstd_safe::InBuffer::around(input_buf);
loop {
let mut output = zstd_safe::OutBuffer::around_pos(output_buf, output_buf.len());
let end_op = zstd_safe::zstd_sys::ZSTD_EndDirective::ZSTD_e_continue;
let to_flush = cctx.compress_stream2(&mut output, &mut input, end_op)?;
if input.pos == input.src.len() {
break; // let the end_stream loop below call reserve_exact with the finalized amount
}
output_buf.reserve(to_flush);
}

loop {
let mut output = zstd_safe::OutBuffer::around_pos(output_buf, output_buf.len());
let to_flush = cctx.end_stream(&mut output)?;
if to_flush == 0 {
break;
}
output_buf.reserve_exact(to_flush);

@mapleFU mapleFU Jul 30, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Emmm calling reserve_exact in a loop is a bit weird, does this guranteed calling once?

@MassivePizza MassivePizza Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

It's not ideal but that's what the Zstd API looks like. end_stream will essentially always produce N followed by 0 bytes. Only multi-threaded encode would need multiple flushes as workers drain.

Just reserve works but I tried to avoid need for future shrinks.
From testing, this is often the first and last alloc too (output_buf Vecs don't seem to be reused much).

}
Ok(())
}

impl Codec for ZSTDCodec {
Expand All @@ -535,35 +561,32 @@ mod zstd_codec {
output_buf: &mut Vec<u8>,
uncompress_size: Option<usize>,
) -> Result<usize> {
let offset = output_buf.len();
let len = uncompress_size
.or_else(|| {
// Get the decompressed size from the zstd frame header
zstd::zstd_safe::get_frame_content_size(input_buf)
.ok()
.flatten()
.map(|size| size as usize)
})
.unwrap_or(input_buf.len().saturating_mul(4));
output_buf.reserve(len);

let mut cursor = Cursor::new(output_buf);
cursor.set_position(offset as u64);
let len = self
.decompressor
.decompress_to_buffer(input_buf, &mut cursor)?;
if let Some(len) = uncompress_size.or_else(|| {
// Get the decompressed size from the zstd frame header
zstd::zstd_safe::get_frame_content_size(input_buf)
.ok()
.flatten()
.map(|size| size as usize)
}) {
output_buf.reserve(len);
}

let mut decoder =
zstd::stream::Decoder::with_context(input_buf, &mut self.dctx).single_frame();

// The default Read::read_to_end impl is acceptable;
// it reads directly into the Vec most of the time.
// Using raw DCtx here would be more annoying than the compress path,
// but could be done if better Vec reserve control is desired in the future.
let len = decoder.read_to_end(output_buf)?;
Ok(len)
}

fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
let offset = output_buf.len();
let len = zstd::zstd_safe::compress_bound(input_buf.len());
output_buf.reserve(len);

let mut cursor = Cursor::new(output_buf);
cursor.set_position(offset as u64);
let _written = self.compressor.compress_to_buffer(input_buf, &mut cursor)?;
Ok(())
compress_to_vec(&mut self.cctx, input_buf, output_buf).map_err(|code| {
let msg = zstd_safe::get_error_name(code);
std::io::Error::other(msg.to_string()).into()
})
}
}
}
Expand Down
Loading