polars_io/utils/
compression.rs

1use std::cmp;
2use std::io::{BufRead, Cursor, Read};
3
4use polars_buffer::Buffer;
5use polars_core::prelude::*;
6use polars_error::{feature_gated, to_compute_err};
7
8use crate::utils::stream_buf_reader::ReaderSource;
9
10/// Represents the compression algorithms that we have decoders for
11#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
12pub enum SupportedCompression {
13    GZIP,
14    ZLIB,
15    ZSTD,
16}
17
18impl SupportedCompression {
19    /// If the given byte slice starts with the "magic" bytes for a supported compression family, return
20    /// that family, for unsupported/uncompressed slices, return None.
21    /// Based on <https://en.wikipedia.org/wiki/List_of_file_signatures>.
22    pub fn check(bytes: &[u8]) -> Option<Self> {
23        if bytes.len() < 4 {
24            // not enough bytes to perform prefix checks
25            return None;
26        }
27        match bytes[..4] {
28            [0x1f, 0x8b, _, _] => Some(Self::GZIP),
29            // Different zlib compression levels without preset dictionary.
30            [0x78, 0x01, _, _] => Some(Self::ZLIB),
31            [0x78, 0x5e, _, _] => Some(Self::ZLIB),
32            [0x78, 0x9c, _, _] => Some(Self::ZLIB),
33            [0x78, 0xda, _, _] => Some(Self::ZLIB),
34            [0x28, 0xb5, 0x2f, 0xfd] => Some(Self::ZSTD),
35            _ => None,
36        }
37    }
38}
39
40/// Decompress `bytes` if compression is detected, otherwise simply return it.
41/// An `out` vec must be given for ownership of the decompressed data.
42#[allow(clippy::ptr_arg)]
43#[deprecated(note = "may cause OOM, use CompressedReader instead")]
44pub fn maybe_decompress_bytes<'a>(bytes: &'a [u8], out: &'a mut Vec<u8>) -> PolarsResult<&'a [u8]> {
45    assert!(out.is_empty());
46
47    let Some(algo) = SupportedCompression::check(bytes) else {
48        return Ok(bytes);
49    };
50
51    feature_gated!("decompress", {
52        match algo {
53            SupportedCompression::GZIP => {
54                flate2::read::MultiGzDecoder::new(bytes)
55                    .read_to_end(out)
56                    .map_err(to_compute_err)?;
57            },
58            SupportedCompression::ZLIB => {
59                flate2::read::ZlibDecoder::new(bytes)
60                    .read_to_end(out)
61                    .map_err(to_compute_err)?;
62            },
63            SupportedCompression::ZSTD => {
64                zstd::Decoder::with_buffer(bytes)?.read_to_end(out)?;
65            },
66        }
67
68        Ok(out)
69    })
70}
71
72/// Reader that implements a streaming read trait for uncompressed, gzip, zlib and zstd
73/// compression.
74///
75/// This allows handling decompression transparently in a streaming fashion.
76pub enum CompressedReader {
77    Uncompressed {
78        slice: Buffer<u8>,
79        offset: usize,
80    },
81    #[cfg(feature = "decompress")]
82    Gzip(flate2::bufread::MultiGzDecoder<Cursor<Buffer<u8>>>),
83    #[cfg(feature = "decompress")]
84    Zlib(flate2::bufread::ZlibDecoder<Cursor<Buffer<u8>>>),
85    #[cfg(feature = "decompress")]
86    Zstd(zstd::Decoder<'static, Cursor<Buffer<u8>>>),
87}
88
89impl CompressedReader {
90    pub fn try_new(slice: Buffer<u8>) -> PolarsResult<Self> {
91        let algo = SupportedCompression::check(&slice);
92
93        Ok(match algo {
94            None => CompressedReader::Uncompressed { slice, offset: 0 },
95            #[cfg(feature = "decompress")]
96            Some(SupportedCompression::GZIP) => {
97                CompressedReader::Gzip(flate2::bufread::MultiGzDecoder::new(Cursor::new(slice)))
98            },
99            #[cfg(feature = "decompress")]
100            Some(SupportedCompression::ZLIB) => {
101                CompressedReader::Zlib(flate2::bufread::ZlibDecoder::new(Cursor::new(slice)))
102            },
103            #[cfg(feature = "decompress")]
104            Some(SupportedCompression::ZSTD) => {
105                CompressedReader::Zstd(zstd::Decoder::with_buffer(Cursor::new(slice))?)
106            },
107            #[cfg(not(feature = "decompress"))]
108            _ => panic!("activate 'decompress' feature"),
109        })
110    }
111
112    pub fn is_compressed(&self) -> bool {
113        !matches!(&self, CompressedReader::Uncompressed { .. })
114    }
115
116    pub const fn initial_read_size() -> usize {
117        // We don't want to read too much at the beginning to keep decompression to a minimum if for
118        // example only the schema is needed or a slice op is used. Keep in sync with
119        // `ideal_read_size` so that `initial_read_size * N * 4 == ideal_read_size`.
120        32 * 1024
121    }
122
123    pub const fn ideal_read_size() -> usize {
124        // Somewhat conservative guess for L2 size, which performs the best on most machines and is
125        // nearly always core exclusive. The loss of going larger and accidentally hitting L3 is not
126        // recouped by amortizing the block processing cost even further.
127        //
128        // It's possible that callers use or need a larger `read_size` if for example a single row
129        // doesn't fit in the 512KB.
130        512 * 1024
131    }
132
133    /// If possible returns the total number of bytes that will be produced by reading from the
134    /// start to finish.
135    pub fn total_len_estimate(&self) -> usize {
136        const ESTIMATED_DEFLATE_RATIO: usize = 3;
137        const ESTIMATED_ZSTD_RATIO: usize = 5;
138
139        match self {
140            CompressedReader::Uncompressed { slice, .. } => slice.len(),
141            #[cfg(feature = "decompress")]
142            CompressedReader::Gzip(reader) => {
143                reader.get_ref().get_ref().len() * ESTIMATED_DEFLATE_RATIO
144            },
145            #[cfg(feature = "decompress")]
146            CompressedReader::Zlib(reader) => {
147                reader.get_ref().get_ref().len() * ESTIMATED_DEFLATE_RATIO
148            },
149            #[cfg(feature = "decompress")]
150            CompressedReader::Zstd(reader) => {
151                reader.get_ref().get_ref().len() * ESTIMATED_ZSTD_RATIO
152            },
153        }
154    }
155
156    /// Reads exactly `read_size` bytes if possible from the internal readers and creates a new
157    /// [`Buffer`] with the content `concat(prev_leftover, new_bytes)`.
158    ///
159    /// Returns the new slice and the number of bytes read, which will be 0 when eof is reached and
160    /// this function is called again.
161    ///
162    /// If the underlying reader is uncompressed the operation is a cheap zero-copy
163    /// [`Buffer::sliced`] operation.
164    ///
165    /// By handling slice concatenation at this level we can implement zero-copy reading *and* make
166    /// the interface easier to use.
167    ///
168    /// It's a logic bug if `prev_leftover` is neither empty nor the last slice returned by this
169    /// function.
170    pub fn read_next_slice(
171        &mut self,
172        prev_leftover: &Buffer<u8>,
173        read_size: usize,
174    ) -> std::io::Result<(Buffer<u8>, usize)> {
175        // Assuming that callers of this function correctly handle re-trying, by continuously growing
176        // prev_leftover if it doesn't contain a single row, this abstraction supports arbitrarily
177        // sized rows.
178        let prev_len = prev_leftover.len();
179
180        let mut buf = Vec::new();
181        if self.is_compressed() {
182            let reserve_size = cmp::min(
183                prev_len.saturating_add(read_size),
184                self.total_len_estimate().saturating_mul(2),
185            );
186            buf.reserve_exact(reserve_size);
187            buf.extend_from_slice(prev_leftover);
188        }
189
190        let new_slice_from_read =
191            |bytes_read: usize, mut buf: Vec<u8>| -> std::io::Result<(Buffer<u8>, usize)> {
192                buf.truncate(prev_len + bytes_read);
193                Ok((Buffer::from_vec(buf), bytes_read))
194            };
195
196        match self {
197            CompressedReader::Uncompressed { slice, offset, .. } => {
198                let bytes_read = cmp::min(read_size, slice.len() - *offset);
199                let new_slice = slice
200                    .clone()
201                    .sliced(*offset - prev_len..*offset + bytes_read);
202                *offset += bytes_read;
203                Ok((new_slice, bytes_read))
204            },
205            #[cfg(feature = "decompress")]
206            CompressedReader::Gzip(decoder) => {
207                new_slice_from_read(decoder.take(read_size as u64).read_to_end(&mut buf)?, buf)
208            },
209            #[cfg(feature = "decompress")]
210            CompressedReader::Zlib(decoder) => {
211                new_slice_from_read(decoder.take(read_size as u64).read_to_end(&mut buf)?, buf)
212            },
213            #[cfg(feature = "decompress")]
214            CompressedReader::Zstd(decoder) => {
215                new_slice_from_read(decoder.take(read_size as u64).read_to_end(&mut buf)?, buf)
216            },
217        }
218    }
219}
220
221/// This implementation is meant for compatibility. Use [`Self::read_next_slice`] for best
222/// performance.
223impl Read for CompressedReader {
224    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
225        match self {
226            CompressedReader::Uncompressed { slice, offset, .. } => {
227                let bytes_read = cmp::min(buf.len(), slice.len() - *offset);
228                buf[..bytes_read].copy_from_slice(&slice[*offset..(*offset + bytes_read)]);
229                *offset += bytes_read;
230                Ok(bytes_read)
231            },
232            #[cfg(feature = "decompress")]
233            CompressedReader::Gzip(decoder) => decoder.read(buf),
234            #[cfg(feature = "decompress")]
235            CompressedReader::Zlib(decoder) => decoder.read(buf),
236            #[cfg(feature = "decompress")]
237            CompressedReader::Zstd(decoder) => decoder.read(buf),
238        }
239    }
240}
241
242/// A byte source that abstracts over in-memory buffers and streaming
243/// readers, with optional transparent decompression and buffering.
244///
245/// Implements `BufRead`, allowing uniform access regardless of whether
246/// the underlying data is an in-memory slice, a raw stream, or a
247/// compressed stream (gzip/zlib/zstd).
248///
249/// This is the generic successor to [`CompressedReader`], which only
250/// supports in-memory (`Buffer<u8>`) sources.
251pub enum ByteSourceReader<R: BufRead> {
252    UncompressedMemory {
253        slice: Buffer<u8>,
254        offset: usize,
255    },
256    UncompressedStream(R),
257    #[cfg(feature = "decompress")]
258    Gzip(flate2::bufread::MultiGzDecoder<R>),
259    #[cfg(feature = "decompress")]
260    Zlib(flate2::bufread::ZlibDecoder<R>),
261    #[cfg(feature = "decompress")]
262    Zstd(zstd::Decoder<'static, R>),
263}
264
265impl<R: BufRead> ByteSourceReader<R> {
266    pub fn try_new(reader: R, compression: Option<SupportedCompression>) -> PolarsResult<Self> {
267        Ok(match compression {
268            None => Self::UncompressedStream(reader),
269            #[cfg(feature = "decompress")]
270            Some(SupportedCompression::GZIP) => {
271                Self::Gzip(flate2::bufread::MultiGzDecoder::new(reader))
272            },
273            #[cfg(feature = "decompress")]
274            Some(SupportedCompression::ZLIB) => {
275                Self::Zlib(flate2::bufread::ZlibDecoder::new(reader))
276            },
277            #[cfg(feature = "decompress")]
278            Some(SupportedCompression::ZSTD) => Self::Zstd(zstd::Decoder::with_buffer(reader)?),
279            #[cfg(not(feature = "decompress"))]
280            _ => panic!("activate 'decompress' feature"),
281        })
282    }
283
284    pub fn is_compressed(&self) -> bool {
285        !matches!(
286            &self,
287            Self::UncompressedMemory { .. } | Self::UncompressedStream(_)
288        )
289    }
290
291    pub fn compression(&self) -> Option<SupportedCompression> {
292        match self {
293            Self::UncompressedMemory { .. } => None,
294            Self::UncompressedStream(_) => None,
295            #[cfg(feature = "decompress")]
296            Self::Gzip(_) => Some(SupportedCompression::GZIP),
297            #[cfg(feature = "decompress")]
298            Self::Zlib(_) => Some(SupportedCompression::ZLIB),
299            #[cfg(feature = "decompress")]
300            Self::Zstd(_) => Some(SupportedCompression::ZSTD),
301        }
302    }
303
304    pub const fn initial_read_size() -> usize {
305        // We don't want to read too much at the beginning to keep decompression to a minimum if for
306        // example only the schema is needed or a slice op is used. Keep in sync with
307        // `ideal_read_size` so that `initial_read_size * N * 4 == ideal_read_size`.
308        32 * 1024
309    }
310
311    pub const fn ideal_read_size() -> usize {
312        // Somewhat conservative guess for L2 size, which performs the best on most machines and is
313        // nearly always core exclusive. The loss of going larger and accidentally hitting L3 is not
314        // recouped by amortizing the block processing cost even further.
315        //
316        // It's possible that callers use or need a larger `read_size` if for example a single row
317        // doesn't fit in the 512KB.
318        512 * 1024
319    }
320
321    /// Reads exactly `read_size` bytes if possible from the internal readers and creates a new
322    /// [`Buffer`] with the content `concat(prev_leftover, new_bytes)`.
323    ///
324    /// Returns the new slice and the number of bytes read, which will be 0 when eof is reached and
325    /// this function is called again.
326    ///
327    /// If the underlying reader is uncompressed the operation is a cheap zero-copy
328    /// [`Buffer::sliced`] operation.
329    ///
330    /// By handling slice concatenation at this level we can implement zero-copy reading *and* make
331    /// the interface easier to use.
332    ///
333    /// It's a logic bug if `prev_leftover` is neither empty nor the last slice returned by this
334    /// function.
335    pub fn read_next_slice(
336        &mut self,
337        prev_leftover: &Buffer<u8>,
338        read_size: usize,
339        uncompressed_size_hint: Option<usize>,
340    ) -> std::io::Result<(Buffer<u8>, usize)> {
341        // Assuming that callers of this function correctly handle re-trying, by continuously growing
342        // prev_leftover if it doesn't contain a single row, this abstraction supports arbitrarily
343        // sized rows.
344        let prev_len = prev_leftover.len();
345
346        let reader: &mut dyn Read = match self {
347            // Zero-copy fast-path — no allocation required
348            Self::UncompressedMemory { slice, offset } => {
349                let bytes_read = cmp::min(read_size, slice.len() - *offset);
350                let new_slice = slice
351                    .clone()
352                    .sliced(*offset - prev_len..*offset + bytes_read);
353                *offset += bytes_read;
354                return Ok((new_slice, bytes_read));
355            },
356            Self::UncompressedStream(reader) => reader,
357            #[cfg(feature = "decompress")]
358            Self::Gzip(reader) => reader,
359            #[cfg(feature = "decompress")]
360            Self::Zlib(reader) => reader,
361            #[cfg(feature = "decompress")]
362            Self::Zstd(reader) => reader,
363        };
364
365        let mut buf = Vec::new();
366
367        // Cap the reserve_size, for the scenario where read_size == usize::MAX
368        let max_reserve_size = uncompressed_size_hint.unwrap_or(4 * 1024 * 1024);
369        let reserve_size = cmp::min(prev_len.saturating_add(read_size), max_reserve_size);
370        buf.reserve_exact(reserve_size);
371        buf.extend_from_slice(prev_leftover);
372
373        let bytes_read = reader.take(read_size as u64).read_to_end(&mut buf)?;
374        buf.truncate(prev_len + bytes_read);
375        Ok((Buffer::from_vec(buf), bytes_read))
376    }
377}
378
379impl ByteSourceReader<ReaderSource> {
380    pub fn from_memory(slice: Buffer<u8>) -> PolarsResult<Self> {
381        let compression = SupportedCompression::check(&slice);
382        match compression {
383            None => Ok(Self::UncompressedMemory { slice, offset: 0 }),
384            _ => Self::try_new(ReaderSource::Memory(Cursor::new(slice)), compression),
385        }
386    }
387}
388
389#[cfg(feature = "decompress")]
390pub use compressed_writer::CompressedWriter;
391
392#[cfg(feature = "decompress")]
393mod compressed_writer {
394    use std::io;
395
396    /// Constructor for `WritableTrait` compressed encoders.
397    pub enum CompressedWriter<'a, W: io::Write> {
398        Gzip(Option<flate2::write::GzEncoder<&'a mut W>>),
399        Zstd(Option<zstd::Encoder<'static, &'a mut W>>),
400    }
401
402    impl<'a, W: io::Write> CompressedWriter<'a, W> {
403        pub fn gzip(writer: &'a mut W, level: Option<u32>) -> Self {
404            Self::Gzip(Some(flate2::write::GzEncoder::new(
405                writer,
406                level.map(flate2::Compression::new).unwrap_or_default(),
407            )))
408        }
409
410        pub fn zstd(writer: &'a mut W, level: Option<u32>) -> io::Result<Self> {
411            zstd::Encoder::new(writer, level.unwrap_or(3) as i32)
412                .map(Some)
413                .map(Self::Zstd)
414        }
415    }
416
417    impl<'a, W: io::Write> io::Write for CompressedWriter<'a, W> {
418        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
419            match self {
420                Self::Gzip(encoder) => encoder.as_mut().unwrap().write(buf),
421                Self::Zstd(encoder) => encoder.as_mut().unwrap().write(buf),
422            }
423        }
424
425        fn flush(&mut self) -> io::Result<()> {
426            match self {
427                Self::Gzip(encoder) => encoder.as_mut().unwrap().flush(),
428                Self::Zstd(encoder) => encoder.as_mut().unwrap().flush(),
429            }
430        }
431    }
432
433    impl<'a, W: io::Write> CompressedWriter<'a, W> {
434        pub fn finish(&mut self) -> io::Result<&'a mut W> {
435            match self {
436                Self::Gzip(encoder) => encoder.take().unwrap().finish(),
437                Self::Zstd(encoder) => encoder.take().unwrap().finish(),
438            }
439        }
440    }
441}
polars_io/utils/compression.rs

polars_io/utils/
compression.rs
