redb_extras/roaring/
value.rs

1//! Roaring bitmap value type for partitioned tables.
2//!
3//! Provides encoding, decoding, and size information for RoaringTreemap values
4//! stored in partitioned segments.
5
6use super::RoaringError;
7use crate::{MergeableValue, Result};
8use redb::Value as RedbValue;
9use roaring::RoaringTreemap;
10
11/// Value type for RoaringTreemap in partitioned tables.
12///
13/// This struct provides the bridge between the generic partitioned storage
14/// and roaring-specific value operations. It handles:
15/// - Serialization/deserialization of RoaringTreemap
16/// - Size queries for segment rolling decisions
17/// - Version management for future migrations
18#[derive(Debug, Clone, PartialEq)]
19pub struct RoaringValue {
20    bitmap: RoaringTreemap,
21}
22
23impl RoaringValue {
24    /// Creates a new RoaringValue from an existing bitmap.
25    pub fn new(bitmap: RoaringTreemap) -> Self {
26        Self { bitmap }
27    }
28
29    /// Creates an empty RoaringValue.
30    pub fn empty() -> Self {
31        Self {
32            bitmap: RoaringTreemap::new(),
33        }
34    }
35
36    /// Returns a reference to the underlying bitmap.
37    pub fn bitmap(&self) -> &RoaringTreemap {
38        &self.bitmap
39    }
40
41    /// Returns a mutable reference to the underlying bitmap.
42    pub fn bitmap_mut(&mut self) -> &mut RoaringTreemap {
43        &mut self.bitmap
44    }
45
46    /// Consumes the value and returns the underlying bitmap.
47    pub fn into_bitmap(self) -> RoaringTreemap {
48        self.bitmap
49    }
50
51    /// Encodes a RoaringTreemap into storage format.
52    ///
53    /// # Arguments
54    /// * `bitmap` - The roaring bitmap to encode
55    ///
56    /// # Returns
57    /// Encoded bytes ready for storage
58    pub fn encode(&self) -> Result<Vec<u8>> {
59        Self::encode_bitmap(&self.bitmap)
60    }
61
62    /// Encodes a RoaringTreemap into storage format.
63    ///
64    /// # Arguments
65    /// * `bitmap` - The roaring bitmap to encode
66    ///
67    /// # Returns
68    /// Encoded bytes ready for storage
69    pub fn encode_bitmap(bitmap: &RoaringTreemap) -> Result<Vec<u8>> {
70        let mut buf = Vec::new();
71        bitmap
72            .serialize_into(&mut buf)
73            .map_err(|e| RoaringError::SerializationFailed(e.to_string()))?;
74
75        // Add version prefix (current version = 1)
76        let mut result = Vec::with_capacity(1 + buf.len());
77        result.push(1u8); // Version byte
78        result.extend_from_slice(&buf);
79
80        Ok(result)
81    }
82
83    /// Decodes storage bytes into a RoaringValue.
84    ///
85    /// # Arguments
86    /// * `data` - The encoded value bytes
87    ///
88    /// # Returns
89    /// Decoded RoaringValue
90    pub fn decode(data: &[u8]) -> Result<Self> {
91        if data.is_empty() {
92            return Err(RoaringError::InvalidBitmap("Empty data".to_string()).into());
93        }
94
95        let version = data[0];
96        let bitmap_bytes = &data[1..];
97
98        if version != 1 {
99            return Err(
100                RoaringError::InvalidBitmap(format!("Unsupported version: {}", version)).into(),
101            );
102        }
103
104        let bitmap = RoaringTreemap::deserialize_from(bitmap_bytes)
105            .map_err(|e| RoaringError::SerializationFailed(e.to_string()))?;
106        Ok(Self { bitmap })
107    }
108
109    /// Gets the serialized size of a RoaringTreemap.
110    ///
111    /// This size is used by the partition layer to determine when to roll
112    /// segments based on the configured maximum segment size.
113    ///
114    /// # Arguments
115    /// * `bitmap` - The roaring bitmap to measure
116    ///
117    /// # Returns
118    /// Serialized size in bytes (including version prefix)
119    pub fn get_serialized_size(&self) -> Result<usize> {
120        Self::get_serialized_size_for(&self.bitmap)
121    }
122
123    /// Gets the serialized size of a RoaringTreemap.
124    ///
125    /// This size is used by the partition layer to determine when to roll
126    /// segments based on the configured maximum segment size.
127    ///
128    /// # Arguments
129    /// * `bitmap` - The roaring bitmap to measure
130    ///
131    /// # Returns
132    /// Serialized size in bytes (including version prefix)
133    pub fn get_serialized_size_for(bitmap: &RoaringTreemap) -> Result<usize> {
134        let mut buf = Vec::new();
135        bitmap
136            .serialize_into(&mut buf)
137            .map_err(|e| RoaringError::SerializationFailed(e.to_string()))?;
138
139        // Include 1 byte for version prefix
140        Ok(1 + buf.len())
141    }
142
143    /// Creates a RoaringValue from a single value.
144    pub fn from_single(value: u64) -> Self {
145        let mut bitmap = RoaringTreemap::new();
146        bitmap.insert(value);
147        Self { bitmap }
148    }
149
150    /// Creates a RoaringValue from an iterator of values.
151    pub fn from_iter<I>(iter: I) -> Self
152    where
153        I: IntoIterator<Item = u64>,
154    {
155        let values: Vec<u64> = iter.into_iter().collect();
156        let bitmap =
157            RoaringTreemap::from_sorted_iter(values.iter().cloned()).unwrap_or_else(|_| {
158                let mut bitmap = RoaringTreemap::new();
159                for value in &values {
160                    bitmap.insert(*value);
161                }
162                bitmap
163            });
164        Self { bitmap }
165    }
166
167    /// Returns the number of members in the bitmap.
168    pub fn len(&self) -> u64 {
169        self.bitmap.len()
170    }
171
172    /// Returns true if the bitmap is empty.
173    pub fn is_empty(&self) -> bool {
174        self.bitmap.is_empty()
175    }
176}
177
178impl From<RoaringTreemap> for RoaringValue {
179    fn from(value: RoaringTreemap) -> Self {
180        Self { bitmap: value }
181    }
182}
183
184impl Default for RoaringValue {
185    fn default() -> Self {
186        Self::empty()
187    }
188}
189
190impl MergeableValue for RoaringValue {
191    fn merge(existing: Option<Self>, incoming: Self) -> Self {
192        match existing {
193            Some(mut existing) => {
194                existing.bitmap.extend(incoming.bitmap.into_iter());
195                existing
196            }
197            None => incoming,
198        }
199    }
200}
201
202impl RedbValue for RoaringValue {
203    type SelfType<'a>
204        = RoaringValue
205    where
206        Self: 'a;
207    type AsBytes<'a>
208        = Vec<u8>
209    where
210        Self: 'a;
211
212    fn fixed_width() -> Option<usize> {
213        None // Variable width serialization
214    }
215
216    fn from_bytes<'a>(data: &'a [u8]) -> Self::SelfType<'a>
217    where
218        Self: 'a,
219    {
220        RoaringValue::decode(data).unwrap_or_else(|_| RoaringValue::empty())
221    }
222
223    fn as_bytes<'a, 'b: 'a>(value: &'a Self::SelfType<'b>) -> Self::AsBytes<'a>
224    where
225        Self: 'b,
226    {
227        value.encode().unwrap_or_else(|_| Vec::new())
228    }
229
230    fn type_name() -> redb::TypeName {
231        redb::TypeName::new("RoaringTreemap")
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn test_encode_decode_roundtrip() {
241        let mut bitmap = RoaringTreemap::new();
242        bitmap.insert(1);
243        bitmap.insert(100);
244        bitmap.insert(1000);
245        let value = RoaringValue::from(bitmap);
246
247        let encoded = value.encode().unwrap();
248        let decoded = RoaringValue::decode(&encoded).unwrap();
249
250        assert_eq!(value, decoded);
251    }
252
253    #[test]
254    fn test_empty_bitmap() {
255        let value = RoaringValue::empty();
256
257        let encoded = value.encode().unwrap();
258        let decoded = RoaringValue::decode(&encoded).unwrap();
259
260        assert_eq!(value, decoded);
261        assert_eq!(decoded.len(), 0);
262    }
263
264    #[test]
265    fn test_serialized_size() {
266        let mut bitmap = RoaringTreemap::new();
267        bitmap.insert(1);
268        bitmap.insert(2);
269
270        let value = RoaringValue::from(bitmap);
271
272        let size = value.get_serialized_size().unwrap();
273        assert!(size > 1); // At least version byte
274        assert!(size < 1000); // Should be reasonably small
275
276        let encoded = value.encode().unwrap();
277        assert_eq!(size, encoded.len());
278    }
279
280    #[test]
281    fn test_single_value() {
282        let value = RoaringValue::from_single(42);
283
284        assert_eq!(value.len(), 1);
285        assert!(value.bitmap().contains(42));
286
287        let encoded = value.encode().unwrap();
288        let decoded = RoaringValue::decode(&encoded).unwrap();
289
290        assert_eq!(value, decoded);
291    }
292
293    #[test]
294    fn test_from_iter() {
295        let values = vec![1, 5, 10, 100];
296        let value = RoaringValue::from_iter(values.clone());
297
298        assert_eq!(value.len(), values.len() as u64);
299        for member in &values {
300            assert!(value.bitmap().contains(*member));
301        }
302
303        let encoded = value.encode().unwrap();
304        let decoded = RoaringValue::decode(&encoded).unwrap();
305
306        assert_eq!(value, decoded);
307    }
308
309    #[test]
310    fn test_invalid_version() {
311        let mut invalid_data = vec![99]; // Invalid version
312        invalid_data.extend_from_slice(b"fake_data");
313
314        let result = RoaringValue::decode(&invalid_data);
315        assert!(result.is_err());
316    }
317}
redb_extras/roaring/value.rs

redb_extras/roaring/
value.rs
