Skip to main content

vortex_array/arrays/bool/
compact.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBufferMeta;
5use vortex_error::VortexResult;
6
7use crate::arrays::BoolArray;
8use crate::buffer::BufferHandle;
9
10impl BoolArray {
11    /// Trims the packed bit buffer to the bytes backing the array's visible bits.
12    ///
13    /// Arrays built over a shared buffer, such as one decoded from a file segment, can hold
14    /// trailing bytes that no read can see but that `nbytes` and serialization still pay for. The
15    /// trim is zero-copy and keeps the leading bit offset, so only whole bytes are dropped.
16    pub fn trim_bits(&self) -> VortexResult<Self> {
17        let byte_len = BitBufferMeta::new(self.meta.offset(), self.len()).byte_len();
18
19        let bits = match self.bits.as_host_opt() {
20            Some(host) => BufferHandle::new_host(host.slice_unaligned(..byte_len)),
21            None => self.bits.slice(0..byte_len),
22        };
23
24        Self::try_new_from_handle(bits, self.meta.offset(), self.len(), self.validity()?)
25    }
26}
27
28#[cfg(test)]
29mod tests {
30    use rstest::rstest;
31    use vortex_buffer::BitBuffer;
32    use vortex_buffer::ByteBuffer;
33    use vortex_error::VortexResult;
34
35    use crate::Canonical;
36    use crate::VortexSessionExecute;
37    use crate::array_session;
38    use crate::arrays::BoolArray;
39    use crate::assert_arrays_eq;
40    use crate::buffer::BufferHandle;
41    use crate::validity::Validity;
42
43    #[rstest]
44    #[case(0, 16, Validity::NonNullable)]
45    #[case(0, 13, Validity::NonNullable)]
46    #[case(3, 10, Validity::AllValid)]
47    #[case(7, 1, Validity::NonNullable)]
48    #[case(0, 0, Validity::NonNullable)]
49    fn trims_to_visible_bytes(
50        #[case] offset: usize,
51        #[case] len: usize,
52        #[case] validity: Validity,
53    ) -> VortexResult<()> {
54        let mut ctx = array_session().create_execution_ctx();
55        let bytes: ByteBuffer = (0..64u8).collect();
56        let handle = BufferHandle::new_host(bytes.clone());
57
58        let trimmed =
59            BoolArray::try_new_from_handle(handle, offset, len, validity.clone())?.trim_bits()?;
60
61        assert_eq!(trimmed.bits.len(), (offset + len).div_ceil(8));
62        let expected = BoolArray::new(BitBuffer::new_with_offset(bytes, len, offset), validity);
63        assert_arrays_eq!(trimmed, expected, &mut ctx);
64        Ok(())
65    }
66
67    #[test]
68    fn canonical_compact_trims_bits() -> VortexResult<()> {
69        let mut ctx = array_session().create_execution_ctx();
70        let handle = BufferHandle::new_host((0..64u8).collect());
71        let array = BoolArray::try_new_from_handle(handle, 0, 20, Validity::NonNullable)?;
72
73        let trimmed = Canonical::Bool(array).compact(&mut ctx)?.into_bool();
74
75        assert_eq!(trimmed.bits.len(), 3);
76        Ok(())
77    }
78}