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