diff --git a/parquet/src/compression.rs b/parquet/src/compression.rs index fe2fb59c5b8c..6f3a18fe60b0 100644 --- a/parquet/src/compression.rs +++ b/parquet/src/compression.rs @@ -505,6 +505,7 @@ pub use lz4_codec::*; mod zstd_codec { use crate::compression::{Codec, ZstdLevel}; use crate::errors::Result; + use std::io::Cursor; /// Codec for Zstandard compression algorithm. /// @@ -534,22 +535,34 @@ mod zstd_codec { output_buf: &mut Vec, uncompress_size: Option, ) -> Result { - let capacity = uncompress_size.unwrap_or_else(|| { - // Get the decompressed size from the zstd frame header - zstd::zstd_safe::get_frame_content_size(input_buf) - .ok() - .flatten() - .unwrap_or(input_buf.len() as u64 * 4) as usize - }); - let decompressed = self.decompressor.decompress(input_buf, capacity)?; - let len = decompressed.len(); - output_buf.extend_from_slice(&decompressed); + 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)?; Ok(len) } fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()> { - let compressed = self.compressor.compress(input_buf)?; - output_buf.extend_from_slice(&compressed); + 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(()) } }