arrow_ipc/
compression.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::CompressionType;
19use arrow_buffer::Buffer;
20use arrow_schema::ArrowError;
21use flatbuffers::FlatBufferBuilder;
22
23const LENGTH_NO_COMPRESSED_DATA: i64 = -1;
24const LENGTH_OF_PREFIX_DATA: i64 = 8;
25const DEFAULT_ZSTD_COMPRESSION_LEVEL: i32 = 3;
26
27/// Additional context that may be needed for compression.
28///
29/// In the case of zstd, this will contain the zstd context, which can be reused between subsequent
30/// compression calls to avoid the performance overhead of initialising a new context for every
31/// compression. Also holds a [`FlatBufferBuilder`] that is reused across IPC writes.
32#[derive(Default)]
33pub struct IpcWriteContext {
34    scratch: Vec<u8>,
35    reserve_scratch: bool,
36    fbb: FlatBufferBuilder<'static>,
37    #[cfg(feature = "zstd")]
38    compressor: Option<zstd::bulk::Compressor<'static>>,
39}
40
41impl IpcWriteContext {
42    /// Get a mutable reference to the [`FlatBufferBuilder`] that is reused across IPC writes.
43    pub(crate) fn mut_fbb(&mut self) -> &mut FlatBufferBuilder<'static> {
44        &mut self.fbb
45    }
46
47    /// Set whether the scratch buffer capacity should be reserved after each encode for reuse
48    /// on the next call. Set to `false` for the final batch in a sequence to avoid a
49    /// pointless allocation. by default, this is set to `false`.
50    pub fn set_reserve_scratch(&mut self, reserve: bool) {
51        self.reserve_scratch = reserve;
52    }
53    /// Reserve the scratch buffer capacity for reuse on the next call. This is a no-op if
54    /// `reserve_scratch` is set to `false`.
55    pub(crate) fn reserve_scratch_with_capacity(&mut self, additional: usize) {
56        if self.reserve_scratch {
57            self.scratch.reserve(additional);
58        }
59    }
60    pub(crate) fn scratch(&mut self) -> Vec<u8> {
61        std::mem::take(&mut self.scratch)
62    }
63
64    #[cfg(feature = "zstd")]
65    fn zstd_compressor(&mut self, level: i32) -> &mut zstd::bulk::Compressor<'static> {
66        self.compressor.get_or_insert_with(|| {
67            zstd::bulk::Compressor::new(level).expect("can use default compression level")
68        })
69    }
70}
71
72impl std::fmt::Debug for IpcWriteContext {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        let mut ds = f.debug_struct("IpcWriteContext");
75
76        #[cfg(feature = "zstd")]
77        ds.field(
78            "compressor",
79            &self.compressor.as_ref().map(|_| "zstd::bulk::Compressor"),
80        );
81
82        ds.finish()
83    }
84}
85
86/// Deprecated alias for [`IpcWriteContext`].
87#[deprecated(since = "59.1.0", note = "Use IpcWriteContext instead")]
88pub type CompressionContext = IpcWriteContext;
89
90/// Additional context that may be needed for decompression.
91///
92/// In the case of zstd, this will contain the zstd decompression context, which can be reused
93/// between subsequent decompression calls to avoid the performance overhead of initialising a new
94/// context for every decompression.
95pub struct DecompressionContext {
96    #[cfg(feature = "zstd")]
97    decompressor: Option<zstd::bulk::Decompressor<'static>>,
98}
99
100impl DecompressionContext {
101    pub(crate) fn new() -> Self {
102        Default::default()
103    }
104
105    #[cfg(feature = "zstd")]
106    fn zstd_decompressor(&mut self) -> &mut zstd::bulk::Decompressor<'static> {
107        self.decompressor.get_or_insert_with(|| {
108            zstd::bulk::Decompressor::new().expect("can create zstd decompressor")
109        })
110    }
111}
112
113#[allow(clippy::derivable_impls)]
114impl Default for DecompressionContext {
115    fn default() -> Self {
116        DecompressionContext {
117            #[cfg(feature = "zstd")]
118            decompressor: None,
119        }
120    }
121}
122
123impl std::fmt::Debug for DecompressionContext {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        let mut ds = f.debug_struct("DecompressionContext");
126
127        #[cfg(feature = "zstd")]
128        ds.field(
129            "decompressor",
130            &self
131                .decompressor
132                .as_ref()
133                .map(|_| "zstd::bulk::Decompressor"),
134        );
135
136        ds.finish()
137    }
138}
139
140/// Represents compressing a ipc stream using a particular compression algorithm
141#[derive(Debug, Clone, Copy, PartialEq, Eq)]
142pub enum CompressionCodec {
143    Lz4Frame,
144    Zstd(i32),
145}
146
147impl TryFrom<CompressionType> for CompressionCodec {
148    type Error = ArrowError;
149
150    fn try_from(compression_type: CompressionType) -> Result<Self, ArrowError> {
151        match compression_type {
152            CompressionType::ZSTD => Ok(CompressionCodec::Zstd(DEFAULT_ZSTD_COMPRESSION_LEVEL)),
153            CompressionType::LZ4_FRAME => Ok(CompressionCodec::Lz4Frame),
154            other_type => Err(ArrowError::NotYetImplemented(format!(
155                "compression type {other_type:?} not supported "
156            ))),
157        }
158    }
159}
160
161impl CompressionCodec {
162    /// Creates a [`CompressionCodec`] with an explicit compression level.
163    ///
164    /// The level is used for [`CompressionType::ZSTD`].
165    /// [`CompressionType::LZ4_FRAME`] does not yet support compression levels
166    /// and ignores this value. Returns an error for unsupported compression
167    /// types.
168    pub(crate) fn try_new_with_compression_level(
169        compression_type: CompressionType,
170        compression_level: i32,
171    ) -> Result<Self, ArrowError> {
172        match compression_type {
173            CompressionType::ZSTD => Ok(CompressionCodec::Zstd(compression_level)),
174            CompressionType::LZ4_FRAME => Ok(CompressionCodec::Lz4Frame),
175            other_type => Err(ArrowError::NotYetImplemented(format!(
176                "compression type {other_type:?} not supported "
177            ))),
178        }
179    }
180
181    /// Compresses the data in `input` to `output` and appends the
182    /// data using the specified compression mechanism.
183    ///
184    /// returns the number of bytes written to the stream
185    ///
186    /// Writes this format to output:
187    /// ```text
188    /// [8 bytes]:         uncompressed length
189    /// [remaining bytes]: compressed data stream
190    /// ```
191    pub(crate) fn compress_to_vec(
192        &self,
193        input: &[u8],
194        output: &mut Vec<u8>,
195        context: &mut IpcWriteContext,
196    ) -> Result<usize, ArrowError> {
197        let uncompressed_data_len = input.len();
198        let original_output_len = output.len();
199
200        if input.is_empty() {
201            // empty input, nothing to do
202        } else {
203            // write compressed data directly into the output buffer
204            output.extend_from_slice(&uncompressed_data_len.to_le_bytes());
205            self.compress(input, output, context)?;
206
207            let compression_len = output.len() - original_output_len;
208            if compression_len > uncompressed_data_len {
209                // length of compressed data was larger than
210                // uncompressed data, use the uncompressed data with
211                // length -1 to indicate that we don't compress the
212                // data
213                output.truncate(original_output_len);
214                output.extend_from_slice(&LENGTH_NO_COMPRESSED_DATA.to_le_bytes());
215                output.extend_from_slice(input);
216            }
217        }
218        Ok(output.len() - original_output_len)
219    }
220
221    /// Decompresses the input into a [`Buffer`]
222    ///
223    /// The input should look like:
224    /// ```text
225    /// [8 bytes]:         uncompressed length
226    /// [remaining bytes]: compressed data stream
227    /// ```
228    pub(crate) fn decompress_to_buffer(
229        &self,
230        input: &Buffer,
231        context: &mut DecompressionContext,
232    ) -> Result<Buffer, ArrowError> {
233        // read the first 8 bytes to determine if the data is
234        // compressed
235        let decompressed_length = read_uncompressed_size(input)?;
236        let buffer = if decompressed_length == 0 {
237            // empty
238            Buffer::from([])
239        } else if decompressed_length == LENGTH_NO_COMPRESSED_DATA {
240            // no compression
241            input.slice(LENGTH_OF_PREFIX_DATA as usize)
242        } else if let Ok(decompressed_length) = usize::try_from(decompressed_length) {
243            // decompress data using the codec
244            let input_data = &input[(LENGTH_OF_PREFIX_DATA as usize)..];
245            let v = self.decompress(input_data, decompressed_length as _, context)?;
246            Buffer::from_vec(v)
247        } else {
248            return Err(ArrowError::IpcError(format!(
249                "Invalid uncompressed length: {decompressed_length}"
250            )));
251        };
252        Ok(buffer)
253    }
254
255    /// Compress the data in input buffer and write to output buffer
256    /// using the specified compression
257    fn compress(
258        &self,
259        input: &[u8],
260        output: &mut Vec<u8>,
261        context: &mut IpcWriteContext,
262    ) -> Result<(), ArrowError> {
263        match self {
264            CompressionCodec::Lz4Frame => compress_lz4(input, output),
265            CompressionCodec::Zstd(level) => compress_zstd(input, output, context, *level),
266        }
267    }
268
269    /// Decompress the data in input buffer and write to output buffer
270    /// using the specified compression
271    fn decompress(
272        &self,
273        input: &[u8],
274        decompressed_size: usize,
275        context: &mut DecompressionContext,
276    ) -> Result<Vec<u8>, ArrowError> {
277        let ret = match self {
278            CompressionCodec::Lz4Frame => decompress_lz4(input, decompressed_size)?,
279            CompressionCodec::Zstd(_) => decompress_zstd(input, decompressed_size, context)?,
280        };
281        if ret.len() != decompressed_size {
282            return Err(ArrowError::IpcError(format!(
283                "Expected compressed length of {decompressed_size} got {}",
284                ret.len()
285            )));
286        }
287        Ok(ret)
288    }
289}
290
291#[cfg(feature = "lz4")]
292fn compress_lz4(input: &[u8], output: &mut Vec<u8>) -> Result<(), ArrowError> {
293    use std::io::Write;
294    let mut encoder = lz4_flex::frame::FrameEncoder::new(output);
295    encoder.write_all(input)?;
296    encoder
297        .finish()
298        .map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
299    Ok(())
300}
301
302#[cfg(not(feature = "lz4"))]
303#[allow(clippy::ptr_arg)]
304fn compress_lz4(_input: &[u8], _output: &mut Vec<u8>) -> Result<(), ArrowError> {
305    Err(ArrowError::InvalidArgumentError(
306        "lz4 IPC compression requires the lz4 feature".to_string(),
307    ))
308}
309
310#[cfg(feature = "lz4")]
311fn decompress_lz4(input: &[u8], decompressed_size: usize) -> Result<Vec<u8>, ArrowError> {
312    use std::io::Read;
313    let mut output = Vec::with_capacity(decompressed_size);
314    lz4_flex::frame::FrameDecoder::new(input).read_to_end(&mut output)?;
315    Ok(output)
316}
317
318#[cfg(not(feature = "lz4"))]
319#[allow(clippy::ptr_arg)]
320fn decompress_lz4(_input: &[u8], _decompressed_size: usize) -> Result<Vec<u8>, ArrowError> {
321    Err(ArrowError::InvalidArgumentError(
322        "lz4 IPC decompression requires the lz4 feature".to_string(),
323    ))
324}
325
326#[cfg(feature = "zstd")]
327fn compress_zstd(
328    input: &[u8],
329    output: &mut Vec<u8>,
330    context: &mut IpcWriteContext,
331    level: i32,
332) -> Result<(), ArrowError> {
333    let result = context.zstd_compressor(level).compress(input)?;
334    output.extend_from_slice(&result);
335    Ok(())
336}
337
338#[cfg(not(feature = "zstd"))]
339#[allow(clippy::ptr_arg)]
340fn compress_zstd(
341    _input: &[u8],
342    _output: &mut Vec<u8>,
343    _context: &mut IpcWriteContext,
344    _level: i32,
345) -> Result<(), ArrowError> {
346    Err(ArrowError::InvalidArgumentError(
347        "zstd IPC compression requires the zstd feature".to_string(),
348    ))
349}
350
351#[cfg(feature = "zstd")]
352fn decompress_zstd(
353    input: &[u8],
354    decompressed_size: usize,
355    context: &mut DecompressionContext,
356) -> Result<Vec<u8>, ArrowError> {
357    let output = context
358        .zstd_decompressor()
359        .decompress(input, decompressed_size)?;
360    Ok(output)
361}
362
363#[cfg(not(feature = "zstd"))]
364#[allow(clippy::ptr_arg)]
365fn decompress_zstd(
366    _input: &[u8],
367    _decompressed_size: usize,
368    _context: &mut DecompressionContext,
369) -> Result<Vec<u8>, ArrowError> {
370    Err(ArrowError::InvalidArgumentError(
371        "zstd IPC decompression requires the zstd feature".to_string(),
372    ))
373}
374
375/// Get the uncompressed length
376/// Notes:
377///   LENGTH_NO_COMPRESSED_DATA: indicate that the data that follows is not compressed
378///    0: indicate that there is no data
379///   positive number: indicate the uncompressed length for the following data
380/// Returns an error if the input buffer is shorter than 8 bytes
381#[inline]
382fn read_uncompressed_size(buffer: &[u8]) -> Result<i64, ArrowError> {
383    let len_buffer = buffer.get(..LENGTH_OF_PREFIX_DATA as usize).ok_or_else(|| {
384        ArrowError::IpcError(format!(
385            "Compressed IPC buffer is too short: expected at least {LENGTH_OF_PREFIX_DATA} bytes, got {}",
386            buffer.len()
387        ))
388    })?;
389    Ok(i64::from_le_bytes(len_buffer.try_into().unwrap()))
390}
391
392#[cfg(test)]
393mod tests {
394    #[test]
395    #[cfg(feature = "lz4")]
396    fn test_lz4_compression() {
397        let input_bytes = b"hello lz4";
398        let codec = super::CompressionCodec::Lz4Frame;
399        let mut output_bytes: Vec<u8> = Vec::new();
400        codec
401            .compress(input_bytes, &mut output_bytes, &mut Default::default())
402            .unwrap();
403        let result = codec
404            .decompress(
405                output_bytes.as_slice(),
406                input_bytes.len(),
407                &mut Default::default(),
408            )
409            .unwrap();
410        assert_eq!(input_bytes, result.as_slice());
411    }
412
413    #[test]
414    #[cfg(feature = "zstd")]
415    fn test_zstd_compression() {
416        let input_bytes = b"hello zstd";
417        let codec = super::CompressionCodec::Zstd(super::DEFAULT_ZSTD_COMPRESSION_LEVEL);
418        let mut output_bytes: Vec<u8> = Vec::new();
419        codec
420            .compress(input_bytes, &mut output_bytes, &mut Default::default())
421            .unwrap();
422        let result = codec
423            .decompress(
424                output_bytes.as_slice(),
425                input_bytes.len(),
426                &mut Default::default(),
427            )
428            .unwrap();
429        assert_eq!(input_bytes, result.as_slice());
430    }
431
432    #[test]
433    fn test_read_uncompressed_size_rejects_short_prefix() {
434        let err = super::read_uncompressed_size(&[1, 2, 3, 4, 5, 6, 7])
435            .expect_err("short compressed IPC prefix should return an error");
436
437        assert!(
438            err.to_string()
439                .contains("Compressed IPC buffer is too short"),
440            "unexpected error: {err}"
441        );
442    }
443}
arrow_ipc/compression.rs

arrow_ipc/
compression.rs
