Skip to main content

polars_utils/
compression.rs

1use polars_error::{PolarsResult, polars_bail};
2#[cfg(feature = "serde")]
3use serde::{Deserialize, Serialize};
4
5#[cfg(feature = "zstd")]
6pub mod zstd;
7
8/// Defines valid compression levels.
9pub trait CompressionLevel<T: std::fmt::Display + std::cmp::PartialOrd> {
10    const MINIMUM_LEVEL: T;
11    const MAXIMUM_LEVEL: T;
12
13    /// Tests if the provided compression level is valid.
14    fn is_valid_level(level: T) -> PolarsResult<()> {
15        let compression_range = Self::MINIMUM_LEVEL..=Self::MAXIMUM_LEVEL;
16        if compression_range.contains(&level) {
17            Ok(())
18        } else {
19            polars_bail!(InvalidOperation: "valid compression range {}..={} exceeded.",
20                compression_range.start(),
21                compression_range.end()
22            )
23        }
24    }
25}
26
27/// Represents a valid brotli compression level.
28#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
29#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
30#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
31pub struct BrotliLevel(u32);
32
33impl Default for BrotliLevel {
34    fn default() -> Self {
35        Self(1)
36    }
37}
38
39impl CompressionLevel<u32> for BrotliLevel {
40    const MINIMUM_LEVEL: u32 = 0;
41    const MAXIMUM_LEVEL: u32 = 11;
42}
43
44impl BrotliLevel {
45    /// Attempts to create a brotli compression level.
46    ///
47    /// Compression levels must be valid.
48    pub fn try_new(level: u32) -> PolarsResult<Self> {
49        Self::is_valid_level(level).map(|_| Self(level))
50    }
51
52    /// Returns the compression level.
53    pub fn compression_level(&self) -> u32 {
54        self.0
55    }
56}
57
58/// Represents a valid gzip compression level.
59#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
60#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
61#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
62pub struct GzipLevel(u8);
63
64impl Default for GzipLevel {
65    fn default() -> Self {
66        // The default as of miniz_oxide 0.5.1 is 6 for compression level
67        // (miniz_oxide::deflate::CompressionLevel::DefaultLevel)
68        Self(6)
69    }
70}
71
72impl CompressionLevel<u8> for GzipLevel {
73    const MINIMUM_LEVEL: u8 = 0;
74    const MAXIMUM_LEVEL: u8 = 9;
75}
76
77impl GzipLevel {
78    /// Attempts to create a gzip compression level.
79    ///
80    /// Compression levels must be valid (i.e. be acceptable for `flate2::Compression`).
81    pub fn try_new(level: u8) -> PolarsResult<Self> {
82        Self::is_valid_level(level).map(|_| Self(level))
83    }
84
85    /// Returns the compression level.
86    pub fn compression_level(&self) -> u8 {
87        self.0
88    }
89}
90
91/// Represents a valid zstd compression level.
92#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
93#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
94#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
95pub struct ZstdLevel(i32);
96
97impl CompressionLevel<i32> for ZstdLevel {
98    // zstd binds to C, and hence zstd::compression_level_range() is not const as this calls the
99    // underlying C library.
100    const MINIMUM_LEVEL: i32 = 1;
101    const MAXIMUM_LEVEL: i32 = 22;
102}
103
104impl ZstdLevel {
105    /// Attempts to create a zstd compression level from a given compression level.
106    ///
107    /// Compression levels must be valid (i.e. be acceptable for `zstd::compression_level_range`).
108    pub fn try_new(level: i32) -> PolarsResult<Self> {
109        Self::is_valid_level(level).map(|_| Self(level))
110    }
111
112    /// Returns the compression level.
113    pub fn compression_level(&self) -> i32 {
114        self.0
115    }
116}
117
118impl Default for ZstdLevel {
119    fn default() -> Self {
120        Self(3)
121    }
122}