vortex_compute/comparison/
collection.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBuffer;
5
6use crate::comparison::Compare;
7use crate::comparison::ComparisonOperator;
8
9/// Adapter to implement `Compare` for any `ComparableCollection`.
10pub(crate) struct ComparableCollectionAdapter<C>(pub C);
11
12impl<Op, C> Compare<Op> for ComparableCollectionAdapter<C>
13where
14    C: ComparableCollection,
15    Op: ComparisonOperator<C::Item>,
16{
17    type Output = BitBuffer;
18
19    fn compare(self, rhs: Self) -> Self::Output {
20        assert_eq!(self.0.len(), rhs.0.len());
21
22        BitBuffer::from_iter((0..self.0.len()).map(|i| {
23            let left = unsafe { self.0.item_unchecked(i) };
24            let right = unsafe { rhs.0.item_unchecked(i) };
25            Op::apply(&left, &right)
26        }))
27    }
28}
29
30/// Marker trait for comparable collections.
31pub trait ComparableCollection {
32    /// The item type that can be compared.
33    type Item;
34
35    /// Get the length of the comparable collection.
36    fn len(&self) -> usize;
37
38    /// Get the item at the specified index without bounds checking.
39    unsafe fn item_unchecked(&self, index: usize) -> Self::Item;
40}
41
42impl<T: Copy> ComparableCollection for &[T] {
43    type Item = T;
44
45    fn len(&self) -> usize {
46        <[T]>::len(self)
47    }
48
49    unsafe fn item_unchecked(&self, index: usize) -> Self::Item {
50        unsafe { *self.get_unchecked(index) }
51    }
52}
53
54impl<Op, T> Compare<Op> for &[T]
55where
56    T: Copy,
57    Op: ComparisonOperator<T>,
58{
59    type Output = BitBuffer;
60
61    fn compare(self, rhs: Self) -> Self::Output {
62        Compare::<Op>::compare(
63            ComparableCollectionAdapter(self),
64            ComparableCollectionAdapter(rhs),
65        )
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use vortex_buffer::bitbuffer;
72
73    use super::*;
74    use crate::comparison::Equal;
75    use crate::comparison::GreaterThan;
76    use crate::comparison::LessThan;
77    use crate::comparison::NotEqual;
78
79    #[test]
80    fn test_slice_equal() {
81        let left: &[u32] = &[1, 2, 3, 4];
82        let right: &[u32] = &[1, 2, 5, 4];
83
84        let result = Compare::<Equal>::compare(left, right);
85        assert_eq!(result, bitbuffer![1 1 0 1]);
86    }
87
88    #[test]
89    fn test_slice_not_equal() {
90        let left: &[u32] = &[1, 2, 3, 4];
91        let right: &[u32] = &[1, 2, 5, 4];
92
93        let result = Compare::<NotEqual>::compare(left, right);
94        assert_eq!(result, bitbuffer![0 0 1 0]);
95    }
96
97    #[test]
98    fn test_slice_less_than() {
99        let left: &[u32] = &[1, 2, 3, 4];
100        let right: &[u32] = &[2, 2, 1, 5];
101
102        let result = Compare::<LessThan>::compare(left, right);
103        assert_eq!(result, bitbuffer![1 0 0 1]);
104    }
105
106    #[test]
107    fn test_slice_greater_than() {
108        let left: &[u32] = &[3, 2, 1, 5];
109        let right: &[u32] = &[1, 2, 3, 4];
110
111        let result = Compare::<GreaterThan>::compare(left, right);
112        assert_eq!(result, bitbuffer![1 0 0 1]);
113    }
114}