vortex_compute/comparison/
bool.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::BitAnd;
5
6use vortex_buffer::BitBuffer;
7use vortex_buffer::BufferMut;
8use vortex_vector::VectorOps;
9use vortex_vector::bool::BoolScalar;
10use vortex_vector::bool::BoolVector;
11
12use crate::comparison::Compare;
13use crate::comparison::Equal;
14use crate::comparison::GreaterThan;
15use crate::comparison::GreaterThanOrEqual;
16use crate::comparison::LessThan;
17use crate::comparison::LessThanOrEqual;
18use crate::comparison::NotEqual;
19
20impl Compare<Equal> for BoolScalar {
21    type Output = BoolScalar;
22
23    fn compare(self, rhs: Self) -> Self::Output {
24        BoolScalar::new(self.value().zip(rhs.value()).map(|(l, r)| l == r))
25    }
26}
27
28impl Compare<NotEqual> for BoolScalar {
29    type Output = BoolScalar;
30
31    fn compare(self, rhs: Self) -> Self::Output {
32        BoolScalar::new(self.value().zip(rhs.value()).map(|(l, r)| l != r))
33    }
34}
35
36impl Compare<LessThan> for BoolScalar {
37    type Output = BoolScalar;
38
39    fn compare(self, rhs: Self) -> Self::Output {
40        BoolScalar::new(self.value().zip(rhs.value()).map(|(l, r)| !l && r))
41    }
42}
43
44impl Compare<LessThanOrEqual> for BoolScalar {
45    type Output = BoolScalar;
46
47    fn compare(self, rhs: Self) -> Self::Output {
48        BoolScalar::new(self.value().zip(rhs.value()).map(|(l, r)| !l || r))
49    }
50}
51
52impl Compare<GreaterThan> for BoolScalar {
53    type Output = BoolScalar;
54
55    fn compare(self, rhs: Self) -> Self::Output {
56        BoolScalar::new(self.value().zip(rhs.value()).map(|(l, r)| l && !r))
57    }
58}
59
60impl Compare<GreaterThanOrEqual> for BoolScalar {
61    type Output = BoolScalar;
62
63    fn compare(self, rhs: Self) -> Self::Output {
64        BoolScalar::new(self.value().zip(rhs.value()).map(|(l, r)| l || !r))
65    }
66}
67
68impl<Op> Compare<Op> for BoolVector
69where
70    Op: BitComparisonOperator,
71{
72    type Output = BoolVector;
73
74    fn compare(self, rhs: Self) -> Self::Output {
75        let validity = self.validity().bitand(rhs.validity());
76
77        let lhs = self.bits().chunks();
78        let rhs = rhs.bits().chunks();
79
80        // Reserve one extra chunk to account for partial padding chunk at the end.
81        let mut buffer = BufferMut::<u64>::with_capacity(lhs.chunk_len() + 1);
82        buffer.extend(
83            lhs.iter_padded()
84                .zip(rhs.iter_padded())
85                .map(|(a_chunk, b_chunk)| Op::apply(&a_chunk, &b_chunk)),
86        );
87        let bits = BitBuffer::new(buffer.freeze().into_byte_buffer(), self.len());
88
89        BoolVector::new(bits, validity)
90    }
91}
92
93pub trait BitComparisonOperator {
94    fn apply(a: &u64, b: &u64) -> u64;
95}
96
97impl BitComparisonOperator for Equal {
98    fn apply(a: &u64, b: &u64) -> u64 {
99        !(a ^ b)
100    }
101}
102impl BitComparisonOperator for NotEqual {
103    fn apply(a: &u64, b: &u64) -> u64 {
104        a ^ b
105    }
106}
107impl BitComparisonOperator for LessThan {
108    fn apply(a: &u64, b: &u64) -> u64 {
109        (!a) & b
110    }
111}
112impl BitComparisonOperator for LessThanOrEqual {
113    fn apply(a: &u64, b: &u64) -> u64 {
114        !(a & (!b))
115    }
116}
117impl BitComparisonOperator for GreaterThan {
118    fn apply(a: &u64, b: &u64) -> u64 {
119        a & (!b)
120    }
121}
122impl BitComparisonOperator for GreaterThanOrEqual {
123    fn apply(a: &u64, b: &u64) -> u64 {
124        !((!a) & b)
125    }
126}