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

vortex_array/arrays/bool/compute/
zip.rs
