diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index 0ae0182263ed..303e7a0e730c 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -57,7 +57,7 @@ brotli = { version = "8.0", default-features = false, features = ["std"], option # To use `flate2` you must enable either the `flate2-zlib-rs` or `flate2-rust_backend` backends flate2 = { version = "1.1", default-features = false, optional = true } lz4_flex = { version = "0.14", default-features = false, features = ["std", "frame"], optional = true } -zstd = { version = "0.13", optional = true, default-features = false } +zstd = { version = "0.13", optional = true, default-features = false, features = ["experimental"] } chrono = { workspace = true } num-bigint = { version = "0.5", default-features = false } num-integer = { version = "0.1.46", default-features = false, features = ["std"] } diff --git a/parquet/src/compression.rs b/parquet/src/compression.rs index 6f3a18fe60b0..c28521e9a65f 100644 --- a/parquet/src/compression.rs +++ b/parquet/src/compression.rs @@ -281,8 +281,11 @@ mod gzip_codec { &mut self, input_buf: &[u8], output_buf: &mut Vec, - _uncompress_size: Option, + uncompress_size: Option, ) -> Result { + 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()) } @@ -389,8 +392,10 @@ mod brotli_codec { output_buf: &mut Vec, uncompress_size: Option, ) -> Result { - 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()) } @@ -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 {} @@ -463,33 +466,19 @@ mod lz4_codec { &mut self, input_buf: &[u8], output_buf: &mut Vec, - _uncompress_size: Option, + uncompress_size: Option, ) -> Result { - 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) -> 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))), @@ -503,6 +492,8 @@ 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; @@ -512,20 +503,64 @@ mod zstd_codec { /// 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 } + } + } + + fn map_error_code(code: zstd_safe::ErrorCode) -> crate::errors::ParquetError { + let msg = zstd_safe::get_error_name(code); + std::io::Error::other(msg.to_string()).into() + } + + /// 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, + ) -> std::result::Result<(), zstd_safe::ErrorCode> { + cctx.reset(zstd_safe::ResetDirective::SessionOnly)?; + // causes frame header to include size + cctx.set_pledged_src_size(Some(input_buf.len() as u64))?; + // zstd can avoid allocating and copying into an input window + cctx.set_parameter(zstd_safe::CParameter::StableInBuffer(true))?; + + 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 end_op = zstd_safe::zstd_sys::ZSTD_EndDirective::ZSTD_e_end; + let to_flush = cctx.compress_stream2(&mut output, &mut input, end_op)?; + if to_flush == 0 { + break; + } + output_buf.reserve_exact(to_flush); + } + Ok(()) } impl Codec for ZSTDCodec { @@ -535,35 +570,30 @@ mod zstd_codec { output_buf: &mut Vec, uncompress_size: Option, ) -> Result { - 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); + if let Some(len) = uncompress_size { + output_buf.reserve_exact(len); + } else if let Some(len) = zstd_safe::find_decompressed_size(input_buf) + .map_err(|err| std::io::Error::other(err.to_string()))? + { + output_buf.reserve_exact(len as usize); + } else { + let len = zstd_safe::decompress_bound(input_buf).map_err(map_error_code)?; + output_buf.reserve(len as usize); + } + let offset = output_buf.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)?; + .dctx + .decompress(&mut cursor, input_buf) + .map_err(map_error_code)?; Ok(len) } fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> 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(map_error_code) } } }