vortex_array/arrays/bool/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hash;
5use std::hash::Hasher;
6
7use kernel::PARENT_KERNELS;
8use prost::Message;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_ensure;
13use vortex_error::vortex_panic;
14use vortex_session::VortexSession;
15
16use crate::ArrayRef;
17use crate::ExecutionCtx;
18use crate::ExecutionResult;
19use crate::array::Array;
20use crate::array::ArrayParts;
21use crate::array::ArrayView;
22use crate::array::VTable;
23use crate::array::child_to_validity;
24use crate::arrays::bool::BoolData;
25use crate::arrays::bool::array::SLOT_NAMES;
26use crate::buffer::BufferHandle;
27use crate::dtype::DType;
28use crate::serde::ArrayChildren;
29use crate::validity::Validity;
30mod canonical;
31mod kernel;
32mod operations;
33mod validity;
34
35use vortex_session::registry::CachedId;
36
37use crate::Precision;
38use crate::array::ArrayId;
39use crate::arrays::bool::compute::rules::RULES;
40use crate::hash::ArrayEq;
41use crate::hash::ArrayHash;
42
43/// A [`Bool`]-encoded Vortex array.
44pub type BoolArray = Array<Bool>;
45
46#[derive(prost::Message)]
47pub struct BoolMetadata {
48    // The offset in bits must be <8
49    #[prost(uint32, tag = "1")]
50    pub offset: u32,
51}
52
53impl ArrayHash for BoolData {
54    fn array_hash<H: Hasher>(&self, state: &mut H, precision: Precision) {
55        self.bits.array_hash(state, precision);
56        self.offset.hash(state);
57    }
58}
59
60impl ArrayEq for BoolData {
61    fn array_eq(&self, other: &Self, precision: Precision) -> bool {
62        self.offset == other.offset && self.bits.array_eq(&other.bits, precision)
63    }
64}
65
66impl VTable for Bool {
67    type ArrayData = BoolData;
68
69    type OperationsVTable = Self;
70    type ValidityVTable = Self;
71
72    fn id(&self) -> ArrayId {
73        static ID: CachedId = CachedId::new("vortex.bool");
74        *ID
75    }
76
77    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
78        1
79    }
80
81    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
82        match idx {
83            0 => array.bits.clone(),
84            _ => vortex_panic!("BoolArray buffer index {idx} out of bounds"),
85        }
86    }
87
88    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
89        match idx {
90            0 => Some("bits".to_string()),
91            _ => None,
92        }
93    }
94
95    fn serialize(
96        array: ArrayView<'_, Self>,
97        _session: &VortexSession,
98    ) -> VortexResult<Option<Vec<u8>>> {
99        assert!(array.offset < 8, "Offset must be <8, got {}", array.offset);
100        Ok(Some(
101            BoolMetadata {
102                offset: u32::try_from(array.offset).vortex_expect("checked"),
103            }
104            .encode_to_vec(),
105        ))
106    }
107
108    fn validate(
109        &self,
110        data: &BoolData,
111        dtype: &DType,
112        len: usize,
113        slots: &[Option<ArrayRef>],
114    ) -> VortexResult<()> {
115        let DType::Bool(nullability) = dtype else {
116            vortex_bail!("Expected bool dtype, got {dtype:?}");
117        };
118        vortex_ensure!(
119            data.bits.len() * 8 >= data.offset + len,
120            "BoolArray buffer with offset {} cannot back outer length {} (buffer bits = {})",
121            data.offset,
122            len,
123            data.bits.len() * 8
124        );
125
126        let validity = child_to_validity(slots[0].as_ref(), *nullability);
127        if let Some(validity_len) = validity.maybe_len() {
128            vortex_ensure!(
129                validity_len == len,
130                "BoolArray validity len {} does not match outer length {}",
131                validity_len,
132                len
133            );
134        }
135
136        Ok(())
137    }
138
139    fn deserialize(
140        &self,
141        dtype: &DType,
142        len: usize,
143        metadata: &[u8],
144        buffers: &[BufferHandle],
145        children: &dyn ArrayChildren,
146        _session: &VortexSession,
147    ) -> VortexResult<ArrayParts<Self>> {
148        let metadata = BoolMetadata::decode(metadata)?;
149        if buffers.len() != 1 {
150            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
151        }
152
153        let validity = if children.is_empty() {
154            Validity::from(dtype.nullability())
155        } else if children.len() == 1 {
156            let validity = children.get(0, &Validity::DTYPE, len)?;
157            Validity::Array(validity)
158        } else {
159            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
160        };
161
162        let buffer = buffers[0].clone();
163        let slots = BoolData::make_slots(&validity, len);
164        let data = BoolData::try_new_from_handle(buffer, metadata.offset as usize, len, validity)?;
165        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
166    }
167
168    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
169        SLOT_NAMES[idx].to_string()
170    }
171
172    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
173        Ok(ExecutionResult::done(array))
174    }
175
176    fn execute_parent(
177        array: ArrayView<'_, Self>,
178        parent: &ArrayRef,
179        child_idx: usize,
180        ctx: &mut ExecutionCtx,
181    ) -> VortexResult<Option<ArrayRef>> {
182        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
183    }
184
185    fn reduce_parent(
186        array: ArrayView<'_, Self>,
187        parent: &ArrayRef,
188        child_idx: usize,
189    ) -> VortexResult<Option<ArrayRef>> {
190        RULES.evaluate(array, parent, child_idx)
191    }
192}
193
194#[derive(Clone, Debug)]
195pub struct Bool;
196
197#[cfg(test)]
198mod tests {
199    use vortex_buffer::ByteBufferMut;
200    use vortex_session::registry::ReadContext;
201
202    use crate::ArrayContext;
203    use crate::IntoArray;
204    use crate::LEGACY_SESSION;
205    use crate::arrays::BoolArray;
206    use crate::assert_arrays_eq;
207    use crate::serde::SerializeOptions;
208    use crate::serde::SerializedArray;
209
210    #[test]
211    fn test_nullable_bool_serde_roundtrip() {
212        let array = BoolArray::from_iter([Some(true), None, Some(false), None]);
213        let dtype = array.dtype().clone();
214        let len = array.len();
215
216        let ctx = ArrayContext::empty();
217        let serialized = array
218            .clone()
219            .into_array()
220            .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default())
221            .unwrap();
222
223        let mut concat = ByteBufferMut::empty();
224        for buf in serialized {
225            concat.extend_from_slice(buf.as_ref());
226        }
227        let parts = SerializedArray::try_from(concat.freeze()).unwrap();
228        let decoded = parts
229            .decode(
230                &dtype,
231                len,
232                &ReadContext::new(ctx.to_ids()),
233                &LEGACY_SESSION,
234            )
235            .unwrap();
236
237        assert_arrays_eq!(decoded, array);
238    }
239}
vortex_array/arrays/bool/vtable/mod.rs

vortex_array/arrays/bool/vtable/
mod.rs
