mapleFU commented on code in PR #10483:
URL: https://github.com/apache/arrow-rs/pull/10483#discussion_r3679474776
##########
parquet/src/compression.rs:
##########
@@ -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);
Review Comment:
Emmm calling reserve_exact in a loop is a bit weird, does this guranteed
calling once?
--
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]
