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)]
26pub struct BrotliLevel(u32);
27
28impl Default for BrotliLevel {
29    fn default() -> Self {
30        Self(1)
31    }
32}
33
34impl CompressionLevel<u32> for BrotliLevel {
35    const MINIMUM_LEVEL: u32 = 0;
36    const MAXIMUM_LEVEL: u32 = 11;
37}
38
39impl BrotliLevel {
40    /// Attempts to create a brotli compression level.
41    ///
42    /// Compression levels must be valid.
43    pub fn try_new(level: u32) -> PolarsResult<Self> {
44        Self::is_valid_level(level).map(|_| Self(level))
45    }
46
47    /// Returns the compression level.
48    pub fn compression_level(&self) -> u32 {
49        self.0
50    }
51}
52
53/// Represents a valid gzip compression level.
54#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
55pub struct GzipLevel(u8);
56
57impl Default for GzipLevel {
58    fn default() -> Self {
59        // The default as of miniz_oxide 0.5.1 is 6 for compression level
60        // (miniz_oxide::deflate::CompressionLevel::DefaultLevel)
61        Self(6)
62    }
63}
64
65impl CompressionLevel<u8> for GzipLevel {
66    const MINIMUM_LEVEL: u8 = 0;
67    const MAXIMUM_LEVEL: u8 = 9;
68}
69
70impl GzipLevel {
71    /// Attempts to create a gzip compression level.
72    ///
73    /// Compression levels must be valid (i.e. be acceptable for `flate2::Compression`).
74    pub fn try_new(level: u8) -> PolarsResult<Self> {
75        Self::is_valid_level(level).map(|_| Self(level))
76    }
77
78    /// Returns the compression level.
79    pub fn compression_level(&self) -> u8 {
80        self.0
81    }
82}
83
84/// Represents a valid zstd compression level.
85#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
86#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
87#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
88pub struct ZstdLevel(i32);
89
90impl CompressionLevel<i32> for ZstdLevel {
91    // zstd binds to C, and hence zstd::compression_level_range() is not const as this calls the
92    // underlying C library.
93    const MINIMUM_LEVEL: i32 = 1;
94    const MAXIMUM_LEVEL: i32 = 22;
95}
96
97impl ZstdLevel {
98    /// Attempts to create a zstd compression level from a given compression level.
99    ///
100    /// Compression levels must be valid (i.e. be acceptable for `zstd::compression_level_range`).
101    pub fn try_new(level: i32) -> PolarsResult<Self> {
102        Self::is_valid_level(level).map(|_| Self(level))
103    }
104
105    /// Returns the compression level.
106    pub fn compression_level(&self) -> i32 {
107        self.0
108    }
109}
110
111impl Default for ZstdLevel {
112    fn default() -> Self {
113        Self(3)
114    }
115}