arrow_ipc/
compression.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::CompressionType;
6use arrow_buffer::Buffer;
7use arrow_schema::ArrowError;
8use flatbuffers::FlatBufferBuilder;
9
10const LENGTH_NO_COMPRESSED_DATA: i64 = -1;
11const LENGTH_OF_PREFIX_DATA: i64 = 8;
12const DEFAULT_ZSTD_COMPRESSION_LEVEL: i32 = 3;
13
14/// Additional context that may be needed for compression.
15///
16/// In the case of zstd, this will contain the zstd context, which can be reused between subsequent
17/// compression calls to avoid the performance overhead of initialising a new context for every
18/// compression. In the case of lz4, this will contain the lz4 frame encoder, whose hash table and
19/// block buffers are likewise reused. Also holds a [`FlatBufferBuilder`] that is reused across
20/// IPC writes.
21#[derive(Default)]
22pub struct IpcWriteContext {
23    scratch: Vec<u8>,
24    reserve_scratch: bool,
25    fbb: FlatBufferBuilder<'static>,
26    #[cfg(feature = "zstd")]
27    compressor: Option<zstd::bulk::Compressor<'static>>,
28    #[cfg(feature = "lz4")]
29    lz4_compressor: Option<Lz4Compressor>,
30}
31
32impl IpcWriteContext {
33    /// Get a mutable reference to the [`FlatBufferBuilder`] that is reused across IPC writes.
34    pub(crate) fn mut_fbb(&mut self) -> &mut FlatBufferBuilder<'static> {
35        &mut self.fbb
36    }
37
38    /// Set whether the scratch buffer capacity should be reserved after each encode for reuse
39    /// on the next call. Set to `false` for the final batch in a sequence to avoid a
40    /// pointless allocation. by default, this is set to `false`.
41    pub fn set_reserve_scratch(&mut self, reserve: bool) {
42        self.reserve_scratch = reserve;
43    }
44    /// Reserve the scratch buffer capacity for reuse on the next call. This is a no-op if
45    /// `reserve_scratch` is set to `false`.
46    pub(crate) fn reserve_scratch_with_capacity(&mut self, additional: usize) {
47        if self.reserve_scratch {
48            self.scratch.reserve(additional);
49        }
50    }
51    pub(crate) fn scratch(&mut self) -> Vec<u8> {
52        std::mem::take(&mut self.scratch)
53    }
54
55    #[cfg(feature = "zstd")]
56    fn zstd_compressor(&mut self, level: i32) -> &mut zstd::bulk::Compressor<'static> {
57        self.compressor.get_or_insert_with(|| {
58            zstd::bulk::Compressor::new(level).expect("can use default compression level")
59        })
60    }
61
62    #[cfg(feature = "lz4")]
63    fn lz4_compressor(&mut self) -> &mut Lz4Compressor {
64        self.lz4_compressor.get_or_insert_with(Lz4Compressor::new)
65    }
66}
67
68impl std::fmt::Debug for IpcWriteContext {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        let mut ds = f.debug_struct("IpcWriteContext");
71
72        #[cfg(feature = "zstd")]
73        ds.field(
74            "compressor",
75            &self.compressor.as_ref().map(|_| "zstd::bulk::Compressor"),
76        );
77
78        #[cfg(feature = "lz4")]
79        ds.field(
80            "lz4_compressor",
81            &self.lz4_compressor.as_ref().map(|_| "Lz4Compressor"),
82        );
83
84        ds.finish()
85    }
86}
87
88/// Deprecated alias for [`IpcWriteContext`].
89#[deprecated(since = "59.1.0", note = "Use IpcWriteContext instead")]
90pub type CompressionContext = IpcWriteContext;
91
92/// Additional context that may be needed for decompression.
93///
94/// In the case of zstd, this will contain the zstd decompression context, which can be reused
95/// between subsequent decompression calls to avoid the performance overhead of initialising a new
96/// context for every decompression.
97pub struct DecompressionContext {
98    #[cfg(feature = "zstd")]
99    decompressor: Option<zstd::bulk::Decompressor<'static>>,
100}
101
102impl DecompressionContext {
103    pub(crate) fn new() -> Self {
104        Default::default()
105    }
106
107    #[cfg(feature = "zstd")]
108    fn zstd_decompressor(&mut self) -> &mut zstd::bulk::Decompressor<'static> {
109        self.decompressor.get_or_insert_with(|| {
110            zstd::bulk::Decompressor::new().expect("can create zstd decompressor")
111        })
112    }
113}
114
115#[expect(clippy::derivable_impls)]
116impl Default for DecompressionContext {
117    fn default() -> Self {
118        DecompressionContext {
119            #[cfg(feature = "zstd")]
120            decompressor: None,
121        }
122    }
123}
124
125impl std::fmt::Debug for DecompressionContext {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        let mut ds = f.debug_struct("DecompressionContext");
128
129        #[cfg(feature = "zstd")]
130        ds.field(
131            "decompressor",
132            &self
133                .decompressor
134                .as_ref()
135                .map(|_| "zstd::bulk::Decompressor"),
136        );
137
138        ds.finish()
139    }
140}
141
142/// Represents compressing a ipc stream using a particular compression algorithm
143#[derive(Debug, Clone, Copy, PartialEq, Eq)]
144pub enum CompressionCodec {
145    Lz4Frame,
146    Zstd(i32),
147}
148
149impl TryFrom<CompressionType> for CompressionCodec {
150    type Error = ArrowError;
151
152    fn try_from(compression_type: CompressionType) -> Result<Self, ArrowError> {
153        match compression_type {
154            CompressionType::ZSTD => Ok(CompressionCodec::Zstd(DEFAULT_ZSTD_COMPRESSION_LEVEL)),
155            CompressionType::LZ4_FRAME => Ok(CompressionCodec::Lz4Frame),
156            other_type => Err(ArrowError::NotYetImplemented(format!(
157                "compression type {other_type:?} not supported "
158            ))),
159        }
160    }
161}
162
163impl CompressionCodec {
164    /// Creates a [`CompressionCodec`] with an explicit compression level.
165    ///
166    /// The level is used for [`CompressionType::ZSTD`].
167    /// [`CompressionType::LZ4_FRAME`] does not yet support compression levels
168    /// and ignores this value. Returns an error for unsupported compression
169    /// types.
170    pub(crate) fn try_new_with_compression_level(
171        compression_type: CompressionType,
172        compression_level: i32,
173    ) -> Result<Self, ArrowError> {
174        match compression_type {
175            CompressionType::ZSTD => Ok(CompressionCodec::Zstd(compression_level)),
176            CompressionType::LZ4_FRAME => Ok(CompressionCodec::Lz4Frame),
177            other_type => Err(ArrowError::NotYetImplemented(format!(
178                "compression type {other_type:?} not supported "
179            ))),
180        }
181    }
182
183    /// Compresses the data in `input` to `output` and appends the
184    /// data using the specified compression mechanism.
185    ///
186    /// returns the number of bytes written to the stream
187    ///
188    /// Writes this format to output:
189    /// ```text
190    /// [8 bytes]:         uncompressed length
191    /// [remaining bytes]: compressed data stream
192    /// ```
193    pub(crate) fn compress_to_vec(
194        &self,
195        input: &[u8],
196        output: &mut Vec<u8>,
197        context: &mut IpcWriteContext,
198    ) -> Result<usize, ArrowError> {
199        let uncompressed_data_len = input.len();
200        let original_output_len = output.len();
201
202        if input.is_empty() {
203            // empty input, nothing to do
204        } else {
205            // write compressed data directly into the output buffer
206            output.extend_from_slice(&(uncompressed_data_len as i64).to_le_bytes());
207            self.compress(input, output, context)?;
208
209            let compression_len = output.len() - original_output_len;
210            if compression_len > uncompressed_data_len {
211                // length of compressed data was larger than
212                // uncompressed data, use the uncompressed data with
213                // length -1 to indicate that we don't compress the
214                // data
215                output.truncate(original_output_len);
216                output.extend_from_slice(&LENGTH_NO_COMPRESSED_DATA.to_le_bytes());
217                output.extend_from_slice(input);
218            }
219        }
220        Ok(output.len() - original_output_len)
221    }
222
223    /// Decompresses the input into a [`Buffer`]
224    ///
225    /// The input should look like:
226    /// ```text
227    /// [8 bytes]:         uncompressed length
228    /// [remaining bytes]: compressed data stream
229    /// ```
230    pub(crate) fn decompress_to_buffer(
231        &self,
232        input: &Buffer,
233        context: &mut DecompressionContext,
234    ) -> Result<Buffer, ArrowError> {
235        // read the first 8 bytes to determine if the data is
236        // compressed
237        let decompressed_length = read_uncompressed_size(input)?;
238        let buffer = if decompressed_length == 0 {
239            // empty
240            Buffer::from([])
241        } else if decompressed_length == LENGTH_NO_COMPRESSED_DATA {
242            // no compression
243            input.slice(LENGTH_OF_PREFIX_DATA as usize)
244        } else if let Ok(decompressed_length) = usize::try_from(decompressed_length) {
245            // decompress data using the codec
246            let input_data = &input[(LENGTH_OF_PREFIX_DATA as usize)..];
247            let v = self.decompress(input_data, decompressed_length, context)?;
248            Buffer::from_vec(v)
249        } else {
250            return Err(ArrowError::IpcError(format!(
251                "Invalid uncompressed length: {decompressed_length}"
252            )));
253        };
254        Ok(buffer)
255    }
256
257    /// Compress the data in input buffer and write to output buffer
258    /// using the specified compression
259    fn compress(
260        &self,
261        input: &[u8],
262        output: &mut Vec<u8>,
263        context: &mut IpcWriteContext,
264    ) -> Result<(), ArrowError> {
265        match self {
266            CompressionCodec::Lz4Frame => compress_lz4(input, output, context),
267            CompressionCodec::Zstd(level) => compress_zstd(input, output, context, *level),
268        }
269    }
270
271    /// Decompress the data in input buffer and write to output buffer
272    /// using the specified compression
273    fn decompress(
274        &self,
275        input: &[u8],
276        decompressed_size: usize,
277        context: &mut DecompressionContext,
278    ) -> Result<Vec<u8>, ArrowError> {
279        let ret = match self {
280            CompressionCodec::Lz4Frame => decompress_lz4(input, decompressed_size)?,
281            CompressionCodec::Zstd(_) => decompress_zstd(input, decompressed_size, context)?,
282        };
283        if ret.len() != decompressed_size {
284            return Err(ArrowError::IpcError(format!(
285                "Expected decompressed length of {decompressed_size} got {}",
286                ret.len()
287            )));
288        }
289        Ok(ret)
290    }
291}
292
293/// Reusable lz4 frame compressor, the lz4 counterpart of `zstd::bulk::Compressor`.
294///
295/// `lz4_flex` has no bulk frame API, so this wraps a [`lz4_flex::frame::FrameEncoder`] whose hash
296/// table and block buffers are retained across calls to [`Self::compress_to_buffer`].
297#[cfg(feature = "lz4")]
298struct Lz4Compressor {
299    encoder: lz4_flex::frame::FrameEncoder<Vec<u8>>,
300}
301
302#[cfg(feature = "lz4")]
303impl Lz4Compressor {
304    fn new() -> Self {
305        // Pin the block size: with `BlockSize::Auto` the first buffer written would fix it for
306        // every subsequent frame of the reused encoder. 256KB keeps large buffers as fast as
307        // with bigger blocks while the encoder retains only ~0.5MB of buffers.
308        let frame_info =
309            lz4_flex::frame::FrameInfo::new().block_size(lz4_flex::frame::BlockSize::Max256KB);
310        Self {
311            encoder: lz4_flex::frame::FrameEncoder::with_frame_info(frame_info, Vec::new()),
312        }
313    }
314
315    /// Compresses `input` as a single lz4 frame appended to `output`.
316    fn compress_to_buffer(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), ArrowError> {
317        use std::io::Write;
318        // The encoder skips the frame header on empty input once it has written a frame.
319        debug_assert!(!input.is_empty());
320        // Lend `output` to the encoder so that the frame is written straight into it.
321        std::mem::swap(self.encoder.get_mut(), output);
322        let result = self
323            .encoder
324            .write_all(input)
325            .map_err(ArrowError::from)
326            .and_then(|()| {
327                self.encoder
328                    .try_finish()
329                    .map_err(|e| ArrowError::ExternalError(Box::new(e)))
330            });
331        std::mem::swap(self.encoder.get_mut(), output);
332        if result.is_err() {
333            // The encoder may be left mid-frame, start from a fresh one next time.
334            *self = Self::new();
335        }
336        result
337    }
338}
339
340#[cfg(feature = "lz4")]
341fn compress_lz4(
342    input: &[u8],
343    output: &mut Vec<u8>,
344    context: &mut IpcWriteContext,
345) -> Result<(), ArrowError> {
346    context.lz4_compressor().compress_to_buffer(input, output)
347}
348
349#[cfg(not(feature = "lz4"))]
350fn compress_lz4(
351    _input: &[u8],
352    _output: &mut Vec<u8>,
353    _context: &mut IpcWriteContext,
354) -> Result<(), ArrowError> {
355    Err(ArrowError::InvalidArgumentError(
356        "lz4 IPC compression requires the lz4 feature".to_string(),
357    ))
358}
359
360#[cfg(feature = "lz4")]
361fn decompress_lz4(input: &[u8], decompressed_size: usize) -> Result<Vec<u8>, ArrowError> {
362    use std::io::Read;
363    let mut output = Vec::with_capacity(decompressed_size.min(crate::reader::MAX_PREALLOC_BYTES));
364    let mut decoder = lz4_flex::frame::FrameDecoder::new(input);
365    decoder
366        .by_ref()
367        .take(decompressed_size as u64)
368        .read_to_end(&mut output)?;
369
370    // Probe without growing `output` to reject data exceeding the advertised size.
371    if decoder.read(&mut [0])? != 0 {
372        return Err(ArrowError::IpcError(format!(
373            "LZ4 decompressed buffer exceeds advertised size of {decompressed_size}"
374        )));
375    }
376    Ok(output)
377}
378
379#[cfg(not(feature = "lz4"))]
380fn decompress_lz4(_input: &[u8], _decompressed_size: usize) -> Result<Vec<u8>, ArrowError> {
381    Err(ArrowError::InvalidArgumentError(
382        "lz4 IPC decompression requires the lz4 feature".to_string(),
383    ))
384}
385
386#[cfg(feature = "zstd")]
387fn compress_zstd(
388    input: &[u8],
389    output: &mut Vec<u8>,
390    context: &mut IpcWriteContext,
391    level: i32,
392) -> Result<(), ArrowError> {
393    let start = output.len();
394    output.reserve(zstd::zstd_safe::compress_bound(input.len()));
395
396    let mut cursor = std::io::Cursor::new(output);
397    cursor.set_position(start as u64);
398    context
399        .zstd_compressor(level)
400        .compress_to_buffer(input, &mut cursor)?;
401
402    Ok(())
403}
404
405#[cfg(not(feature = "zstd"))]
406fn compress_zstd(
407    _input: &[u8],
408    _output: &mut Vec<u8>,
409    _context: &mut IpcWriteContext,
410    _level: i32,
411) -> Result<(), ArrowError> {
412    Err(ArrowError::InvalidArgumentError(
413        "zstd IPC compression requires the zstd feature".to_string(),
414    ))
415}
416
417#[cfg(feature = "zstd")]
418fn decompress_zstd(
419    input: &[u8],
420    decompressed_size: usize,
421    context: &mut DecompressionContext,
422) -> Result<Vec<u8>, ArrowError> {
423    if decompressed_size > crate::reader::MAX_PREALLOC_BYTES {
424        // The bulk decompressor writes into a buffer reserved up front, so stream instead.
425        use std::io::Read;
426        let mut output =
427            Vec::with_capacity(decompressed_size.min(crate::reader::MAX_PREALLOC_BYTES));
428        let mut decoder = zstd::stream::read::Decoder::with_buffer(input)?;
429        decoder
430            .by_ref()
431            .take(decompressed_size as u64)
432            .read_to_end(&mut output)?;
433
434        // Probe without growing `output` to reject data exceeding the advertised size.
435        if decoder.read(&mut [0])? != 0 {
436            return Err(ArrowError::IpcError(format!(
437                "ZSTD decompressed buffer exceeds advertised size of {decompressed_size}"
438            )));
439        }
440        return Ok(output);
441    }
442    let output = context
443        .zstd_decompressor()
444        .decompress(input, decompressed_size)?;
445    Ok(output)
446}
447
448#[cfg(not(feature = "zstd"))]
449fn decompress_zstd(
450    _input: &[u8],
451    _decompressed_size: usize,
452    _context: &mut DecompressionContext,
453) -> Result<Vec<u8>, ArrowError> {
454    Err(ArrowError::InvalidArgumentError(
455        "zstd IPC decompression requires the zstd feature".to_string(),
456    ))
457}
458
459/// Get the uncompressed length
460/// Notes:
461///   LENGTH_NO_COMPRESSED_DATA: indicate that the data that follows is not compressed
462///    0: indicate that there is no data
463///   positive number: indicate the uncompressed length for the following data
464/// Returns an error if the input buffer is shorter than 8 bytes
465#[inline]
466fn read_uncompressed_size(buffer: &[u8]) -> Result<i64, ArrowError> {
467    let len_buffer = buffer.get(..LENGTH_OF_PREFIX_DATA as usize).ok_or_else(|| {
468        ArrowError::IpcError(format!(
469            "Compressed IPC buffer is too short: expected at least {LENGTH_OF_PREFIX_DATA} bytes, got {}",
470            buffer.len()
471        ))
472    })?;
473    Ok(i64::from_le_bytes(len_buffer.try_into().unwrap()))
474}
475
476#[cfg(test)]
477mod tests {
478    #[test]
479    #[cfg(feature = "lz4")]
480    fn test_lz4_compression() {
481        let input_bytes = b"hello lz4";
482        let codec = super::CompressionCodec::Lz4Frame;
483        let mut output_bytes: Vec<u8> = Vec::new();
484        codec
485            .compress(input_bytes, &mut output_bytes, &mut Default::default())
486            .unwrap();
487        let result = codec
488            .decompress(
489                output_bytes.as_slice(),
490                input_bytes.len(),
491                &mut Default::default(),
492            )
493            .unwrap();
494        assert_eq!(input_bytes, result.as_slice());
495    }
496
497    #[test]
498    #[cfg(feature = "lz4")]
499    fn test_lz4_decompression_rejects_output_exceeding_advertised_size() {
500        let input_bytes = b"hello lz4";
501        let codec = super::CompressionCodec::Lz4Frame;
502        let mut compressed = Vec::new();
503        codec
504            .compress(input_bytes, &mut compressed, &mut Default::default())
505            .unwrap();
506
507        let err = codec
508            .decompress(&compressed, input_bytes.len() - 1, &mut Default::default())
509            .expect_err("output larger than the advertised size should fail");
510
511        assert!(
512            err.to_string().contains("exceeds advertised size"),
513            "unexpected error: {err}"
514        );
515    }
516
517    #[test]
518    #[cfg(feature = "lz4")]
519    fn test_lz4_compression_multiple_frames() {
520        let codec = super::CompressionCodec::Lz4Frame;
521        let mut context = super::IpcWriteContext::default();
522        // Spans several 256KB blocks, followed by a smaller frame.
523        let large: Vec<u8> = (0..200_000u32)
524            .flat_map(|i| (i % 1000).to_le_bytes())
525            .collect();
526        let small = b"hello lz4".to_vec();
527        for input in [&large, &small, &large] {
528            let mut output = b"prefix".to_vec();
529            codec.compress(input, &mut output, &mut context).unwrap();
530            assert_eq!(&output[..6], b"prefix");
531            let result = codec
532                .decompress(&output[6..], input.len(), &mut Default::default())
533                .unwrap();
534            assert_eq!(input, &result);
535        }
536    }
537
538    #[test]
539    #[cfg(feature = "zstd")]
540    fn test_zstd_compression() {
541        let input_bytes = b"hello zstd";
542        let codec = super::CompressionCodec::Zstd(super::DEFAULT_ZSTD_COMPRESSION_LEVEL);
543        let mut output_bytes: Vec<u8> = Vec::new();
544        codec
545            .compress(input_bytes, &mut output_bytes, &mut Default::default())
546            .unwrap();
547        let result = codec
548            .decompress(
549                output_bytes.as_slice(),
550                input_bytes.len(),
551                &mut Default::default(),
552            )
553            .unwrap();
554        assert_eq!(input_bytes, result.as_slice());
555    }
556
557    /// A buffer larger than `MAX_PREALLOC_BYTES` round-trips, and a size its data does not back fails.
558    #[test]
559    #[cfg(any(feature = "lz4", feature = "zstd"))]
560    fn test_decompression_beyond_prealloc() {
561        let input = vec![7u8; crate::reader::MAX_PREALLOC_BYTES + 2];
562        for codec in [
563            #[cfg(feature = "lz4")]
564            super::CompressionCodec::Lz4Frame,
565            #[cfg(feature = "zstd")]
566            super::CompressionCodec::Zstd(super::DEFAULT_ZSTD_COMPRESSION_LEVEL),
567        ] {
568            let mut compressed = Vec::new();
569            codec
570                .compress(&input, &mut compressed, &mut Default::default())
571                .unwrap();
572            let decompress = |len| codec.decompress(&compressed, len, &mut Default::default());
573            assert_eq!(decompress(input.len()).unwrap(), input, "{codec:?}");
574            assert!(decompress(input.len() - 1).is_err(), "{codec:?}");
575            assert!(decompress(usize::MAX).is_err(), "{codec:?}");
576        }
577    }
578
579    #[test]
580    fn test_read_uncompressed_size_rejects_short_prefix() {
581        let err = super::read_uncompressed_size(&[1, 2, 3, 4, 5, 6, 7])
582            .expect_err("short compressed IPC prefix should return an error");
583
584        assert!(
585            err.to_string()
586                .contains("Compressed IPC buffer is too short"),
587            "unexpected error: {err}"
588        );
589    }
590
591    #[test]
592    #[cfg(feature = "lz4")]
593    fn test_compress_to_vec_writes_8_byte_length_prefix() {
594        // The length prefix must always be 8 bytes (i64),
595        // even on platforms where `usize` is narrower (e.g. wasm32).
596        let input_bytes = vec![42u8; 132];
597        let codec = super::CompressionCodec::Lz4Frame;
598        let mut output_bytes: Vec<u8> = Vec::new();
599        codec
600            .compress_to_vec(&input_bytes, &mut output_bytes, &mut Default::default())
601            .unwrap();
602
603        let prefix: [u8; 8] = output_bytes[..8].try_into().unwrap();
604        assert_eq!(i64::from_le_bytes(prefix), input_bytes.len() as i64);
605    }
606}
arrow_ipc/compression.rs

arrow_ipc/
compression.rs
