polars_utils/
compression.rs

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