Vector comparison compute (#5100)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-compute/Cargo.toml

@@ -28,8 +28,9 @@ vortex-vector = { workspace = true }
 num-traits = { workspace = true }
 
 [features]
-default = ["arithmetic", "filter", "logical"]
+default = ["arithmetic", "comparison", "filter", "logical"]
 
 arithmetic = []
+comparison = []
 filter = []
 logical = []

vortex-compute/src/comparison/bool.rs

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::BitAnd;
+
+use vortex_buffer::{BitBuffer, BufferMut};
+use vortex_vector::{BoolVector, VectorOps};
+
+use crate::comparison::{
+    Compare, Equal, GreaterThan, GreaterThanOrEqual, LessThan, LessThanOrEqual, NotEqual,
+};
+
+impl<Op> Compare<Op> for &BoolVector
+where
+    Op: BitComparisonOperator,
+{
+    type Output = BoolVector;
+
+    fn compare(self, rhs: Self) -> Self::Output {
+        let validity = self.validity().bitand(rhs.validity());
+
+        let lhs = self.bits().chunks();
+        let rhs = rhs.bits().chunks();
+
+        // Reserve one extra chunk to account for partial padding chunk at the end.
+        let mut buffer = BufferMut::<u64>::with_capacity(lhs.chunk_len() + 1);
+        buffer.extend(
+            lhs.iter_padded()
+                .zip(rhs.iter_padded())
+                .map(|(a_chunk, b_chunk)| Op::apply(&a_chunk, &b_chunk)),
+        );
+        let bits = BitBuffer::new(buffer.freeze().into_byte_buffer(), self.len());
+
+        BoolVector::new(bits, validity)
+    }
+}
+
+pub trait BitComparisonOperator {
+    fn apply(a: &u64, b: &u64) -> u64;
+}
+
+impl BitComparisonOperator for Equal {
+    fn apply(a: &u64, b: &u64) -> u64 {
+        !(a ^ b)
+    }
+}
+impl BitComparisonOperator for NotEqual {
+    fn apply(a: &u64, b: &u64) -> u64 {
+        a ^ b
+    }
+}
+impl BitComparisonOperator for LessThan {
+    fn apply(a: &u64, b: &u64) -> u64 {
+        (!a) & b
+    }
+}
+impl BitComparisonOperator for LessThanOrEqual {
+    fn apply(a: &u64, b: &u64) -> u64 {
+        !(a & (!b))
+    }
+}
+impl BitComparisonOperator for GreaterThan {
+    fn apply(a: &u64, b: &u64) -> u64 {
+        a & (!b)
+    }
+}
+impl BitComparisonOperator for GreaterThanOrEqual {
+    fn apply(a: &u64, b: &u64) -> u64 {
+        !((!a) & b)
+    }
+}

vortex-compute/src/comparison/collection.rs

@@ -0,0 +1,110 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_buffer::BitBuffer;
+
+use crate::comparison::{Compare, ComparisonOperator};
+
+/// Adapter to implement `Compare` for any `ComparableCollection`.
+pub(crate) struct ComparableCollectionAdapter<C>(pub C);
+
+impl<Op, C> Compare<Op> for ComparableCollectionAdapter<C>
+where
+    C: ComparableCollection,
+    Op: ComparisonOperator<C::Item>,
+{
+    type Output = BitBuffer;
+
+    fn compare(self, rhs: Self) -> Self::Output {
+        assert_eq!(self.0.len(), rhs.0.len());
+
+        BitBuffer::from_iter((0..self.0.len()).map(|i| {
+            let left = unsafe { self.0.item_unchecked(i) };
+            let right = unsafe { rhs.0.item_unchecked(i) };
+            Op::apply(&left, &right)
+        }))
+    }
+}
+
+/// Marker trait for comparable collections.
+pub trait ComparableCollection {
+    /// The item type that can be compared.
+    type Item;
+
+    /// Get the length of the comparable collection.
+    fn len(&self) -> usize;
+
+    /// Get the item at the specified index without bounds checking.
+    unsafe fn item_unchecked(&self, index: usize) -> Self::Item;
+}
+
+impl<T: Copy> ComparableCollection for &[T] {
+    type Item = T;
+
+    fn len(&self) -> usize {
+        <[T]>::len(self)
+    }
+
+    unsafe fn item_unchecked(&self, index: usize) -> Self::Item {
+        unsafe { *self.get_unchecked(index) }
+    }
+}
+
+impl<Op, T> Compare<Op> for &[T]
+where
+    T: Copy,
+    Op: ComparisonOperator<T>,
+{
+    type Output = BitBuffer;
+
+    fn compare(self, rhs: Self) -> Self::Output {
+        Compare::<Op>::compare(
+            ComparableCollectionAdapter(self),
+            ComparableCollectionAdapter(rhs),
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::bitbuffer;
+
+    use super::*;
+    use crate::comparison::{Equal, GreaterThan, LessThan, NotEqual};
+
+    #[test]
+    fn test_slice_equal() {
+        let left: &[u32] = &[1, 2, 3, 4];
+        let right: &[u32] = &[1, 2, 5, 4];
+
+        let result = Compare::<Equal>::compare(left, right);
+        assert_eq!(result, bitbuffer![1 1 0 1]);
+    }
+
+    #[test]
+    fn test_slice_not_equal() {
+        let left: &[u32] = &[1, 2, 3, 4];
+        let right: &[u32] = &[1, 2, 5, 4];
+
+        let result = Compare::<NotEqual>::compare(left, right);
+        assert_eq!(result, bitbuffer![0 0 1 0]);
+    }
+
+    #[test]
+    fn test_slice_less_than() {
+        let left: &[u32] = &[1, 2, 3, 4];
+        let right: &[u32] = &[2, 2, 1, 5];
+
+        let result = Compare::<LessThan>::compare(left, right);
+        assert_eq!(result, bitbuffer![1 0 0 1]);
+    }
+
+    #[test]
+    fn test_slice_greater_than() {
+        let left: &[u32] = &[3, 2, 1, 5];
+        let right: &[u32] = &[1, 2, 3, 4];
+
+        let result = Compare::<GreaterThan>::compare(left, right);
+        assert_eq!(result, bitbuffer![1 0 0 1]);
+    }
+}

vortex-compute/src/comparison/mod.rs

@@ -0,0 +1,130 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Comparison operations for Vortex vectors.
+
+use vortex_dtype::half::f16;
+
+mod bool;
+mod collection;
+mod pvector;
+
+/// Trait for comparison operations.
+pub trait Compare<Op, Rhs = Self> {
+    /// The result type after performing the operation.
+    type Output;
+
+    /// Perform the comparison operation.
+    fn compare(self, rhs: Rhs) -> Self::Output;
+}
+
+/// Trait for comparison operators.
+pub trait ComparisonOperator<T> {
+    /// Apply the operator to the two operands.
+    fn apply(a: &T, b: &T) -> bool;
+}
+
+/// A marker type for equality comparison operations.
+pub struct Equal;
+/// A marker type for inequality comparison operations.
+pub struct NotEqual;
+/// A marker type for less-than comparison operations.
+pub struct LessThan;
+/// A marker type for less-than-or-equal comparison operations.
+pub struct LessThanOrEqual;
+/// A marker type for greater-than comparison operations.
+pub struct GreaterThan;
+/// A marker type for greater-than-or-equal comparison operations.
+pub struct GreaterThanOrEqual;
+
+/// Marker trait for comparable items.
+pub trait ComparableItem {
+    /// Check if two items are equal.
+    fn is_equal(lhs: &Self, rhs: &Self) -> bool;
+
+    /// Check if the `lhs` item is less than the `rhs` item.
+    fn is_less_than(lhs: &Self, rhs: &Self) -> bool;
+}
+
+impl<T: ComparableItem> ComparisonOperator<T> for Equal {
+    fn apply(a: &T, b: &T) -> bool {
+        T::is_equal(a, b)
+    }
+}
+
+impl<T: ComparableItem> ComparisonOperator<T> for NotEqual {
+    fn apply(a: &T, b: &T) -> bool {
+        !T::is_equal(a, b)
+    }
+}
+
+impl<T: ComparableItem> ComparisonOperator<T> for LessThan {
+    fn apply(a: &T, b: &T) -> bool {
+        T::is_less_than(a, b)
+    }
+}
+
+impl<T: ComparableItem> ComparisonOperator<T> for GreaterThanOrEqual {
+    fn apply(a: &T, b: &T) -> bool {
+        !T::is_less_than(a, b)
+    }
+}
+
+impl<T: ComparableItem> ComparisonOperator<T> for GreaterThan {
+    fn apply(a: &T, b: &T) -> bool {
+        T::is_less_than(b, a)
+    }
+}
+
+impl<T: ComparableItem> ComparisonOperator<T> for LessThanOrEqual {
+    fn apply(a: &T, b: &T) -> bool {
+        !T::is_less_than(b, a)
+    }
+}
+
+macro_rules! impl_integer {
+    ($T:ty) => {
+        impl ComparableItem for $T {
+            #[inline(always)]
+            fn is_equal(lhs: &Self, rhs: &Self) -> bool {
+                lhs == rhs
+            }
+
+            #[inline(always)]
+            fn is_less_than(lhs: &Self, rhs: &Self) -> bool {
+                lhs < rhs
+            }
+        }
+    };
+}
+
+impl_integer!(i8);
+impl_integer!(i16);
+impl_integer!(i32);
+impl_integer!(i64);
+impl_integer!(i128);
+impl_integer!(u8);
+impl_integer!(u16);
+impl_integer!(u32);
+impl_integer!(u64);
+impl_integer!(u128);
+
+macro_rules! impl_float {
+    ($T:ty) => {
+        impl ComparableItem for $T {
+            #[inline(always)]
+            fn is_equal(lhs: &Self, rhs: &Self) -> bool {
+                lhs.to_bits().eq(&rhs.to_bits())
+            }
+
+            #[inline(always)]
+            fn is_less_than(lhs: &Self, rhs: &Self) -> bool {
+                lhs.total_cmp(rhs).is_lt()
+            }
+        }
+    };
+}
+
+impl_float!(f16);
+impl_float!(f32);
+impl_float!(f64);