use rudb_common::{Error, Result};
use crate::keymap::{Form, KeyMap, Observed};
pub const HEADER_BYTES: usize = 40;
const LAYOUT: u8 = 1;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Payload {
pub flags: u32,
pub header_bytes: u32,
pub bytes: Vec<u8>,
}
pub fn encode(map: &KeyMap, type_tag: u8) -> Result<Payload> {
let observed = *map.observed();
let mut bytes = Vec::with_capacity(HEADER_BYTES + map.bytes());
bytes.extend_from_slice(&map.base().to_le_bytes());
bytes.extend_from_slice(&observed.rows.to_le_bytes());
bytes.extend_from_slice(&observed.nulls.to_le_bytes());
bytes.push(map.form().tag());
bytes.push(type_tag);
bytes.push(u8::from(observed.distinct));
bytes.push(u8::from(observed.sorted));
bytes.push(LAYOUT);
bytes.extend_from_slice(&[0; 3]);
debug_assert_eq!(bytes.len(), HEADER_BYTES, "the key map header is forty bytes");
map.write_body(&mut bytes)?;
Ok(Payload {
flags: u32::from(map.form().tag()),
header_bytes: u32::try_from(HEADER_BYTES).map_err(|_| malformed("header overflow"))?,
bytes,
})
}
pub fn decode(bytes: &[u8]) -> Result<(KeyMap, u8)> {
if bytes.len() < HEADER_BYTES {
return Err(malformed("a key map payload is shorter than its header"));
}
let base = i128::from_le_bytes(bytes[0..16].try_into().map_err(|_| torn())?);
let rows = u64::from_le_bytes(bytes[16..24].try_into().map_err(|_| torn())?);
let nulls = u64::from_le_bytes(bytes[24..32].try_into().map_err(|_| torn())?);
let form = Form::from_tag(bytes[32])?;
let type_tag = bytes[33];
let distinct = flag(bytes[34])?;
let sorted = flag(bytes[35])?;
if bytes[36] != LAYOUT {
return Err(malformed(format!("key map layout {} is not one this build knows", bytes[36])));
}
let observed = Observed {
rows,
nulls,
distinct,
sorted,
min: (rows > 0).then_some(base),
max: None,
};
let map = KeyMap::read_body(form, base, observed, &bytes[HEADER_BYTES..])?;
Ok((map, type_tag))
}
fn flag(byte: u8) -> Result<bool> {
match byte {
0 => Ok(false),
1 => Ok(true),
_ => Err(malformed("a flag byte in a key map header is neither zero nor one")),
}
}
fn torn() -> Error {
malformed("a key map header is torn")
}
fn malformed(message: impl Into<String>) -> Error {
Error::invalid_input(format!("invalid rudb key map payload: {}", message.into()))
}
#[cfg(test)]
mod tests {
use super::*;
const INTEGER: u8 = 4;
fn keys(values: &[i128]) -> Vec<Option<i128>> {
values.iter().copied().map(Some).collect()
}
fn survives(column: &[Option<i128>]) -> KeyMap {
let built = KeyMap::build(column).expect("build");
let payload = encode(&built, INTEGER).expect("encode");
assert_eq!(payload.header_bytes as usize, HEADER_BYTES);
assert_eq!(payload.flags, u32::from(built.form().tag()));
let (read, type_tag) = decode(&payload.bytes).expect("decode");
assert_eq!(type_tag, INTEGER, "the type tag is carried, not interpreted");
assert_eq!(read.form(), built.form(), "the form survives");
assert_eq!(read.observed(), built.observed(), "the observed facts survive");
for (rid, key) in column.iter().enumerate() {
let Some(key) = *key else { continue };
assert_eq!(
read.lookup(key).expect("lookup"),
Some(rid as u64),
"key {key} did not survive the round trip"
);
}
read
}
#[test]
fn the_identity_form_round_trips_and_is_header_only() {
let map = survives(&keys(&(1..=1000).collect::<Vec<i128>>()));
assert_eq!(map.form(), Form::Identity);
let payload = encode(&map, INTEGER).expect("encode");
assert_eq!(payload.bytes.len(), HEADER_BYTES, "section 3.3: no extents beyond the header");
}
#[test]
fn the_dense_form_round_trips_with_its_rank_index() {
let map = survives(&keys(&(0..20_000).map(|value| value * 2).collect::<Vec<i128>>()));
assert_eq!(map.form(), Form::Dense);
let payload = encode(&map, INTEGER).expect("encode");
let bitmap = 40_000 / 8;
let body = payload.bytes.len() - HEADER_BYTES;
assert!(body > bitmap, "the body is {body} bytes and the bitmap alone is {bitmap}");
assert!(body < bitmap * 5 / 4, "the index costs about an eighth, not {body} over {bitmap}");
}
#[test]
fn the_sorted_form_round_trips_with_both_of_its_bit_packed_arrays() {
let map = survives(&keys(&[500, 3, 9000, 12, 7, 88, 41, 6]));
assert_eq!(map.form(), Form::Sorted);
}
#[test]
fn the_permuted_form_round_trips_with_its_bitmap_index_and_rids() {
let column: Vec<Option<i128>> =
(0..5_000_i128).map(|at| Some((at * 7_919 % 5_000) * 3)).collect();
let map = survives(&column);
assert_eq!(map.form(), Form::Permuted);
assert_eq!(map.lookup(1).expect("lookup"), None);
}
#[test]
fn a_column_with_nulls_round_trips_and_keeps_its_null_count() {
let column = vec![Some(10), None, Some(20), None, Some(30)];
let map = survives(&column);
assert_eq!(map.observed().nulls, 2);
assert_eq!(map.observed().rows, 3);
}
#[test]
fn a_column_of_one_key_round_trips() {
survives(&keys(&[42]));
}
#[test]
fn negative_keys_round_trip_because_the_base_is_an_i128() {
survives(&keys(&[i128::MIN + 1, i128::MIN + 9, i128::MIN + 4]));
}
#[test]
fn an_empty_key_map_round_trips_and_resolves_nothing() {
let built = KeyMap::build(&[]).expect("build");
let payload = encode(&built, INTEGER).expect("encode");
let (read, _) = decode(&payload.bytes).expect("decode");
assert!(read.is_empty());
assert_eq!(read.observed().min, None, "an empty map has no minimum, not a minimum of zero");
assert_eq!(read.lookup(0).expect("lookup"), None);
}
#[test]
fn a_non_distinct_column_carries_that_fact_through_the_round_trip() {
let built = KeyMap::build(&keys(&[5, 7, 5, 9])).expect("build");
assert!(!built.observed().distinct);
let payload = encode(&built, INTEGER).expect("encode");
let (read, _) = decode(&payload.bytes).expect("decode");
assert!(!read.observed().distinct);
assert!(!read.observed().usable_as_parent());
}
#[test]
fn a_payload_shorter_than_its_header_is_refused() {
let built = KeyMap::build(&keys(&[1, 2, 3])).expect("build");
let payload = encode(&built, INTEGER).expect("encode");
for cut in [0, 1, HEADER_BYTES - 1] {
assert!(decode(&payload.bytes[..cut]).is_err(), "a payload of {cut} bytes is refused");
}
}
#[test]
fn a_form_this_build_does_not_know_is_refused_rather_than_guessed() {
let built = KeyMap::build(&keys(&[1, 2, 3])).expect("build");
let mut payload = encode(&built, INTEGER).expect("encode");
payload.bytes[32] = 9;
let error = decode(&payload.bytes).expect_err("refused");
assert!(error.to_string().contains("form 9"), "{error}");
}
#[test]
fn a_layout_this_build_does_not_know_is_refused() {
let built = KeyMap::build(&keys(&[1, 2, 3])).expect("build");
let mut payload = encode(&built, INTEGER).expect("encode");
payload.bytes[36] = LAYOUT + 1;
let error = decode(&payload.bytes).expect_err("refused");
assert!(error.to_string().contains("layout"), "{error}");
}
#[test]
fn a_flag_byte_that_is_neither_zero_nor_one_is_refused() {
let built = KeyMap::build(&keys(&[1, 2, 3])).expect("build");
let mut payload = encode(&built, INTEGER).expect("encode");
payload.bytes[34] = 2;
assert!(decode(&payload.bytes).is_err(), "a torn distinct flag is refused");
let mut payload = encode(&built, INTEGER).expect("encode");
payload.bytes[35] = 0xff;
assert!(decode(&payload.bytes).is_err(), "a torn sorted flag is refused");
}
#[test]
fn a_truncated_body_is_refused_rather_than_read_past() {
for column in [
keys(&(0..2000).map(|value| value * 2).collect::<Vec<i128>>()),
keys(&[500, 3, 9000, 12, 7, 88, 41, 6]),
] {
let built = KeyMap::build(&column).expect("build");
let payload = encode(&built, INTEGER).expect("encode");
let short = &payload.bytes[..payload.bytes.len() - 1];
assert!(decode(short).is_err(), "a truncated {:?} body is refused", built.form());
}
}
#[test]
fn a_body_where_the_header_expects_none_is_refused() {
let built = KeyMap::build(&keys(&(1..=10).collect::<Vec<i128>>())).expect("build");
let mut payload = encode(&built, INTEGER).expect("encode");
assert_eq!(payload.bytes.len(), HEADER_BYTES);
payload.bytes.push(0);
assert!(decode(&payload.bytes).is_err());
}
}