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#[expect(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 as i64).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, 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 decompressed 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"))]
303fn compress_lz4(_input: &[u8], _output: &mut Vec<u8>) -> Result<(), ArrowError> {
304    Err(ArrowError::InvalidArgumentError(
305        "lz4 IPC compression requires the lz4 feature".to_string(),
306    ))
307}
308
309#[cfg(feature = "lz4")]
310fn decompress_lz4(input: &[u8], decompressed_size: usize) -> Result<Vec<u8>, ArrowError> {
311    use std::io::Read;
312    let mut output = Vec::with_capacity(decompressed_size);
313    let mut decoder = lz4_flex::frame::FrameDecoder::new(input);
314    decoder
315        .by_ref()
316        .take(decompressed_size as u64)
317        .read_to_end(&mut output)?;
318
319    // Probe without growing `output` to reject data exceeding the advertised size.
320    if decoder.read(&mut [0])? != 0 {
321        return Err(ArrowError::IpcError(format!(
322            "LZ4 decompressed buffer exceeds advertised size of {decompressed_size}"
323        )));
324    }
325    Ok(output)
326}
327
328#[cfg(not(feature = "lz4"))]
329fn decompress_lz4(_input: &[u8], _decompressed_size: usize) -> Result<Vec<u8>, ArrowError> {
330    Err(ArrowError::InvalidArgumentError(
331        "lz4 IPC decompression requires the lz4 feature".to_string(),
332    ))
333}
334
335#[cfg(feature = "zstd")]
336fn compress_zstd(
337    input: &[u8],
338    output: &mut Vec<u8>,
339    context: &mut IpcWriteContext,
340    level: i32,
341) -> Result<(), ArrowError> {
342    let start = output.len();
343    output.reserve(zstd::zstd_safe::compress_bound(input.len()));
344
345    let mut cursor = std::io::Cursor::new(output);
346    cursor.set_position(start as u64);
347    context
348        .zstd_compressor(level)
349        .compress_to_buffer(input, &mut cursor)?;
350
351    Ok(())
352}
353
354#[cfg(not(feature = "zstd"))]
355fn compress_zstd(
356    _input: &[u8],
357    _output: &mut Vec<u8>,
358    _context: &mut IpcWriteContext,
359    _level: i32,
360) -> Result<(), ArrowError> {
361    Err(ArrowError::InvalidArgumentError(
362        "zstd IPC compression requires the zstd feature".to_string(),
363    ))
364}
365
366#[cfg(feature = "zstd")]
367fn decompress_zstd(
368    input: &[u8],
369    decompressed_size: usize,
370    context: &mut DecompressionContext,
371) -> Result<Vec<u8>, ArrowError> {
372    let output = context
373        .zstd_decompressor()
374        .decompress(input, decompressed_size)?;
375    Ok(output)
376}
377
378#[cfg(not(feature = "zstd"))]
379fn decompress_zstd(
380    _input: &[u8],
381    _decompressed_size: usize,
382    _context: &mut DecompressionContext,
383) -> Result<Vec<u8>, ArrowError> {
384    Err(ArrowError::InvalidArgumentError(
385        "zstd IPC decompression requires the zstd feature".to_string(),
386    ))
387}
388
389/// Get the uncompressed length
390/// Notes:
391///   LENGTH_NO_COMPRESSED_DATA: indicate that the data that follows is not compressed
392///    0: indicate that there is no data
393///   positive number: indicate the uncompressed length for the following data
394/// Returns an error if the input buffer is shorter than 8 bytes
395#[inline]
396fn read_uncompressed_size(buffer: &[u8]) -> Result<i64, ArrowError> {
397    let len_buffer = buffer.get(..LENGTH_OF_PREFIX_DATA as usize).ok_or_else(|| {
398        ArrowError::IpcError(format!(
399            "Compressed IPC buffer is too short: expected at least {LENGTH_OF_PREFIX_DATA} bytes, got {}",
400            buffer.len()
401        ))
402    })?;
403    Ok(i64::from_le_bytes(len_buffer.try_into().unwrap()))
404}
405
406#[cfg(test)]
407mod tests {
408    #[test]
409    #[cfg(feature = "lz4")]
410    fn test_lz4_compression() {
411        let input_bytes = b"hello lz4";
412        let codec = super::CompressionCodec::Lz4Frame;
413        let mut output_bytes: Vec<u8> = Vec::new();
414        codec
415            .compress(input_bytes, &mut output_bytes, &mut Default::default())
416            .unwrap();
417        let result = codec
418            .decompress(
419                output_bytes.as_slice(),
420                input_bytes.len(),
421                &mut Default::default(),
422            )
423            .unwrap();
424        assert_eq!(input_bytes, result.as_slice());
425    }
426
427    #[test]
428    #[cfg(feature = "lz4")]
429    fn test_lz4_decompression_rejects_output_exceeding_advertised_size() {
430        let input_bytes = b"hello lz4";
431        let codec = super::CompressionCodec::Lz4Frame;
432        let mut compressed = Vec::new();
433        codec
434            .compress(input_bytes, &mut compressed, &mut Default::default())
435            .unwrap();
436
437        let err = codec
438            .decompress(&compressed, input_bytes.len() - 1, &mut Default::default())
439            .expect_err("output larger than the advertised size should fail");
440
441        assert!(
442            err.to_string().contains("exceeds advertised size"),
443            "unexpected error: {err}"
444        );
445    }
446
447    #[test]
448    #[cfg(feature = "zstd")]
449    fn test_zstd_compression() {
450        let input_bytes = b"hello zstd";
451        let codec = super::CompressionCodec::Zstd(super::DEFAULT_ZSTD_COMPRESSION_LEVEL);
452        let mut output_bytes: Vec<u8> = Vec::new();
453        codec
454            .compress(input_bytes, &mut output_bytes, &mut Default::default())
455            .unwrap();
456        let result = codec
457            .decompress(
458                output_bytes.as_slice(),
459                input_bytes.len(),
460                &mut Default::default(),
461            )
462            .unwrap();
463        assert_eq!(input_bytes, result.as_slice());
464    }
465
466    #[test]
467    fn test_read_uncompressed_size_rejects_short_prefix() {
468        let err = super::read_uncompressed_size(&[1, 2, 3, 4, 5, 6, 7])
469            .expect_err("short compressed IPC prefix should return an error");
470
471        assert!(
472            err.to_string()
473                .contains("Compressed IPC buffer is too short"),
474            "unexpected error: {err}"
475        );
476    }
477
478    #[test]
479    #[cfg(feature = "lz4")]
480    fn test_compress_to_vec_writes_8_byte_length_prefix() {
481        // The length prefix must always be 8 bytes (i64),
482        // even on platforms where `usize` is narrower (e.g. wasm32).
483        let input_bytes = vec![42u8; 132];
484        let codec = super::CompressionCodec::Lz4Frame;
485        let mut output_bytes: Vec<u8> = Vec::new();
486        codec
487            .compress_to_vec(&input_bytes, &mut output_bytes, &mut Default::default())
488            .unwrap();
489
490        let prefix: [u8; 8] = output_bytes[..8].try_into().unwrap();
491        assert_eq!(i64::from_le_bytes(prefix), input_bytes.len() as i64);
492    }
493}
arrow_ipc/compression.rs

arrow_ipc/
compression.rs
