vortex_vector/bool/
vector_mut.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Definition and implementation of [`BoolVectorMut`].
5
6use vortex_buffer::BitBufferMut;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_ensure;
10use vortex_mask::MaskMut;
11
12use crate::VectorMutOps;
13use crate::VectorOps;
14use crate::bool::BoolScalar;
15use crate::bool::BoolVector;
16
17/// A mutable vector of boolean values.
18///
19/// Internally, this `BoolVectorMut` is a wrapper around a [`BitBufferMut`] and a validity mask.
20#[derive(Debug, Clone)]
21pub struct BoolVectorMut {
22    /// The mutable bits that we use to represent booleans.
23    pub(super) bits: BitBufferMut,
24    /// The validity mask (where `true` represents an element is **not** null).
25    pub(super) validity: MaskMut,
26}
27
28impl BoolVectorMut {
29    /// Creates a new [`BoolVectorMut`] from the given bits and validity mask.
30    ///
31    /// # Panics
32    ///
33    /// Panics if the length of the validity mask does not match the length of the bits.
34    pub fn new(bits: BitBufferMut, validity: MaskMut) -> Self {
35        Self::try_new(bits, validity).vortex_expect("Failed to create `BoolVectorMut`")
36    }
37
38    /// Tries to create a new [`BoolVectorMut`] from the given bits and validity mask.
39    ///
40    /// # Errors
41    ///
42    /// Returns an error if the length of the validity mask does not match the length of the bits.
43    pub fn try_new(bits: BitBufferMut, validity: MaskMut) -> VortexResult<Self> {
44        vortex_ensure!(
45            validity.len() == bits.len(),
46            "`BoolVector` validity mask must have the same length as bits"
47        );
48
49        Ok(Self { bits, validity })
50    }
51
52    /// Creates a new [`BoolVectorMut`] from the given bits and validity mask without validation.
53    ///
54    /// # Safety
55    ///
56    /// The caller must ensure that the validity mask has the same length as the bits.
57    ///
58    /// Ideally, they are taken from `into_parts`, mutated in a way that doesn't re-allocate, and
59    /// then passed back to this function.
60    pub unsafe fn new_unchecked(bits: BitBufferMut, validity: MaskMut) -> Self {
61        if cfg!(debug_assertions) {
62            Self::new(bits, validity)
63        } else {
64            Self { bits, validity }
65        }
66    }
67
68    /// Creates a new mutable boolean vector with the given `capacity`.
69    pub fn with_capacity(capacity: usize) -> Self {
70        Self {
71            bits: BitBufferMut::with_capacity(capacity),
72            validity: MaskMut::with_capacity(capacity),
73        }
74    }
75
76    /// Decomposes the boolean vector into its constituent parts (bit buffer and validity).
77    pub fn into_parts(self) -> (BitBufferMut, MaskMut) {
78        (self.bits, self.validity)
79    }
80
81    /// Append n values to the vector.
82    pub fn append_values(&mut self, value: bool, n: usize) {
83        self.bits.append_n(value, n);
84        self.validity.append_n(true, n);
85    }
86
87    /// Returns a readonly handle to the bits backing the vector.
88    pub fn bits(&self) -> &BitBufferMut {
89        &self.bits
90    }
91
92    /// Returns a mutable handle to the bits backing the vector.
93    ///
94    /// # Safety
95    ///
96    /// Caller must ensure that bits and validity always have same length.
97    pub unsafe fn bits_mut(&mut self) -> &mut BitBufferMut {
98        &mut self.bits
99    }
100
101    /// Get a mutable handle to the validity mask of the vector.
102    ///
103    /// # Safety
104    ///
105    /// Caller must ensure that length of the validity always matches
106    /// length of the bits.
107    pub unsafe fn validity_mut(&mut self) -> &mut MaskMut {
108        &mut self.validity
109    }
110}
111
112impl VectorMutOps for BoolVectorMut {
113    type Immutable = BoolVector;
114
115    fn len(&self) -> usize {
116        debug_assert!(self.validity.len() == self.bits.len());
117
118        self.bits.len()
119    }
120
121    fn validity(&self) -> &MaskMut {
122        &self.validity
123    }
124
125    fn capacity(&self) -> usize {
126        self.bits.capacity()
127    }
128
129    fn reserve(&mut self, additional: usize) {
130        self.bits.reserve(additional);
131        self.validity.reserve(additional);
132    }
133
134    fn clear(&mut self) {
135        self.bits.clear();
136        self.validity.clear();
137    }
138
139    fn truncate(&mut self, len: usize) {
140        self.bits.truncate(len);
141        self.validity.truncate(len);
142    }
143
144    fn extend_from_vector(&mut self, other: &BoolVector) {
145        self.bits.append_buffer(&other.bits);
146        self.validity.append_mask(other.validity());
147    }
148
149    fn append_nulls(&mut self, n: usize) {
150        self.bits.append_n(false, n); // Note that the value we push doesn't actually matter.
151        self.validity.append_n(false, n);
152    }
153
154    fn append_zeros(&mut self, n: usize) {
155        self.bits.append_n(false, n);
156        self.validity.append_n(true, n);
157    }
158
159    fn append_scalars(&mut self, scalar: &BoolScalar, n: usize) {
160        match scalar.value() {
161            None => self.append_nulls(n),
162            Some(value) => self.append_values(value, n),
163        }
164    }
165
166    fn freeze(self) -> BoolVector {
167        BoolVector {
168            bits: self.bits.freeze(),
169            validity: self.validity.freeze(),
170        }
171    }
172
173    fn split_off(&mut self, at: usize) -> Self {
174        Self {
175            bits: self.bits.split_off(at),
176            validity: self.validity.split_off(at),
177        }
178    }
179
180    fn unsplit(&mut self, other: Self) {
181        if self.is_empty() {
182            *self = other;
183            return;
184        }
185        self.bits.unsplit(other.bits);
186        self.validity.unsplit(other.validity);
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[test]
195    fn test_from_iter_with_options() {
196        // Test FromIterator<Option<bool>> with nulls and empty.
197        let vec_empty = BoolVectorMut::from_iter(std::iter::empty::<Option<bool>>());
198        assert_eq!(vec_empty.len(), 0);
199
200        let vec = BoolVectorMut::from_iter([Some(true), None, Some(false), None, Some(true)]);
201        assert_eq!(vec.len(), 5);
202        let frozen = vec.freeze();
203        assert_eq!(frozen.validity().true_count(), 3);
204    }
205
206    #[test]
207    fn test_from_iter_non_null() {
208        // Test FromIterator<bool> creates all-valid vector.
209        let vec = BoolVectorMut::from_iter([true, false, true, true, false]);
210        assert_eq!(vec.len(), 5);
211        let frozen = vec.freeze();
212        assert_eq!(frozen.validity().true_count(), 5);
213    }
214
215    #[test]
216    fn test_operations_preserve_validity() {
217        // Comprehensive test for split/unsplit/extend preserving validity.
218        let mut vec = BoolVectorMut::from_iter([Some(true), None, Some(false), None, Some(true)]);
219
220        // Test split.
221        let second_half = vec.split_off(2);
222        assert_eq!(vec.len(), 2);
223        assert_eq!(second_half.len(), 3);
224
225        // Test validity after split.
226        let frozen_first = vec.freeze();
227        assert_eq!(frozen_first.validity().true_count(), 1);
228        let frozen_second = second_half.freeze();
229        assert_eq!(frozen_second.validity().true_count(), 2);
230
231        // Test unsplit.
232        let mut vec1 = BoolVectorMut::from_iter([Some(true), None]);
233        let vec2 = BoolVectorMut::from_iter([Some(false), Some(true)]);
234        vec1.unsplit(vec2);
235        assert_eq!(vec1.len(), 4);
236        let frozen = vec1.freeze();
237        assert_eq!(frozen.validity().true_count(), 3);
238    }
239
240    #[test]
241    fn test_into_iter_roundtrip() {
242        // Test that from_iter followed by into_iter preserves the data.
243        let original_data = vec![
244            Some(true),
245            None,
246            Some(false),
247            Some(true),
248            None,
249            Some(false),
250            None,
251            Some(true),
252        ];
253
254        // Create vector from iterator.
255        let vec = BoolVectorMut::from_iter(original_data.clone());
256
257        // Convert back to iterator and collect.
258        let roundtrip: Vec<_> = vec.into_iter().collect();
259
260        // Should be identical.
261        assert_eq!(roundtrip, original_data);
262
263        // Also test with all valid values.
264        let all_valid = vec![true, false, true, false, true];
265        let vec = BoolVectorMut::from_iter(all_valid.clone());
266        let roundtrip: Vec<_> = vec.into_iter().collect();
267        let expected: Vec<_> = all_valid.into_iter().map(Some).collect();
268        assert_eq!(roundtrip, expected);
269
270        // Test with empty.
271        let empty: Vec<Option<bool>> = vec![];
272        let vec = BoolVectorMut::from_iter(empty.clone());
273        let roundtrip: Vec<_> = vec.into_iter().collect();
274        assert_eq!(roundtrip, empty);
275    }
276}
vortex_vector/bool/vector_mut.rs

vortex_vector/bool/
vector_mut.rs
