mod any;
mod fixed;
mod varlen;
use std::{str, sync::Arc};
use arrow_array::{ArrayRef, LargeStringArray, RecordBatch, RecordBatchOptions, make_array};
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
use arrow_schema::{FieldRef, Schema};
use reifydb_value::{
encoding::LeBytes,
reifydb_assertions,
value::{
container::varlen_array::blob_array,
diff_type::DiffType,
frame::frame::Frame,
value_type::{
ValueType,
field::{FieldType, to_field},
},
},
};
use crate::{
error::DecodeError,
frame::{
encoding::dict::{decode_dict_blob, decode_dict_utf8},
format::{
COL_FLAG_HAS_NONES, COLUMN_DESCRIPTOR_SIZE, Encoding, FRAME_HEADER_SIZE, MESSAGE_HEADER_SIZE,
RBCF_MAGIC, RBCF_VERSION, dict_index_width_from_flags,
},
},
tag::{TypeTag, ValueKind},
};
pub fn decode_frames(data: &[u8]) -> Result<Vec<Frame>, DecodeError> {
let mut pos = 0;
check_len(data, pos, MESSAGE_HEADER_SIZE)?;
let magic = read_u32(data, pos);
pos += 4;
if magic != RBCF_MAGIC {
return Err(DecodeError::InvalidMagic(magic));
}
let version = read_u16(data, pos);
pos += 2;
if version != RBCF_VERSION {
return Err(DecodeError::UnsupportedVersion(version));
}
let _flags = read_u16(data, pos);
pos += 2;
let frame_count = read_u32(data, pos) as usize;
pos += 4;
let _total_size = read_u32(data, pos) as usize;
pos += 4;
let mut frames = Vec::with_capacity(frame_count);
for _ in 0..frame_count {
let (frame, new_pos) = decode_frame(data, pos)?;
frames.push(frame);
pos = new_pos;
}
reifydb_assertions! {
assert!(
pos == _total_size,
"the RBCF message header declared a total size that disagrees with the bytes consumed while decoding, so the message is truncated or carries trailing bytes and a peer would mis-frame the following message (declared={} consumed={})",
_total_size,
pos
);
}
Ok(frames)
}
struct FrameHeader {
row_count: usize,
column_count: usize,
op: Option<DiffType>,
}
fn decode_frame(data: &[u8], start: usize) -> Result<(Frame, usize), DecodeError> {
let (header, pos) = read_frame_header(data, start)?;
let (columns, pos) = read_frame_columns(data, pos, header.column_count)?;
Ok((
Frame {
batch: frame_batch(columns, header.row_count)?,
op: header.op,
},
pos,
))
}
pub(crate) fn frame_batch(columns: Vec<(FieldRef, ArrayRef)>, row_count: usize) -> Result<RecordBatch, DecodeError> {
let (fields, arrays): (Vec<FieldRef>, Vec<ArrayRef>) = columns.into_iter().unzip();
RecordBatch::try_new_with_options(
Arc::new(Schema::new(fields)),
arrays,
&RecordBatchOptions::new().with_row_count(Some(row_count)),
)
.map_err(|error| DecodeError::InvalidData(format!("frame columns do not form a batch: {error}")))
}
#[inline]
fn read_frame_header(data: &[u8], start: usize) -> Result<(FrameHeader, usize), DecodeError> {
let mut pos = start;
check_len(data, pos, FRAME_HEADER_SIZE)?;
let row_count = read_u32(data, pos) as usize;
pos += 4;
let column_count = read_u16(data, pos) as usize;
pos += 2;
let meta_flags = data[pos];
if meta_flags != 0 {
return Err(DecodeError::InvalidData(format!("frame meta byte must be 0, found 0x{meta_flags:02X}")));
}
pos += 1;
let op = match data[pos] {
0 => None,
raw => Some(DiffType::from_u8(raw)
.ok_or_else(|| DecodeError::InvalidData(format!("unknown frame op {raw}")))?),
};
pos += 1;
let _frame_size = read_u32(data, pos);
pos += 4;
Ok((
FrameHeader {
row_count,
column_count,
op,
},
pos,
))
}
#[inline]
fn read_frame_columns(
data: &[u8],
mut pos: usize,
column_count: usize,
) -> Result<(Vec<(FieldRef, ArrayRef)>, usize), DecodeError> {
let mut columns = Vec::with_capacity(column_count);
for _ in 0..column_count {
let (col, new_pos) = decode_column(data, pos)?;
columns.push(col);
pos = new_pos;
}
Ok((columns, pos))
}
fn decode_column(data: &[u8], start: usize) -> Result<((FieldRef, ArrayRef), usize), DecodeError> {
let mut pos = start;
check_len(data, pos, COLUMN_DESCRIPTOR_SIZE)?;
let type_code = data[pos];
pos += 1;
let encoding_byte = data[pos];
pos += 1;
let flags = data[pos];
pos += 1;
let _reserved = data[pos];
pos += 1;
let name_len = read_u16(data, pos) as usize;
pos += 2;
let _reserved2 = read_u16(data, pos);
pos += 2;
let row_count = read_u32(data, pos) as usize;
pos += 4;
let nones_len = read_u32(data, pos) as usize;
pos += 4;
let data_len = read_u32(data, pos) as usize;
pos += 4;
let offsets_len = read_u32(data, pos) as usize;
pos += 4;
let extra_len = read_u32(data, pos) as usize;
pos += 4;
let encoding = Encoding::from_u8(encoding_byte).ok_or(DecodeError::UnknownEncoding(encoding_byte))?;
let has_nones = flags & COL_FLAG_HAS_NONES != 0;
let tag = TypeTag::from_byte(type_code)?;
let depth = tag.depth() as usize;
let base_code = tag.kind_bits();
check_len(data, pos, name_len)?;
let name = str::from_utf8(&data[pos..pos + name_len])
.map_err(|e| DecodeError::InvalidData(format!("invalid column name: {}", e)))?
.to_string();
pos += name_len;
let name_pad = (4 - (name_len % 4)) % 4;
pos += name_pad;
let result = (|| -> Result<((FieldRef, ArrayRef), usize), DecodeError> {
let mut pos = pos;
if has_nones != (depth > 0) {
return Err(DecodeError::InvalidData(format!(
"column type code 0x{type_code:02X} has option depth {depth} but the has-nones flag is {}",
if has_nones {
"set"
} else {
"clear"
}
)));
}
if depth > 1 {
return Err(DecodeError::InvalidData(format!(
"column type code 0x{type_code:02X} has option depth {depth}, but a column holds at most one option layer"
)));
}
let layer_len = row_count.div_ceil(8);
if nones_len != depth * layer_len {
return Err(DecodeError::InvalidData(format!(
"nones length {nones_len} disagrees with option depth {depth} and row count {row_count} (expected {})",
depth * layer_len
)));
}
check_len(data, pos, nones_len)?;
let layers: Vec<BooleanBuffer> = (0..depth)
.map(|layer| {
let start = pos + layer * layer_len;
decode_bitvec(&data[start..start + layer_len], row_count)
})
.collect();
pos += nones_len;
check_len(data, pos, data_len)?;
let data_bytes = &data[pos..pos + data_len];
pos += data_len;
check_len(data, pos, offsets_len)?;
let offsets_bytes = &data[pos..pos + offsets_len];
pos += offsets_len;
check_len(data, pos, extra_len)?;
let extra_bytes = &data[pos..pos + extra_len];
pos += extra_len;
let (value_type, array) = decode_column_dispatch(
base_code,
encoding,
flags,
row_count,
data_bytes,
offsets_bytes,
extra_bytes,
layers.last(),
)?;
Ok((decoded_column(&name, value_type, array, layers.into_iter().next())?, pos))
})()
.map_err(|e| DecodeError::ColumnDecodeFailed {
column_name: name.clone(),
row_index: None,
source: Box::new(e),
})?;
Ok(result)
}
pub(crate) fn decoded_column(
name: &str,
value_type: ValueType,
array: ArrayRef,
layer: Option<BooleanBuffer>,
) -> Result<(FieldRef, ArrayRef), DecodeError> {
let declared_type = matches!(value_type, ValueType::List(_) | ValueType::Record(_) | ValueType::Tuple(_))
.then(|| value_type.clone());
let (value_type, array) = match layer {
Some(layer) => (ValueType::Option(Box::new(value_type)), attach_layer(array, layer)?),
None if array.logical_null_count() == 0 => (value_type, array),
None if matches!(value_type, ValueType::Digest { .. }) => {
(ValueType::Option(Box::new(value_type)), array)
}
None => {
return Err(DecodeError::InvalidData(format!(
"column of type {value_type:?} holds {} none cells but has no option layer",
array.logical_null_count()
)));
}
};
let field_type = FieldType {
value_type: Some(value_type),
declared_type,
..FieldType::default()
};
Ok((Arc::new(to_field(name, &field_type)), array))
}
fn attach_layer(array: ArrayRef, layer: BooleanBuffer) -> Result<ArrayRef, DecodeError> {
let nulls = NullBuffer::union(Some(&NullBuffer::new(layer)), array.logical_nulls().as_ref());
let data =
array.to_data().into_builder().nulls(nulls).build().map_err(|error| {
DecodeError::InvalidData(format!("option layer does not fit the column: {error}"))
})?;
Ok(make_array(data))
}
pub(crate) fn column_type_from_code(type_code: u8) -> Result<ValueType, DecodeError> {
TypeTag::from_byte(type_code)?.to_type()
}
#[allow(clippy::too_many_arguments)]
fn decode_column_dispatch(
type_code: u8,
encoding: Encoding,
flags: u8,
row_count: usize,
data: &[u8],
offsets: &[u8],
extra: &[u8],
defined: Option<&BooleanBuffer>,
) -> Result<(ValueType, ArrayRef), DecodeError> {
if type_code == ValueKind::Digest.byte() {
if encoding != Encoding::Plain {
return Err(DecodeError::InvalidData(format!(
"digest column must use plain encoding, found {encoding:?}"
)));
}
return varlen::decode_digest_plain(row_count, data, offsets, extra);
}
if let Some(kind @ ValueKind::Decimal) = ValueKind::from_byte(type_code) {
return fixed::decode_unscaled_column(kind, encoding, row_count, data, extra);
}
let ty = column_type_from_code(type_code)?;
let array: ArrayRef = match encoding {
Encoding::Plain | Encoding::BitPack => {
if ty == ValueType::Any {
any::decode_any_column(row_count, data, defined)?
} else if let Some(result) = fixed::decode_fixed_plain(type_code, row_count, data) {
result?
} else if let Some(result) = varlen::decode_varlen_plain(type_code, row_count, data, offsets) {
result?
} else {
return Err(DecodeError::UnsupportedType(format!("{:?}", ty)));
}
}
Encoding::Dict => match ty {
ValueType::Utf8 => {
let index_width = dict_index_width_from_flags(flags);
let strings = decode_dict_utf8(data, extra, row_count, index_width)?;
Arc::new(LargeStringArray::from(strings))
}
ValueType::Blob => {
let index_width = dict_index_width_from_flags(flags);
let blobs = decode_dict_blob(data, extra, row_count, index_width)?;
Arc::new(blob_array(&blobs))
}
_ => {
return Err(DecodeError::InvalidData(format!(
"Dict encoding not supported for type {:?}",
ty
)));
}
},
Encoding::Rle => fixed::decode_rle_column(type_code, row_count, data)?,
Encoding::Delta => fixed::decode_delta_column(type_code, row_count, data)?,
Encoding::DeltaRle => fixed::decode_delta_rle_column(type_code, row_count, data)?,
};
Ok((ty, array))
}
fn decode_bitvec(data: &[u8], len: usize) -> BooleanBuffer {
BooleanBuffer::new(Buffer::from_vec(data.to_vec()), 0, len)
}
#[inline]
fn read_u16(data: &[u8], pos: usize) -> u16 {
u16::read_le(&data[pos..])
}
#[inline]
fn read_u32(data: &[u8], pos: usize) -> u32 {
u32::read_le(&data[pos..])
}
fn check_len(data: &[u8], pos: usize, needed: usize) -> Result<(), DecodeError> {
if pos + needed > data.len() {
Err(DecodeError::UnexpectedEof {
expected: needed,
available: data.len().saturating_sub(pos),
})
} else {
Ok(())
}
}