use arrow_schema::{DataType, Field, Schema};
use iris_abi::Node;
use crate::error::{Invariant, Result, Violation};
use crate::layout::{Layout, layout, offset_width, slot_bits};
pub const MAX_DEPTH: usize = 64;
pub fn check_schema(schema: &Schema) -> Result<()> {
let mut work: Vec<(&Field, usize, String)> = schema
.fields()
.iter()
.map(|field| (field.as_ref(), 1, field.name().clone()))
.collect();
while let Some((field, depth, path)) = work.pop() {
if depth > MAX_DEPTH {
return Err(Violation::at(
Invariant::Depth,
&path,
format!("this build walks {MAX_DEPTH} levels of nesting and this is deeper"),
));
}
let found = layout(field.data_type(), &path)?;
for child in found.children {
let child_path = format!("{path}.{}", child.name());
work.push((child, depth + 1, child_path));
}
}
Ok(())
}
pub fn check<B: AsRef<[u8]>>(
schema: &Schema,
rows: u64,
nodes: &[Node],
buffers: &[B],
) -> Result<()> {
check_schema(schema)?;
let mut cursor = Cursor {
nodes,
buffers,
node: 0,
buffer: 0,
};
for field in schema.fields() {
let len = cursor.array(field, field.name())?;
if len != rows {
return Err(Violation::at(
Invariant::Rows,
field.name(),
format!("the batch says {rows} rows and this column has {len}"),
));
}
}
cursor.finish()
}
struct Cursor<'a, B> {
nodes: &'a [Node],
buffers: &'a [B],
node: usize,
buffer: usize,
}
impl<'a, B: AsRef<[u8]>> Cursor<'a, B> {
fn array(&mut self, field: &Field, path: &str) -> Result<u64> {
let node = self.next_node(path)?;
let data_type = field.data_type();
let Layout {
validity,
values,
children,
} = layout(data_type, path)?;
if validity {
let bitmap = self.next_buffer(path)?;
check_validity(bitmap, node.length, node.null_count, path)?;
} else if node.null_count != 0 {
return Err(Violation::at(
Invariant::NullCount,
path,
format!(
"a {data_type} column has no validity buffer and this one says it has {} nulls",
node.null_count
),
));
}
let mut taken = Vec::with_capacity(values);
for _ in 0..values {
taken.push(self.next_buffer(path)?);
}
let mut child_lengths = Vec::with_capacity(children.len());
for child in &children {
let child_path = format!("{path}.{}", child.name());
child_lengths.push(self.array(child, &child_path)?);
}
check_values(data_type, node.length, &taken, &child_lengths, path)?;
Ok(node.length)
}
fn next_node(&mut self, path: &str) -> Result<Node> {
let node = self.nodes.get(self.node).copied().ok_or_else(|| {
Violation::at(
Invariant::Arrays,
path,
format!(
"the schema calls for more arrays than the batch has, which is {}",
self.nodes.len()
),
)
})?;
self.node += 1;
Ok(node)
}
fn next_buffer(&mut self, path: &str) -> Result<&'a [u8]> {
let bytes = self.buffers.get(self.buffer).ok_or_else(|| {
Violation::at(
Invariant::Buffers,
path,
format!(
"the schema calls for more buffers than the batch has, which is {}",
self.buffers.len()
),
)
})?;
self.buffer += 1;
Ok(bytes.as_ref())
}
fn finish(&self) -> Result<()> {
if self.node != self.nodes.len() {
return Err(Violation::at(
Invariant::Arrays,
"",
format!(
"the batch has {} arrays and the schema accounts for {}",
self.nodes.len(),
self.node
),
));
}
if self.buffer != self.buffers.len() {
return Err(Violation::at(
Invariant::Buffers,
"",
format!(
"the batch has {} buffers and the schema accounts for {}",
self.buffers.len(),
self.buffer
),
));
}
Ok(())
}
}
fn check_validity(bitmap: &[u8], len: u64, null_count: u64, path: &str) -> Result<()> {
if bitmap.is_empty() {
if null_count != 0 {
return Err(Violation::at(
Invariant::NullCount,
path,
format!(
"there is no validity buffer and this array says it has {null_count} nulls"
),
));
}
return Ok(());
}
let needed = len.div_ceil(8);
let have = as_u64(bitmap.len());
if have < needed {
return Err(Violation::at(
Invariant::Validity,
path,
format!("{len} slots need {needed} bytes of validity and there are {have}"),
));
}
let counted = count_nulls(bitmap, len);
if counted != null_count {
return Err(Violation::at(
Invariant::NullCount,
path,
format!("this array says it has {null_count} nulls and its bitmap has {counted}"),
));
}
Ok(())
}
fn count_nulls(bitmap: &[u8], len: u64) -> u64 {
let available = as_u64(bitmap.len()).saturating_mul(8);
let considered = len.min(available);
let whole = bitmap
.len()
.min(usize::try_from(considered / 8).unwrap_or(usize::MAX));
let mut set: u64 = bitmap[..whole]
.iter()
.map(|byte| u64::from(byte.count_ones()))
.sum();
let spare = u32::try_from(considered % 8).unwrap_or(0);
if spare != 0 {
let mask = (1u8 << spare) - 1;
set += u64::from((bitmap.get(whole).copied().unwrap_or(0) & mask).count_ones());
}
considered - set
}
fn check_values(
data_type: &DataType,
len: u64,
buffers: &[&[u8]],
child_lengths: &[u64],
path: &str,
) -> Result<()> {
let child = child_lengths.first().copied().unwrap_or(0);
match data_type {
DataType::Null => Ok(()),
DataType::Utf8 | DataType::Binary | DataType::LargeUtf8 | DataType::LargeBinary => {
check_variable(data_type, len, buffers, path)
}
DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _) => {
check_list(data_type, len, buffers, child, path)
}
DataType::FixedSizeList(_, size) => check_fixed_size_list(len, *size, child, path),
DataType::Struct(fields) => {
for (field, child) in fields.iter().zip(child_lengths) {
if *child < len {
return Err(Violation::at(
Invariant::ChildLength,
path,
format!(
"this struct has {len} rows and its {} field has {child}",
field.name()
),
));
}
}
Ok(())
}
other => check_fixed_width(other, len, buffers, path),
}
}
fn check_variable(data_type: &DataType, len: u64, buffers: &[&[u8]], path: &str) -> Result<()> {
let [offsets, data] = buffers else {
return Err(counted_wrong(path, 2, buffers.len()));
};
let width = offset_width(data_type).expect("a variable length type has offsets");
let last = check_offsets(offsets, len, width, path)?;
let have = as_u64(data.len());
if last > have {
return Err(Violation::at(
Invariant::OffsetRange,
path,
format!("the last offset is {last} and the values buffer is {have} bytes"),
));
}
Ok(())
}
fn check_list(
data_type: &DataType,
len: u64,
buffers: &[&[u8]],
child: u64,
path: &str,
) -> Result<()> {
let [offsets] = buffers else {
return Err(counted_wrong(path, 1, buffers.len()));
};
let width = offset_width(data_type).expect("a list has offsets");
let last = check_offsets(offsets, len, width, path)?;
if last > child {
return Err(Violation::at(
Invariant::OffsetRange,
path,
format!("the last offset is {last} and the child array has {child} slots"),
));
}
Ok(())
}
fn check_fixed_size_list(len: u64, size: i32, child: u64, path: &str) -> Result<()> {
let size = u64::try_from(size).map_err(|_| {
Violation::at(
Invariant::ChildLength,
path,
format!("a fixed size list cannot hold {size} values a row"),
)
})?;
let needed = len.checked_mul(size).ok_or_else(|| {
Violation::at(
Invariant::Size,
path,
format!("{len} rows of {size} values is more than this host can address"),
)
})?;
if child < needed {
return Err(Violation::at(
Invariant::ChildLength,
path,
format!("{len} rows of {size} values need {needed} slots and the child has {child}"),
));
}
Ok(())
}
fn check_fixed_width(data_type: &DataType, len: u64, buffers: &[&[u8]], path: &str) -> Result<()> {
let [values] = buffers else {
return Err(counted_wrong(path, 1, buffers.len()));
};
let bits = slot_bits(data_type).ok_or_else(|| {
Violation::at(
Invariant::Unsupported,
path,
format!("this build does not know how wide a {data_type} slot is"),
)
})?;
let needed = len
.checked_mul(bits)
.map(|total| total.div_ceil(8))
.ok_or_else(|| {
Violation::at(
Invariant::Size,
path,
format!("{len} slots of {bits} bits is more than this host can address"),
)
})?;
let have = as_u64(values.len());
if have < needed {
return Err(Violation::at(
Invariant::BufferLength,
path,
format!("{len} slots of {bits} bits need {needed} bytes and there are {have}"),
));
}
Ok(())
}
fn check_offsets(offsets: &[u8], len: u64, width: u64, path: &str) -> Result<u64> {
if len == 0 && offsets.is_empty() {
return Ok(0);
}
let entries = len.checked_add(1).ok_or_else(|| {
Violation::at(
Invariant::Size,
path,
format!("{len} slots need one more offset than that, which does not fit in a count"),
)
})?;
let needed = entries.checked_mul(width).ok_or_else(|| {
Violation::at(
Invariant::Size,
path,
format!("{entries} offsets of {width} bytes is more than this host can address"),
)
})?;
let have = as_u64(offsets.len());
if have < needed {
return Err(Violation::at(
Invariant::BufferLength,
path,
format!("{len} slots need {needed} bytes of offsets and there are {have}"),
));
}
let span = usize::try_from(needed).map_err(|_| {
Violation::at(
Invariant::Size,
path,
"the offsets run past what this host can address".to_owned(),
)
})?;
let run = offsets.get(..span).unwrap_or(offsets);
let previous = if width == 8 {
scan_offsets::<8>(run, path, i64::from_le_bytes)?
} else {
scan_offsets::<4>(run, path, |raw| i64::from(i32::from_le_bytes(raw)))?
};
u64::try_from(previous).map_err(|_| {
Violation::at(
Invariant::OffsetRange,
path,
"the last offset is negative".to_owned(),
)
})
}
fn scan_offsets<const W: usize>(
offsets: &[u8],
path: &str,
read: fn([u8; W]) -> i64,
) -> Result<i64> {
let mut previous: i64 = 0;
for (index, raw) in offsets.as_chunks::<W>().0.iter().enumerate() {
let offset = read(*raw);
if offset < 0 {
return Err(Violation::at(
Invariant::OffsetRange,
path,
format!("offset {index} is {offset}, and an offset is a position"),
));
}
if index > 0 && offset < previous {
return Err(Violation::at(
Invariant::OffsetOrder,
path,
format!("offset {index} is {offset} and the one before it is {previous}"),
));
}
previous = offset;
}
Ok(previous)
}
fn as_u64(len: usize) -> u64 {
u64::try_from(len).unwrap_or(u64::MAX)
}
fn counted_wrong(path: &str, wanted: usize, found: usize) -> Violation {
Violation::at(
Invariant::Buffers,
path,
format!("this column takes {wanted} buffers after its validity buffer and got {found}"),
)
}
#[cfg(test)]
mod tests {
use arrow_schema::{DataType, Field, Fields, Schema};
use iris_abi::Node;
use super::{MAX_DEPTH, check, check_schema, count_nulls};
use crate::error::Invariant;
fn node(length: u64, null_count: u64) -> Node {
Node { length, null_count }
}
fn i64s(values: &[i64]) -> Vec<u8> {
values.iter().flat_map(|v| v.to_le_bytes()).collect()
}
fn i32s(values: &[i32]) -> Vec<u8> {
values.iter().flat_map(|v| v.to_le_bytes()).collect()
}
#[test]
fn a_sound_batch_passes() {
let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
let buffers = vec![Vec::new(), i64s(&[1, 2, 3])];
check(&schema, 3, &[node(3, 0)], &buffers).expect("this batch is sound");
}
#[test]
fn a_column_shorter_than_the_batch_is_caught() {
let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
let buffers = vec![Vec::new(), i64s(&[1, 2])];
let err = check(&schema, 3, &[node(2, 0)], &buffers).expect_err("two is not three");
assert_eq!(err.invariant, Invariant::Rows);
}
#[test]
fn a_values_buffer_one_slot_short_is_caught() {
let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
let buffers = vec![Vec::new(), i64s(&[1, 2])];
let err =
check(&schema, 3, &[node(3, 0)], &buffers).expect_err("three slots need 24 bytes");
assert_eq!(err.invariant, Invariant::BufferLength);
assert!(err.to_string().contains("24 bytes"), "{err}");
}
#[test]
fn a_bitmap_with_too_few_bits_is_caught() {
let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
let buffers = vec![Vec::new(), i64s(&[1, 2, 3])];
let short = vec![vec![0xffu8], i64s(&[1; 100])];
let wide = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
let err = check(&wide, 100, &[node(100, 0)], &short).expect_err("100 slots need 13 bytes");
assert_eq!(err.invariant, Invariant::Validity);
check(&schema, 3, &[node(3, 0)], &buffers).expect("the empty bitmap case still passes");
}
#[test]
fn an_offset_one_past_the_end_is_caught() {
let schema = Schema::new(vec![Field::new("s", DataType::Utf8, false)]);
let buffers = vec![Vec::new(), i32s(&[0, 2, 6]), b"hoyea".to_vec()];
let err = check(&schema, 2, &[node(2, 0)], &buffers).expect_err("six is past five");
assert_eq!(err.invariant, Invariant::OffsetRange);
}
#[test]
fn offsets_that_run_backwards_are_caught() {
let schema = Schema::new(vec![Field::new("s", DataType::Utf8, false)]);
let buffers = vec![Vec::new(), i32s(&[0, 4, 2]), b"hoyea".to_vec()];
let err = check(&schema, 2, &[node(2, 0)], &buffers).expect_err("two is less than four");
assert_eq!(err.invariant, Invariant::OffsetOrder);
}
#[test]
fn a_child_one_row_short_of_its_parent_is_caught() {
let children = Fields::from(vec![Field::new("x", DataType::Int64, false)]);
let schema = Schema::new(vec![Field::new("p", DataType::Struct(children), false)]);
let buffers = vec![Vec::new(), Vec::new(), i64s(&[1, 2])];
let err = check(&schema, 3, &[node(3, 0), node(2, 0)], &buffers)
.expect_err("a struct's child cannot be shorter than the struct");
assert_eq!(err.invariant, Invariant::ChildLength);
}
#[test]
fn a_length_that_overflows_a_width_is_caught_rather_than_wrapped() {
let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
let buffers = vec![Vec::new(), i64s(&[1])];
let err = check(&schema, u64::MAX, &[node(u64::MAX, 0)], &buffers)
.expect_err("that many slots is not addressable");
assert_eq!(err.invariant, Invariant::Size);
}
#[test]
fn a_length_that_wraps_the_count_of_offsets_is_caught() {
for data_type in [DataType::Binary, DataType::LargeBinary] {
let schema = Schema::new(vec![Field::new("a", data_type, false)]);
let buffers = vec![Vec::new(), Vec::new(), Vec::new()];
let err = check(&schema, u64::MAX, &[node(u64::MAX, 0)], &buffers)
.expect_err("one more offset than that does not fit in a count");
assert_eq!(err.invariant, Invariant::Size);
}
}
#[test]
fn a_schema_nested_past_the_bound_is_refused_without_recursing_into_it() {
let mut data_type = DataType::Int64;
for _ in 0..MAX_DEPTH + 10 {
data_type = DataType::List(std::sync::Arc::new(Field::new("item", data_type, false)));
}
let schema = Schema::new(vec![Field::new("deep", data_type, false)]);
let err = check_schema(&schema).expect_err("that is deeper than this build walks");
assert_eq!(err.invariant, Invariant::Depth);
}
#[test]
fn a_schema_at_the_bound_is_still_walked() {
let mut data_type = DataType::Int64;
for _ in 0..MAX_DEPTH - 1 {
data_type = DataType::List(std::sync::Arc::new(Field::new("item", data_type, false)));
}
let schema = Schema::new(vec![Field::new("deep", data_type, false)]);
check_schema(&schema).expect("this is exactly as deep as the bound allows");
}
#[test]
fn spare_buffers_are_an_error_rather_than_something_ignored() {
let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
let buffers = vec![Vec::new(), i64s(&[1, 2, 3]), i64s(&[4])];
let err =
check(&schema, 3, &[node(3, 0)], &buffers).expect_err("a spare buffer is not fine");
assert_eq!(err.invariant, Invariant::Buffers);
}
#[test]
fn counting_nulls_stops_at_the_length_rather_than_the_byte() {
assert_eq!(count_nulls(&[0b0001_1111], 5), 0);
assert_eq!(count_nulls(&[0b0001_1110], 5), 1);
assert_eq!(count_nulls(&[0x00, 0xff], 9), 8);
}
#[test]
fn counting_nulls_agrees_with_reading_the_bits_one_at_a_time() {
fn one_at_a_time(bitmap: &[u8], len: u64) -> u64 {
let mut nulls = 0;
for bit in 0..len {
let byte = usize::try_from(bit / 8).expect("this test is small");
let Some(value) = bitmap.get(byte) else {
break;
};
let shift = u32::try_from(bit % 8).expect("a bit in a byte");
if value >> shift & 1 == 0 {
nulls += 1;
}
}
nulls
}
for bitmap in [
[0x00u8, 0x00, 0x00],
[0xff, 0xff, 0xff],
[0b1010_1010, 0b0000_1111, 0b1100_0011],
] {
for len in 0..40 {
assert_eq!(
count_nulls(&bitmap, len),
one_at_a_time(&bitmap, len),
"bitmap {bitmap:?} at length {len}"
);
}
}
}
}