1use arrow_schema::{DataType, Field, Schema};
4use iris_abi::Node;
5
6use crate::error::{Invariant, Result, Violation};
7use crate::layout::{Layout, layout, offset_width, slot_bits};
8
9pub const MAX_DEPTH: usize = 64;
17
18pub fn check_schema(schema: &Schema) -> Result<()> {
29 let mut work: Vec<(&Field, usize, String)> = schema
33 .fields()
34 .iter()
35 .map(|field| (field.as_ref(), 1, field.name().clone()))
36 .collect();
37
38 while let Some((field, depth, path)) = work.pop() {
39 if depth > MAX_DEPTH {
40 return Err(Violation::at(
41 Invariant::Depth,
42 &path,
43 format!("this build walks {MAX_DEPTH} levels of nesting and this is deeper"),
44 ));
45 }
46 let found = layout(field.data_type(), &path)?;
47 for child in found.children {
48 let child_path = format!("{path}.{}", child.name());
49 work.push((child, depth + 1, child_path));
50 }
51 }
52
53 Ok(())
54}
55
56pub fn check<B: AsRef<[u8]>>(
73 schema: &Schema,
74 rows: u64,
75 nodes: &[Node],
76 buffers: &[B],
77) -> Result<()> {
78 check_schema(schema)?;
79
80 let mut cursor = Cursor {
81 nodes,
82 buffers,
83 node: 0,
84 buffer: 0,
85 };
86
87 for field in schema.fields() {
88 let len = cursor.array(field, field.name())?;
89 if len != rows {
90 return Err(Violation::at(
91 Invariant::Rows,
92 field.name(),
93 format!("the batch says {rows} rows and this column has {len}"),
94 ));
95 }
96 }
97
98 cursor.finish()
99}
100
101struct Cursor<'a, B> {
103 nodes: &'a [Node],
104 buffers: &'a [B],
105 node: usize,
106 buffer: usize,
107}
108
109impl<'a, B: AsRef<[u8]>> Cursor<'a, B> {
110 fn array(&mut self, field: &Field, path: &str) -> Result<u64> {
112 let node = self.next_node(path)?;
113 let data_type = field.data_type();
114 let Layout {
115 validity,
116 values,
117 children,
118 } = layout(data_type, path)?;
119
120 if validity {
121 let bitmap = self.next_buffer(path)?;
122 check_validity(bitmap, node.length, node.null_count, path)?;
123 } else if node.null_count != 0 {
124 return Err(Violation::at(
125 Invariant::NullCount,
126 path,
127 format!(
128 "a {data_type} column has no validity buffer and this one says it has {} nulls",
129 node.null_count
130 ),
131 ));
132 }
133
134 let mut taken = Vec::with_capacity(values);
135 for _ in 0..values {
136 taken.push(self.next_buffer(path)?);
137 }
138
139 let mut child_lengths = Vec::with_capacity(children.len());
142 for child in &children {
143 let child_path = format!("{path}.{}", child.name());
144 child_lengths.push(self.array(child, &child_path)?);
145 }
146
147 check_values(data_type, node.length, &taken, &child_lengths, path)?;
148 Ok(node.length)
149 }
150
151 fn next_node(&mut self, path: &str) -> Result<Node> {
152 let node = self.nodes.get(self.node).copied().ok_or_else(|| {
153 Violation::at(
154 Invariant::Arrays,
155 path,
156 format!(
157 "the schema calls for more arrays than the batch has, which is {}",
158 self.nodes.len()
159 ),
160 )
161 })?;
162 self.node += 1;
163 Ok(node)
164 }
165
166 fn next_buffer(&mut self, path: &str) -> Result<&'a [u8]> {
167 let bytes = self.buffers.get(self.buffer).ok_or_else(|| {
168 Violation::at(
169 Invariant::Buffers,
170 path,
171 format!(
172 "the schema calls for more buffers than the batch has, which is {}",
173 self.buffers.len()
174 ),
175 )
176 })?;
177 self.buffer += 1;
178 Ok(bytes.as_ref())
179 }
180
181 fn finish(&self) -> Result<()> {
187 if self.node != self.nodes.len() {
188 return Err(Violation::at(
189 Invariant::Arrays,
190 "",
191 format!(
192 "the batch has {} arrays and the schema accounts for {}",
193 self.nodes.len(),
194 self.node
195 ),
196 ));
197 }
198 if self.buffer != self.buffers.len() {
199 return Err(Violation::at(
200 Invariant::Buffers,
201 "",
202 format!(
203 "the batch has {} buffers and the schema accounts for {}",
204 self.buffers.len(),
205 self.buffer
206 ),
207 ));
208 }
209 Ok(())
210 }
211}
212
213fn check_validity(bitmap: &[u8], len: u64, null_count: u64, path: &str) -> Result<()> {
218 if bitmap.is_empty() {
219 if null_count != 0 {
220 return Err(Violation::at(
221 Invariant::NullCount,
222 path,
223 format!(
224 "there is no validity buffer and this array says it has {null_count} nulls"
225 ),
226 ));
227 }
228 return Ok(());
229 }
230
231 let needed = len.div_ceil(8);
232 let have = as_u64(bitmap.len());
233 if have < needed {
234 return Err(Violation::at(
235 Invariant::Validity,
236 path,
237 format!("{len} slots need {needed} bytes of validity and there are {have}"),
238 ));
239 }
240
241 let counted = count_nulls(bitmap, len);
245 if counted != null_count {
246 return Err(Violation::at(
247 Invariant::NullCount,
248 path,
249 format!("this array says it has {null_count} nulls and its bitmap has {counted}"),
250 ));
251 }
252
253 Ok(())
254}
255
256fn count_nulls(bitmap: &[u8], len: u64) -> u64 {
265 let available = as_u64(bitmap.len()).saturating_mul(8);
268 let considered = len.min(available);
269
270 let whole = bitmap
271 .len()
272 .min(usize::try_from(considered / 8).unwrap_or(usize::MAX));
273 let mut set: u64 = bitmap[..whole]
274 .iter()
275 .map(|byte| u64::from(byte.count_ones()))
276 .sum();
277
278 let spare = u32::try_from(considered % 8).unwrap_or(0);
279 if spare != 0 {
280 let mask = (1u8 << spare) - 1;
281 set += u64::from((bitmap.get(whole).copied().unwrap_or(0) & mask).count_ones());
282 }
283
284 considered - set
285}
286
287fn check_values(
292 data_type: &DataType,
293 len: u64,
294 buffers: &[&[u8]],
295 child_lengths: &[u64],
296 path: &str,
297) -> Result<()> {
298 let child = child_lengths.first().copied().unwrap_or(0);
299
300 match data_type {
301 DataType::Null => Ok(()),
302 DataType::Utf8 | DataType::Binary | DataType::LargeUtf8 | DataType::LargeBinary => {
303 check_variable(data_type, len, buffers, path)
304 }
305 DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _) => {
306 check_list(data_type, len, buffers, child, path)
307 }
308 DataType::FixedSizeList(_, size) => check_fixed_size_list(len, *size, child, path),
309 DataType::Struct(fields) => {
310 for (field, child) in fields.iter().zip(child_lengths) {
311 if *child < len {
312 return Err(Violation::at(
313 Invariant::ChildLength,
314 path,
315 format!(
316 "this struct has {len} rows and its {} field has {child}",
317 field.name()
318 ),
319 ));
320 }
321 }
322 Ok(())
323 }
324 other => check_fixed_width(other, len, buffers, path),
326 }
327}
328
329fn check_variable(data_type: &DataType, len: u64, buffers: &[&[u8]], path: &str) -> Result<()> {
331 let [offsets, data] = buffers else {
332 return Err(counted_wrong(path, 2, buffers.len()));
333 };
334 let width = offset_width(data_type).expect("a variable length type has offsets");
335 let last = check_offsets(offsets, len, width, path)?;
336 let have = as_u64(data.len());
337 if last > have {
338 return Err(Violation::at(
339 Invariant::OffsetRange,
340 path,
341 format!("the last offset is {last} and the values buffer is {have} bytes"),
342 ));
343 }
344 Ok(())
345}
346
347fn check_list(
349 data_type: &DataType,
350 len: u64,
351 buffers: &[&[u8]],
352 child: u64,
353 path: &str,
354) -> Result<()> {
355 let [offsets] = buffers else {
356 return Err(counted_wrong(path, 1, buffers.len()));
357 };
358 let width = offset_width(data_type).expect("a list has offsets");
359 let last = check_offsets(offsets, len, width, path)?;
360 if last > child {
361 return Err(Violation::at(
362 Invariant::OffsetRange,
363 path,
364 format!("the last offset is {last} and the child array has {child} slots"),
365 ));
366 }
367 Ok(())
368}
369
370fn check_fixed_size_list(len: u64, size: i32, child: u64, path: &str) -> Result<()> {
372 let size = u64::try_from(size).map_err(|_| {
373 Violation::at(
374 Invariant::ChildLength,
375 path,
376 format!("a fixed size list cannot hold {size} values a row"),
377 )
378 })?;
379 let needed = len.checked_mul(size).ok_or_else(|| {
380 Violation::at(
381 Invariant::Size,
382 path,
383 format!("{len} rows of {size} values is more than this host can address"),
384 )
385 })?;
386 if child < needed {
387 return Err(Violation::at(
388 Invariant::ChildLength,
389 path,
390 format!("{len} rows of {size} values need {needed} slots and the child has {child}"),
391 ));
392 }
393 Ok(())
394}
395
396fn check_fixed_width(data_type: &DataType, len: u64, buffers: &[&[u8]], path: &str) -> Result<()> {
398 let [values] = buffers else {
399 return Err(counted_wrong(path, 1, buffers.len()));
400 };
401 let bits = slot_bits(data_type).ok_or_else(|| {
402 Violation::at(
403 Invariant::Unsupported,
404 path,
405 format!("this build does not know how wide a {data_type} slot is"),
406 )
407 })?;
408 let needed = len
409 .checked_mul(bits)
410 .map(|total| total.div_ceil(8))
411 .ok_or_else(|| {
412 Violation::at(
413 Invariant::Size,
414 path,
415 format!("{len} slots of {bits} bits is more than this host can address"),
416 )
417 })?;
418 let have = as_u64(values.len());
419 if have < needed {
420 return Err(Violation::at(
421 Invariant::BufferLength,
422 path,
423 format!("{len} slots of {bits} bits need {needed} bytes and there are {have}"),
424 ));
425 }
426 Ok(())
427}
428
429fn check_offsets(offsets: &[u8], len: u64, width: u64, path: &str) -> Result<u64> {
434 if len == 0 && offsets.is_empty() {
437 return Ok(0);
438 }
439
440 let entries = len.checked_add(1).ok_or_else(|| {
445 Violation::at(
446 Invariant::Size,
447 path,
448 format!("{len} slots need one more offset than that, which does not fit in a count"),
449 )
450 })?;
451 let needed = entries.checked_mul(width).ok_or_else(|| {
452 Violation::at(
453 Invariant::Size,
454 path,
455 format!("{entries} offsets of {width} bytes is more than this host can address"),
456 )
457 })?;
458 let have = as_u64(offsets.len());
459 if have < needed {
460 return Err(Violation::at(
461 Invariant::BufferLength,
462 path,
463 format!("{len} slots need {needed} bytes of offsets and there are {have}"),
464 ));
465 }
466
467 let span = usize::try_from(needed).map_err(|_| {
472 Violation::at(
473 Invariant::Size,
474 path,
475 "the offsets run past what this host can address".to_owned(),
476 )
477 })?;
478 let run = offsets.get(..span).unwrap_or(offsets);
479 let previous = if width == 8 {
480 scan_offsets::<8>(run, path, i64::from_le_bytes)?
481 } else {
482 scan_offsets::<4>(run, path, |raw| i64::from(i32::from_le_bytes(raw)))?
483 };
484
485 u64::try_from(previous).map_err(|_| {
486 Violation::at(
487 Invariant::OffsetRange,
488 path,
489 "the last offset is negative".to_owned(),
490 )
491 })
492}
493
494fn scan_offsets<const W: usize>(
501 offsets: &[u8],
502 path: &str,
503 read: fn([u8; W]) -> i64,
504) -> Result<i64> {
505 let mut previous: i64 = 0;
506 for (index, raw) in offsets.as_chunks::<W>().0.iter().enumerate() {
507 let offset = read(*raw);
508
509 if offset < 0 {
510 return Err(Violation::at(
511 Invariant::OffsetRange,
512 path,
513 format!("offset {index} is {offset}, and an offset is a position"),
514 ));
515 }
516 if index > 0 && offset < previous {
517 return Err(Violation::at(
518 Invariant::OffsetOrder,
519 path,
520 format!("offset {index} is {offset} and the one before it is {previous}"),
521 ));
522 }
523 previous = offset;
524 }
525 Ok(previous)
526}
527
528fn as_u64(len: usize) -> u64 {
530 u64::try_from(len).unwrap_or(u64::MAX)
531}
532
533fn counted_wrong(path: &str, wanted: usize, found: usize) -> Violation {
534 Violation::at(
535 Invariant::Buffers,
536 path,
537 format!("this column takes {wanted} buffers after its validity buffer and got {found}"),
538 )
539}
540
541#[cfg(test)]
542mod tests {
543 use arrow_schema::{DataType, Field, Fields, Schema};
544 use iris_abi::Node;
545
546 use super::{MAX_DEPTH, check, check_schema, count_nulls};
547 use crate::error::Invariant;
548
549 fn node(length: u64, null_count: u64) -> Node {
550 Node { length, null_count }
551 }
552
553 fn i64s(values: &[i64]) -> Vec<u8> {
554 values.iter().flat_map(|v| v.to_le_bytes()).collect()
555 }
556
557 fn i32s(values: &[i32]) -> Vec<u8> {
558 values.iter().flat_map(|v| v.to_le_bytes()).collect()
559 }
560
561 #[test]
562 fn a_sound_batch_passes() {
563 let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
564 let buffers = vec![Vec::new(), i64s(&[1, 2, 3])];
565 check(&schema, 3, &[node(3, 0)], &buffers).expect("this batch is sound");
566 }
567
568 #[test]
569 fn a_column_shorter_than_the_batch_is_caught() {
570 let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
571 let buffers = vec![Vec::new(), i64s(&[1, 2])];
572 let err = check(&schema, 3, &[node(2, 0)], &buffers).expect_err("two is not three");
573 assert_eq!(err.invariant, Invariant::Rows);
574 }
575
576 #[test]
577 fn a_values_buffer_one_slot_short_is_caught() {
578 let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
579 let buffers = vec![Vec::new(), i64s(&[1, 2])];
580 let err =
581 check(&schema, 3, &[node(3, 0)], &buffers).expect_err("three slots need 24 bytes");
582 assert_eq!(err.invariant, Invariant::BufferLength);
583 assert!(err.to_string().contains("24 bytes"), "{err}");
584 }
585
586 #[test]
587 fn a_bitmap_with_too_few_bits_is_caught() {
588 let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
589 let buffers = vec![Vec::new(), i64s(&[1, 2, 3])];
590 let short = vec![vec![0xffu8], i64s(&[1; 100])];
592 let wide = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
593 let err = check(&wide, 100, &[node(100, 0)], &short).expect_err("100 slots need 13 bytes");
594 assert_eq!(err.invariant, Invariant::Validity);
595 check(&schema, 3, &[node(3, 0)], &buffers).expect("the empty bitmap case still passes");
596 }
597
598 #[test]
599 fn an_offset_one_past_the_end_is_caught() {
600 let schema = Schema::new(vec![Field::new("s", DataType::Utf8, false)]);
601 let buffers = vec![Vec::new(), i32s(&[0, 2, 6]), b"hoyea".to_vec()];
602 let err = check(&schema, 2, &[node(2, 0)], &buffers).expect_err("six is past five");
603 assert_eq!(err.invariant, Invariant::OffsetRange);
604 }
605
606 #[test]
607 fn offsets_that_run_backwards_are_caught() {
608 let schema = Schema::new(vec![Field::new("s", DataType::Utf8, false)]);
609 let buffers = vec![Vec::new(), i32s(&[0, 4, 2]), b"hoyea".to_vec()];
610 let err = check(&schema, 2, &[node(2, 0)], &buffers).expect_err("two is less than four");
611 assert_eq!(err.invariant, Invariant::OffsetOrder);
612 }
613
614 #[test]
615 fn a_child_one_row_short_of_its_parent_is_caught() {
616 let children = Fields::from(vec![Field::new("x", DataType::Int64, false)]);
617 let schema = Schema::new(vec![Field::new("p", DataType::Struct(children), false)]);
618 let buffers = vec![Vec::new(), Vec::new(), i64s(&[1, 2])];
619 let err = check(&schema, 3, &[node(3, 0), node(2, 0)], &buffers)
620 .expect_err("a struct's child cannot be shorter than the struct");
621 assert_eq!(err.invariant, Invariant::ChildLength);
622 }
623
624 #[test]
625 fn a_length_that_overflows_a_width_is_caught_rather_than_wrapped() {
626 let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
627 let buffers = vec![Vec::new(), i64s(&[1])];
628 let err = check(&schema, u64::MAX, &[node(u64::MAX, 0)], &buffers)
629 .expect_err("that many slots is not addressable");
630 assert_eq!(err.invariant, Invariant::Size);
631 }
632
633 #[test]
640 fn a_length_that_wraps_the_count_of_offsets_is_caught() {
641 for data_type in [DataType::Binary, DataType::LargeBinary] {
642 let schema = Schema::new(vec![Field::new("a", data_type, false)]);
643 let buffers = vec![Vec::new(), Vec::new(), Vec::new()];
644 let err = check(&schema, u64::MAX, &[node(u64::MAX, 0)], &buffers)
645 .expect_err("one more offset than that does not fit in a count");
646 assert_eq!(err.invariant, Invariant::Size);
647 }
648 }
649
650 #[test]
651 fn a_schema_nested_past_the_bound_is_refused_without_recursing_into_it() {
652 let mut data_type = DataType::Int64;
653 for _ in 0..MAX_DEPTH + 10 {
654 data_type = DataType::List(std::sync::Arc::new(Field::new("item", data_type, false)));
655 }
656 let schema = Schema::new(vec![Field::new("deep", data_type, false)]);
657 let err = check_schema(&schema).expect_err("that is deeper than this build walks");
658 assert_eq!(err.invariant, Invariant::Depth);
659 }
660
661 #[test]
662 fn a_schema_at_the_bound_is_still_walked() {
663 let mut data_type = DataType::Int64;
664 for _ in 0..MAX_DEPTH - 1 {
665 data_type = DataType::List(std::sync::Arc::new(Field::new("item", data_type, false)));
666 }
667 let schema = Schema::new(vec![Field::new("deep", data_type, false)]);
668 check_schema(&schema).expect("this is exactly as deep as the bound allows");
669 }
670
671 #[test]
672 fn spare_buffers_are_an_error_rather_than_something_ignored() {
673 let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
674 let buffers = vec![Vec::new(), i64s(&[1, 2, 3]), i64s(&[4])];
675 let err =
676 check(&schema, 3, &[node(3, 0)], &buffers).expect_err("a spare buffer is not fine");
677 assert_eq!(err.invariant, Invariant::Buffers);
678 }
679
680 #[test]
681 fn counting_nulls_stops_at_the_length_rather_than_the_byte() {
682 assert_eq!(count_nulls(&[0b0001_1111], 5), 0);
684 assert_eq!(count_nulls(&[0b0001_1110], 5), 1);
685 assert_eq!(count_nulls(&[0x00, 0xff], 9), 8);
686 }
687
688 #[test]
689 fn counting_nulls_agrees_with_reading_the_bits_one_at_a_time() {
690 fn one_at_a_time(bitmap: &[u8], len: u64) -> u64 {
695 let mut nulls = 0;
696 for bit in 0..len {
697 let byte = usize::try_from(bit / 8).expect("this test is small");
698 let Some(value) = bitmap.get(byte) else {
699 break;
700 };
701 let shift = u32::try_from(bit % 8).expect("a bit in a byte");
702 if value >> shift & 1 == 0 {
703 nulls += 1;
704 }
705 }
706 nulls
707 }
708
709 for bitmap in [
710 [0x00u8, 0x00, 0x00],
711 [0xff, 0xff, 0xff],
712 [0b1010_1010, 0b0000_1111, 0b1100_0011],
713 ] {
714 for len in 0..40 {
715 assert_eq!(
716 count_nulls(&bitmap, len),
717 one_at_a_time(&bitmap, len),
718 "bitmap {bitmap:?} at length {len}"
719 );
720 }
721 }
722 }
723}