risingwave_common/util/
compress.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::iter_util::ZipEqFast;
16
17/// This function compresses sequential repeated data in a vector. The compression result contains
18/// two vectors, one for the last indices of sequential repeated elements, and another for the
19/// repeated data. For example, [14, 14, 14, 27, 27] will be compressed to [2, 4], [14, 27].
20pub fn compress_data<T>(original_data: &[T]) -> (Vec<u32>, Vec<T>)
21where
22    T: PartialEq + Copy,
23{
24    let mut original_indices = Vec::new();
25    let mut data = Vec::new();
26
27    for i in 1..original_data.len() {
28        if original_data[i - 1] != original_data[i] {
29            original_indices.push(i as u32 - 1);
30            data.push(original_data[i - 1]);
31        }
32    }
33
34    if let Some(&last) = original_data.last() {
35        original_indices.push(original_data.len() as u32 - 1);
36        data.push(last);
37    }
38
39    (original_indices, data)
40}
41
42/// Works in a reversed way as `compress_data`.
43pub fn decompress_data<T>(original_indices: &[u32], data: &[T]) -> Vec<T>
44where
45    T: Copy,
46{
47    match original_indices.last() {
48        Some(last_idx) => {
49            let mut original_data = Vec::with_capacity(*last_idx as usize + 1);
50            original_indices
51                .iter()
52                .zip_eq_fast(data)
53                .for_each(|(&idx, &x)| {
54                    original_data.resize(idx as usize + 1, x);
55                });
56            original_data
57        }
58        None => Vec::new(),
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::{compress_data, decompress_data};
65
66    #[test]
67    fn test_compress() {
68        // Simple
69        let original_data = [3u32, 3, 3, 3, 3, 4, 4, 5, 5, 6, 7, 8, 8, 8, 9];
70        let (compressed_original_indices, compressed_data) = compress_data(&original_data);
71        let expect_original_indices = Vec::from([4u32, 6, 8, 9, 10, 13, 14]);
72        let expect_data = Vec::from([3u32, 4, 5, 6, 7, 8, 9]);
73        assert_eq!(compressed_original_indices, expect_original_indices);
74        assert_eq!(compressed_data, expect_data);
75        let decompressed_data = decompress_data(&compressed_original_indices, &compressed_data);
76        assert_eq!(decompressed_data, original_data);
77
78        // Complex
79        let mut long_original_data = Vec::new();
80        long_original_data.resize(512, 1);
81        long_original_data.resize(1024, 2);
82        long_original_data.resize(1536, 3);
83        long_original_data.resize(2048, 4);
84        long_original_data[0] = 5;
85        long_original_data[2046] = 5;
86        let (compressed_original_indices, compressed_data) = compress_data(&long_original_data);
87        let expect_original_indices = Vec::from([0u32, 511, 1023, 1535, 2045, 2046, 2047]);
88        let expect_data = Vec::from([5u32, 1, 2, 3, 4, 5, 4]);
89        assert_eq!(compressed_original_indices, expect_original_indices);
90        assert_eq!(compressed_data, expect_data);
91        let decompressed_data = decompress_data(&compressed_original_indices, &compressed_data);
92        assert_eq!(decompressed_data, long_original_data);
93
94        // Empty
95        let (compressed_original_indices, compressed_data) = compress_data::<u8>(&[]);
96        assert!(compressed_original_indices.is_empty());
97        assert!(compressed_data.is_empty());
98        let decompressed_data = decompress_data(&compressed_original_indices, &compressed_data);
99        assert!(decompressed_data.is_empty());
100    }
101}
risingwave_common/util/compress.rs

risingwave_common/util/
compress.rs
