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    /// A reader over bytes that are known to be uncompressed, e.g. already decompressed. Unlike
389    /// [`Self::from_memory`], it does not look for compression magic bytes, which text may start
390    /// with.
391    pub fn from_uncompressed_memory(slice: Buffer<u8>) -> Self {
392        Self::UncompressedMemory { slice, offset: 0 }
393    }
394}
395
396#[cfg(feature = "decompress")]
397pub use compressed_writer::CompressedWriter;
398
399#[cfg(feature = "decompress")]
400mod compressed_writer {
401    use std::io;
402
403    /// Constructor for `WritableTrait` compressed encoders.
404    pub enum CompressedWriter<'a, W: io::Write> {
405        Gzip(Option<flate2::write::GzEncoder<GzipTarget<'a, W>>>),
406        Zstd(Option<zstd::Encoder<'static, &'a mut W>>),
407    }
408
409    /// Discards writes once `discard` is set.
410    pub struct GzipTarget<'a, W> {
411        inner: &'a mut W,
412        discard: bool,
413    }
414
415    impl<W: io::Write> io::Write for GzipTarget<'_, W> {
416        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
417            if self.discard {
418                Ok(buf.len())
419            } else {
420                self.inner.write(buf)
421            }
422        }
423
424        fn flush(&mut self) -> io::Result<()> {
425            if self.discard {
426                Ok(())
427            } else {
428                self.inner.flush()
429            }
430        }
431    }
432
433    impl<'a, W: io::Write> CompressedWriter<'a, W> {
434        pub fn gzip(writer: &'a mut W, level: Option<u32>) -> Self {
435            Self::Gzip(Some(flate2::write::GzEncoder::new(
436                GzipTarget {
437                    inner: writer,
438                    discard: false,
439                },
440                level.map(flate2::Compression::new).unwrap_or_default(),
441            )))
442        }
443
444        pub fn zstd(writer: &'a mut W, level: Option<u32>) -> io::Result<Self> {
445            zstd::Encoder::new(writer, level.unwrap_or(3) as i32)
446                .map(Some)
447                .map(Self::Zstd)
448        }
449    }
450
451    impl<'a, W: io::Write> io::Write for CompressedWriter<'a, W> {
452        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
453            match self {
454                Self::Gzip(encoder) => encoder.as_mut().unwrap().write(buf),
455                Self::Zstd(encoder) => encoder.as_mut().unwrap().write(buf),
456            }
457        }
458
459        fn flush(&mut self) -> io::Result<()> {
460            match self {
461                Self::Gzip(encoder) => encoder.as_mut().unwrap().flush(),
462                Self::Zstd(encoder) => encoder.as_mut().unwrap().flush(),
463            }
464        }
465    }
466
467    impl<'a, W: io::Write> CompressedWriter<'a, W> {
468        pub fn finish(&mut self) -> io::Result<&'a mut W> {
469            match self {
470                Self::Gzip(encoder) => {
471                    // On failure the encoder stays, so that its drop can't write the trailer.
472                    encoder.as_mut().unwrap().try_finish()?;
473                    encoder.take().unwrap().finish().map(|t| t.inner)
474                },
475                Self::Zstd(encoder) => encoder.take().unwrap().finish(),
476            }
477        }
478    }
479
480    impl<W: io::Write> Drop for CompressedWriter<'_, W> {
481        fn drop(&mut self) {
482            // `GzEncoder` writes its trailer on drop, which would make an unfinished file look complete.
483            if let Self::Gzip(Some(encoder)) = self {
484                encoder.get_mut().discard = true;
485            }
486        }
487    }
488
489    #[cfg(test)]
490    mod tests {
491        use std::io::{self, Read, Write};
492
493        use super::CompressedWriter;
494
495        fn gzip_written(finish: bool) -> Vec<u8> {
496            let mut out = Vec::new();
497            let mut writer = CompressedWriter::gzip(&mut out, None);
498            writer.write_all(b"x\n1\n").unwrap();
499            writer.flush().unwrap();
500            if finish {
501                writer.finish().unwrap();
502            }
503            drop(writer);
504            out
505        }
506
507        #[test]
508        fn dropped_gzip_writer_leaves_no_trailer() {
509            let out = gzip_written(false);
510            assert!(!out.is_empty());
511            let err = flate2::read::GzDecoder::new(out.as_slice())
512                .read_to_end(&mut Vec::new())
513                .unwrap_err();
514            assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
515        }
516
517        struct FailOnce<'a> {
518            out: Vec<u8>,
519            fail: &'a std::cell::Cell<bool>,
520        }
521
522        impl Write for FailOnce<'_> {
523            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
524                if self.fail.replace(false) {
525                    return Err(io::Error::other("write failed"));
526                }
527                self.out.extend_from_slice(buf);
528                Ok(buf.len())
529            }
530
531            fn flush(&mut self) -> io::Result<()> {
532                Ok(())
533            }
534        }
535
536        #[test]
537        fn failed_gzip_finish_leaves_no_trailer() {
538            let fail = std::cell::Cell::new(false);
539            let mut target = FailOnce {
540                out: Vec::new(),
541                fail: &fail,
542            };
543            let mut writer = CompressedWriter::gzip(&mut target, None);
544            writer.write_all(b"x\n1\n").unwrap();
545            writer.flush().unwrap();
546            fail.set(true);
547            assert!(writer.finish().is_err());
548            drop(writer);
549            let err = flate2::read::GzDecoder::new(target.out.as_slice())
550                .read_to_end(&mut Vec::new())
551                .unwrap_err();
552            assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof);
553        }
554
555        #[test]
556        fn finished_gzip_writer_is_complete() {
557            let out = gzip_written(true);
558            let mut decoded = Vec::new();
559            flate2::read::GzDecoder::new(out.as_slice())
560                .read_to_end(&mut decoded)
561                .unwrap();
562            assert_eq!(decoded, b"x\n1\n");
563        }
564    }
565}
polars_io/utils/compression.rs

polars_io/utils/
compression.rs
