use rudb_common::{Error, Result};
use crate::reader::Reader;
use crate::bitpack::{self, VALUES};
const MAX_DEPTH: u8 = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Kind {
Constant = 0,
Packed = 1,
Delta = 2,
Rle = 3,
Dict = 4,
Sparse = 5,
}
impl Kind {
fn tag(self) -> u8 {
self as u8
}
fn from_tag(tag: u8) -> Result<Self> {
match tag {
0 => Ok(Self::Constant),
1 => Ok(Self::Packed),
2 => Ok(Self::Delta),
3 => Ok(Self::Rle),
4 => Ok(Self::Dict),
5 => Ok(Self::Sparse),
other => Err(Error::internal(format!("unknown encoding tag {other}"))),
}
}
#[must_use]
pub fn name(self) -> &'static str {
match self {
Self::Constant => "CONSTANT",
Self::Packed => "FOR+BITPACK",
Self::Delta => "DELTA",
Self::Rle => "RLE",
Self::Dict => "DICT",
Self::Sparse => "SPARSE",
}
}
}
pub fn encode(values: &[i64]) -> Result<Vec<u8>> {
encode_at(values, 0)
}
pub fn decode(bytes: &[u8]) -> Result<Vec<i64>> {
let mut reader = Reader::new(bytes);
let values = decode_chunk(&mut reader)?;
if reader.remaining() != 0 {
return Err(Error::internal(format!(
"{} bytes left over after decoding a chunk",
reader.remaining()
)));
}
Ok(values)
}
pub fn decode_prefix(bytes: &[u8]) -> Result<(Vec<i64>, usize)> {
let mut reader = Reader::new(bytes);
let values = decode_chunk(&mut reader)?;
Ok((values, reader.used()))
}
pub fn describe_prefix(bytes: &[u8]) -> Result<(String, usize)> {
let mut reader = Reader::new(bytes);
let text = describe_chunk(&mut reader)?;
Ok((text, reader.used()))
}
pub fn candidate_sizes(values: &[i64]) -> Result<Vec<(Kind, usize)>> {
let mut sizes = Vec::new();
for kind in candidates(values, 0) {
if let Some(bytes) = encode_as(kind, values, 0)? {
sizes.push((kind, bytes.len()));
}
}
Ok(sizes)
}
pub fn describe(bytes: &[u8]) -> Result<String> {
let mut reader = Reader::new(bytes);
describe_chunk(&mut reader)
}
fn encode_at(values: &[i64], depth: u8) -> Result<Vec<u8>> {
let mut best: Option<Vec<u8>> = None;
for kind in candidates(values, depth) {
let Some(bytes) = encode_as(kind, values, depth)? else {
continue;
};
if best.as_ref().is_none_or(|current| bytes.len() < current.len()) {
best = Some(bytes);
}
}
best.ok_or_else(|| Error::internal("no encoding applied to the chunk"))
}
fn candidates(values: &[i64], depth: u8) -> Vec<Kind> {
let mut kinds = vec![Kind::Packed];
if depth >= MAX_DEPTH || values.is_empty() {
return kinds;
}
if values.iter().all(|value| *value == values[0]) {
return vec![Kind::Constant];
}
if values.len() >= 2 && deltas(values).is_some() {
kinds.push(Kind::Delta);
}
if run_count(values) * 4 <= values.len() * 3 {
kinds.push(Kind::Rle);
}
let distinct = distinct_values(values);
if distinct.len() * 2 <= values.len() {
kinds.push(Kind::Dict);
}
match dominant_value(values) {
Some((_, count)) if count * 10 >= values.len() * 8 => kinds.push(Kind::Sparse),
_ => {}
}
kinds
}
fn encode_as(kind: Kind, values: &[i64], depth: u8) -> Result<Option<Vec<u8>>> {
let mut out = Vec::new();
put_u8(&mut out, kind.tag());
put_u32(&mut out, u32::try_from(values.len()).map_err(|_| too_long(values.len()))?);
match kind {
Kind::Constant => {
let Some(first) = values.first() else {
return Ok(None);
};
if values.iter().any(|value| value != first) {
return Ok(None);
}
put_i64(&mut out, *first);
}
Kind::Packed => encode_packed(values, &mut out)?,
Kind::Delta => {
let Some(deltas) = deltas(values) else {
return Ok(None);
};
put_i64(&mut out, values[0]);
out.extend_from_slice(&encode_at(&deltas, depth + 1)?);
}
Kind::Rle => {
let (run_values, run_lengths) = runs(values);
if run_values.is_empty() {
return Ok(None);
}
out.extend_from_slice(&encode_at(&run_values, depth + 1)?);
out.extend_from_slice(&encode_at(&run_lengths, depth + 1)?);
}
Kind::Dict => {
let dictionary = distinct_values(values);
if dictionary.is_empty() {
return Ok(None);
}
let codes = codes_over(values, &dictionary);
out.extend_from_slice(&encode_at(&dictionary, depth + 1)?);
out.extend_from_slice(&encode_at(&codes, depth + 1)?);
}
Kind::Sparse => {
let Some((value, _)) = dominant_value(values) else {
return Ok(None);
};
let mut positions = Vec::new();
let mut exceptions = Vec::new();
for (index, other) in values.iter().enumerate() {
if *other != value {
positions.push(index as i64);
exceptions.push(*other);
}
}
put_i64(&mut out, value);
put_u32(
&mut out,
u32::try_from(positions.len()).map_err(|_| too_long(positions.len()))?,
);
out.extend_from_slice(&encode_at(&positions, depth + 1)?);
out.extend_from_slice(&encode_at(&exceptions, depth + 1)?);
}
}
Ok(Some(out))
}
fn encode_packed(values: &[i64], out: &mut Vec<u8>) -> Result<()> {
for unit in values.chunks(VALUES) {
let base = unit.iter().copied().min().unwrap_or(0);
let offsets: Vec<u64> = unit.iter().map(|value| offset_from(*value, base)).collect();
let width = bitpack::required_width(&offsets);
put_i64(out, base);
put_u8(out, u8::try_from(width).map_err(|_| Error::internal("impossible width"))?);
if unit.len() == VALUES {
let mut packed = vec![0u64; bitpack::packed_len::<u64>(width)];
bitpack::pack(&offsets, width, &mut packed)?;
for word in packed {
put_u64(out, word);
}
} else {
bitpack::pack_tail(&offsets, width, out)?;
}
}
Ok(())
}
fn decode_chunk(reader: &mut Reader<'_>) -> Result<Vec<i64>> {
let kind = Kind::from_tag(reader.u8()?)?;
let count = reader.u32()? as usize;
match kind {
Kind::Constant => Ok(vec![reader.i64()?; count]),
Kind::Packed => {
let mut values = Vec::with_capacity(count);
while values.len() < count {
let base = reader.i64()?;
let width = reader.u8()? as usize;
let wanted = (count - values.len()).min(VALUES);
if wanted == VALUES {
let mut packed = vec![0u64; bitpack::packed_len::<u64>(width)];
for word in &mut packed {
*word = reader.u64()?;
}
let mut unit = vec![0u64; VALUES];
bitpack::unpack(&packed, width, &mut unit)?;
values.extend(unit.iter().map(|offset| value_from(*offset, base)));
} else {
let bytes = reader.bytes(bitpack::tail_len(wanted, width))?;
let unit = bitpack::unpack_tail(bytes, width, wanted)?;
values.extend(unit.iter().map(|offset| value_from(*offset, base)));
}
}
Ok(values)
}
Kind::Delta => {
let first = reader.i64()?;
let deltas = decode_chunk(reader)?;
let mut values = Vec::with_capacity(count);
values.push(first);
let mut current = first;
for delta in deltas {
current = current.wrapping_add(unzigzag(delta as u64));
values.push(current);
}
check_count(values.len(), count)?;
Ok(values)
}
Kind::Rle => {
let run_values = decode_chunk(reader)?;
let run_lengths = decode_chunk(reader)?;
if run_values.len() != run_lengths.len() {
return Err(Error::internal("an RLE chunk has more runs than run lengths"));
}
let mut values = Vec::with_capacity(count);
for (value, length) in run_values.into_iter().zip(run_lengths) {
let length = usize::try_from(length)
.map_err(|_| Error::internal("a negative RLE run length"))?;
values.extend(std::iter::repeat_n(value, length));
}
check_count(values.len(), count)?;
Ok(values)
}
Kind::Dict => {
let dictionary = decode_chunk(reader)?;
let codes = decode_chunk(reader)?;
let mut values = Vec::with_capacity(count);
for code in codes {
let index =
usize::try_from(code).ok().and_then(|index| dictionary.get(index)).ok_or_else(
|| Error::internal(format!("code {code} is not in the dictionary")),
)?;
values.push(*index);
}
check_count(values.len(), count)?;
Ok(values)
}
Kind::Sparse => {
let value = reader.i64()?;
let exception_count = reader.u32()? as usize;
let positions = decode_chunk(reader)?;
let exceptions = decode_chunk(reader)?;
if positions.len() != exception_count || exceptions.len() != exception_count {
return Err(Error::internal("a sparse chunk disagrees about its exception count"));
}
let mut values = vec![value; count];
for (position, exception) in positions.into_iter().zip(exceptions) {
let position = usize::try_from(position)
.ok()
.filter(|position| *position < count)
.ok_or_else(|| {
Error::internal(format!("exception at {position} is outside the chunk"))
})?;
values[position] = exception;
}
Ok(values)
}
}
}
fn describe_chunk(reader: &mut Reader<'_>) -> Result<String> {
let kind = Kind::from_tag(reader.u8()?)?;
let count = reader.u32()? as usize;
Ok(match kind {
Kind::Constant => {
reader.i64()?;
"CONSTANT".to_string()
}
Kind::Packed => {
let mut widths = Vec::new();
let mut seen = 0;
while seen < count {
reader.i64()?;
let width = reader.u8()? as usize;
let wanted = (count - seen).min(VALUES);
if wanted == VALUES {
for _ in 0..bitpack::packed_len::<u64>(width) {
reader.u64()?;
}
} else {
reader.bytes(bitpack::tail_len(wanted, width))?;
}
widths.push(width);
seen += wanted;
}
let low = widths.iter().copied().min().unwrap_or(0);
let high = widths.iter().copied().max().unwrap_or(0);
if low == high {
format!("FOR+BITPACK[{low}]")
} else {
format!("FOR+BITPACK[{low}..{high}]")
}
}
Kind::Delta => {
reader.i64()?;
format!("DELTA({})", describe_chunk(reader)?)
}
Kind::Rle => {
let values = describe_chunk(reader)?;
let lengths = describe_chunk(reader)?;
format!("RLE({values}, {lengths})")
}
Kind::Dict => {
let dictionary = describe_chunk(reader)?;
let codes = describe_chunk(reader)?;
format!("DICT({dictionary}, {codes})")
}
Kind::Sparse => {
reader.i64()?;
reader.u32()?;
let positions = describe_chunk(reader)?;
let exceptions = describe_chunk(reader)?;
format!("SPARSE({positions}, {exceptions})")
}
})
}
fn offset_from(value: i64, base: i64) -> u64 {
(i128::from(value) - i128::from(base)) as u64
}
fn value_from(offset: u64, base: i64) -> i64 {
(i128::from(base) + i128::from(offset)) as i64
}
fn zigzag(value: i64) -> u64 {
((value << 1) ^ (value >> 63)) as u64
}
fn unzigzag(value: u64) -> i64 {
((value >> 1) as i64) ^ -((value & 1) as i64)
}
fn deltas(values: &[i64]) -> Option<Vec<i64>> {
let mut deltas = Vec::with_capacity(values.len().saturating_sub(1));
for pair in values.windows(2) {
let difference = i128::from(pair[1]) - i128::from(pair[0]);
let difference = i64::try_from(difference).ok()?;
deltas.push(zigzag(difference) as i64);
}
Some(deltas)
}
fn run_count(values: &[i64]) -> usize {
let mut runs = 0;
let mut previous = None;
for value in values {
if previous != Some(value) {
runs += 1;
previous = Some(value);
}
}
runs
}
fn runs(values: &[i64]) -> (Vec<i64>, Vec<i64>) {
let mut run_values: Vec<i64> = Vec::new();
let mut run_lengths: Vec<i64> = Vec::new();
for value in values {
if run_values.last() == Some(value) {
*run_lengths.last_mut().expect("a run length exists beside every run value") += 1;
} else {
run_values.push(*value);
run_lengths.push(1);
}
}
(run_values, run_lengths)
}
fn distinct_values(values: &[i64]) -> Vec<i64> {
let mut distinct = values.to_vec();
distinct.sort_unstable();
distinct.dedup();
distinct
}
fn codes_over(values: &[i64], dictionary: &[i64]) -> Vec<i64> {
values
.iter()
.map(|value| {
dictionary
.binary_search(value)
.expect("the dictionary is the distinct values of this chunk") as i64
})
.collect()
}
fn dominant_value(values: &[i64]) -> Option<(i64, usize)> {
let mut sorted = values.to_vec();
sorted.sort_unstable();
let mut best: Option<(i64, usize)> = None;
let mut index = 0;
while index < sorted.len() {
let value = sorted[index];
let mut end = index;
while end < sorted.len() && sorted[end] == value {
end += 1;
}
let count = end - index;
if best.is_none_or(|(_, seen)| count > seen) {
best = Some((value, count));
}
index = end;
}
best
}
fn check_count(actual: usize, expected: usize) -> Result<()> {
if actual == expected {
Ok(())
} else {
Err(Error::internal(format!(
"a chunk says it holds {expected} values and decoded to {actual}"
)))
}
}
fn too_long(len: usize) -> Error {
Error::internal(format!("a chunk of {len} values is longer than the format allows"))
}
fn put_u8(out: &mut Vec<u8>, value: u8) {
out.push(value);
}
fn put_u32(out: &mut Vec<u8>, value: u32) {
out.extend_from_slice(&value.to_le_bytes());
}
fn put_u64(out: &mut Vec<u8>, value: u64) {
out.extend_from_slice(&value.to_le_bytes());
}
fn put_i64(out: &mut Vec<u8>, value: i64) {
out.extend_from_slice(&value.to_le_bytes());
}
#[cfg(test)]
mod tests {
use super::*;
fn round_trip(values: &[i64]) -> Vec<u8> {
let bytes = encode(values).unwrap();
assert_eq!(decode(&bytes).unwrap(), values, "{}", describe(&bytes).unwrap());
bytes
}
fn kind_of(bytes: &[u8]) -> Kind {
Kind::from_tag(bytes[0]).unwrap()
}
struct Random(u64);
impl Random {
fn new() -> Self {
Self(0x9e37_79b9_7f4a_7c15)
}
fn next(&mut self) -> u64 {
self.0 ^= self.0 << 13;
self.0 ^= self.0 >> 7;
self.0 ^= self.0 << 17;
self.0
}
}
#[test]
fn an_empty_chunk_round_trips() {
let bytes = round_trip(&[]);
assert_eq!(bytes.len(), 5);
}
#[test]
fn a_constant_column_costs_thirteen_bytes_however_long_it_is() {
let bytes = round_trip(&vec![42; 1_000_000]);
assert_eq!(kind_of(&bytes), Kind::Constant);
assert_eq!(bytes.len(), 13);
}
#[test]
fn a_narrow_range_is_packed_at_the_width_of_the_range_and_not_of_the_type() {
let mut random = Random::new();
let values: Vec<i64> = (0..100_000).map(|_| 1000 + (random.next() % 64) as i64).collect();
let bytes = round_trip(&values);
assert_eq!(kind_of(&bytes), Kind::Packed);
let packed = 100_000 * 6 / 8;
assert!(bytes.len() < packed + 2000, "{} bytes for {packed} of payload", bytes.len());
assert!(bytes.len() > packed, "{} bytes cannot hold {packed}", bytes.len());
}
#[test]
fn a_counter_becomes_deltas_and_then_a_constant() {
let values: Vec<i64> = (0..1_000_000).collect();
let bytes = round_trip(&values);
assert_eq!(kind_of(&bytes), Kind::Delta);
assert_eq!(describe(&bytes).unwrap(), "DELTA(CONSTANT)");
assert!(bytes.len() < 40, "{} bytes for a counter", bytes.len());
}
#[test]
fn a_column_that_counts_down_is_as_cheap_as_one_that_counts_up() {
let up: Vec<i64> = (0..100_000).collect();
let down: Vec<i64> = (0..100_000).rev().collect();
assert_eq!(round_trip(&up).len(), round_trip(&down).len());
}
#[test]
fn long_runs_become_rle() {
let mut values = Vec::new();
for run in 0..1000 {
values.extend(std::iter::repeat_n(run % 7, 200));
}
let bytes = round_trip(&values);
assert_eq!(kind_of(&bytes), Kind::Rle);
assert!(bytes.len() < 2000, "{} bytes for 1000 runs", bytes.len());
}
#[test]
fn a_low_cardinality_column_becomes_a_dictionary() {
let mut random = Random::new();
let dictionary: Vec<i64> = (0..40).map(|index| 1_000_000_000 + index * 7919).collect();
let values: Vec<i64> =
(0..100_000).map(|_| dictionary[(random.next() % 40) as usize]).collect();
let bytes = round_trip(&values);
assert_eq!(kind_of(&bytes), Kind::Dict);
assert!(bytes.len() < 100_000, "{} bytes", bytes.len());
}
#[test]
fn a_nearly_constant_column_becomes_sparse() {
let mut values = vec![0i64; 100_000];
for index in 0..300 {
values[index * 331] = 1 << 40;
}
let bytes = round_trip(&values);
assert_eq!(kind_of(&bytes), Kind::Sparse);
assert!(bytes.len() < 3000, "{} bytes for 300 exceptions", bytes.len());
}
#[test]
fn the_cascade_goes_more_than_one_level_deep() {
let mut values = Vec::new();
for index in 0..2000i64 {
values.extend(std::iter::repeat_n(1_000_000 + (index % 5) * 104_729, 100));
}
let bytes = round_trip(&values);
let shape = describe(&bytes).unwrap();
assert!(shape.contains('('), "{shape} is not a cascade");
assert!(bytes.len() < 4000, "{} bytes: {shape}", bytes.len());
}
#[test]
fn random_data_is_packed_at_full_width_and_costs_what_it_costs() {
let mut random = Random::new();
let values: Vec<i64> = (0..10_000).map(|_| random.next() as i64).collect();
let bytes = round_trip(&values);
assert_eq!(kind_of(&bytes), Kind::Packed);
assert!(bytes.len() < 10_000 * 8 + 1000, "{} bytes", bytes.len());
}
#[test]
fn the_extremes_of_the_type_survive() {
let values = vec![i64::MIN, i64::MAX, 0, -1, i64::MIN, i64::MAX];
round_trip(&values);
round_trip(&[i64::MIN; 3]);
round_trip(&[i64::MIN, i64::MIN + 1]);
}
#[test]
fn a_chunk_that_is_not_a_multiple_of_the_unit_round_trips() {
for len in [1, 2, 1023, 1024, 1025, 2047, 2049] {
let values: Vec<i64> = (0..len).map(|index| (index * 31 % 97) as i64).collect();
round_trip(&values);
}
}
#[test]
fn a_partial_unit_costs_its_own_values_and_not_a_whole_unit() {
let values = vec![1i64 << 39, (1 << 39) + 7, 1 << 38];
let bytes = encode_as(Kind::Packed, &values, 0).unwrap().unwrap();
assert_eq!(bytes.len(), 5 + 9 + 15);
assert_eq!(decode(&bytes).unwrap(), values);
}
#[test]
fn the_frame_of_reference_is_per_unit_and_not_per_chunk() {
let values: Vec<i64> =
(0..4096i64).map(|index| (index / 1024) * 1_000_000 + (index % 1024)).collect();
let bytes = encode_as(Kind::Packed, &values, 0).unwrap().unwrap();
assert_eq!(describe(&bytes).unwrap(), "FOR+BITPACK[10]");
assert_eq!(decode(&bytes).unwrap(), values);
}
#[test]
fn every_candidate_that_applies_decodes_to_the_input() {
let mut values = vec![5i64; 3000];
for (index, value) in values.iter_mut().enumerate() {
if index % 500 == 0 {
*value = index as i64;
}
}
let applicable = candidates(&values, 0);
assert!(applicable.len() >= 4, "{applicable:?}");
for kind in applicable {
let bytes = encode_as(kind, &values, 0).unwrap().unwrap();
assert_eq!(decode(&bytes).unwrap(), values, "{}", kind.name());
}
}
#[test]
fn the_chooser_picks_the_smallest_candidate_rather_than_the_first_that_applies() {
let mut values = vec![5i64; 3000];
values[1500] = 9;
let chosen = encode(&values).unwrap();
for (_, size) in candidate_sizes(&values).unwrap() {
assert!(chosen.len() <= size);
}
}
#[test]
fn a_truncated_chunk_is_an_error_and_not_a_panic() {
let bytes = encode(&[1, 2, 3, 4, 5]).unwrap();
for len in 0..bytes.len() {
let error = decode(&bytes[..len]).unwrap_err();
assert!(error.message().contains("chunk"), "{error}");
}
}
#[test]
fn trailing_bytes_are_an_error() {
let mut bytes = encode(&[1, 2, 3]).unwrap();
bytes.push(0);
let error = decode(&bytes).unwrap_err();
assert!(error.message().contains("left over"), "{error}");
}
#[test]
fn an_unknown_tag_is_an_error() {
let error = decode(&[99, 0, 0, 0, 0]).unwrap_err();
assert!(error.message().contains("unknown encoding tag"), "{error}");
}
#[test]
fn a_dictionary_code_outside_the_dictionary_is_an_error() {
let mut bytes = vec![Kind::Dict.tag()];
put_u32(&mut bytes, 1);
bytes.extend_from_slice(&encode(&[10]).unwrap());
bytes.extend_from_slice(&encode(&[5]).unwrap());
let error = decode(&bytes).unwrap_err();
assert!(error.message().contains("not in the dictionary"), "{error}");
}
#[test]
fn a_negative_run_length_is_an_error() {
let mut bytes = vec![Kind::Rle.tag()];
put_u32(&mut bytes, 4);
bytes.extend_from_slice(&encode(&[7]).unwrap());
bytes.extend_from_slice(&encode(&[-4]).unwrap());
let error = decode(&bytes).unwrap_err();
assert!(error.message().contains("negative"), "{error}");
}
#[test]
fn the_cascade_depth_is_bounded() {
let values: Vec<i64> = (0..50_000).map(|index| (index / 100) % 250).collect();
let bytes = round_trip(&values);
let shape = describe(&bytes).unwrap();
let depth = shape.matches('(').count();
assert!(depth <= MAX_DEPTH as usize, "{shape} is {depth} deep");
}
#[test]
fn candidate_sizes_reports_what_the_chooser_looked_at() {
let values: Vec<i64> = (0..5000).map(|index| index % 17).collect();
let sizes = candidate_sizes(&values).unwrap();
assert!(sizes.iter().any(|(kind, _)| *kind == Kind::Dict));
assert!(sizes.iter().any(|(kind, _)| *kind == Kind::Packed));
assert!(sizes.iter().all(|(_, size)| *size > 0));
}
#[test]
fn a_chunk_can_be_read_from_the_front_of_a_longer_buffer() {
let first = encode(&[1, 2, 3]).unwrap();
let second: Vec<i64> = (0..3000).map(|index| index % 11).collect();
let second_bytes = encode(&second).unwrap();
let mut joined = first.clone();
joined.extend_from_slice(&second_bytes);
joined.extend_from_slice(b"and then something else");
let (values, used) = decode_prefix(&joined).unwrap();
assert_eq!(values, vec![1, 2, 3]);
assert_eq!(used, first.len());
let (more, used_again) = decode_prefix(&joined[used..]).unwrap();
assert_eq!(more, second);
assert_eq!(used_again, second_bytes.len());
let (text, described) = describe_prefix(&joined).unwrap();
assert_eq!(described, first.len());
assert_eq!(text, describe(&first).unwrap());
}
#[test]
fn a_truncated_chunk_is_still_an_error_when_read_as_a_prefix() {
let bytes = encode(&(0..2000).collect::<Vec<i64>>()).unwrap();
for len in 0..bytes.len() {
assert!(decode_prefix(&bytes[..len]).is_err(), "{len} bytes decoded");
}
}
}