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

arrow_ipc/
compression.rs
