vortex_array/arrays/bool/compute/
mask.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::IntoArray;
8use crate::array::ArrayView;
9use crate::arrays::Bool;
10use crate::arrays::BoolArray;
11use crate::scalar_fn::fns::mask::MaskReduce;
12use crate::validity::Validity;
13
14impl MaskReduce for Bool {
15    const VALIDITY_IS_METADATA_ONLY: bool = true;
16
17    fn mask(array: ArrayView<'_, Bool>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
18        Ok(Some(
19            BoolArray::try_new_from_handle(
20                array.bits.clone(),
21                array.meta.offset(),
22                array.len(),
23                array.validity()?.and(Validity::Array(mask.clone()))?,
24            )?
25            .into_array(),
26        ))
27    }
28}
29
30#[cfg(test)]
31mod tests {
32    use std::any::Any;
33    use std::ops::Range;
34    use std::sync::Arc;
35
36    use futures::future::BoxFuture;
37    use rstest::rstest;
38    use vortex_buffer::Alignment;
39    use vortex_buffer::ByteBuffer;
40    use vortex_error::VortexResult;
41    use vortex_error::vortex_bail;
42
43    use crate::ArrayEq;
44    use crate::ArrayRef;
45    use crate::EqMode;
46    use crate::IntoArray;
47    use crate::VortexSessionExecute;
48    use crate::array_session;
49    use crate::arrays::Bool;
50    use crate::arrays::BoolArray;
51    use crate::buffer::BufferHandle;
52    use crate::buffer::DeviceBuffer;
53    use crate::compute::conformance::mask::test_mask_conformance;
54    use crate::scalar_fn::fns::mask::MaskReduce;
55    use crate::validity::Validity;
56
57    // Any access beyond buffer metadata would invalidate the metadata-only reduction.
58    #[derive(Debug, PartialEq, Eq, Hash)]
59    struct MetadataOnlyDeviceBuffer(usize);
60
61    impl DeviceBuffer for MetadataOnlyDeviceBuffer {
62        fn as_any(&self) -> &dyn Any {
63            self
64        }
65
66        fn len(&self) -> usize {
67            self.0
68        }
69
70        fn alignment(&self) -> Alignment {
71            Alignment::of::<u8>()
72        }
73
74        fn copy_to_host_sync(&self, _alignment: Alignment) -> VortexResult<ByteBuffer> {
75            vortex_bail!("mask reduction must not copy device values")
76        }
77
78        fn copy_to_host(
79            &self,
80            _alignment: Alignment,
81        ) -> VortexResult<BoxFuture<'static, VortexResult<ByteBuffer>>> {
82            vortex_bail!("mask reduction must not copy device values")
83        }
84
85        fn slice(&self, _range: Range<usize>) -> Arc<dyn DeviceBuffer> {
86            panic!("mask reduction must preserve the original buffer");
87        }
88
89        fn aligned(self: Arc<Self>, alignment: Alignment) -> VortexResult<Arc<dyn DeviceBuffer>> {
90            assert_eq!(self.alignment(), alignment);
91            Ok(self)
92        }
93    }
94
95    #[test]
96    fn mask_preserves_device_buffer_and_bit_offset() -> VortexResult<()> {
97        let handle = BufferHandle::new_device(Arc::new(MetadataOnlyDeviceBuffer(1)));
98        let input = BoolArray::try_new_from_handle(handle, 3, 5, Validity::AllValid)?;
99        let mask = BoolArray::from_iter([true, false, true, false, true]).into_array();
100        let masked = <Bool as MaskReduce>::mask(input.as_ref().as_::<Bool>(), &mask)?
101            .expect("Boolean masking must reduce");
102        let output = masked.as_::<Bool>();
103
104        assert!(output.bits.array_eq(&input.bits, EqMode::Ptr));
105        assert_eq!(output.meta.offset(), input.meta.offset());
106        assert_eq!(output.len(), input.len());
107        let Validity::Array(validity) = masked.validity()? else {
108            panic!("mask must remain array-backed");
109        };
110        assert!(ArrayRef::ptr_eq(&validity, &mask));
111
112        Ok(())
113    }
114
115    #[rstest]
116    #[case(BoolArray::from_iter([true, false, true, true, false]))]
117    #[case(BoolArray::from_iter([Some(true), None, Some(false), Some(true), None]))]
118    #[case(BoolArray::from_iter([true]))]
119    #[case(BoolArray::from_iter([false, false]))]
120    #[case(BoolArray::from_iter((0..100).map(|i| i % 2 == 0)))]
121    fn test_mask_bool_conformance(#[case] array: BoolArray) {
122        test_mask_conformance(
123            &array.into_array(),
124            &mut array_session().create_execution_ctx(),
125        );
126    }
127}
vortex_array/arrays/bool/compute/mask.rs

vortex_array/arrays/bool/compute/
mask.rs
