datafusion_roaring/
codec.rs

1//! Versioned serialization for roaring values.
2
3use std::io::{self, Cursor};
4
5use datafusion_common::{DataFusionError, Result};
6use roaring::RoaringBitmap;
7
8const MAGIC: &[u8; 4] = b"DFRB";
9const VERSION: u8 = 1;
10const BITMAP_KIND: u8 = 1;
11const HEADER_LEN: usize = 6;
12
13fn codec_error(action: &str, message: impl std::fmt::Display) -> DataFusionError {
14    DataFusionError::Execution(format!("could not {action} roaring value: {message}"))
15}
16
17fn payload<'a>(bytes: &'a [u8], expected_kind: u8, type_name: &str) -> Result<&'a [u8]> {
18    if bytes.len() < HEADER_LEN {
19        return Err(codec_error(
20            "decode",
21            "value is shorter than the format header",
22        ));
23    }
24    if &bytes[..4] != MAGIC {
25        return Err(codec_error("decode", "invalid format magic"));
26    }
27    if bytes[4] != VERSION {
28        return Err(codec_error(
29            "decode",
30            format!("unsupported format version {}", bytes[4]),
31        ));
32    }
33    if bytes[5] != expected_kind {
34        return Err(codec_error("decode", format!("value is not a {type_name}")));
35    }
36    Ok(&bytes[HEADER_LEN..])
37}
38
39fn encode_header(output: &mut Vec<u8>, kind: u8) {
40    output.extend_from_slice(MAGIC);
41    output.push(VERSION);
42    output.push(kind);
43}
44
45/// Serializes a 32-bit roaring bitmap using this crate's versioned format.
46pub fn encode_bitmap(bitmap: &RoaringBitmap) -> Result<Vec<u8>> {
47    let mut output = Vec::with_capacity(HEADER_LEN + bitmap.serialized_size());
48    encode_header(&mut output, BITMAP_KIND);
49    bitmap
50        .serialize_into(&mut output)
51        .map_err(|error| codec_error("encode UInt32", error))?;
52    Ok(output)
53}
54
55/// Deserializes a 32-bit roaring bitmap encoded by [`encode_bitmap`].
56pub fn decode_bitmap(bytes: &[u8]) -> Result<RoaringBitmap> {
57    let payload = payload(bytes, BITMAP_KIND, "UInt32 bitmap")?;
58    let mut input = Cursor::new(payload);
59    let bitmap = RoaringBitmap::deserialize_from(&mut input)
60        .map_err(|error: io::Error| codec_error("decode UInt32", error))?;
61    if input.position() != payload.len() as u64 {
62        return Err(codec_error("decode", "trailing bytes after bitmap payload"));
63    }
64    Ok(bitmap)
65}
datafusion_roaring/codec.rs

datafusion_roaring/
codec.rs
