ailake_vec/
compress.rs

1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2pub enum CompressionCodec {
3    None,
4    Lz4,
5    Zstd,
6}
7
8pub struct BlockCompressor {
9    codec: CompressionCodec,
10    zstd_level: i32,
11}
12
13impl BlockCompressor {
14    pub fn none() -> Self {
15        Self {
16            codec: CompressionCodec::None,
17            zstd_level: 3,
18        }
19    }
20
21    pub fn lz4() -> Self {
22        Self {
23            codec: CompressionCodec::Lz4,
24            zstd_level: 3,
25        }
26    }
27
28    pub fn zstd(level: i32) -> Self {
29        Self {
30            codec: CompressionCodec::Zstd,
31            zstd_level: level,
32        }
33    }
34
35    pub fn codec(&self) -> CompressionCodec {
36        self.codec
37    }
38
39    pub fn compress(&self, data: &[u8]) -> Vec<u8> {
40        match self.codec {
41            CompressionCodec::None => data.to_vec(),
42            CompressionCodec::Lz4 => lz4_flex::compress_prepend_size(data),
43            CompressionCodec::Zstd => {
44                zstd::bulk::compress(data, self.zstd_level).unwrap_or_else(|_| data.to_vec())
45            }
46        }
47    }
48
49    pub fn decompress(&self, data: &[u8]) -> Vec<u8> {
50        match self.codec {
51            CompressionCodec::None => data.to_vec(),
52            CompressionCodec::Lz4 => {
53                lz4_flex::decompress_size_prepended(data).unwrap_or_else(|_| data.to_vec())
54            }
55            CompressionCodec::Zstd => {
56                zstd::bulk::decompress(data, 64 * 1024 * 1024).unwrap_or_else(|_| data.to_vec())
57            }
58        }
59    }
60}
61
62impl Default for BlockCompressor {
63    fn default() -> Self {
64        Self::zstd(3)
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    fn roundtrip(codec: BlockCompressor, data: &[u8]) {
73        let compressed = codec.compress(data);
74        let decompressed = codec.decompress(&compressed);
75        assert_eq!(decompressed, data);
76    }
77
78    #[test]
79    fn lz4_roundtrip() {
80        let data: Vec<u8> = (0u8..200).cycle().take(4096).collect();
81        roundtrip(BlockCompressor::lz4(), &data);
82    }
83
84    #[test]
85    fn zstd_roundtrip() {
86        let data: Vec<u8> = (0u8..200).cycle().take(4096).collect();
87        roundtrip(BlockCompressor::zstd(3), &data);
88    }
89
90    #[test]
91    fn none_passthrough() {
92        let data = b"hello ailake";
93        roundtrip(BlockCompressor::none(), data);
94    }
95
96    #[test]
97    fn zstd_compresses_repetitive_data() {
98        // Repetitive float data (like zero vectors) should compress well
99        let data = vec![0u8; 8192];
100        let c = BlockCompressor::zstd(3);
101        let compressed = c.compress(&data);
102        assert!(
103            compressed.len() < data.len() / 4,
104            "expected >4x compression ratio"
105        );
106    }
107}
ailake_vec/compress.rs

ailake_vec/
compress.rs
