ailake_vec/
compress.rs

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

ailake_vec/
compress.rs
