use rudb_common::{Error, Result};
use rudb_encoding::bitpack;
use crate::bits::Rank;
use crate::rid::Rid;
pub const DENSE_THRESHOLD: u64 = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Form {
Identity,
Dense,
Sorted,
Permuted,
}
impl Form {
#[must_use]
pub fn tag(self) -> u8 {
match self {
Self::Identity => 0,
Self::Dense => 1,
Self::Sorted => 2,
Self::Permuted => 3,
}
}
#[must_use]
pub fn label(self) -> &'static str {
match self {
Self::Identity => "identity",
Self::Dense => "dense",
Self::Sorted => "sorted",
Self::Permuted => "permuted",
}
}
pub fn from_tag(tag: u8) -> Result<Self> {
match tag {
0 => Ok(Self::Identity),
1 => Ok(Self::Dense),
2 => Ok(Self::Sorted),
3 => Ok(Self::Permuted),
_ => Err(malformed(format!("key map form {tag} is not one this build knows"))),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Observed {
pub rows: u64,
pub nulls: u64,
pub distinct: bool,
pub sorted: bool,
pub min: Option<i128>,
pub max: Option<i128>,
}
impl Observed {
#[must_use]
pub fn usable_as_parent(&self) -> bool {
self.distinct
}
}
#[derive(Debug, Clone)]
enum Body {
Identity {
base: i128,
count: u64,
},
Dense {
base: i128,
range: u64,
bits: Vec<u64>,
rank: Rank,
},
Sorted {
base: i128,
key_width: usize,
keys: Vec<u8>,
rid_width: usize,
perm: Vec<u8>,
count: u64,
},
Permuted {
base: i128,
range: u64,
bits: Vec<u64>,
rank: Rank,
rid_width: usize,
perm: Vec<u8>,
},
}
#[derive(Debug, Clone)]
pub struct KeyMap {
body: Body,
observed: Observed,
}
impl KeyMap {
pub fn build(keys: &[Option<i128>]) -> Result<Self> {
let mut observed = observe(keys);
if !observed.distinct {
return Ok(Self { body: Body::Identity { base: 0, count: 0 }, observed });
}
Ok(match plan(&observed)? {
Plan::Empty => Self { body: Body::Identity { base: 0, count: 0 }, observed },
Plan::Identity { base, count } => {
Self { body: Body::Identity { base, count }, observed }
}
Plan::Dense { base, range } => Self { body: dense(keys, base, range)?, observed },
Plan::Permuted { base, range } => {
let mut bits = DenseBits::new(base, range);
for key in keys.iter().flatten() {
if !bits.mark(*key)? {
observed.distinct = false;
return Ok(Self { body: Body::Identity { base: 0, count: 0 }, observed });
}
}
let mut perm = Permutation::new(bits, keys.len())?;
for (rid, key) in keys.iter().enumerate() {
if let Some(key) = *key {
perm.place(key, rid)?;
}
}
Self { body: perm.finish()?, observed }
}
Plan::Sorted { base } => {
let (body, distinct) = sorted(keys, base, observed.rows)?;
observed.distinct = distinct;
Self { body, observed }
}
})
}
pub fn build_from<K: Keys + ?Sized>(keys: &K) -> Result<Self> {
let mut observer = Observer::new();
keys.scan(&mut |key| {
observer.push(key);
Ok(())
})?;
let mut observed = observer.observed;
if !observed.distinct {
return Ok(Self { body: Body::Identity { base: 0, count: 0 }, observed });
}
Ok(match plan(&observed)? {
Plan::Empty => Self { body: Body::Identity { base: 0, count: 0 }, observed },
Plan::Identity { base, count } => {
Self { body: Body::Identity { base, count }, observed }
}
Plan::Dense { base, range } => {
let mut bits = DenseBits::new(base, range);
keys.scan(&mut |key| match key {
Some(key) => bits.push(key),
None => Ok(()),
})?;
Self { body: bits.finish(), observed }
}
Plan::Permuted { base, range } => {
let mut bits = DenseBits::new(base, range);
let mut repeated = false;
keys.scan(&mut |key| {
if let Some(key) = key {
repeated |= !bits.mark(key)?;
}
Ok(())
})?;
if repeated {
observed.distinct = false;
return Ok(Self { body: Body::Identity { base: 0, count: 0 }, observed });
}
let rows = observed.rows + observed.nulls;
let rows = usize::try_from(rows)
.map_err(|_| malformed("the column is too long for this machine"))?;
let mut perm = Permutation::new(bits, rows)?;
let mut rid = 0_usize;
keys.scan(&mut |key| {
if let Some(key) = key {
perm.place(key, rid)?;
}
rid += 1;
Ok(())
})?;
Self { body: perm.finish()?, observed }
}
Plan::Sorted { base } => {
let mut held = Vec::with_capacity(
usize::try_from(observed.rows + observed.nulls).unwrap_or_default(),
);
keys.scan(&mut |key| {
held.push(key);
Ok(())
})?;
let (body, distinct) = sorted(&held, base, observed.rows)?;
observed.distinct = distinct;
Self { body, observed }
}
})
}
#[must_use]
pub fn form(&self) -> Form {
match self.body {
Body::Identity { .. } => Form::Identity,
Body::Dense { .. } => Form::Dense,
Body::Sorted { .. } => Form::Sorted,
Body::Permuted { .. } => Form::Permuted,
}
}
#[must_use]
pub fn span(&self) -> Option<(i128, u64)> {
match self.body {
Body::Identity { base, count } => Some((base, count)),
Body::Dense { base, range, .. } | Body::Permuted { base, range, .. } => {
Some((base, range))
}
Body::Sorted { .. } => None,
}
}
#[must_use]
pub fn observed(&self) -> &Observed {
&self.observed
}
pub(crate) fn base(&self) -> i128 {
match &self.body {
Body::Identity { base, .. }
| Body::Dense { base, .. }
| Body::Sorted { base, .. }
| Body::Permuted { base, .. } => *base,
}
}
pub(crate) fn write_body(&self, out: &mut Vec<u8>) -> Result<()> {
match &self.body {
Body::Identity { .. } => Ok(()),
Body::Dense { range, bits, rank, .. } => {
out.extend_from_slice(&range.to_le_bytes());
for word in bits {
out.extend_from_slice(&word.to_le_bytes());
}
rank.write(out);
Ok(())
}
Body::Sorted { key_width, keys, rid_width, perm, .. } => {
let widths = [*key_width, *rid_width];
for width in widths {
let width = u8::try_from(width)
.map_err(|_| malformed("a sorted key map's width does not fit a byte"))?;
out.push(width);
}
out.extend_from_slice(keys);
out.extend_from_slice(perm);
Ok(())
}
Body::Permuted { range, bits, rank, rid_width, perm, .. } => {
out.extend_from_slice(&range.to_le_bytes());
for word in bits {
out.extend_from_slice(&word.to_le_bytes());
}
rank.write(out);
let width = u8::try_from(*rid_width)
.map_err(|_| malformed("a permuted key map's width does not fit a byte"))?;
out.push(width);
out.extend_from_slice(perm);
Ok(())
}
}
}
pub(crate) fn read_body(
form: Form,
base: i128,
mut observed: Observed,
body: &[u8],
) -> Result<Self> {
match form {
Form::Identity => {
if !body.is_empty() {
return Err(malformed("an identity key map has no body"));
}
if observed.rows > 0 {
observed.max = Some(
base.checked_add(i128::from(observed.rows) - 1)
.ok_or_else(|| malformed("an identity key map's range overflows"))?,
);
}
Ok(Self { body: Body::Identity { base, count: observed.rows }, observed })
}
Form::Dense => {
let Some(head) = body.get(..size_of::<u64>()) else {
return Err(malformed("a dense key map has no range"));
};
let range = u64::from_le_bytes(head.try_into().expect("eight bytes"));
let Ok(range_usize) = usize::try_from(range) else {
return Err(malformed("a dense key map's range does not fit this machine"));
};
let words = range_usize.div_ceil(64);
let bitmap = words * size_of::<u64>();
let rest = &body[size_of::<u64>()..];
if rest.len() < bitmap {
return Err(malformed("a dense key map's bitmap is shorter than its range"));
}
let bits: Vec<u64> = rest[..bitmap]
.chunks_exact(size_of::<u64>())
.map(|word| u64::from_le_bytes(word.try_into().expect("eight bytes")))
.collect();
let rank = Rank::read(&rest[bitmap..], words)?;
observed.max = Some(
base.checked_add(i128::from(range) - 1)
.ok_or_else(|| malformed("a dense key map's range overflows"))?,
);
Ok(Self { body: Body::Dense { base, range, bits, rank }, observed })
}
Form::Sorted => {
if body.len() < 2 {
return Err(malformed("a sorted key map has no widths"));
}
let key_width = usize::from(body[0]);
let rid_width = usize::from(body[1]);
if key_width == 0 || key_width > 64 || rid_width == 0 || rid_width > 64 {
return Err(malformed("a sorted key map's width is not one a u64 can take"));
}
let count = observed.rows;
let Ok(count_usize) = usize::try_from(count) else {
return Err(malformed(
"a sorted key map holds more keys than this machine can",
));
};
let key_bytes = (count_usize * key_width).div_ceil(8);
let perm_bytes = (count_usize * rid_width).div_ceil(8);
let rest = &body[2..];
if rest.len() != key_bytes + perm_bytes {
return Err(malformed(
"a sorted key map's arrays are not the size its widths and count imply",
));
}
let keys = rest[..key_bytes].to_vec();
let perm = rest[key_bytes..].to_vec();
if count > 0 {
let largest = bitpack::tail_at(&keys, key_width, count_usize - 1)?;
observed.max =
Some(base.checked_add(i128::from(largest)).ok_or_else(|| {
malformed("a sorted key map's largest key overflows")
})?);
}
Ok(Self {
body: Body::Sorted { base, key_width, keys, rid_width, perm, count },
observed,
})
}
Form::Permuted => {
let Some(head) = body.get(..size_of::<u64>()) else {
return Err(malformed("a permuted key map has no range"));
};
let range = u64::from_le_bytes(head.try_into().expect("eight bytes"));
let Ok(range_usize) = usize::try_from(range) else {
return Err(malformed("a permuted key map's range does not fit this machine"));
};
let words = range_usize.div_ceil(64);
let bitmap = words * size_of::<u64>();
let (blocks, superblocks) = Rank::shape(words);
let ranks = superblocks * size_of::<u32>() + blocks * size_of::<u16>();
let rest = &body[size_of::<u64>()..];
if rest.len() < bitmap + ranks + 1 {
return Err(malformed("a permuted key map is shorter than its range implies"));
}
let bits: Vec<u64> = rest[..bitmap]
.chunks_exact(size_of::<u64>())
.map(|word| u64::from_le_bytes(word.try_into().expect("eight bytes")))
.collect();
let rank = Rank::read(&rest[bitmap..bitmap + ranks], words)?;
let rid_width = usize::from(rest[bitmap + ranks]);
if rid_width == 0 || rid_width > 64 {
return Err(malformed("a permuted key map's width is not one a u64 can take"));
}
let Ok(count) = usize::try_from(observed.rows) else {
return Err(malformed(
"a permuted key map holds more keys than this machine can",
));
};
let perm = rest[bitmap + ranks + 1..].to_vec();
if perm.len() != (count * rid_width).div_ceil(8) {
return Err(malformed(
"a permuted key map's permutation is not the size its width and count imply",
));
}
observed.max = Some(
base.checked_add(i128::from(range) - 1)
.ok_or_else(|| malformed("a permuted key map's range overflows"))?,
);
Ok(Self {
body: Body::Permuted { base, range, bits, rank, rid_width, perm },
observed,
})
}
}
}
#[must_use]
pub fn len(&self) -> u64 {
match &self.body {
Body::Identity { count, .. } | Body::Sorted { count, .. } => *count,
Body::Dense { .. } | Body::Permuted { .. } => self.observed.rows,
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn bytes(&self) -> usize {
match &self.body {
Body::Identity { .. } => size_of::<i128>() + size_of::<u64>(),
Body::Dense { bits, rank, .. } => bits.len() * size_of::<u64>() + rank.bytes(),
Body::Sorted { keys, perm, .. } => keys.len() + perm.len(),
Body::Permuted { bits, rank, perm, .. } => {
bits.len() * size_of::<u64>() + rank.bytes() + perm.len()
}
}
}
pub fn rows_of_span(&self, held: &[u64], rows: u64) -> Result<Option<Vec<u64>>> {
let words = usize::try_from(rows.div_ceil(64)).unwrap_or(usize::MAX);
let Some((_, range)) = self.span() else { return Ok(None) };
let span_words = usize::try_from(range.div_ceil(64)).unwrap_or(usize::MAX);
if held.len() < span_words || held[span_words..].iter().any(|&word| word != 0) {
return Ok(None);
}
if range % 64 != 0 && held[span_words - 1] >> (range % 64) != 0 {
return Ok(None);
}
let held = &held[..span_words];
let mut out = vec![0_u64; words];
match &self.body {
Body::Identity { count, .. } => {
if *count > rows {
return Ok(None);
}
out[..span_words].copy_from_slice(held);
}
Body::Dense { bits, .. } => {
if bits.len() < held.len() {
return Ok(None);
}
let mut before = 0_u64;
for (&keys, &present) in held.iter().zip(bits) {
if keys & !present != 0 {
return Ok(None);
}
let mut left = keys;
while left != 0 {
let below = (1_u64 << left.trailing_zeros()) - 1;
let rid = before + u64::from((present & below).count_ones());
let Some(word) = out.get_mut((rid / 64) as usize) else { return Ok(None) };
*word |= 1 << (rid % 64);
left &= left - 1;
}
before += u64::from(present.count_ones());
}
}
Body::Permuted { bits, rid_width, perm, .. } => {
if bits.len() < held.len() {
return Ok(None);
}
let mut before = 0_u64;
for (&keys, &present) in held.iter().zip(bits) {
if keys & !present != 0 {
return Ok(None);
}
let mut left = keys;
while left != 0 {
let below = (1_u64 << left.trailing_zeros()) - 1;
let place = before + u64::from((present & below).count_ones());
#[expect(
clippy::cast_possible_truncation,
reason = "a rank is below the key count, which the build checked fits a usize"
)]
let rid = bitpack::tail_at(perm, *rid_width, place as usize)?;
let Some(word) = out.get_mut((rid / 64) as usize) else { return Ok(None) };
*word |= 1 << (rid % 64);
left &= left - 1;
}
before += u64::from(present.count_ones());
}
}
Body::Sorted { .. } => return Ok(None),
}
if !rows.is_multiple_of(64) && out.last().is_some_and(|&word| word >> (rows % 64) != 0) {
return Ok(None);
}
Ok(Some(out))
}
pub fn lookup(&self, key: i128) -> Result<Option<Rid>> {
match &self.body {
Body::Identity { base, count } => {
let Some(offset) = key.checked_sub(*base) else {
return Ok(None);
};
match u64::try_from(offset) {
Ok(rid) if rid < *count => Ok(Some(rid)),
_ => Ok(None),
}
}
Body::Dense { base, range, bits, rank } => {
let Some(offset) = key.checked_sub(*base) else {
return Ok(None);
};
let Ok(offset) = u64::try_from(offset) else {
return Ok(None);
};
if offset >= *range {
return Ok(None);
}
#[expect(
clippy::cast_possible_truncation,
reason = "the build checked the range fits a usize"
)]
let at = offset as usize;
if bits[at / 64] >> (at % 64) & 1 == 0 {
return Ok(None);
}
Ok(Some(rank.rank(bits, at)))
}
Body::Permuted { base, range, bits, rank, rid_width, perm } => {
let Some(offset) = key.checked_sub(*base) else {
return Ok(None);
};
let Ok(offset) = u64::try_from(offset) else {
return Ok(None);
};
if offset >= *range {
return Ok(None);
}
#[expect(
clippy::cast_possible_truncation,
reason = "the build checked the range fits a usize"
)]
let at = offset as usize;
if bits[at / 64] >> (at % 64) & 1 == 0 {
return Ok(None);
}
#[expect(
clippy::cast_possible_truncation,
reason = "a rank is below the key count, which the build checked fits a usize"
)]
let place = rank.rank(bits, at) as usize;
Ok(Some(bitpack::tail_at(perm, *rid_width, place)?))
}
Body::Sorted { base, key_width, keys, rid_width, perm, count } => {
let Some(offset) = key.checked_sub(*base) else {
return Ok(None);
};
let Ok(wanted) = u64::try_from(offset) else {
return Ok(None);
};
#[expect(
clippy::cast_possible_truncation,
reason = "the build refused a column wider than a usize of rows"
)]
let len = *count as usize;
let mut low = 0_usize;
let mut high = len;
while low < high {
let mid = low + (high - low) / 2;
let at = bitpack::tail_at(keys, *key_width, mid)?;
if at < wanted {
low = mid + 1;
} else {
high = mid;
}
}
if low >= len || bitpack::tail_at(keys, *key_width, low)? != wanted {
return Ok(None);
}
Ok(Some(bitpack::tail_at(perm, *rid_width, low)?))
}
}
}
}
pub trait Keys {
fn scan(&self, each: &mut dyn FnMut(Option<i128>) -> Result<()>) -> Result<()>;
}
impl Keys for [Option<i128>] {
fn scan(&self, each: &mut dyn FnMut(Option<i128>) -> Result<()>) -> Result<()> {
for key in self {
each(*key)?;
}
Ok(())
}
}
enum Plan {
Empty,
Identity { base: i128, count: u64 },
Dense { base: i128, range: u64 },
Permuted { base: i128, range: u64 },
Sorted { base: i128 },
}
fn plan(observed: &Observed) -> Result<Plan> {
if observed.rows == 0 {
return Ok(Plan::Empty);
}
let (Some(min), Some(max)) = (observed.min, observed.max) else {
return Err(malformed("a column with keys in it reported no minimum"));
};
let range = range_of(min, max)?;
let positional = observed.distinct && observed.sorted && observed.nulls == 0;
if positional && range == observed.rows {
return Ok(Plan::Identity { base: min, count: observed.rows });
}
let compact = usize::try_from(range).is_ok() && range / observed.rows < DENSE_THRESHOLD;
if positional && compact {
return Ok(Plan::Dense { base: min, range });
}
if observed.distinct && compact {
return Ok(Plan::Permuted { base: min, range });
}
Ok(Plan::Sorted { base: min })
}
struct Observer {
observed: Observed,
previous: Option<i128>,
}
impl Observer {
fn new() -> Self {
Self {
observed: Observed {
rows: 0,
nulls: 0,
distinct: true,
sorted: true,
min: None,
max: None,
},
previous: None,
}
}
fn push(&mut self, key: Option<i128>) {
let Some(key) = key else {
self.observed.nulls += 1;
return;
};
self.observed.rows += 1;
self.observed.min = Some(self.observed.min.map_or(key, |held| held.min(key)));
self.observed.max = Some(self.observed.max.map_or(key, |held| held.max(key)));
if let Some(previous) = self.previous {
if key < previous {
self.observed.sorted = false;
} else if key == previous {
self.observed.distinct = false;
}
}
self.previous = Some(key);
}
}
fn observe(keys: &[Option<i128>]) -> Observed {
let mut observer = Observer::new();
for key in keys {
observer.push(*key);
}
observer.observed
}
fn range_of(min: i128, max: i128) -> Result<u64> {
let span = max.wrapping_sub(min) as u128;
u64::try_from(span)
.ok()
.and_then(|span| span.checked_add(1))
.ok_or_else(|| malformed("the key column spans more than a u64 of values"))
}
fn offset_of(key: i128, base: i128) -> Result<u64> {
let offset = key
.checked_sub(base)
.ok_or_else(|| malformed("a key is further from the base than an i128 holds"))?;
u64::try_from(offset)
.map_err(|_| malformed("a key is below the base or further from it than a u64 holds"))
}
struct DenseBits {
base: i128,
range: u64,
bits: Vec<u64>,
previous: Option<i128>,
}
impl DenseBits {
fn new(base: i128, range: u64) -> Self {
#[expect(
clippy::cast_possible_truncation,
reason = "the caller checked the range fits a usize"
)]
let range_usize = range as usize;
Self { base, range, bits: vec![0_u64; range_usize.div_ceil(64)], previous: None }
}
fn push(&mut self, key: i128) -> Result<()> {
debug_assert!(
self.previous.is_none_or(|held| key > held),
"the bitmap form needs a distinct ascending column, because a rank is a count of keys below a value and that is a rid only there"
);
self.previous = Some(key);
let offset = offset_of(key, self.base)?;
#[expect(
clippy::cast_possible_truncation,
reason = "the caller checked the range fits a usize and the offset is inside it"
)]
let at = offset as usize;
self.bits[at / 64] |= 1 << (at % 64);
Ok(())
}
fn finish(self) -> Body {
let rank = Rank::build(&self.bits);
Body::Dense { base: self.base, range: self.range, bits: self.bits, rank }
}
fn mark(&mut self, key: i128) -> Result<bool> {
let offset = offset_of(key, self.base)?;
#[expect(
clippy::cast_possible_truncation,
reason = "the caller checked the range fits a usize and the offset is inside it"
)]
let at = offset as usize;
let bit = 1 << (at % 64);
let fresh = self.bits[at / 64] & bit == 0;
self.bits[at / 64] |= bit;
Ok(fresh)
}
}
struct Permutation {
base: i128,
range: u64,
bits: Vec<u64>,
rank: Rank,
places: Vec<u64>,
rid_width: usize,
}
impl Permutation {
fn new(bits: DenseBits, rows: usize) -> Result<Self> {
let keys: u64 = bits.bits.iter().map(|word| u64::from(word.count_ones())).sum();
let keys =
usize::try_from(keys).map_err(|_| malformed("too many keys for this machine"))?;
let largest = u64::try_from(rows.saturating_sub(1))
.map_err(|_| malformed("the column is too long for a rid"))?;
let rank = Rank::build(&bits.bits);
Ok(Self {
base: bits.base,
range: bits.range,
bits: bits.bits,
rank,
places: vec![0; keys],
rid_width: width_for(largest),
})
}
fn place(&mut self, key: i128, rid: usize) -> Result<()> {
let offset = offset_of(key, self.base)?;
#[expect(
clippy::cast_possible_truncation,
reason = "the offset is inside a range the caller checked fits a usize"
)]
let at = offset as usize;
let place = usize::try_from(self.rank.rank(&self.bits, at))
.map_err(|_| malformed("a rank past what this machine can index"))?;
let rid = u64::try_from(rid).map_err(|_| malformed("the column is too long for a rid"))?;
let slot = self
.places
.get_mut(place)
.ok_or_else(|| malformed("a key was placed that the bitmap does not hold"))?;
*slot = rid;
Ok(())
}
fn finish(self) -> Result<Body> {
let mut perm = Vec::new();
bitpack::pack_linear(&self.places, self.rid_width, &mut perm)?;
Ok(Body::Permuted {
base: self.base,
range: self.range,
bits: self.bits,
rank: self.rank,
rid_width: self.rid_width,
perm,
})
}
}
fn dense(keys: &[Option<i128>], base: i128, range: u64) -> Result<Body> {
let mut bits = DenseBits::new(base, range);
for key in keys.iter().flatten() {
bits.push(*key)?;
}
Ok(bits.finish())
}
fn sorted(keys: &[Option<i128>], base: i128, rows: u64) -> Result<(Body, bool)> {
let mut pairs: Vec<(u64, u64)> = Vec::with_capacity(keys.len());
for (rid, key) in keys.iter().enumerate() {
let Some(key) = *key else { continue };
let offset = offset_of(key, base)?;
let rid = u64::try_from(rid).map_err(|_| malformed("the column is too long for a rid"))?;
pairs.push((offset, rid));
}
pairs.sort_unstable();
let distinct = pairs.windows(2).all(|pair| pair[0].0 != pair[1].0);
debug_assert_eq!(
u64::try_from(pairs.len()).ok(),
Some(rows),
"the pair list is the non-null column"
);
let key_width = width_for(pairs.last().map_or(0, |pair| pair.0));
let rows_width = u64::try_from(keys.len().saturating_sub(1))
.map_err(|_| malformed("the column is too long for a rid"))?;
let rid_width = width_for(rows_width);
let mut key_bytes = Vec::new();
let mut rid_bytes = Vec::new();
let key_values: Vec<u64> = pairs.iter().map(|pair| pair.0).collect();
let rid_values: Vec<u64> = pairs.iter().map(|pair| pair.1).collect();
bitpack::pack_linear(&key_values, key_width, &mut key_bytes)?;
bitpack::pack_linear(&rid_values, rid_width, &mut rid_bytes)?;
Ok((
Body::Sorted { base, key_width, keys: key_bytes, rid_width, perm: rid_bytes, count: rows },
distinct,
))
}
fn width_for(largest: u64) -> usize {
let bits = u64::BITS - largest.leading_zeros();
bits.max(1) as usize
}
fn malformed(message: impl Into<String>) -> Error {
Error::invalid_input(format!("invalid rudb key map: {}", message.into()))
}
#[cfg(test)]
mod tests {
use super::*;
fn keys(values: &[i128]) -> Vec<Option<i128>> {
values.iter().copied().map(Some).collect()
}
fn resolves(column: &[Option<i128>], map: &KeyMap) {
for (rid, key) in column.iter().enumerate() {
let Some(key) = *key else { continue };
let found = map.lookup(key).expect("lookup").expect("a key in the column resolves");
assert_eq!(found, rid as u64, "key {key} resolved to {found} rather than {rid}");
}
}
#[test]
fn a_sequence_from_one_is_the_identity_form_and_stores_two_numbers() {
let column = keys(&(1..=1000).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Identity);
assert_eq!(map.bytes(), 24, "section 4.2 says identity is twenty four bytes");
assert_eq!(map.len(), 1000);
resolves(&column, &map);
assert_eq!(map.lookup(0).expect("lookup"), None, "below the base");
assert_eq!(map.lookup(1001).expect("lookup"), None, "past the end");
assert_eq!(map.span(), Some((1, 1000)));
}
#[test]
fn the_rows_of_a_span_are_the_rows_a_lookup_of_each_key_finds() {
let identity = keys(&(5..1005).collect::<Vec<i128>>());
let dense = keys(&(0..3000).filter(|value| value % 3 != 1).collect::<Vec<i128>>());
let mut shuffled = (0..3000).filter(|value| value % 3 != 1).collect::<Vec<i128>>();
shuffled.sort_by_key(|value| (value * 7919) % 3001);
let permuted = keys(&shuffled);
for (column, form) in
[(identity, Form::Identity), (dense, Form::Dense), (permuted, Form::Permuted)]
{
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), form);
let rows = column.len() as u64;
let (base, range) = map.span().expect("a span");
let mut held = vec![0_u64; (range / 64 + 1) as usize];
let mut wanted = vec![0_u64; rows.div_ceil(64) as usize];
for key in column.iter().flatten().filter(|key| *key % 5 == 0 || *key % 7 == 3) {
let offset = (key - base) as u64;
held[(offset / 64) as usize] |= 1 << (offset % 64);
let rid = map.lookup(*key).expect("lookup").expect("a key in the column");
wanted[(rid / 64) as usize] |= 1 << (rid % 64);
}
assert_eq!(map.rows_of_span(&held, rows).expect("rows"), Some(wanted), "{form:?}");
if form != Form::Identity {
let missing = (0..range).find(|offset| {
map.lookup(base + i128::from(*offset)).expect("lookup").is_none()
});
let offset = missing.expect("a hole in the span");
held[(offset / 64) as usize] |= 1 << (offset % 64);
assert_eq!(map.rows_of_span(&held, rows).expect("rows"), None, "{form:?}");
}
let last = held.len() - 1;
held[last] |= 1 << 63;
assert_eq!(map.rows_of_span(&held, rows).expect("rows"), None, "past the span");
}
let sorted = keys(&(0..1000).map(|value| value * 1000).collect::<Vec<i128>>());
let map = KeyMap::build(&sorted).expect("build");
assert_eq!(map.rows_of_span(&[0; 4], 1000).expect("rows"), None, "no span");
}
#[test]
fn a_sequence_from_zero_is_also_the_identity_form() {
let column = keys(&(0..64).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Identity);
resolves(&column, &map);
}
#[test]
fn a_sequence_with_a_gap_in_it_is_the_dense_form() {
let column = keys(&(0..1000).map(|value| value * 2).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Dense);
resolves(&column, &map);
assert_eq!(
map.lookup(1).expect("lookup"),
None,
"a value in the range and not in the column"
);
assert_eq!(map.lookup(2001).expect("lookup"), None, "past the range");
assert_eq!(map.span(), Some((0, 1999)), "from the smallest key to the largest");
}
#[test]
fn a_range_too_sparse_for_a_bitmap_is_the_sorted_form() {
let column = keys(&(0..1000).map(|value| value * 1000).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Sorted);
resolves(&column, &map);
assert_eq!(map.lookup(500).expect("lookup"), None);
assert_eq!(map.span(), None, "too sparse for a bitmap over the range");
}
#[test]
fn the_sorted_form_is_not_bounded_by_a_packed_unit() {
let column = keys(&(0..5000).map(|value| (value * 7919) % 100_003).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Sorted);
resolves(&column, &map);
}
#[test]
fn keys_in_no_order_at_all_resolve_to_the_rows_that_hold_them() {
let column = keys(&[500, 3, 9000, 12, 7, 88, 41, 6]);
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Sorted);
resolves(&column, &map);
}
#[test]
fn a_descending_column_dense_enough_for_a_bitmap_still_resolves_correctly() {
let column = keys(&(0..500).rev().collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Permuted, "a descending column cannot take the bare bitmap");
resolves(&column, &map);
}
#[test]
fn nulls_are_not_keys_and_do_not_shift_the_rows_around_them() {
let column = vec![Some(10), None, Some(20), None, Some(30)];
let map = KeyMap::build(&column).expect("build");
assert_eq!(
map.form(),
Form::Permuted,
"a null before a key shifts it out of the positional forms"
);
resolves(&column, &map);
assert_eq!(map.observed().nulls, 2);
assert_eq!(map.observed().rows, 3);
assert_eq!(
map.lookup(20).expect("lookup"),
Some(2),
"the rid is the position in the column"
);
}
#[test]
fn a_leading_null_keeps_an_otherwise_perfect_sequence_out_of_the_identity_form() {
let mut column = vec![None];
column.extend((1..=1000).map(Some));
let map = KeyMap::build(&column).expect("build");
assert_ne!(map.form(), Form::Identity);
resolves(&column, &map);
assert_eq!(map.lookup(1).expect("lookup"), Some(1), "row zero is the null, not key one");
}
#[test]
fn a_null_only_column_builds_and_resolves_nothing() {
let column = vec![None, None, None];
let map = KeyMap::build(&column).expect("build");
assert!(map.is_empty());
assert_eq!(map.observed().nulls, 3);
assert_eq!(map.lookup(0).expect("lookup"), None);
}
#[test]
fn an_empty_column_builds_and_resolves_nothing() {
let map = KeyMap::build(&[]).expect("build");
assert!(map.is_empty());
assert_eq!(map.lookup(0).expect("lookup"), None);
assert!(map.observed().usable_as_parent(), "an empty parent is unique, vacuously");
}
#[test]
fn a_duplicated_key_is_reported_rather_than_resolved_to_one_of_its_rows() {
let column = keys(&[5, 7, 5, 9]);
let map = KeyMap::build(&column).expect("build");
assert!(!map.observed().distinct);
assert!(!map.observed().usable_as_parent(), "a non-unique parent side takes no link");
}
#[test]
fn a_column_that_arrives_with_its_repeats_together_is_not_sorted_into_a_map() {
let column = keys(&[1, 1, 2, 2, 2, 90_000, 90_000]);
let map = KeyMap::build_from(&column[..]).expect("build");
assert!(!map.observed().distinct);
assert_eq!(map.observed().rows, 7, "the column was still counted");
assert_eq!(map.observed().max, Some(90_000));
assert_eq!(map.bytes(), KeyMap::build(&keys(&[])).expect("build").bytes());
assert_eq!(map.lookup(2).expect("lookup"), None, "and it answers nothing, as it must");
}
#[test]
fn a_single_key_column_resolves_it() {
let column = keys(&[42]);
let map = KeyMap::build(&column).expect("build");
resolves(&column, &map);
assert_eq!(map.lookup(41).expect("lookup"), None);
assert_eq!(map.lookup(43).expect("lookup"), None);
}
#[test]
fn negative_keys_resolve_because_the_base_is_the_minimum_and_not_zero() {
let column = keys(&[-9000, -3, -1, 0, 7]);
let map = KeyMap::build(&column).expect("build");
resolves(&column, &map);
assert_eq!(map.lookup(-9001).expect("lookup"), None);
}
#[test]
fn a_column_spanning_more_than_a_u64_of_values_is_refused_and_not_panicked_over() {
let column = keys(&[i128::MIN, 0, i128::MAX]);
let error = KeyMap::build(&column).expect_err("refused");
assert!(error.to_string().contains("spans more than a u64"), "{error}");
}
#[test]
fn keys_at_the_far_end_of_the_integer_type_resolve_when_their_range_is_narrow() {
let column = keys(&[i128::MIN, i128::MIN + 5, i128::MIN + 2]);
let map = KeyMap::build(&column).expect("build");
resolves(&column, &map);
assert_eq!(map.lookup(i128::MAX).expect("lookup"), None);
assert_eq!(map.lookup(0).expect("lookup"), None);
}
#[test]
fn the_rank_index_agrees_with_counting_the_bits_by_hand() {
let column = keys(&(0..10_000).map(|value| value * 2).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Dense);
resolves(&column, &map);
}
#[test]
fn a_string_key_arrives_as_dictionary_codes_and_never_as_text() {
let codes = keys(&[7, 1, 4, 9, 2]);
let map = KeyMap::build(&codes).expect("build");
resolves(&codes, &map);
}
#[test]
fn the_form_tag_round_trips_and_an_unknown_one_is_refused() {
for form in [Form::Identity, Form::Dense, Form::Sorted, Form::Permuted] {
assert_eq!(Form::from_tag(form.tag()).expect("a known tag"), form);
}
assert!(Form::from_tag(4).is_err(), "an unfamiliar form is refused rather than guessed");
}
#[test]
fn the_dense_form_costs_a_bitmap_and_about_an_eighth_again() {
let column = keys(&(0..10_000).map(|value| value * 8).collect::<Vec<i128>>());
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Dense);
let bitmap = 80_000 / 8;
let bytes = map.bytes();
assert!(bytes > bitmap, "the map took {bytes} bytes and the bitmap alone is {bitmap}");
assert!(
bytes < bitmap * 5 / 4,
"the map took {bytes} bytes, more than a quarter over the bitmap's {bitmap}"
);
}
struct Counted {
column: Vec<Option<i128>>,
scans: std::cell::Cell<usize>,
}
impl Keys for Counted {
fn scan(&self, each: &mut dyn FnMut(Option<i128>) -> Result<()>) -> Result<()> {
self.scans.set(self.scans.get() + 1);
self.column.scan(each)
}
}
#[test]
fn a_build_from_a_scan_is_the_same_map_as_a_build_from_a_slice() {
let columns: Vec<Vec<Option<i128>>> = vec![
Vec::new(),
keys(&[]),
keys(&(1..=1000).collect::<Vec<i128>>()),
keys(&(0..500).map(|value| value * 4).collect::<Vec<i128>>()),
keys(&[100, 3, 40, 7, 9000]),
keys(&[5, 5, 9]),
vec![Some(10), None, Some(20), None, Some(30)],
vec![None, None],
shuffled(1_000, 4),
keys(&[1, 3, 2, 1]),
];
for column in &columns {
let held = KeyMap::build(column).expect("build from a slice");
let read = KeyMap::build_from(&column[..]).expect("build from a scan");
assert_eq!(read.form(), held.form(), "{column:?}");
assert_eq!(read.observed(), held.observed(), "{column:?}");
assert_eq!(read.len(), held.len(), "{column:?}");
assert_eq!(read.bytes(), held.bytes(), "{column:?}");
if read.observed().usable_as_parent() {
resolves(column, &read);
}
}
}
#[test]
fn the_identity_form_is_built_without_reading_the_column_twice() {
let identity =
Counted { column: keys(&(1..=1000).collect::<Vec<i128>>()), scans: 0.into() };
assert_eq!(KeyMap::build_from(&identity).expect("build").form(), Form::Identity);
assert_eq!(
identity.scans.get(),
1,
"the identity form is the observation and nothing more"
);
let dense = Counted {
column: keys(&(0..500).map(|v| v * 4).collect::<Vec<i128>>()),
scans: 0.into(),
};
assert_eq!(KeyMap::build_from(&dense).expect("build").form(), Form::Dense);
assert_eq!(dense.scans.get(), 2);
let sorted = Counted { column: keys(&[100, 3, 40, 7, 9000]), scans: 0.into() };
assert_eq!(KeyMap::build_from(&sorted).expect("build").form(), Form::Sorted);
assert_eq!(sorted.scans.get(), 2);
let permuted = Counted { column: shuffled(1_000, 4), scans: 0.into() };
assert_eq!(KeyMap::build_from(&permuted).expect("build").form(), Form::Permuted);
assert_eq!(permuted.scans.get(), 3);
}
fn shuffled(count: i128, step: i128) -> Vec<Option<i128>> {
(0..count).map(|at| Some((at * 7_919 % count) * step)).collect()
}
#[test]
fn dense_keys_stored_out_of_key_order_take_the_permuted_form() {
let column = shuffled(10_000, 4);
let map = KeyMap::build(&column).expect("build");
assert_eq!(map.form(), Form::Permuted);
assert!(map.observed().distinct);
assert!(!map.observed().sorted);
resolves(&column, &map);
assert_eq!(map.lookup(1).expect("lookup"), None, "a key between two keys is not a key");
assert_eq!(map.lookup(40_000).expect("lookup"), None, "past the end");
assert_eq!(map.span(), Some((0, 39_997)), "the span is about the keys and not the rows");
let sorted = 10_000 * (16 + 14) / 8;
assert!(map.bytes() * 10 < sorted * 7, "{} bytes against {sorted}", map.bytes());
}
#[test]
fn a_repeat_that_is_not_adjacent_keeps_a_dense_column_out_of_every_form() {
let column = keys(&[4, 1, 3, 2, 4]);
let map = KeyMap::build(&column).expect("build");
assert!(!map.observed().distinct);
assert!(!map.observed().usable_as_parent());
assert_eq!(map.lookup(4).expect("lookup"), None);
let read = KeyMap::build_from(&column[..]).expect("build");
assert_eq!(read.observed(), map.observed());
}
#[test]
fn a_scan_that_fails_stops_the_build_rather_than_half_finishing_it() {
struct Broken;
impl Keys for Broken {
fn scan(&self, _: &mut dyn FnMut(Option<i128>) -> Result<()>) -> Result<()> {
Err(malformed("the column could not be read"))
}
}
let error = KeyMap::build_from(&Broken).expect_err("a build over an unreadable column");
assert!(error.to_string().contains("could not be read"), "{error}");
}
}