vortex_compute/comparison/
pvector.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::BitAnd;
5
6use vortex_dtype::NativePType;
7use vortex_vector::VectorOps;
8use vortex_vector::bool::BoolVector;
9use vortex_vector::primitive::PVector;
10
11use crate::comparison::Compare;
12use crate::comparison::ComparisonOperator;
13use crate::comparison::collection::ComparableCollectionAdapter;
14
15impl<Op, T> Compare<Op> for &PVector<T>
16where
17    T: NativePType,
18    Op: ComparisonOperator<T>,
19{
20    type Output = BoolVector;
21
22    fn compare(self, rhs: &PVector<T>) -> Self::Output {
23        let validity = self.validity().bitand(rhs.validity());
24
25        // TODO(ngates): match on density of validity mask to choose optimal implementation
26
27        let bits = Compare::<Op>::compare(
28            ComparableCollectionAdapter(self.elements().as_slice()),
29            ComparableCollectionAdapter(rhs.elements().as_slice()),
30        );
31
32        BoolVector::new(bits, validity)
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use vortex_buffer::bitbuffer;
39    use vortex_buffer::buffer;
40    use vortex_mask::Mask;
41    use vortex_vector::bool::BoolVector;
42
43    use super::*;
44    use crate::comparison::Equal;
45    use crate::comparison::GreaterThan;
46    use crate::comparison::GreaterThanOrEqual;
47    use crate::comparison::LessThan;
48    use crate::comparison::LessThanOrEqual;
49    use crate::comparison::NotEqual;
50
51    #[test]
52    fn test_equal() {
53        let left = PVector::new(buffer![1u32, 2, 3, 4], Mask::new_true(4));
54        let right = PVector::new(buffer![1u32, 2, 5, 4], Mask::new_true(4));
55
56        let result = Compare::<Equal>::compare(&left, &right);
57        let expected = BoolVector::new(bitbuffer![1 1 0 1], Mask::new_true(4));
58        assert_eq!(result, expected);
59    }
60
61    #[test]
62    fn test_not_equal() {
63        let left = PVector::new(buffer![1u32, 2, 3, 4], Mask::new_true(4));
64        let right = PVector::new(buffer![1u32, 2, 5, 4], Mask::new_true(4));
65
66        let result = Compare::<NotEqual>::compare(&left, &right);
67        let expected = BoolVector::new(bitbuffer![0 0 1 0], Mask::new_true(4));
68        assert_eq!(result, expected);
69    }
70
71    #[test]
72    fn test_less_than() {
73        let left = PVector::new(buffer![1u32, 2, 3, 4], Mask::new_true(4));
74        let right = PVector::new(buffer![2u32, 2, 1, 5], Mask::new_true(4));
75
76        let result = Compare::<LessThan>::compare(&left, &right);
77        let expected = BoolVector::new(bitbuffer![1 0 0 1], Mask::new_true(4));
78        assert_eq!(result, expected);
79    }
80
81    #[test]
82    fn test_less_than_or_equal() {
83        let left = PVector::new(buffer![1u32, 2, 3, 4], Mask::new_true(4));
84        let right = PVector::new(buffer![2u32, 2, 1, 5], Mask::new_true(4));
85
86        let result = Compare::<LessThanOrEqual>::compare(&left, &right);
87        let expected = BoolVector::new(bitbuffer![1 1 0 1], Mask::new_true(4));
88        assert_eq!(result, expected);
89    }
90
91    #[test]
92    fn test_greater_than() {
93        let left = PVector::new(buffer![3u32, 2, 1, 5], Mask::new_true(4));
94        let right = PVector::new(buffer![1u32, 2, 3, 4], Mask::new_true(4));
95
96        let result = Compare::<GreaterThan>::compare(&left, &right);
97        let expected = BoolVector::new(bitbuffer![1 0 0 1], Mask::new_true(4));
98        assert_eq!(result, expected);
99    }
100
101    #[test]
102    fn test_greater_than_or_equal() {
103        let left = PVector::new(buffer![3u32, 2, 1, 5], Mask::new_true(4));
104        let right = PVector::new(buffer![1u32, 2, 3, 4], Mask::new_true(4));
105
106        let result = Compare::<GreaterThanOrEqual>::compare(&left, &right);
107        let expected = BoolVector::new(bitbuffer![1 1 0 1], Mask::new_true(4));
108        assert_eq!(result, expected);
109    }
110
111    #[test]
112    fn test_compare_with_nulls() {
113        let left = PVector::new(buffer![1u32, 2, 3], Mask::from_iter([true, false, true]));
114        let right = PVector::new(buffer![1u32, 2, 3], Mask::new_true(3));
115
116        let result = Compare::<Equal>::compare(&left, &right);
117        // Validity is AND'd, so if either side is null, result validity is null
118        let expected = BoolVector::new(bitbuffer![1 1 1], Mask::from_iter([true, false, true]));
119        assert_eq!(result, expected);
120    }
121}