vortex_array/pipeline/operators/
binary_bool.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::hash::Hash;
6use std::rc::Rc;
7
8use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_panic};
9
10use crate::pipeline::bits::BitView;
11use crate::pipeline::operators::{BindContext, Operator};
12use crate::pipeline::types::VType;
13use crate::pipeline::vec::VectorId;
14use crate::pipeline::view::ViewMut;
15use crate::pipeline::{Kernel, KernelContext};
16
17/// Boolean operations supported by the binary boolean operator.
18#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
19pub enum BoolOp {
20    And,
21    Or,
22}
23
24/// Pipeline operator for binary boolean operations (AND, OR) on two boolean arrays.
25#[derive(Debug, Hash)]
26pub struct BinaryBoolOpOperator {
27    children: [Rc<dyn Operator>; 2],
28    op: BoolOp,
29}
30
31impl BinaryBoolOpOperator {
32    fn new(left: Rc<dyn Operator>, right: Rc<dyn Operator>, op: BoolOp) -> Self {
33        // Verify both children are boolean type
34        let VType::Bool = left.vtype() else {
35            vortex_panic!("BinaryBoolOpOperator left child must be boolean type");
36        };
37        let VType::Bool = right.vtype() else {
38            vortex_panic!("BinaryBoolOpOperator right child must be boolean type");
39        };
40
41        Self {
42            children: [left, right],
43            op,
44        }
45    }
46
47    pub fn and(left: Rc<dyn Operator>, right: Rc<dyn Operator>) -> Self {
48        Self::new(left, right, BoolOp::And)
49    }
50
51    pub fn or(left: Rc<dyn Operator>, right: Rc<dyn Operator>) -> Self {
52        Self::new(left, right, BoolOp::Or)
53    }
54}
55
56impl Operator for BinaryBoolOpOperator {
57    fn as_any(&self) -> &dyn Any {
58        self
59    }
60
61    fn vtype(&self) -> VType {
62        VType::Bool
63    }
64
65    fn children(&self) -> &[Rc<dyn Operator>] {
66        &self.children
67    }
68
69    fn with_children(&self, children: Vec<Rc<dyn Operator>>) -> Rc<dyn Operator> {
70        let [left, right] = children
71            .try_into()
72            .ok()
73            .vortex_expect("Expected 2 children");
74        Rc::new(BinaryBoolOpOperator::new(left, right, self.op))
75    }
76
77    fn bind(&self, ctx: &dyn BindContext) -> VortexResult<Box<dyn Kernel>> {
78        if self.vtype() != VType::Bool {
79            vortex_bail!("BinaryBoolOpOperator only supports boolean types");
80        }
81
82        Ok(Box::new(BinaryBoolOpKernel {
83            left: ctx.children()[0],
84            right: ctx.children()[1],
85            op: self.op,
86        }))
87    }
88}
89
90/// Kernel that performs binary boolean operations on two input vectors.
91pub struct BinaryBoolOpKernel {
92    left: VectorId,
93    right: VectorId,
94    op: BoolOp,
95}
96
97impl Kernel for BinaryBoolOpKernel {
98    fn seek(&mut self, _chunk_idx: usize) -> VortexResult<()> {
99        Ok(())
100    }
101
102    fn step(
103        &mut self,
104        ctx: &KernelContext,
105        selected: BitView,
106        out: &mut ViewMut,
107    ) -> VortexResult<()> {
108        let left_vec = ctx.vector(self.left);
109        let right_vec = ctx.vector(self.right);
110        let left_values = left_vec.as_slice::<bool>();
111        let right_values = right_vec.as_slice::<bool>();
112        let out_slice = out.as_slice_mut::<bool>();
113
114        assert!(selected.true_count() <= left_values.len());
115        assert!(selected.true_count() <= right_values.len());
116        assert!(selected.true_count() <= out_slice.len());
117
118        match self.op {
119            BoolOp::And => {
120                for i in 0..selected.true_count() {
121                    unsafe {
122                        let left_val = *left_values.get_unchecked(i);
123                        let right_val = *right_values.get_unchecked(i);
124                        *out_slice.get_unchecked_mut(i) = left_val && right_val;
125                    }
126                }
127            }
128            BoolOp::Or => {
129                for i in 0..selected.true_count() {
130                    unsafe {
131                        let left_val = *left_values.get_unchecked(i);
132                        let right_val = *right_values.get_unchecked(i);
133                        *out_slice.get_unchecked_mut(i) = left_val || right_val;
134                    }
135                }
136            }
137        }
138
139        Ok(())
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use std::rc::Rc;
146
147    use vortex_buffer::BufferMut;
148    use vortex_dtype::Nullability;
149    use vortex_scalar::Scalar;
150
151    use super::*;
152    use crate::arrays::{ConstantArray, PrimitiveArray};
153    use crate::compute::Operator as BinaryOperator;
154    use crate::pipeline::bits::BitView;
155    use crate::pipeline::operators::scalar_compare::ScalarCompareOperator;
156    use crate::pipeline::query::QueryPlan;
157    use crate::pipeline::view::ViewMut;
158    use crate::pipeline::{N, N_WORDS};
159
160    #[test]
161    fn test_binary_bool_and_basic() {
162        // Create left data: [1, 0, 1, 0] to generate [true, false, true, false]
163        let size = 4;
164        let left_primitive_array = [1i32, 0, 1, 0].into_iter().collect::<PrimitiveArray>();
165        let left_primitive_op = left_primitive_array
166            .as_ref()
167            .to_operator()
168            .unwrap()
169            .unwrap();
170
171        // Create right data: [1, 1, 0, 0] to generate [true, true, false, false]
172        let right_primitive_array = [1i32, 1, 0, 0].into_iter().collect::<PrimitiveArray>();
173        let right_primitive_op = right_primitive_array.to_operator().unwrap().unwrap();
174
175        let zero_scalar = Scalar::primitive(0i32, Nullability::NonNullable);
176        let left_bool_op = Rc::new(ScalarCompareOperator::new(
177            left_primitive_op,
178            BinaryOperator::Gt,
179            zero_scalar.clone(),
180        ));
181        let right_bool_op = Rc::new(ScalarCompareOperator::new(
182            right_primitive_op,
183            BinaryOperator::Gt,
184            zero_scalar,
185        ));
186
187        // Create binary AND operator: left_bool AND right_bool
188        let binary_and_op = Rc::new(BinaryBoolOpOperator::and(left_bool_op, right_bool_op));
189
190        // Create query plan from the operator
191        let plan = QueryPlan::new(binary_and_op.as_ref()).unwrap();
192        let mut pipeline = plan.executable_plan().unwrap();
193
194        // Create mask for first 4 elements
195        let mut mask_data = [0usize; N_WORDS];
196        mask_data[0] = 0b1111; // First 4 bits set
197        let mask_view = BitView::new(&mask_data);
198
199        // Create output buffer for results
200        let mut output = BufferMut::<bool>::with_capacity(N);
201        unsafe { output.set_len(N) };
202        let mut output_view = ViewMut::new(&mut output[..], None);
203
204        // Execute the pipeline
205        let result = pipeline._step(mask_view, &mut output_view);
206        assert!(result.is_ok());
207
208        // Verify results: left[i] AND right[i]
209        // [true&true, false&true, true&false, false&false] = [true, false, false, false]
210        let expected = [true, false, false, false];
211        for i in 0..size {
212            assert_eq!(
213                output[i], expected[i],
214                "Position {}: expected {}, got {}",
215                i, expected[i], output[i]
216            );
217        }
218    }
219
220    #[test]
221    fn test_binary_bool_or_basic() {
222        // Create left data: [1, 0, 1, 0] to generate [true, false, true, false]
223        let size = 4;
224        let left_primitive_array = [1i32, 0, 1, 0].into_iter().collect::<PrimitiveArray>();
225        let left_primitive_op = left_primitive_array.to_operator().unwrap().unwrap();
226
227        // Create right data: [1, 1, 0, 0] to generate [true, true, false, false]
228        let right_primitive_array = [1i32, 1, 0, 0].into_iter().collect::<PrimitiveArray>();
229        let right_primitive_op = right_primitive_array.to_operator().unwrap().unwrap();
230
231        // Create comparisons to generate boolean values: value > 0
232        let zero_scalar = Scalar::primitive(0i32, Nullability::NonNullable);
233        let left_bool_op = Rc::new(ScalarCompareOperator::new(
234            left_primitive_op,
235            BinaryOperator::Gt,
236            zero_scalar.clone(),
237        ));
238        let right_bool_op = Rc::new(ScalarCompareOperator::new(
239            right_primitive_op,
240            BinaryOperator::Gt,
241            zero_scalar,
242        ));
243
244        // Create binary OR operator: left_bool OR right_bool
245        let binary_or_op = Rc::new(BinaryBoolOpOperator::or(left_bool_op, right_bool_op));
246
247        // Create query plan from the operator
248        let plan = QueryPlan::new(binary_or_op.as_ref()).unwrap();
249        let mut pipeline = plan.executable_plan().unwrap();
250
251        // Create mask for first 4 elements
252        let mut mask_data = [0usize; N_WORDS];
253        mask_data[0] = 0b1111; // First 4 bits set
254        let mask_view = BitView::new(&mask_data);
255
256        // Create output buffer for results
257        let mut output = BufferMut::<bool>::with_capacity(N);
258        unsafe { output.set_len(N) };
259        let mut output_view = ViewMut::new(&mut output[..], None);
260
261        // Execute the pipeline
262        let result = pipeline._step(mask_view, &mut output_view);
263        assert!(result.is_ok());
264
265        // Verify results: left[i] OR right[i]
266        // [true|true, false|true, true|false, false|false] = [true, true, true, false]
267        let expected = [true, true, true, false];
268        for i in 0..size {
269            assert_eq!(
270                output[i], expected[i],
271                "Position {}: expected {}, got {}",
272                i, expected[i], output[i]
273            );
274        }
275    }
276
277    #[test]
278    fn test_binary_bool_with_constant() {
279        // Test combining boolean operators with constants: (x AND y) OR true
280        let size = 4;
281
282        // Create x data: [1, 0, 1, 0] to generate [true, false, true, false]
283        let x_primitive_array = [1i32, 0, 1, 0].into_iter().collect::<PrimitiveArray>();
284        let x_primitive_op = x_primitive_array.to_operator().unwrap().unwrap();
285
286        // Create y data: [0, 0, 1, 1] to generate [false, false, true, true]
287        let y_primitive_array = [0i32, 0, 1, 1].into_iter().collect::<PrimitiveArray>();
288        let y_primitive_op = y_primitive_array.to_operator().unwrap().unwrap();
289
290        // Create comparisons to generate boolean values: value > 0
291        use vortex_dtype::Nullability;
292        use vortex_scalar::Scalar;
293
294        use crate::compute::Operator as BinaryOperator;
295        use crate::pipeline::operators::scalar_compare::ScalarCompareOperator;
296
297        let zero_scalar = Scalar::primitive(0i32, Nullability::NonNullable);
298        let x_bool_op = Rc::new(ScalarCompareOperator::new(
299            x_primitive_op,
300            BinaryOperator::Gt,
301            zero_scalar.clone(),
302        ));
303        let y_bool_op = Rc::new(ScalarCompareOperator::new(
304            y_primitive_op,
305            BinaryOperator::Gt,
306            zero_scalar,
307        ));
308
309        // Create constant true
310        let constant_true = ConstantArray::new(Scalar::from(true), size)
311            .to_operator()
312            .unwrap()
313            .unwrap();
314
315        // Build pipeline: (x AND y) OR true
316        let x_and_y = Rc::new(BinaryBoolOpOperator::and(x_bool_op, y_bool_op));
317        let final_op = Rc::new(BinaryBoolOpOperator::or(x_and_y, constant_true));
318
319        // Create query plan
320        let plan = QueryPlan::new(final_op.as_ref()).unwrap();
321        let mut pipeline = plan.executable_plan().unwrap();
322
323        // Create mask for first 4 elements
324        let mut mask_data = [0usize; N_WORDS];
325        mask_data[0] = 0b1111; // First 4 bits set
326        let mask_view = BitView::new(&mask_data);
327
328        // Create output buffer
329        let mut output = BufferMut::<bool>::with_capacity(N);
330        unsafe { output.set_len(N) };
331        let mut output_view = ViewMut::new(&mut output[..], None);
332
333        // Execute the pipeline
334        let result = pipeline._step(mask_view, &mut output_view);
335        assert!(result.is_ok());
336
337        // Verify results: (x[i] AND y[i]) OR true
338        // [(true&false)|true, (false&false)|true, (true&true)|true, (false&true)|true] = [true, true, true, true]
339        let expected = [true, true, true, true];
340        for i in 0..size {
341            assert_eq!(
342                output[i], expected[i],
343                "Position {}: expected {}, got {}",
344                i, expected[i], output[i]
345            );
346        }
347    }
348}
vortex_array/pipeline/operators/binary_bool.rs

vortex_array/pipeline/operators/
binary_bool.rs
