polars_utils/compression/
zstd.rs

1//! zstd compression and decompression using state that is reused per thread.
2//!
3//! Creating a zstd context allocates and initializes its internal state, which
4//! can cost more than compressing or decompressing a small buffer.
5
6use std::cell::RefCell;
7use std::io::Cursor;
8
9use ::zstd::zstd_safe::{self, CCtx, CParameter, DCtx, InBuffer, OutBuffer, ResetDirective};
10use polars_error::{PolarsError, PolarsResult, polars_ensure, polars_err};
11
12/// Upper bound on the size of each zstd context kept per thread between calls.
13/// Contexts for high compression levels or large windows exceed it and are not
14/// kept, but for those the cost of creating a context is small relative to the
15/// work.
16const MAX_CACHED_CONTEXT_BYTES: usize = 2 * 1000 * 1000;
17
18thread_local! {
19    static CCTX: RefCell<Option<CCtx<'static>>> = const { RefCell::new(None) };
20    static DCTX: RefCell<Option<DCtx<'static>>> = const { RefCell::new(None) };
21}
22
23fn zstd_err(code: zstd_safe::ErrorCode) -> PolarsError {
24    polars_err!(ComputeError: "zstd error: {}", zstd_safe::get_error_name(code))
25}
26
27fn with_cctx<R>(
28    level: i32,
29    f: impl FnOnce(&mut CCtx<'static>) -> PolarsResult<R>,
30) -> PolarsResult<R> {
31    let run = |ctx: &mut CCtx<'static>| {
32        ctx.set_parameter(CParameter::CompressionLevel(level))
33            .map_err(zstd_err)?;
34        f(ctx)
35    };
36    CCTX.with(|cell| {
37        let Ok(mut cached) = cell.try_borrow_mut() else {
38            // Re-entrant use.
39            return run(&mut CCtx::create());
40        };
41        let ctx = cached.get_or_insert_with(CCtx::create);
42        let out = run(ctx);
43        if ctx.sizeof() > MAX_CACHED_CONTEXT_BYTES {
44            *cached = None;
45        }
46        out
47    })
48}
49
50fn with_dctx<R>(f: impl FnOnce(&mut DCtx<'static>) -> PolarsResult<R>) -> PolarsResult<R> {
51    DCTX.with(|cell| {
52        let Ok(mut cached) = cell.try_borrow_mut() else {
53            // Re-entrant use.
54            return f(&mut DCtx::create());
55        };
56        let ctx = cached.get_or_insert_with(DCtx::create);
57        let out = f(ctx);
58        if ctx.sizeof() > MAX_CACHED_CONTEXT_BYTES {
59            *cached = None;
60        }
61        out
62    })
63}
64
65/// Compresses `src` into a single zstd frame appended to `dst`.
66pub fn compress(src: &[u8], dst: &mut Vec<u8>, level: i32) -> PolarsResult<()> {
67    with_cctx(level, |ctx| {
68        dst.reserve(zstd_safe::compress_bound(src.len()));
69        let mut out = Cursor::new(&mut *dst);
70        out.set_position(out.get_ref().len() as u64);
71        ctx.compress2(&mut out, src).map_err(zstd_err)?;
72        Ok(())
73    })
74}
75
76/// Decompresses the zstd frames in `src`, appending them to `dst`, and returns
77/// the number of bytes appended. Fails if they do not fit in the spare
78/// capacity of `dst`.
79pub fn decompress(src: &[u8], dst: &mut Vec<u8>) -> PolarsResult<usize> {
80    with_dctx(|ctx| {
81        let mut out = Cursor::new(&mut *dst);
82        out.set_position(out.get_ref().len() as u64);
83        ctx.decompress(&mut out, src).map_err(zstd_err)
84    })
85}
86
87/// Decompresses the zstd frames in `src`, appending the first `len` decompressed
88/// bytes to `dst`. Fails if they decompress to fewer bytes.
89pub fn decompress_prefix(src: &[u8], dst: &mut Vec<u8>, len: usize) -> PolarsResult<()> {
90    let end = dst.len() + len;
91    dst.reserve_exact(len);
92
93    // A single frame that records exactly `len` bytes is decompressed in one shot.
94    // zstd verifies the recorded size while decompressing.
95    if matches!(zstd_safe::find_frame_compressed_size(src), Ok(n) if n == src.len())
96        && matches!(zstd_safe::get_frame_content_size(src), Ok(Some(n)) if n == len as u64)
97    {
98        decompress(src, dst)?;
99        return Ok(());
100    }
101
102    with_dctx(|ctx| {
103        ctx.reset(ResetDirective::SessionOnly).map_err(zstd_err)?;
104        let mut out = Cursor::new(&mut *dst);
105        out.set_position(out.get_ref().len() as u64);
106        let mut output = OutBuffer::around(&mut out);
107        let mut input = InBuffer::around(src);
108        while output.pos() < len {
109            let progress = (input.pos(), output.pos());
110            ctx.decompress_stream(&mut output, &mut input)
111                .map_err(zstd_err)?;
112            polars_ensure!(
113                (input.pos(), output.pos()) != progress,
114                ComputeError: "zstd data decompresses to fewer than {len} bytes"
115            );
116        }
117        Ok(())
118    })?;
119    // The spare capacity may have received more than `len` bytes.
120    dst.truncate(end);
121    Ok(())
122}
polars_utils/compression/zstd.rs

polars_utils/compression/
zstd.rs
