diff --git a/vortex-array/src/arrays/bool/compact.rs b/vortex-array/src/arrays/bool/compact.rs new file mode 100644 index 00000000000..37727843cb8 --- /dev/null +++ b/vortex-array/src/arrays/bool/compact.rs @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBufferMeta; +use vortex_error::VortexResult; + +use crate::arrays::BoolArray; +use crate::buffer::BufferHandle; + +impl BoolArray { + /// Trims the packed bit buffer to the bytes backing the array's visible bits. + /// + /// Arrays built over a shared buffer, such as one decoded from a file segment, can hold + /// trailing bytes that no read can see but that `nbytes` and serialization still pay for. The + /// trim is zero-copy and keeps the leading bit offset, so only whole bytes are dropped. + pub fn trim_bits(&self) -> VortexResult { + let byte_len = BitBufferMeta::new(self.meta.offset(), self.len()).byte_len(); + + let bits = match self.bits.as_host_opt() { + Some(host) => BufferHandle::new_host(host.slice_unaligned(..byte_len)), + None => self.bits.slice(0..byte_len), + }; + + Self::try_new_from_handle(bits, self.meta.offset(), self.len(), self.validity()?) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::BitBuffer; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + + use crate::Canonical; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::assert_arrays_eq; + use crate::buffer::BufferHandle; + use crate::validity::Validity; + + #[rstest] + #[case(0, 16, Validity::NonNullable)] + #[case(0, 13, Validity::NonNullable)] + #[case(3, 10, Validity::AllValid)] + #[case(7, 1, Validity::NonNullable)] + #[case(0, 0, Validity::NonNullable)] + fn trims_to_visible_bytes( + #[case] offset: usize, + #[case] len: usize, + #[case] validity: Validity, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let bytes: ByteBuffer = (0..64u8).collect(); + let handle = BufferHandle::new_host(bytes.clone()); + + let trimmed = + BoolArray::try_new_from_handle(handle, offset, len, validity.clone())?.trim_bits()?; + + assert_eq!(trimmed.bits.len(), (offset + len).div_ceil(8)); + let expected = BoolArray::new(BitBuffer::new_with_offset(bytes, len, offset), validity); + assert_arrays_eq!(trimmed, expected, &mut ctx); + Ok(()) + } + + #[test] + fn canonical_compact_trims_bits() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let handle = BufferHandle::new_host((0..64u8).collect()); + let array = BoolArray::try_new_from_handle(handle, 0, 20, Validity::NonNullable)?; + + let trimmed = Canonical::Bool(array).compact(&mut ctx)?.into_bool(); + + assert_eq!(trimmed.bits.len(), 3); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/bool/mod.rs b/vortex-array/src/arrays/bool/mod.rs index 96856c66a75..813853b9d2d 100644 --- a/vortex-array/src/arrays/bool/mod.rs +++ b/vortex-array/src/arrays/bool/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod array; +mod compact; mod patch; pub use array::BoolArrayExt; diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 5330aa2f1ed..ede547d95d2 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -282,10 +282,11 @@ impl Canonical { /// This is mostly relevant for the variable-length types such as Utf8, Binary or List where /// they can accumulate wasted space after slicing and taking operations. /// - /// This operation is very expensive and can result in things like allocations, full-scans + /// This operation can be very expensive and can result in things like allocations, full-scans /// and copy operations. pub fn compact(&self, ctx: &mut ExecutionCtx) -> VortexResult { match self { + Canonical::Bool(array) => Ok(Canonical::Bool(array.trim_bits()?)), Canonical::VarBinView(array) => Ok(Canonical::VarBinView(array.compact_buffers(ctx)?)), Canonical::List(array) => Ok(Canonical::List( array.rebuild(ListViewRebuildMode::TrimElements, ctx)?,