mz_repr/
bytes.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::fmt::{self, Display};
11use std::str::FromStr;
12
13use proptest_derive::Arbitrary;
14use serde::{Deserialize, Serialize};
15
16use mz_ore::cast::CastLossy;
17
18/// Struct for postgres compatible size units which is different from
19/// `bytesize::ByteSize`. Instead of MiB or GiB and so on, it uses MB, GB for the sizes
20/// with 1024 multiplier. Valid units are B, kB, MB, GB, TB with multiples of 1024
21/// where 1MB = 1024kB.
22///
23/// In postgres, each setting has a base unit (for eg. B, kB) and the value can either
24/// be integer or float. The base unit serves as the default unit if a number is provided
25/// without a unit and it's also the minimum unit in which values
26/// can be rounded to. For example, with base unit of kB, 30.1kB will be rounded to
27/// 30kB since it can't have a lower unit, but 30.1MB will be rounded to 30822kB.
28/// For [`ByteSize`], the value is an integer and the base unit is bytes (`B`).
29#[derive(
30    Arbitrary,
31    Debug,
32    Clone,
33    PartialEq,
34    Eq,
35    Hash,
36    PartialOrd,
37    Ord,
38    Serialize,
39    Deserialize,
40    Default
41)]
42pub struct ByteSize(u64);
43
44impl ByteSize {
45    pub const fn b(size: u64) -> ByteSize {
46        ByteSize(size)
47    }
48
49    pub const fn kb(size: u64) -> ByteSize {
50        ByteSize(size * BytesUnit::Kb.value())
51    }
52
53    pub const fn mb(size: u64) -> ByteSize {
54        ByteSize(size * BytesUnit::Mb.value())
55    }
56
57    pub const fn gb(size: u64) -> ByteSize {
58        ByteSize(size * BytesUnit::Gb.value())
59    }
60
61    pub const fn tb(size: u64) -> ByteSize {
62        ByteSize(size * BytesUnit::Tb.value())
63    }
64
65    pub fn as_bytes(&self) -> u64 {
66        self.0
67    }
68
69    fn format_string(&self) -> String {
70        match self.0 {
71            zero if zero == 0 => "0".to_string(),
72            tb if tb % BytesUnit::Tb.value() == 0 => {
73                format!("{}{}", tb / BytesUnit::Tb.value(), BytesUnit::Tb)
74            }
75            gb if gb % BytesUnit::Gb.value() == 0 => {
76                format!("{}{}", gb / BytesUnit::Gb.value(), BytesUnit::Gb)
77            }
78            mb if mb % BytesUnit::Mb.value() == 0 => {
79                format!("{}{}", mb / BytesUnit::Mb.value(), BytesUnit::Mb)
80            }
81            kb if kb % BytesUnit::Kb.value() == 0 => {
82                format!("{}{}", kb / BytesUnit::Kb.value(), BytesUnit::Kb)
83            }
84            b => format!("{}{}", b, BytesUnit::B),
85        }
86    }
87}
88
89impl Display for ByteSize {
90    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
91        f.pad(&self.format_string())
92    }
93}
94
95impl FromStr for ByteSize {
96    type Err = String;
97
98    // To behave the same as in postgres, this always
99    // rounds down to the next lower unit if possible.
100    // For example 30.9B, will be rounded to 31B, since there's no
101    // lower unit than B. But 30.1kB will be rounded to
102    // 31642B.
103    fn from_str(value: &str) -> Result<Self, Self::Err> {
104        let number: String = value
105            .chars()
106            .take_while(|c| c.is_digit(10) || c == &'.')
107            .collect();
108
109        let suffix: String = value
110            .chars()
111            .skip_while(|c| c.is_whitespace() || c.is_digit(10) || c == &'.')
112            .collect();
113
114        let unit = if suffix.is_empty() {
115            BytesUnit::B
116        } else {
117            suffix
118            .parse::<BytesUnit>()
119            .map_err(|e| format!("couldn't parse {:?} into a known SI unit, {}. Valid units are B, kB, MB, GB, and TB", suffix, e))?
120        };
121
122        let (size, unit) = if let Ok(integer) = number.parse::<u64>() {
123            (integer, unit)
124        } else {
125            let num = number
126                .parse::<f64>()
127                .map_err(|e| format!("couldn't parse {} as a number, {}", number, e))?;
128
129            // checking if number has no fractional part
130            if num.trunc() == num {
131                let size = u64::cast_lossy(num);
132                (size, unit)
133            } else {
134                match unit {
135                    BytesUnit::B => (u64::cast_lossy(num.round()), BytesUnit::B),
136                    BytesUnit::Kb => (u64::cast_lossy((num * 1024.0).round()), BytesUnit::B),
137                    BytesUnit::Mb => (u64::cast_lossy((num * 1024.0).round()), BytesUnit::Kb),
138                    BytesUnit::Gb => (u64::cast_lossy((num * 1024.0).round()), BytesUnit::Mb),
139                    BytesUnit::Tb => (u64::cast_lossy((num * 1024.0).round()), BytesUnit::Gb),
140                }
141            }
142        };
143
144        let bytes = size
145            .checked_mul(unit.value())
146            .ok_or_else(|| "bytes value exceeds u64 range".to_string())?;
147        Ok(Self(bytes))
148    }
149}
150
151/// Valid units for representing bytes
152#[derive(
153    Arbitrary,
154    Debug,
155    Clone,
156    PartialEq,
157    Eq,
158    Hash,
159    PartialOrd,
160    Ord,
161    Serialize,
162    Deserialize,
163    Default
164)]
165pub enum BytesUnit {
166    #[default]
167    B,
168    Kb,
169    Mb,
170    Gb,
171    Tb,
172}
173
174impl BytesUnit {
175    const fn value(&self) -> u64 {
176        match &self {
177            BytesUnit::B => 1,
178            BytesUnit::Kb => 1_024,
179            BytesUnit::Mb => 1_048_576,
180            BytesUnit::Gb => 1_073_741_824,
181            BytesUnit::Tb => 1_099_511_627_776,
182        }
183    }
184}
185
186impl fmt::Display for BytesUnit {
187    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
188        f.write_str(match self {
189            BytesUnit::B => "B",
190            BytesUnit::Kb => "kB",
191            BytesUnit::Mb => "MB",
192            BytesUnit::Gb => "GB",
193            BytesUnit::Tb => "TB",
194        })
195    }
196}
197
198impl FromStr for BytesUnit {
199    type Err = String;
200
201    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
202        match s {
203            "B" => Ok(Self::B),
204            "kB" => Ok(Self::Kb),
205            "MB" => Ok(Self::Mb),
206            "GB" => Ok(Self::Gb),
207            "TB" => Ok(Self::Tb),
208            _ => Err(format!(
209                "invalid BytesUnit: {}. Valid units are B, kB, MB, GB, and TB",
210                s
211            )),
212        }
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use crate::bytes::ByteSize;
219    use mz_ore::assert_err;
220    use proptest::prelude::*;
221    use proptest::proptest;
222
223    #[mz_ore::test]
224    fn test_to_string() {
225        fn assert_to_string(expected: &str, b: ByteSize) {
226            assert_eq!(expected.to_string(), b.to_string());
227        }
228        assert_to_string("0", ByteSize::gb(0));
229        assert_to_string("1GB", ByteSize::mb(1024));
230        assert_to_string("215B", ByteSize::b(215));
231        assert_to_string("1kB", ByteSize::kb(1));
232        assert_to_string("301kB", ByteSize::kb(301));
233        assert_to_string("419MB", ByteSize::mb(419));
234        assert_to_string("518GB", ByteSize::gb(518));
235        assert_to_string("815TB", ByteSize::tb(815));
236        assert_to_string("10kB", ByteSize::b(10240));
237        assert_to_string("10MB", ByteSize::kb(10240));
238        assert_to_string("10GB", ByteSize::mb(10240));
239        assert_to_string("10TB", ByteSize::gb(10240));
240        assert_to_string("10240TB", ByteSize::tb(10240));
241    }
242
243    #[mz_ore::test]
244    fn test_parse() {
245        // shortcut for writing test cases
246        fn parse(s: &str) -> ByteSize {
247            s.parse::<ByteSize>().unwrap()
248        }
249
250        assert_eq!(parse("0"), ByteSize::b(0));
251        assert_eq!(parse("9.9"), ByteSize::b(10));
252        assert_eq!(parse("0B"), ByteSize::b(0));
253        assert_eq!(parse("0MB"), ByteSize::b(0));
254        assert_eq!(parse("500"), ByteSize::b(500));
255        assert_eq!(parse("1kB"), ByteSize::kb(1));
256        assert_eq!(parse("1.5kB"), ByteSize::b(1536));
257        assert_eq!(parse("1 kB"), ByteSize::kb(1));
258        assert_eq!(parse("3 MB"), ByteSize::mb(3));
259        assert_eq!(parse("6 GB"), ByteSize::gb(6));
260        assert_eq!(parse("4GB"), ByteSize::gb(4));
261        assert_eq!(parse("88TB"), ByteSize::tb(88));
262        assert_eq!(parse("521  TB"), ByteSize::tb(521));
263
264        // parsing errors
265        assert_err!("".parse::<ByteSize>());
266        assert_err!("a124GB".parse::<ByteSize>());
267        assert_err!("1K".parse::<ByteSize>());
268        assert_err!("B".parse::<ByteSize>());
269        // postgres is strict about matching capitalization
270        assert_err!("1gb".parse::<ByteSize>());
271        assert_err!("1KB".parse::<ByteSize>());
272    }
273
274    #[mz_ore::test]
275    fn test_rounding() {
276        // shortcut for writing test cases
277        fn parse(s: &str) -> ByteSize {
278            s.parse::<ByteSize>().unwrap()
279        }
280
281        fn assert_equivalent(v1: &str, v2: &str) {
282            assert_eq!(parse(v1), parse(v2))
283        }
284
285        assert_equivalent("0", "0");
286        assert_equivalent("0 TB", "0");
287        assert_equivalent("0kB", "0");
288        assert_equivalent("13.89", "14B");
289        assert_equivalent("500", "500B");
290        assert_equivalent("1073741824", "1GB");
291        assert_equivalent("1073741824.0", "1GB");
292        assert_equivalent("1073741824.1", "1GB");
293        assert_equivalent("1073741824.9", "1073741825B");
294        assert_equivalent("2147483648", "2GB");
295        assert_equivalent("3221225472", "3GB");
296        assert_equivalent("4294967296", "4GB");
297        assert_equivalent("4294967295", "4294967295B");
298        assert_equivalent("1024.1", "1kB");
299        assert_equivalent("1024.9", "1025B");
300        assert_equivalent("1024.1MB", "1048678kB");
301        assert_equivalent("1024.9MB", "1049498kB");
302        assert_equivalent("1.01B", "1B");
303        assert_equivalent("1.01kB", "1034B");
304        assert_equivalent("1.0kB", "1kB");
305        assert_equivalent("10240B", "10kB");
306        assert_equivalent("1.5kB", "1536B");
307        assert_equivalent("30.1GB", "30822MB");
308        assert_equivalent("30.1MB", "30822kB");
309        assert_equivalent("30.1TB", "30822GB");
310        assert_equivalent("39.9TB", "40858GB");
311        assert_equivalent("30.9B", "31B");
312    }
313
314    proptest! {
315      #[mz_ore::test]
316      fn proptest_bytes_roundtrips_string(og: ByteSize) {
317        // Not all [`ByteSize`] values can successfully roundtrip.
318        // For example, '30.1 MB' will be rounded off to '30822 kB'.
319        // So instead testing roundtrip of string representations to
320        // [`ByteSize`] and vice versa.
321        let og_string = og.to_string();
322        let roundtrip = og_string.parse::<ByteSize>().expect("roundtrip").to_string();
323        prop_assert_eq!(og_string, roundtrip);
324      }
325    }
326}
mz_repr/bytes.rs

mz_repr/
bytes.rs
