vortex_array/arrays/bool/compute/
zip.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBuffer;
5use vortex_buffer::BufferMut;
6use vortex_error::VortexResult;
7use vortex_mask::Mask;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::array::ArrayView;
13use crate::arrays::Bool;
14use crate::arrays::BoolArray;
15use crate::arrays::bool::BoolArrayExt;
16use crate::builtins::ArrayBuiltins;
17use crate::scalar_fn::fns::zip::ZipKernel;
18use crate::scalar_fn::fns::zip::zip_validity;
19
20/// A branchless boolean zip kernel that blends the two value bitmaps with the mask in one pass.
21///
22/// Booleans are bit-packed, so selecting `if_true` where the mask is set and `if_false` where it is
23/// not is a single bitwise blend over the packed words — `(true & mask) | (false & !mask)` — instead
24/// of the generic per-run builder. Validity is combined with the shared `zip_validity`, which itself
25/// reuses this kernel (terminating immediately, since validity bitmaps are non-nullable).
26impl ZipKernel for Bool {
27    fn zip(
28        if_true: ArrayView<'_, Bool>,
29        if_false: &ArrayRef,
30        mask: &ArrayRef,
31        ctx: &mut ExecutionCtx,
32    ) -> VortexResult<Option<ArrayRef>> {
33        let Some(if_false) = if_false.as_opt::<Bool>() else {
34            return Ok(None);
35        };
36
37        // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics.
38        let mask = mask.clone().fill_null(false)?.execute::<Mask>(ctx)?;
39        let mask_values = match &mask {
40            // Defer trivial masks to the generic zip, which just casts the surviving side.
41            Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None),
42            Mask::Values(values) => values,
43        };
44        let mask_bits = mask_values.bit_buffer();
45
46        let values = zip_value_bits(
47            &if_true.to_bit_buffer(),
48            &if_false.to_bit_buffer(),
49            mask_bits,
50        );
51
52        let validity = zip_validity(if_true.validity()?, if_false.validity()?, &mask)?;
53
54        Ok(Some(BoolArray::new(values, validity).into_array()))
55    }
56}
57
58fn zip_value_bits(if_true: &BitBuffer, if_false: &BitBuffer, mask: &BitBuffer) -> BitBuffer {
59    assert_eq!(if_true.len(), if_false.len());
60    assert_eq!(if_true.len(), mask.len());
61
62    let true_chunks = if_true.chunks();
63    let false_chunks = if_false.chunks();
64    let mask_chunks = mask.chunks();
65
66    let mut values = BufferMut::<u64>::with_capacity(true_chunks.num_u64s());
67    for ((true_bits, false_bits), mask_bits) in true_chunks
68        .iter()
69        .zip(false_chunks.iter())
70        .zip(mask_chunks.iter())
71    {
72        values.push((true_bits & mask_bits) | (false_bits & !mask_bits));
73    }
74
75    if true_chunks.remainder_len() != 0 {
76        let true_bits = true_chunks.remainder_bits();
77        let false_bits = false_chunks.remainder_bits();
78        let mask_bits = mask_chunks.remainder_bits();
79        values.push((true_bits & mask_bits) | (false_bits & !mask_bits));
80    }
81
82    BitBuffer::new(values.freeze().into_byte_buffer(), if_true.len())
83}
84
85#[cfg(test)]
86mod tests {
87    use vortex_error::VortexResult;
88    use vortex_mask::Mask;
89
90    use super::zip_value_bits;
91    use crate::ArrayRef;
92    use crate::IntoArray;
93    use crate::VortexSessionExecute;
94    use crate::array_session;
95    use crate::arrays::Bool;
96    use crate::arrays::BoolArray;
97    use crate::assert_arrays_eq;
98    use crate::builtins::ArrayBuiltins;
99
100    #[test]
101    fn blend_value_bits_boundaries() {
102        for len in [0usize, 1, 2, 7, 8, 9, 63, 64, 65, 127, 128] {
103            let if_true = (0..len).map(|i| i.is_multiple_of(2)).collect();
104            let if_false = (0..len).map(|i| i.is_multiple_of(3)).collect();
105            let mask = (0..len).map(|i| i % 3 != 1).collect();
106
107            let values = zip_value_bits(&if_true, &if_false, &mask);
108
109            assert_eq!(values.len(), len);
110            assert_eq!(
111                values.iter().collect::<Vec<_>>(),
112                (0..len)
113                    .map(|i| {
114                        if i % 3 != 1 {
115                            i.is_multiple_of(2)
116                        } else {
117                            i.is_multiple_of(3)
118                        }
119                    })
120                    .collect::<Vec<_>>(),
121                "failed for len {len}",
122            );
123        }
124    }
125
126    /// Blend two non-nullable bool arrays across the 64-bit mask chunk boundary + remainder.
127    #[test]
128    fn zip_nonnull_spans_mask_chunks() -> VortexResult<()> {
129        let len = 150usize;
130        let if_true = BoolArray::from_iter((0..len).map(|i| i.is_multiple_of(2))).into_array();
131        let if_false = BoolArray::from_iter((0..len).map(|i| i.is_multiple_of(3))).into_array();
132
133        let bits: Vec<bool> = (0..len).map(|i| i.is_multiple_of(5) || i == 64).collect();
134        let mask = Mask::from_iter(bits.iter().copied());
135
136        let mut ctx = array_session().create_execution_ctx();
137        let result = mask
138            .into_array()
139            .zip(if_true, if_false)?
140            .execute::<ArrayRef>(&mut ctx)?;
141        assert!(result.is::<Bool>());
142
143        let expected = BoolArray::from_iter((0..len).map(|i| {
144            if bits[i] {
145                i.is_multiple_of(2)
146            } else {
147                i.is_multiple_of(3)
148            }
149        }))
150        .into_array();
151        assert_arrays_eq!(result, expected, &mut ctx);
152        Ok(())
153    }
154
155    /// With `Validity::Array` on both sides, select values and validity from the chosen side.
156    #[test]
157    fn zip_nullable_selects_values_and_validity() -> VortexResult<()> {
158        let len = 130usize;
159        let if_true = BoolArray::from_iter(
160            (0..len).map(|i| (!i.is_multiple_of(4)).then_some(i.is_multiple_of(2))),
161        )
162        .into_array();
163        let if_false = BoolArray::from_iter(
164            (0..len).map(|i| (!i.is_multiple_of(5)).then_some(i.is_multiple_of(3))),
165        )
166        .into_array();
167
168        let bits: Vec<bool> = (0..len).map(|i| i.is_multiple_of(2)).collect();
169        let mask = Mask::from_iter(bits.iter().copied());
170
171        let mut ctx = array_session().create_execution_ctx();
172        let result = mask
173            .into_array()
174            .zip(if_true, if_false)?
175            .execute::<ArrayRef>(&mut ctx)?;
176        assert!(result.is::<Bool>());
177
178        let expected = BoolArray::from_iter((0..len).map(|i| {
179            if bits[i] {
180                (!i.is_multiple_of(4)).then_some(i.is_multiple_of(2))
181            } else {
182                (!i.is_multiple_of(5)).then_some(i.is_multiple_of(3))
183            }
184        }))
185        .into_array();
186        assert_arrays_eq!(result, expected, &mut ctx);
187        Ok(())
188    }
189}
vortex_array/arrays/bool/compute/zip.rs

vortex_array/arrays/bool/compute/
zip.rs
