iris_guard/
layout.rs

1//! How many buffers and children each Arrow type takes.
2//!
3//! This is a transcription of the Arrow specification's layout section, and it lives here rather
4//! than next to the code that assembles arrays because two copies of it would eventually disagree,
5//! and the copy that disagrees silently is the one in the checker.
6//!
7//! | Type | Validity | Then | Children |
8//! | --- | --- | --- | --- |
9//! | Null | no | nothing | none |
10//! | Boolean, the fixed width types, `FixedSizeBinary` | yes | values | none |
11//! | `Utf8`, `Binary` and their large forms | yes | offsets, values | none |
12//! | `List`, `LargeList`, `Map` | yes | offsets | one |
13//! | `FixedSizeList` | yes | nothing | one |
14//! | `Struct` | yes | nothing | one per field |
15//!
16//! Unions, dictionaries, run end encoding and the view types are not here yet. Each of them needs
17//! something the batch this crate checks cannot carry: a union has a type id map and no validity
18//! buffer, a dictionary needs its values to arrive out of band, and a view array has a variable
19//! number of data buffers, which is the one thing that breaks counting buffers against a schema.
20//! They are refused by name rather than skipped, because a column that quietly does not arrive is
21//! worse than one that fails.
22//!
23//! The two checks those types need are written and tested even so, in [`crate::check_dictionary`]
24//! and [`crate::check_views`], because the checks are the hard part and having them ready is what
25//! makes carrying the types a format question rather than a safety question.
26
27use arrow_schema::{DataType, Field};
28
29use crate::error::{Invariant, Result, Violation};
30
31/// What one array takes out of a batch.
32#[derive(Clone, PartialEq, Eq, Debug)]
33pub struct Layout<'a> {
34    /// Whether the array has a validity buffer. Only `Null` does not.
35    pub validity: bool,
36    /// How many buffers follow the validity one.
37    pub values: usize,
38    /// The children, in order.
39    pub children: Vec<&'a Field>,
40}
41
42/// The layout of one type.
43///
44/// # Errors
45///
46/// Returns a [`Invariant::Unsupported`] violation for a type this crate cannot check, which is the
47/// same thing as a type iris cannot carry.
48pub fn layout<'a>(data_type: &'a DataType, path: &str) -> Result<Layout<'a>> {
49    let (validity, values, children) = match data_type {
50        DataType::Null => (false, 0, Vec::new()),
51
52        DataType::Boolean
53        | DataType::Int8
54        | DataType::Int16
55        | DataType::Int32
56        | DataType::Int64
57        | DataType::UInt8
58        | DataType::UInt16
59        | DataType::UInt32
60        | DataType::UInt64
61        | DataType::Float16
62        | DataType::Float32
63        | DataType::Float64
64        | DataType::Date32
65        | DataType::Date64
66        | DataType::Time32(_)
67        | DataType::Time64(_)
68        | DataType::Timestamp(_, _)
69        | DataType::Duration(_)
70        | DataType::Interval(_)
71        | DataType::Decimal32(_, _)
72        | DataType::Decimal64(_, _)
73        | DataType::Decimal128(_, _)
74        | DataType::Decimal256(_, _)
75        | DataType::FixedSizeBinary(_) => (true, 1, Vec::new()),
76
77        DataType::Utf8 | DataType::Binary | DataType::LargeUtf8 | DataType::LargeBinary => {
78            (true, 2, Vec::new())
79        }
80
81        DataType::List(field) | DataType::LargeList(field) | DataType::Map(field, _) => {
82            (true, 1, vec![field.as_ref()])
83        }
84
85        DataType::FixedSizeList(field, _) => (true, 0, vec![field.as_ref()]),
86
87        DataType::Struct(fields) => (true, 0, fields.iter().map(AsRef::as_ref).collect()),
88
89        other => {
90            return Err(Violation::at(
91                Invariant::Unsupported,
92                path,
93                format!("this build does not carry {other} columns"),
94            ));
95        }
96    };
97
98    Ok(Layout {
99        validity,
100        values,
101        children,
102    })
103}
104
105/// How wide one slot of a fixed width type is, in bits.
106///
107/// Bits rather than bytes because of `Boolean`, which is the whole reason this returns a number that
108/// has to be divided rather than multiplied.
109pub(crate) fn slot_bits(data_type: &DataType) -> Option<u64> {
110    let bits = match data_type {
111        DataType::Boolean => 1,
112        DataType::Int8 | DataType::UInt8 => 8,
113        DataType::Int16 | DataType::UInt16 | DataType::Float16 => 16,
114        DataType::Int32
115        | DataType::UInt32
116        | DataType::Float32
117        | DataType::Date32
118        | DataType::Time32(_)
119        | DataType::Decimal32(_, _) => 32,
120        DataType::Int64
121        | DataType::UInt64
122        | DataType::Float64
123        | DataType::Date64
124        | DataType::Time64(_)
125        | DataType::Timestamp(_, _)
126        | DataType::Duration(_)
127        | DataType::Decimal64(_, _) => 64,
128        DataType::Decimal128(_, _) => 128,
129        DataType::Decimal256(_, _) => 256,
130        DataType::Interval(unit) => match unit {
131            arrow_schema::IntervalUnit::YearMonth => 32,
132            arrow_schema::IntervalUnit::DayTime => 64,
133            arrow_schema::IntervalUnit::MonthDayNano => 128,
134        },
135        // A negative width is not a width. Arrow's own type carries an `i32` and nothing stops a
136        // schema from declaring one, so it is turned down here rather than cast.
137        DataType::FixedSizeBinary(width) => u64::try_from(*width).ok()? * 8,
138        _ => return None,
139    };
140    Some(bits)
141}
142
143/// How wide the offsets of a variable length type are, in bytes.
144pub(crate) const fn offset_width(data_type: &DataType) -> Option<u64> {
145    match data_type {
146        DataType::Utf8 | DataType::Binary | DataType::List(_) | DataType::Map(_, _) => Some(4),
147        DataType::LargeUtf8 | DataType::LargeBinary | DataType::LargeList(_) => Some(8),
148        _ => None,
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use arrow_schema::{DataType, Field, Fields};
155
156    use super::{layout, offset_width, slot_bits};
157    use crate::error::Invariant;
158
159    #[test]
160    fn a_flat_column_takes_a_validity_buffer_and_a_values_buffer() {
161        let found = layout(&DataType::Int64, "a").expect("Int64 is a type iris carries");
162        assert!(found.validity);
163        assert_eq!(found.values, 1);
164        assert!(found.children.is_empty());
165    }
166
167    #[test]
168    fn a_struct_takes_no_values_buffer_and_one_child_per_field() {
169        let fields = Fields::from(vec![
170            Field::new("x", DataType::Int64, false),
171            Field::new("y", DataType::Int64, false),
172        ]);
173        let point = DataType::Struct(fields);
174        let found = layout(&point, "p").expect("Struct is a type iris carries");
175        assert_eq!(found.values, 0);
176        assert_eq!(found.children.len(), 2);
177    }
178
179    #[test]
180    fn null_is_the_one_type_with_no_validity_buffer() {
181        let found = layout(&DataType::Null, "n").expect("Null is a type iris carries");
182        assert!(!found.validity);
183        assert_eq!(found.values, 0);
184    }
185
186    #[test]
187    fn a_type_this_build_does_not_carry_is_refused_by_name() {
188        let dictionary = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
189        let err = layout(&dictionary, "d").expect_err("a dictionary is not carried yet");
190        assert_eq!(err.invariant, Invariant::Unsupported);
191        assert!(err.to_string().contains("Dictionary"), "{err}");
192    }
193
194    #[test]
195    fn a_boolean_slot_is_one_bit_and_a_timestamp_is_sixty_four() {
196        assert_eq!(slot_bits(&DataType::Boolean), Some(1));
197        assert_eq!(
198            slot_bits(&DataType::Timestamp(
199                arrow_schema::TimeUnit::Nanosecond,
200                None
201            )),
202            Some(64)
203        );
204        assert_eq!(slot_bits(&DataType::FixedSizeBinary(7)), Some(56));
205    }
206
207    #[test]
208    fn a_negative_fixed_width_is_not_a_width() {
209        assert_eq!(slot_bits(&DataType::FixedSizeBinary(-1)), None);
210    }
211
212    #[test]
213    fn the_large_variants_are_the_ones_with_eight_byte_offsets() {
214        assert_eq!(offset_width(&DataType::Utf8), Some(4));
215        assert_eq!(offset_width(&DataType::LargeUtf8), Some(8));
216        assert_eq!(offset_width(&DataType::Int64), None);
217    }
218}
iris_guard/layout.rs

iris_guard/
layout.rs
