use chronix_core::config::FloatEncoding;
use crate::alp::AlpDecoder;
use crate::bitmap::{BitmapDecoder, BitmapEncoder};
use crate::chimp::{Chimp128Decoder, ChimpDecoder};
use crate::delta::{DeltaOfDeltaDecoder, DeltaOfDeltaEncoder};
use crate::dictionary::{DictionaryDecoder, DictionaryEncoder};
use crate::error::{EncodingError, Result};
use crate::for_encoding::{ForDecoder, ForEncoder};
use crate::gorilla::{GorillaDecoder, GorillaEncoder};
use crate::integer::{IntegerDecoder, IntegerEncoder, VarintDecoder, VarintEncoder};
use crate::patas::PatasDecoder;
use crate::plain::{PlainDecoder, PlainEncoder};
use crate::rle::{RleDecoder, RleEncoder};
macro_rules! encoding_types {
($(
$(#[$meta:meta])*
$variant:ident = $tag:literal, $name:literal;
)+) => {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum EncodingType {
$(
$(#[$meta])*
$variant = $tag,
)+
}
impl EncodingType {
pub const ALL: &'static [Self] = &[$(Self::$variant),+];
pub fn from_tag(tag: u8) -> Result<Self> {
match tag {
$($tag => Ok(Self::$variant),)+
other => Err(EncodingError::UnsupportedEncoding { tag: other }),
}
}
#[must_use]
pub fn tag(self) -> u8 {
self as u8
}
}
impl std::fmt::Display for EncodingType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
$(Self::$variant => $name,)+
})
}
}
};
}
encoding_types! {
DeltaOfDelta = 0, "delta-of-delta";
Chimp = 1, "chimp";
Gorilla = 2, "gorilla";
IntegerI64 = 3, "integer-i64";
IntegerU64 = 4, "integer-u64";
Dictionary = 5, "dictionary";
Bitmap = 6, "bitmap";
PlainF64 = 7, "plain-f64";
PlainI64 = 8, "plain-i64";
PlainU64 = 9, "plain-u64";
PlainBool = 10, "plain-bool";
PlainString = 11, "plain-string";
Nullable = 12, "nullable";
Rle = 13, "rle";
VarintI64 = 14, "varint-i64";
VarintU64 = 15, "varint-u64";
Chimp128 = 16, "chimp128";
ForI64 = 17, "for-i64";
ForU64 = 18, "for-u64";
Patas = 19, "patas";
Alp = 20, "alp";
Pco = 21, "pco";
PcoI64 = 22, "pco_i64";
PcoU64 = 23, "pco_u64";
DecimalI128 = 24, "decimal-i128";
Bytes = 25, "bytes";
}
#[derive(Debug, Clone)]
pub struct EncodedBlock {
pub encoding: EncodingType,
pub payload: Vec<u8>,
}
impl EncodedBlock {
#[must_use]
pub fn to_bytes(&self) -> Vec<u8> {
let mut out = Vec::with_capacity(1 + self.payload.len());
out.push(self.encoding.tag());
out.extend_from_slice(&self.payload);
out
}
pub fn from_bytes(data: &[u8]) -> Result<Self> {
if data.is_empty() {
return Err(EncodingError::CorruptData {
detail: "empty encoded block".to_string(),
});
}
let encoding = EncodingType::from_tag(data[0])?;
let payload = data[1..].to_vec();
Ok(Self { encoding, payload })
}
}
const MIN_COMPRESSION_RATIO: f64 = 1.5;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TimestampEncoding {
#[default]
Auto,
Pco,
DeltaOfDelta,
Plain,
}
#[derive(Debug, Clone, Copy)]
pub struct ColumnEncoder;
impl ColumnEncoder {
pub fn encode_timestamps(values: &[i64]) -> Result<EncodedBlock> {
Self::encode_timestamps_with(values, TimestampEncoding::Auto)
}
pub fn encode_timestamps_with(
values: &[i64],
strategy: TimestampEncoding,
) -> Result<EncodedBlock> {
let raw_size = values.len() * 8;
let (encoding, payload) = match strategy {
TimestampEncoding::Plain => {
return Ok(EncodedBlock {
encoding: EncodingType::PlainI64,
payload: PlainEncoder::encode_i64(values)?,
});
}
TimestampEncoding::DeltaOfDelta => (
EncodingType::DeltaOfDelta,
DeltaOfDeltaEncoder::encode(values)?,
),
TimestampEncoding::Pco => (
EncodingType::PcoI64,
crate::pco::PcoEncoder::encode_timestamps(values)?,
),
TimestampEncoding::Auto => {
let pco = crate::pco::PcoEncoder::encode_timestamps(values)?;
let dod = DeltaOfDeltaEncoder::encode(values)?;
if dod.len() < pco.len() {
(EncodingType::DeltaOfDelta, dod)
} else {
(EncodingType::PcoI64, pco)
}
}
};
if should_use_specialized(&payload, raw_size) {
Ok(EncodedBlock { encoding, payload })
} else {
Ok(EncodedBlock {
encoding: EncodingType::PlainI64,
payload: PlainEncoder::encode_i64(values)?,
})
}
}
pub fn encode_f64(values: &[f64], strategy: FloatEncoding) -> Result<EncodedBlock> {
match strategy {
FloatEncoding::Chimp => {
let selector = crate::adaptive::AdaptiveSelector::new();
selector.encode_f64_adaptive(values)
}
FloatEncoding::Gorilla => {
let raw_size = values.len() * 8;
let gorilla_payload = GorillaEncoder::encode(values)?;
if should_use_specialized(&gorilla_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Gorilla,
payload: gorilla_payload,
});
}
Ok(EncodedBlock {
encoding: EncodingType::PlainF64,
payload: PlainEncoder::encode_f64(values)?,
})
}
FloatEncoding::Plain => Ok(EncodedBlock {
encoding: EncodingType::PlainF64,
payload: PlainEncoder::encode_f64(values)?,
}),
}
}
pub fn encode_i64(values: &[i64]) -> Result<EncodedBlock> {
let selector = crate::adaptive::AdaptiveSelector::new();
let pattern = selector.analyze_integers(values);
if pattern == crate::adaptive::IntegerPattern::Constant {
let raw_size = values.len() * 8;
let rle_payload = RleEncoder::encode_i64(values)?;
if should_use_specialized(&rle_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Rle,
payload: rle_payload,
});
}
}
let raw_size = values.len() * 8;
let fixed_payload = IntegerEncoder::encode_i64(values)?;
let varint_payload = VarintEncoder::encode_i64(values)?;
let for_payload = ForEncoder::encode_i64(values)?;
let pco_payload = crate::pco::PcoEncoder::encode_i64(values)?;
let mut best_enc = EncodingType::IntegerI64;
let mut best_payload = fixed_payload;
if varint_payload.len() < best_payload.len() {
best_enc = EncodingType::VarintI64;
best_payload = varint_payload;
}
if for_payload.len() < best_payload.len() {
best_enc = EncodingType::ForI64;
best_payload = for_payload;
}
if pco_payload.len() < best_payload.len() {
best_enc = EncodingType::PcoI64;
best_payload = pco_payload;
}
if should_use_specialized(&best_payload, raw_size) {
Ok(EncodedBlock {
encoding: best_enc,
payload: best_payload,
})
} else {
Ok(EncodedBlock {
encoding: EncodingType::PlainI64,
payload: PlainEncoder::encode_i64(values)?,
})
}
}
pub fn encode_u64(values: &[u64]) -> Result<EncodedBlock> {
let as_i64: Vec<i64> = values.iter().map(|&v| v as i64).collect();
let selector = crate::adaptive::AdaptiveSelector::new();
if selector.analyze_integers(&as_i64) == crate::adaptive::IntegerPattern::Constant {
let raw_size = values.len() * 8;
let rle_payload = RleEncoder::encode_u64(values)?;
if should_use_specialized(&rle_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Rle,
payload: rle_payload,
});
}
}
let raw_size = values.len() * 8;
let fixed_payload = IntegerEncoder::encode_u64(values)?;
let varint_payload = VarintEncoder::encode_u64(values)?;
let for_payload = ForEncoder::encode_u64(values)?;
let pco_payload = crate::pco::PcoEncoder::encode_u64(values)?;
let mut best_enc = EncodingType::IntegerU64;
let mut best_payload = fixed_payload;
if varint_payload.len() < best_payload.len() {
best_enc = EncodingType::VarintU64;
best_payload = varint_payload;
}
if for_payload.len() < best_payload.len() {
best_enc = EncodingType::ForU64;
best_payload = for_payload;
}
if pco_payload.len() < best_payload.len() {
best_enc = EncodingType::PcoU64;
best_payload = pco_payload;
}
if should_use_specialized(&best_payload, raw_size) {
Ok(EncodedBlock {
encoding: best_enc,
payload: best_payload,
})
} else {
Ok(EncodedBlock {
encoding: EncodingType::PlainU64,
payload: PlainEncoder::encode_u64(values)?,
})
}
}
pub fn encode_string(values: &[&str]) -> Result<EncodedBlock> {
let raw_size: usize = values.iter().map(|s| 4 + s.len()).sum();
match DictionaryEncoder::encode(values) {
Ok(payload) if should_use_specialized(&payload, raw_size) => Ok(EncodedBlock {
encoding: EncodingType::Dictionary,
payload,
}),
Ok(_) | Err(EncodingError::DictionaryOverflow { .. }) => Ok(EncodedBlock {
encoding: EncodingType::PlainString,
payload: PlainEncoder::encode_string(values)?,
}),
Err(e) => Err(e),
}
}
pub fn encode_bool(values: &[bool]) -> Result<EncodedBlock> {
if !values.is_empty() {
let first = values[0];
if values.iter().all(|&v| v == first) {
let raw_size = values.len(); let rle_payload = RleEncoder::encode_bool(values)?;
if should_use_specialized(&rle_payload, raw_size) {
return Ok(EncodedBlock {
encoding: EncodingType::Rle,
payload: rle_payload,
});
}
}
}
Ok(EncodedBlock {
encoding: EncodingType::Bitmap,
payload: BitmapEncoder::encode(values)?,
})
}
pub fn encode_decimal(values: &[i128]) -> Result<EncodedBlock> {
Ok(EncodedBlock {
encoding: EncodingType::DecimalI128,
payload: crate::decimal::DecimalEncoder::encode(values)?,
})
}
pub fn encode_f64_nullable(
values: &[Option<f64>],
strategy: FloatEncoding,
) -> Result<EncodedBlock> {
if values.is_empty() {
return Err(EncodingError::EmptyInput {
context: "nullable f64 encoder",
});
}
let (non_null, validity) = extract_non_null_f64(values);
if non_null.is_empty() {
return Ok(wrap_nullable(
EncodingType::PlainF64,
values.len(),
&validity,
&[],
));
}
let inner = Self::encode_f64(&non_null, strategy)?;
Ok(wrap_nullable(
inner.encoding,
values.len(),
&validity,
&inner.payload,
))
}
pub fn encode_i64_nullable(values: &[Option<i64>]) -> Result<EncodedBlock> {
if values.is_empty() {
return Err(EncodingError::EmptyInput {
context: "nullable i64 encoder",
});
}
let (non_null, validity) = extract_non_null_i64(values);
if non_null.is_empty() {
return Ok(wrap_nullable(
EncodingType::PlainI64,
values.len(),
&validity,
&[],
));
}
let inner = Self::encode_i64(&non_null)?;
Ok(wrap_nullable(
inner.encoding,
values.len(),
&validity,
&inner.payload,
))
}
pub fn encode_u64_nullable(values: &[Option<u64>]) -> Result<EncodedBlock> {
if values.is_empty() {
return Err(EncodingError::EmptyInput {
context: "nullable u64 encoder",
});
}
let (non_null, validity) = extract_non_null_u64(values);
if non_null.is_empty() {
return Ok(wrap_nullable(
EncodingType::PlainU64,
values.len(),
&validity,
&[],
));
}
let inner = Self::encode_u64(&non_null)?;
Ok(wrap_nullable(
inner.encoding,
values.len(),
&validity,
&inner.payload,
))
}
pub fn encode_string_nullable(values: &[Option<&str>]) -> Result<EncodedBlock> {
if values.is_empty() {
return Err(EncodingError::EmptyInput {
context: "nullable string encoder",
});
}
let (non_null, validity): (Vec<&str>, Vec<bool>) = values
.iter()
.map(|v| match v {
Some(s) => (*s, true),
None => ("", false),
})
.unzip();
let non_null_only: Vec<&str> = non_null
.into_iter()
.zip(validity.iter())
.filter(|(_, v)| **v)
.map(|(s, _)| s)
.collect();
if non_null_only.is_empty() {
return Ok(wrap_nullable(
EncodingType::PlainString,
values.len(),
&validity,
&[],
));
}
let inner = Self::encode_string(&non_null_only)?;
Ok(wrap_nullable(
inner.encoding,
values.len(),
&validity,
&inner.payload,
))
}
pub fn encode_decimal_nullable(values: &[Option<i128>]) -> Result<EncodedBlock> {
if values.is_empty() {
return Err(EncodingError::EmptyInput {
context: "nullable decimal encoder",
});
}
let validity: Vec<bool> = values.iter().map(std::option::Option::is_some).collect();
let non_null: Vec<i128> = values.iter().filter_map(|v| *v).collect();
if non_null.is_empty() {
return Ok(wrap_nullable(
EncodingType::DecimalI128,
values.len(),
&validity,
&[],
));
}
let inner = Self::encode_decimal(&non_null)?;
Ok(wrap_nullable(
inner.encoding,
values.len(),
&validity,
&inner.payload,
))
}
pub fn encode_bool_nullable(values: &[Option<bool>]) -> Result<EncodedBlock> {
if values.is_empty() {
return Err(EncodingError::EmptyInput {
context: "nullable bool encoder",
});
}
let validity: Vec<bool> = values.iter().map(std::option::Option::is_some).collect();
let non_null: Vec<bool> = values.iter().filter_map(|v| *v).collect();
if non_null.is_empty() {
return Ok(wrap_nullable(
EncodingType::Bitmap,
values.len(),
&validity,
&[],
));
}
let inner = Self::encode_bool(&non_null)?;
Ok(wrap_nullable(
inner.encoding,
values.len(),
&validity,
&inner.payload,
))
}
}
#[derive(Debug, Clone)]
pub enum DecodedColumn {
I64(Vec<i64>),
U64(Vec<u64>),
F64(Vec<f64>),
Bool(Vec<bool>),
String(Vec<String>),
Decimal(Vec<i128>),
Bytes(Vec<Vec<u8>>),
NullableI64(Vec<Option<i64>>),
NullableU64(Vec<Option<u64>>),
NullableF64(Vec<Option<f64>>),
NullableBool(Vec<Option<bool>>),
NullableString(Vec<Option<String>>),
NullableDecimal(Vec<Option<i128>>),
NullableBytes(Vec<Option<Vec<u8>>>),
}
fn f64_slice_eq(a: &[f64], b: &[f64]) -> bool {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
}
fn opt_f64_slice_eq(a: &[Option<f64>], b: &[Option<f64>]) -> bool {
a.len() == b.len()
&& a.iter().zip(b).all(|(x, y)| match (x, y) {
(Some(x), Some(y)) => x.to_bits() == y.to_bits(),
(None, None) => true,
_ => false,
})
}
impl PartialEq for DecodedColumn {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(Self::I64(a), Self::I64(b)) => a == b,
(Self::U64(a), Self::U64(b)) => a == b,
(Self::F64(a), Self::F64(b)) => f64_slice_eq(a, b),
(Self::Bool(a), Self::Bool(b)) => a == b,
(Self::String(a), Self::String(b)) => a == b,
(Self::NullableI64(a), Self::NullableI64(b)) => a == b,
(Self::NullableU64(a), Self::NullableU64(b)) => a == b,
(Self::NullableF64(a), Self::NullableF64(b)) => opt_f64_slice_eq(a, b),
(Self::NullableBool(a), Self::NullableBool(b)) => a == b,
(Self::NullableString(a), Self::NullableString(b)) => a == b,
(Self::Decimal(a), Self::Decimal(b)) => a == b,
(Self::NullableDecimal(a), Self::NullableDecimal(b)) => a == b,
_ => false,
}
}
}
impl DecodedColumn {
#[must_use]
pub fn len(&self) -> usize {
match self {
Self::I64(v) => v.len(),
Self::U64(v) => v.len(),
Self::F64(v) => v.len(),
Self::Bool(v) => v.len(),
Self::String(v) => v.len(),
Self::NullableI64(v) => v.len(),
Self::NullableU64(v) => v.len(),
Self::NullableF64(v) => v.len(),
Self::NullableBool(v) => v.len(),
Self::NullableString(v) => v.len(),
Self::Decimal(v) => v.len(),
Self::NullableDecimal(v) => v.len(),
Self::Bytes(v) => v.len(),
Self::NullableBytes(v) => v.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[derive(Debug, Clone, Copy)]
pub struct ColumnDecoder;
impl ColumnDecoder {
pub fn decode(block: &EncodedBlock) -> Result<DecodedColumn> {
match block.encoding {
EncodingType::DeltaOfDelta => {
let values = DeltaOfDeltaDecoder::decode(&block.payload)?;
Ok(DecodedColumn::I64(values))
}
EncodingType::Chimp => {
let values = ChimpDecoder::decode(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::Chimp128 => {
let values = Chimp128Decoder::decode(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::Patas => {
let values = PatasDecoder::decode(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::Alp => {
let values = AlpDecoder::decode(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::Pco => {
let values = crate::pco::PcoDecoder::decode_f64(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::PcoI64 => {
let values = crate::pco::PcoDecoder::decode_i64(&block.payload)?;
Ok(DecodedColumn::I64(values))
}
EncodingType::PcoU64 => {
let values = crate::pco::PcoDecoder::decode_u64(&block.payload)?;
Ok(DecodedColumn::U64(values))
}
EncodingType::Gorilla => {
let values = GorillaDecoder::decode(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::IntegerI64 => {
let values = IntegerDecoder::decode_i64(&block.payload)?;
Ok(DecodedColumn::I64(values))
}
EncodingType::IntegerU64 => {
let values = IntegerDecoder::decode_u64(&block.payload)?;
Ok(DecodedColumn::U64(values))
}
EncodingType::Dictionary => {
let values = DictionaryDecoder::decode(&block.payload)?;
Ok(DecodedColumn::String(values))
}
EncodingType::Bitmap => {
let values = BitmapDecoder::decode(&block.payload)?;
Ok(DecodedColumn::Bool(values))
}
EncodingType::PlainF64 => {
let values = PlainDecoder::decode_f64(&block.payload)?;
Ok(DecodedColumn::F64(values))
}
EncodingType::PlainI64 => {
let values = PlainDecoder::decode_i64(&block.payload)?;
Ok(DecodedColumn::I64(values))
}
EncodingType::PlainU64 => {
let values = PlainDecoder::decode_u64(&block.payload)?;
Ok(DecodedColumn::U64(values))
}
EncodingType::PlainBool => {
let values = PlainDecoder::decode_bool(&block.payload)?;
Ok(DecodedColumn::Bool(values))
}
EncodingType::PlainString => {
let values = PlainDecoder::decode_string(&block.payload)?;
Ok(DecodedColumn::String(values))
}
EncodingType::Bytes => {
let values = PlainDecoder::decode_bytes(&block.payload)?;
Ok(DecodedColumn::Bytes(values))
}
EncodingType::Nullable => Self::decode_nullable(&block.payload),
EncodingType::Rle => Self::decode_rle(&block.payload),
EncodingType::VarintI64 => {
let values = VarintDecoder::decode_i64(&block.payload)?;
Ok(DecodedColumn::I64(values))
}
EncodingType::VarintU64 => {
let values = VarintDecoder::decode_u64(&block.payload)?;
Ok(DecodedColumn::U64(values))
}
EncodingType::ForI64 => {
let values = ForDecoder::decode_i64(&block.payload)?;
Ok(DecodedColumn::I64(values))
}
EncodingType::ForU64 => {
let values = ForDecoder::decode_u64(&block.payload)?;
Ok(DecodedColumn::U64(values))
}
EncodingType::DecimalI128 => {
let values = crate::decimal::DecimalDecoder::decode(&block.payload)?;
Ok(DecodedColumn::Decimal(values))
}
}
}
fn decode_nullable(payload: &[u8]) -> Result<DecodedColumn> {
if payload.len() < 5 {
return Err(EncodingError::CorruptData {
detail: "nullable payload too short".to_string(),
});
}
let inner_enc = EncodingType::from_tag(payload[0])?;
let total_count = crate::coding::checked_decode_count(
u32::from_le_bytes([payload[1], payload[2], payload[3], payload[4]]) as usize,
"nullable",
)?;
let bitmap_bytes = total_count.div_ceil(8);
let header_len = 5 + bitmap_bytes;
if payload.len() < header_len {
return Err(EncodingError::CorruptData {
detail: "nullable bitmap truncated".to_string(),
});
}
let bitmap = &payload[5..header_len];
let inner_payload = &payload[header_len..];
let validity: Vec<bool> = (0..total_count)
.map(|i| bitmap[i / 8] & (1 << (i % 8)) != 0)
.collect();
let non_null_count = validity.iter().filter(|v| **v).count();
if non_null_count == 0 {
return match inner_enc {
EncodingType::PlainF64
| EncodingType::Chimp
| EncodingType::Chimp128
| EncodingType::Gorilla
| EncodingType::Patas
| EncodingType::Alp
| EncodingType::Pco => Ok(DecodedColumn::NullableF64(vec![None; total_count])),
EncodingType::PlainI64
| EncodingType::IntegerI64
| EncodingType::VarintI64
| EncodingType::ForI64
| EncodingType::PcoI64
| EncodingType::DeltaOfDelta => {
Ok(DecodedColumn::NullableI64(vec![None; total_count]))
}
EncodingType::PlainU64
| EncodingType::IntegerU64
| EncodingType::VarintU64
| EncodingType::ForU64
| EncodingType::PcoU64 => Ok(DecodedColumn::NullableU64(vec![None; total_count])),
EncodingType::PlainString | EncodingType::Dictionary => {
Ok(DecodedColumn::NullableString(vec![None; total_count]))
}
EncodingType::Bitmap | EncodingType::PlainBool => {
Ok(DecodedColumn::NullableBool(vec![None; total_count]))
}
EncodingType::DecimalI128 => {
Ok(DecodedColumn::NullableDecimal(vec![None; total_count]))
}
_ => Err(EncodingError::CorruptData {
detail: format!("unexpected inner encoding in nullable: {inner_enc}"),
}),
};
}
let inner_block = EncodedBlock {
encoding: inner_enc,
payload: inner_payload.to_vec(),
};
let decoded = Self::decode(&inner_block)?;
match decoded {
DecodedColumn::F64(vals) => Ok(DecodedColumn::NullableF64(scatter(vals, &validity))),
DecodedColumn::I64(vals) => Ok(DecodedColumn::NullableI64(scatter(vals, &validity))),
DecodedColumn::U64(vals) => Ok(DecodedColumn::NullableU64(scatter(vals, &validity))),
DecodedColumn::Bool(vals) => Ok(DecodedColumn::NullableBool(scatter(vals, &validity))),
DecodedColumn::String(vals) => {
Ok(DecodedColumn::NullableString(scatter(vals, &validity)))
}
DecodedColumn::Decimal(vals) => {
Ok(DecodedColumn::NullableDecimal(scatter(vals, &validity)))
}
_ => Err(EncodingError::CorruptData {
detail: "unexpected decoded type in nullable wrapper".to_string(),
}),
}
}
fn decode_rle(payload: &[u8]) -> Result<DecodedColumn> {
let vtype = RleDecoder::value_type(payload)?;
match vtype {
0 => Ok(DecodedColumn::I64(RleDecoder::decode_i64(payload)?)),
1 => Ok(DecodedColumn::U64(RleDecoder::decode_u64(payload)?)),
2 => Ok(DecodedColumn::F64(RleDecoder::decode_f64(payload)?)),
3 => Ok(DecodedColumn::Bool(RleDecoder::decode_bool(payload)?)),
4 => Ok(DecodedColumn::String(RleDecoder::decode_string(payload)?)),
other => Err(EncodingError::CorruptData {
detail: format!("unknown RLE value type: {other}"),
}),
}
}
pub fn decode_bytes(data: &[u8]) -> Result<DecodedColumn> {
let block = EncodedBlock::from_bytes(data)?;
Self::decode(&block)
}
}
pub(crate) fn should_use_specialized(encoded: &[u8], raw_size: usize) -> bool {
if raw_size == 0 {
return false;
}
#[allow(clippy::cast_precision_loss)]
let ratio = raw_size as f64 / encoded.len() as f64;
ratio >= MIN_COMPRESSION_RATIO
}
fn encode_validity_bitmap(validity: &[bool]) -> Vec<u8> {
let byte_count = validity.len().div_ceil(8);
let mut bitmap = vec![0u8; byte_count];
for (i, &valid) in validity.iter().enumerate() {
if valid {
bitmap[i / 8] |= 1 << (i % 8);
}
}
bitmap
}
fn wrap_nullable(
inner_enc: EncodingType,
total_count: usize,
validity: &[bool],
inner_payload: &[u8],
) -> EncodedBlock {
let bitmap = encode_validity_bitmap(validity);
let mut payload = Vec::with_capacity(1 + 4 + bitmap.len() + inner_payload.len());
payload.push(inner_enc.tag());
payload.extend_from_slice(&(total_count as u32).to_le_bytes());
payload.extend_from_slice(&bitmap);
payload.extend_from_slice(inner_payload);
EncodedBlock {
encoding: EncodingType::Nullable,
payload,
}
}
fn extract_non_null_f64(values: &[Option<f64>]) -> (Vec<f64>, Vec<bool>) {
let validity: Vec<bool> = values.iter().map(std::option::Option::is_some).collect();
let non_null: Vec<f64> = values.iter().filter_map(|v| *v).collect();
(non_null, validity)
}
fn extract_non_null_i64(values: &[Option<i64>]) -> (Vec<i64>, Vec<bool>) {
let validity: Vec<bool> = values.iter().map(std::option::Option::is_some).collect();
let non_null: Vec<i64> = values.iter().filter_map(|v| *v).collect();
(non_null, validity)
}
fn extract_non_null_u64(values: &[Option<u64>]) -> (Vec<u64>, Vec<bool>) {
let validity: Vec<bool> = values.iter().map(std::option::Option::is_some).collect();
let non_null: Vec<u64> = values.iter().filter_map(|v| *v).collect();
(non_null, validity)
}
fn scatter<T: Default + Clone>(values: Vec<T>, validity: &[bool]) -> Vec<Option<T>> {
let mut result = Vec::with_capacity(validity.len());
let mut val_iter = values.into_iter();
for &valid in validity {
if valid {
result.push(val_iter.next().map(Some).unwrap_or(None));
} else {
result.push(None);
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn encoding_type_tag_roundtrip() {
let all = [
EncodingType::DeltaOfDelta,
EncodingType::Chimp,
EncodingType::Gorilla,
EncodingType::IntegerI64,
EncodingType::IntegerU64,
EncodingType::Dictionary,
EncodingType::Bitmap,
EncodingType::PlainF64,
EncodingType::PlainI64,
EncodingType::PlainU64,
EncodingType::PlainBool,
EncodingType::PlainString,
EncodingType::Nullable,
EncodingType::Rle,
];
for enc in &all {
let tag = enc.tag();
let recovered = EncodingType::from_tag(tag).unwrap();
assert_eq!(*enc, recovered);
}
}
#[test]
fn encoding_type_unknown_tag() {
assert!(EncodingType::from_tag(255).is_err());
}
#[test]
fn all_lists_every_variant_and_agrees_with_from_tag() {
let scanned: Vec<EncodingType> = (0..=u8::MAX)
.filter_map(|t| EncodingType::from_tag(t).ok())
.collect();
assert_eq!(EncodingType::ALL, scanned.as_slice());
for enc in EncodingType::ALL {
assert_eq!(EncodingType::from_tag(enc.tag()).unwrap(), *enc);
}
for (i, enc) in EncodingType::ALL.iter().enumerate() {
assert_eq!(usize::from(enc.tag()), i, "{enc} has a non-dense tag");
}
}
#[test]
fn encoded_block_roundtrip() {
let block = EncodedBlock {
encoding: EncodingType::Chimp,
payload: vec![1, 2, 3, 4],
};
let bytes = block.to_bytes();
assert_eq!(bytes[0], EncodingType::Chimp.tag());
let recovered = EncodedBlock::from_bytes(&bytes).unwrap();
assert_eq!(recovered.encoding, block.encoding);
assert_eq!(recovered.payload, block.payload);
}
#[test]
fn encoded_block_empty_error() {
assert!(EncodedBlock::from_bytes(&[]).is_err());
}
#[test]
fn encode_timestamps_monotonic() {
let ts: Vec<i64> = (0..1000).map(|i| 1_000_000 + i * 10_000).collect();
let block = ColumnEncoder::encode_timestamps(&ts).unwrap();
assert_eq!(block.encoding, EncodingType::PcoI64);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::I64(ts));
}
#[test]
fn encode_timestamps_honours_an_explicit_codec() {
let ts: Vec<i64> = (0..1000).map(|i| 1_000_000 + i * 10_000).collect();
for (strategy, expected) in [
(TimestampEncoding::Pco, EncodingType::PcoI64),
(TimestampEncoding::DeltaOfDelta, EncodingType::DeltaOfDelta),
(TimestampEncoding::Plain, EncodingType::PlainI64),
(TimestampEncoding::Auto, EncodingType::PcoI64),
] {
let block = ColumnEncoder::encode_timestamps_with(&ts, strategy).unwrap();
assert_eq!(block.encoding, expected, "{strategy:?}");
assert_eq!(
ColumnDecoder::decode(&block).unwrap(),
DecodedColumn::I64(ts.clone())
);
}
}
#[test]
fn auto_timestamps_never_loses_to_delta_of_delta() {
let mut seed = 7u64;
let ts: Vec<i64> = (0..2000)
.map(|_| {
seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
(seed >> 20) as i64
})
.collect();
let auto = ColumnEncoder::encode_timestamps_with(&ts, TimestampEncoding::Auto).unwrap();
let dod =
ColumnEncoder::encode_timestamps_with(&ts, TimestampEncoding::DeltaOfDelta).unwrap();
let pco = ColumnEncoder::encode_timestamps_with(&ts, TimestampEncoding::Pco).unwrap();
assert!(auto.payload.len() <= dod.payload.len().min(pco.payload.len()));
assert_eq!(
ColumnDecoder::decode(&auto).unwrap(),
DecodedColumn::I64(ts)
);
}
#[test]
fn encode_timestamps_bytes_roundtrip() {
let ts: Vec<i64> = (0..100).map(|i| 1_000_000 + i * 10_000).collect();
let block = ColumnEncoder::encode_timestamps(&ts).unwrap();
let raw = block.to_bytes();
let decoded = ColumnDecoder::decode_bytes(&raw).unwrap();
assert_eq!(decoded, DecodedColumn::I64(ts));
}
#[test]
fn encode_f64_adaptive_strategy() {
let values: Vec<f64> = (0..500).map(|i| 20.0 + (i as f64) * 0.01).collect();
let block = ColumnEncoder::encode_f64(&values, FloatEncoding::Chimp).unwrap();
assert_eq!(block.encoding, EncodingType::Pco);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::F64(values));
}
#[test]
fn encode_f64_gorilla_strategy() {
let values: Vec<f64> = (0..500).map(|i| 20.0 + (i as f64) * 0.01).collect();
let block = ColumnEncoder::encode_f64(&values, FloatEncoding::Gorilla).unwrap();
assert!(
block.encoding == EncodingType::Gorilla || block.encoding == EncodingType::PlainF64
);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::F64(values));
}
#[test]
fn encode_f64_plain_strategy() {
let values = vec![1.0, 2.0, 3.0];
let block = ColumnEncoder::encode_f64(&values, FloatEncoding::Plain).unwrap();
assert_eq!(block.encoding, EncodingType::PlainF64);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::F64(values));
}
#[test]
fn encode_i64_with_fallback() {
let values: Vec<i64> = (0..1000).collect();
let block = ColumnEncoder::encode_i64(&values).unwrap();
assert!(
block.encoding == EncodingType::PcoI64
|| block.encoding == EncodingType::IntegerI64
|| block.encoding == EncodingType::PlainI64
);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::I64(values));
}
#[test]
fn encode_u64_with_fallback() {
let values: Vec<u64> = (0..1000).collect();
let block = ColumnEncoder::encode_u64(&values).unwrap();
assert!(
block.encoding == EncodingType::PcoU64
|| block.encoding == EncodingType::IntegerU64
|| block.encoding == EncodingType::PlainU64
);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::U64(values));
}
#[test]
fn encode_string_low_cardinality() {
let raw = vec!["us-east-1"; 500];
let block = ColumnEncoder::encode_string(&raw).unwrap();
assert_eq!(block.encoding, EncodingType::Dictionary);
let decoded = ColumnDecoder::decode(&block).unwrap();
let expected: Vec<String> = raw.iter().map(|s| (*s).to_owned()).collect();
assert_eq!(decoded, DecodedColumn::String(expected));
}
#[test]
fn encode_string_high_cardinality_fallback() {
let raw: Vec<String> = (0..200).map(|i| format!("uuid-{i}")).collect();
let refs: Vec<&str> = raw.iter().map(String::as_str).collect();
let block = ColumnEncoder::encode_string(&refs).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::String(raw));
}
#[test]
fn encode_bool_bitmap() {
let values: Vec<bool> = (0..100).map(|i| i % 3 == 0).collect();
let block = ColumnEncoder::encode_bool(&values).unwrap();
assert_eq!(block.encoding, EncodingType::Bitmap);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::Bool(values));
}
#[test]
fn empty_inputs_error() {
assert!(ColumnEncoder::encode_timestamps(&[]).is_err());
assert!(ColumnEncoder::encode_f64(&[], FloatEncoding::Chimp).is_err());
assert!(ColumnEncoder::encode_i64(&[]).is_err());
assert!(ColumnEncoder::encode_u64(&[]).is_err());
let empty: &[&str] = &[];
assert!(ColumnEncoder::encode_string(empty).is_err());
assert!(ColumnEncoder::encode_bool(&[]).is_err());
}
#[test]
fn should_use_specialized_logic() {
assert!(should_use_specialized(&[0u8; 50], 100));
assert!(!should_use_specialized(&[0u8; 80], 100));
assert!(!should_use_specialized(&[0u8; 10], 0));
}
#[test]
fn nullable_f64_mixed() {
let values = vec![Some(1.0), None, Some(3.0), None, Some(5.0)];
let block = ColumnEncoder::encode_f64_nullable(&values, FloatEncoding::Plain).unwrap();
assert_eq!(block.encoding, EncodingType::Nullable);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::NullableF64(values));
}
#[test]
fn nullable_f64_all_present() {
let values = vec![Some(1.0), Some(2.0), Some(3.0)];
let block = ColumnEncoder::encode_f64_nullable(&values, FloatEncoding::Plain).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::NullableF64(values));
}
#[test]
fn nullable_f64_all_null() {
let values: Vec<Option<f64>> = vec![None, None, None];
let block = ColumnEncoder::encode_f64_nullable(&values, FloatEncoding::Plain).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::NullableF64(values));
}
#[test]
fn nullable_i64_mixed() {
let values = vec![Some(10), None, Some(30), Some(40), None];
let block = ColumnEncoder::encode_i64_nullable(&values).unwrap();
assert_eq!(block.encoding, EncodingType::Nullable);
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::NullableI64(values));
}
#[test]
fn nullable_u64_mixed() {
let values = vec![None, Some(100u64), Some(200), None];
let block = ColumnEncoder::encode_u64_nullable(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
assert_eq!(decoded, DecodedColumn::NullableU64(values));
}
#[test]
fn nullable_string_mixed() {
let values = vec![Some("hello"), None, Some("world"), None];
let block = ColumnEncoder::encode_string_nullable(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
let expected: Vec<Option<String>> =
values.iter().map(|v| v.map(ToString::to_string)).collect();
assert_eq!(decoded, DecodedColumn::NullableString(expected));
}
#[test]
fn nullable_f64_bytes_roundtrip() {
let values = vec![Some(1.0), None, Some(3.0)];
let block = ColumnEncoder::encode_f64_nullable(&values, FloatEncoding::Plain).unwrap();
let raw = block.to_bytes();
let decoded = ColumnDecoder::decode_bytes(&raw).unwrap();
assert_eq!(decoded, DecodedColumn::NullableF64(values));
}
#[test]
fn nullable_empty_errors() {
let empty_f64: &[Option<f64>] = &[];
assert!(ColumnEncoder::encode_f64_nullable(empty_f64, FloatEncoding::Plain).is_err());
let empty_i64: &[Option<i64>] = &[];
assert!(ColumnEncoder::encode_i64_nullable(empty_i64).is_err());
let empty_u64: &[Option<u64>] = &[];
assert!(ColumnEncoder::encode_u64_nullable(empty_u64).is_err());
let empty_str: &[Option<&str>] = &[];
assert!(ColumnEncoder::encode_string_nullable(empty_str).is_err());
let empty_bool: &[Option<bool>] = &[];
assert!(ColumnEncoder::encode_bool_nullable(empty_bool).is_err());
}
#[test]
fn nullable_distinguishes_zero_from_null() {
let values = vec![Some(0.0), None, Some(0.0), None];
let block = ColumnEncoder::encode_f64_nullable(&values, FloatEncoding::Plain).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
match &decoded {
DecodedColumn::NullableF64(vals) => {
assert_eq!(vals[0], Some(0.0));
assert_eq!(vals[1], None); assert_eq!(vals[2], Some(0.0));
assert_eq!(vals[3], None); }
other => panic!("Expected NullableF64, got {other:?}"),
}
}
#[test]
fn nullable_i64_distinguishes_zero_from_null() {
let values = vec![Some(0i64), None, Some(0), None, Some(42)];
let block = ColumnEncoder::encode_i64_nullable(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
match &decoded {
DecodedColumn::NullableI64(vals) => {
assert_eq!(vals[0], Some(0));
assert_eq!(vals[1], None);
assert_eq!(vals[2], Some(0));
assert_eq!(vals[3], None);
assert_eq!(vals[4], Some(42));
}
other => panic!("Expected NullableI64, got {other:?}"),
}
}
}
#[cfg(test)]
mod proptests {
use super::*;
use proptest::prelude::*;
proptest! {
#[test]
fn roundtrip_timestamps(mut values in proptest::collection::vec(any::<i64>(), 1..300)) {
values.sort_unstable();
let block = ColumnEncoder::encode_timestamps(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::I64(values), decoded);
}
#[test]
fn roundtrip_i64(values in proptest::collection::vec(any::<i64>(), 1..300)) {
let block = ColumnEncoder::encode_i64(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::I64(values), decoded);
}
#[test]
fn roundtrip_u64(values in proptest::collection::vec(any::<u64>(), 1..300)) {
let block = ColumnEncoder::encode_u64(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::U64(values), decoded);
}
#[test]
fn roundtrip_f64(values in proptest::collection::vec(any::<f64>(), 1..300)) {
let block = ColumnEncoder::encode_f64(&values, FloatEncoding::Gorilla).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::F64(values), decoded);
}
#[test]
fn roundtrip_bool(values in proptest::collection::vec(any::<bool>(), 1..300)) {
let block = ColumnEncoder::encode_bool(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::Bool(values), decoded);
}
#[test]
fn roundtrip_string(values in proptest::collection::vec("[a-z]{0,8}", 1..100)) {
let refs: Vec<&str> = values.iter().map(String::as_str).collect();
let block = ColumnEncoder::encode_string(&refs).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::String(values), decoded);
}
#[test]
fn roundtrip_nullable_i64(values in proptest::collection::vec(any::<Option<i64>>(), 1..300)) {
let block = ColumnEncoder::encode_i64_nullable(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::NullableI64(values), decoded);
}
#[test]
fn roundtrip_nullable_u64(values in proptest::collection::vec(any::<Option<u64>>(), 1..300)) {
let block = ColumnEncoder::encode_u64_nullable(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::NullableU64(values), decoded);
}
#[test]
fn roundtrip_nullable_f64(values in proptest::collection::vec(any::<Option<f64>>(), 1..300)) {
let block = ColumnEncoder::encode_f64_nullable(&values, FloatEncoding::Gorilla).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::NullableF64(values), decoded);
}
#[test]
fn roundtrip_nullable_bool(values in proptest::collection::vec(any::<Option<bool>>(), 1..300)) {
let block = ColumnEncoder::encode_bool_nullable(&values).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
prop_assert_eq!(DecodedColumn::NullableBool(values), decoded);
}
#[test]
fn roundtrip_nullable_string(values in proptest::collection::vec(
proptest::option::of("[a-z]{0,8}"), 1..100
)) {
let refs: Vec<Option<&str>> = values.iter().map(|o| o.as_deref()).collect();
let block = ColumnEncoder::encode_string_nullable(&refs).unwrap();
let decoded = ColumnDecoder::decode(&block).unwrap();
let expected: Vec<Option<String>> = values.clone();
prop_assert_eq!(DecodedColumn::NullableString(expected), decoded);
}
}
}