vortex_bytebool/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5
6use arrow_buffer::BooleanBuffer;
7use vortex_array::arrays::BoolArray;
8use vortex_array::stats::{ArrayStats, StatsSetRef};
9use vortex_array::validity::Validity;
10use vortex_array::vtable::{
11    ArrayVTable, CanonicalVTable, NotSupported, OperationsVTable, VTable, ValidityHelper,
12    ValidityVTableFromValidityHelper,
13};
14use vortex_array::{ArrayRef, Canonical, EncodingId, EncodingRef, IntoArray, vtable};
15use vortex_buffer::ByteBuffer;
16use vortex_dtype::DType;
17use vortex_error::{VortexResult, vortex_panic};
18use vortex_scalar::Scalar;
19
20vtable!(ByteBool);
21
22impl VTable for ByteBoolVTable {
23    type Array = ByteBoolArray;
24    type Encoding = ByteBoolEncoding;
25
26    type ArrayVTable = Self;
27    type CanonicalVTable = Self;
28    type OperationsVTable = Self;
29    type ValidityVTable = ValidityVTableFromValidityHelper;
30    type VisitorVTable = Self;
31    type ComputeVTable = NotSupported;
32    type EncodeVTable = NotSupported;
33    type SerdeVTable = Self;
34
35    fn id(_encoding: &Self::Encoding) -> EncodingId {
36        EncodingId::new_ref("vortex.bytebool")
37    }
38
39    fn encoding(_array: &Self::Array) -> EncodingRef {
40        EncodingRef::new_ref(ByteBoolEncoding.as_ref())
41    }
42}
43
44#[derive(Clone, Debug)]
45pub struct ByteBoolArray {
46    dtype: DType,
47    buffer: ByteBuffer,
48    validity: Validity,
49    stats_set: ArrayStats,
50}
51
52#[derive(Clone, Debug)]
53pub struct ByteBoolEncoding;
54
55impl ByteBoolArray {
56    pub fn new(buffer: ByteBuffer, validity: Validity) -> Self {
57        let length = buffer.len();
58        if let Some(vlen) = validity.maybe_len()
59            && length != vlen
60        {
61            vortex_panic!(
62                "Buffer length ({}) does not match validity length ({})",
63                length,
64                vlen
65            );
66        }
67        Self {
68            dtype: DType::Bool(validity.nullability()),
69            buffer,
70            validity,
71            stats_set: Default::default(),
72        }
73    }
74
75    // TODO(ngates): deprecate construction from vec
76    pub fn from_vec<V: Into<Validity>>(data: Vec<bool>, validity: V) -> Self {
77        let validity = validity.into();
78        // SAFETY: we are transmuting a Vec<bool> into a Vec<u8>
79        let data: Vec<u8> = unsafe { std::mem::transmute(data) };
80        Self::new(ByteBuffer::from(data), validity)
81    }
82
83    pub fn buffer(&self) -> &ByteBuffer {
84        &self.buffer
85    }
86
87    pub fn as_slice(&self) -> &[bool] {
88        // Safety: The internal buffer contains byte-sized bools
89        unsafe { std::mem::transmute(self.buffer().as_slice()) }
90    }
91}
92
93impl ValidityHelper for ByteBoolArray {
94    fn validity(&self) -> &Validity {
95        &self.validity
96    }
97}
98
99impl ArrayVTable<ByteBoolVTable> for ByteBoolVTable {
100    fn len(array: &ByteBoolArray) -> usize {
101        array.buffer.len()
102    }
103
104    fn dtype(array: &ByteBoolArray) -> &DType {
105        &array.dtype
106    }
107
108    fn stats(array: &ByteBoolArray) -> StatsSetRef<'_> {
109        array.stats_set.to_ref(array.as_ref())
110    }
111}
112
113impl CanonicalVTable<ByteBoolVTable> for ByteBoolVTable {
114    fn canonicalize(array: &ByteBoolArray) -> VortexResult<Canonical> {
115        let boolean_buffer = BooleanBuffer::from(array.as_slice());
116        let validity = array.validity().clone();
117        Ok(Canonical::Bool(BoolArray::new(boolean_buffer, validity)))
118    }
119}
120
121impl OperationsVTable<ByteBoolVTable> for ByteBoolVTable {
122    fn slice(array: &ByteBoolArray, start: usize, stop: usize) -> ArrayRef {
123        ByteBoolArray::new(
124            array.buffer().slice(start..stop),
125            array.validity().slice(start, stop),
126        )
127        .into_array()
128    }
129
130    fn scalar_at(array: &ByteBoolArray, index: usize) -> Scalar {
131        Scalar::bool(array.buffer()[index] == 1, array.dtype().nullability())
132    }
133}
134
135impl From<Vec<bool>> for ByteBoolArray {
136    fn from(value: Vec<bool>) -> Self {
137        Self::from_vec(value, Validity::AllValid)
138    }
139}
140
141impl From<Vec<Option<bool>>> for ByteBoolArray {
142    fn from(value: Vec<Option<bool>>) -> Self {
143        let validity = Validity::from_iter(value.iter().map(|v| v.is_some()));
144
145        // This doesn't reallocate, and the compiler even vectorizes it
146        let data = value.into_iter().map(Option::unwrap_or_default).collect();
147
148        Self::from_vec(data, validity)
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155
156    // #[cfg_attr(miri, ignore)]
157    // #[test]
158    // fn test_bytebool_metadata() {
159    //     check_metadata(
160    //         "bytebool.metadata",
161    //         SerdeMetadata(ByteBoolMetadata {
162    //             validity: ValidityMetadata::AllValid,
163    //         }),
164    //     );
165    // }
166
167    #[test]
168    fn test_validity_construction() {
169        let v = vec![true, false];
170        let v_len = v.len();
171
172        let arr = ByteBoolArray::from(v);
173        assert_eq!(v_len, arr.len());
174
175        for idx in 0..arr.len() {
176            assert!(arr.is_valid(idx).unwrap());
177        }
178
179        let v = vec![Some(true), None, Some(false)];
180        let arr = ByteBoolArray::from(v);
181        assert!(arr.is_valid(0).unwrap());
182        assert!(!arr.is_valid(1).unwrap());
183        assert!(arr.is_valid(2).unwrap());
184        assert_eq!(arr.len(), 3);
185
186        let v: Vec<Option<bool>> = vec![None, None];
187        let v_len = v.len();
188
189        let arr = ByteBoolArray::from(v);
190        assert_eq!(v_len, arr.len());
191
192        for idx in 0..arr.len() {
193            assert!(!arr.is_valid(idx).unwrap());
194        }
195        assert_eq!(arr.len(), 2);
196    }
197}
vortex_bytebool/array.rs

vortex_bytebool/
array.rs
