Array comparison compute function

vortex-array/Cargo.toml

@@ -64,4 +64,8 @@ harness = false
 
 [[bench]]
 name = "filter_indices"
-harness = false
+harness = false
+
+[[bench]]
+name = "compare_arrays"
+harness = false

vortex-array/benches/compare_arrays.rs

@@ -0,0 +1,72 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use itertools::Itertools;
+use rand::distributions::Uniform;
+use rand::{thread_rng, Rng};
+use vortex::array::bool::BoolArray;
+use vortex::IntoArray;
+use vortex_error::VortexError;
+use vortex_expr::operators::Operator;
+
+fn filter_bool_indices(c: &mut Criterion) {
+    let mut group = c.benchmark_group("compare_arrays");
+
+    let mut rng = thread_rng();
+    let range = Uniform::new(0u8, 1);
+    let arr = BoolArray::from(
+        (0..10_000_000)
+            .map(|_| rng.sample(range) == 0)
+            .collect_vec(),
+    )
+    .into_array();
+    let arr2 = BoolArray::from(
+        (0..10_000_000)
+            .map(|_| rng.sample(range) == 0)
+            .collect_vec(),
+    )
+    .into_array();
+
+    group.bench_function("compare_bool", |b| {
+        b.iter(|| {
+            let indices = vortex::compute::compare::compare_arrays(
+                &arr,
+                &arr2,
+                Operator::GreaterThanOrEqualTo,
+            )
+            .unwrap();
+            black_box(indices);
+            Ok::<(), VortexError>(())
+        });
+    });
+}
+
+fn filter_indices(c: &mut Criterion) {
+    let mut group = c.benchmark_group("compare_arrays");
+
+    let mut rng = thread_rng();
+    let range = Uniform::new(0i64, 100_000_000);
+    let arr = (0..10_000_000)
+        .map(|_| rng.sample(range))
+        .collect_vec()
+        .into_array();
+
+    let arr2 = (0..10_000_000)
+        .map(|_| rng.sample(range))
+        .collect_vec()
+        .into_array();
+
+    group.bench_function("compare_int", |b| {
+        b.iter(|| {
+            let indices = vortex::compute::compare::compare_arrays(
+                &arr,
+                &arr2,
+                Operator::GreaterThanOrEqualTo,
+            )
+            .unwrap();
+            black_box(indices);
+            Ok::<(), VortexError>(())
+        });
+    });
+}
+
+criterion_group!(benches, filter_indices, filter_bool_indices);
+criterion_main!(benches);

vortex-array/src/array/bool/compute/compare.rs

@@ -0,0 +1,115 @@
+use std::ops::{BitAnd, BitOr, BitXor, Not};
+
+use vortex_error::VortexResult;
+use vortex_expr::operators::Operator;
+
+use crate::array::bool::BoolArray;
+use crate::compute::compare::CompareArraysFn;
+use crate::{Array, ArrayTrait, IntoArray};
+
+impl CompareArraysFn for BoolArray {
+    fn compare_arrays(&self, other: &Array, op: Operator) -> VortexResult<Array> {
+        let flattened = other.clone().flatten_bool()?;
+        let lhs = self.boolean_buffer();
+        let rhs = flattened.boolean_buffer();
+        let result_buf = match op {
+            Operator::EqualTo => lhs.bitxor(&rhs).not(),
+            Operator::NotEqualTo => lhs.bitxor(&rhs),
+
+            Operator::GreaterThan => lhs.bitand(&rhs).bitxor(&lhs),
+            Operator::GreaterThanOrEqualTo => {
+                let gt = lhs.bitand(&rhs).bitxor(&lhs);
+                let eq = &lhs.bitxor(&rhs).not();
+                gt.bitor(eq)
+            }
+            Operator::LessThan => lhs.bitor(&rhs).bitxor(&lhs),
+            Operator::LessThanOrEqualTo => {
+                let eq = lhs.bitxor(&rhs).not();
+                let lt = lhs.bitor(&rhs).bitxor(&lhs);
+                lt.bitor(&eq)
+            }
+        };
+        let present_buf = self
+            .validity()
+            .to_logical(self.len())
+            .to_present_null_buffer()?
+            .into_inner();
+
+        Ok(BoolArray::from(result_buf.bitand(&present_buf)).into_array())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use itertools::Itertools;
+
+    use super::*;
+    use crate::validity::Validity;
+    use crate::ToArray;
+
+    fn to_int_indices(indices_bits: BoolArray) -> Vec<u64> {
+        let filtered = indices_bits
+            .boolean_buffer()
+            .iter()
+            .enumerate()
+            .flat_map(|(idx, v)| if v { Some(idx as u64) } else { None })
+            .collect_vec();
+        filtered
+    }
+
+    #[test]
+    fn test_basic_comparison() {
+        let arr = BoolArray::from_vec(
+            vec![true, true, false, true, false],
+            Validity::Array(BoolArray::from(vec![false, true, true, true, true]).into_array()),
+        );
+
+        let matches = arr
+            .compare_arrays(&arr.to_array(), Operator::EqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [1u64, 2, 3, 4]);
+
+        let matches = arr
+            .compare_arrays(&arr.to_array(), Operator::NotEqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        let empty: [u64; 0] = [];
+        assert_eq!(to_int_indices(matches), empty);
+
+        let other = BoolArray::from_vec(
+            vec![false, false, false, true, true],
+            Validity::Array(BoolArray::from(vec![false, true, true, true, true]).into_array()),
+        );
+
+        let matches = arr
+            .compare_arrays(&other.to_array(), Operator::LessThanOrEqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [2u64, 3, 4]);
+
+        let matches = arr
+            .compare_arrays(&other.to_array(), Operator::LessThan)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [4u64]);
+
+        let matches = other
+            .compare_arrays(&arr.to_array(), Operator::GreaterThanOrEqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [2u64, 3, 4]);
+
+        let matches = other
+            .compare_arrays(&arr.to_array(), Operator::GreaterThan)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [4u64]);
+    }
+}

vortex-array/src/array/bool/compute/mod.rs

@@ -1,6 +1,7 @@
 use crate::array::bool::BoolArray;
 use crate::compute::as_arrow::AsArrowArray;
 use crate::compute::as_contiguous::AsContiguousFn;
+use crate::compute::compare::CompareArraysFn;
 use crate::compute::fill::FillForwardFn;
 use crate::compute::scalar_at::ScalarAtFn;
 use crate::compute::slice::SliceFn;
@@ -9,6 +10,7 @@ use crate::compute::ArrayCompute;
 
 mod as_arrow;
 mod as_contiguous;
+mod compare;
 mod fill;
 mod flatten;
 mod scalar_at;
@@ -24,6 +26,10 @@ impl ArrayCompute for BoolArray {
         Some(self)
     }
 
+    fn compare_arrays(&self) -> Option<&dyn CompareArraysFn> {
+        Some(self)
+    }
+
     fn fill_forward(&self) -> Option<&dyn FillForwardFn> {
         Some(self)
     }

vortex-array/src/array/primitive/compute/compare.rs

@@ -0,0 +1,138 @@
+use std::ops::BitAnd;
+
+use arrow_buffer::BooleanBuffer;
+use vortex_dtype::{match_each_native_ptype, NativePType};
+use vortex_error::VortexResult;
+use vortex_expr::operators::Operator;
+
+use crate::array::bool::BoolArray;
+use crate::array::primitive::PrimitiveArray;
+use crate::compute::compare::CompareArraysFn;
+use crate::{Array, ArrayTrait, IntoArray};
+
+impl CompareArraysFn for PrimitiveArray {
+    fn compare_arrays(&self, other: &Array, predicate: Operator) -> VortexResult<Array> {
+        let flattened = other.clone().flatten_primitive()?;
+
+        let matching_idxs = match_each_native_ptype!(self.ptype(), |$T| {
+            let predicate_fn = &predicate.to_predicate::<$T>();
+            apply_predicate(self.typed_data::<$T>(), flattened.typed_data::<$T>(), predicate_fn)
+        });
+
+        let present = self
+            .validity()
+            .to_logical(self.len())
+            .to_present_null_buffer()?
+            .into_inner();
+        let present_other = flattened
+            .validity()
+            .to_logical(self.len())
+            .to_present_null_buffer()?
+            .into_inner();
+
+        Ok(BoolArray::from(matching_idxs.bitand(&present).bitand(&present_other)).into_array())
+    }
+}
+
+fn apply_predicate<T: NativePType, F: Fn(&T, &T) -> bool>(
+    lhs: &[T],
+    rhs: &[T],
+    f: F,
+) -> BooleanBuffer {
+    let matches = lhs.iter().zip(rhs.iter()).map(|(lhs, rhs)| f(lhs, rhs));
+    BooleanBuffer::from_iter(matches)
+}
+
+#[cfg(test)]
+mod test {
+    use itertools::Itertools;
+
+    use super::*;
+    use crate::ToArray;
+
+    fn to_int_indices(indices_bits: BoolArray) -> Vec<u64> {
+        let filtered = indices_bits
+            .boolean_buffer()
+            .iter()
+            .enumerate()
+            .flat_map(|(idx, v)| if v { Some(idx as u64) } else { None })
+            .collect_vec();
+        filtered
+    }
+
+    #[test]
+    fn test_basic_comparisons() {
+        let arr = PrimitiveArray::from_nullable_vec(vec![
+            Some(1i32),
+            Some(2),
+            Some(3),
+            Some(4),
+            None,
+            Some(5),
+            Some(6),
+            Some(7),
+            Some(8),
+            None,
+            Some(9),
+            None,
+        ]);
+
+        let matches = arr
+            .compare_arrays(&arr.to_array(), Operator::EqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [0u64, 1, 2, 3, 5, 6, 7, 8, 10]);
+
+        let matches = arr
+            .compare_arrays(&arr.to_array(), Operator::NotEqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        let empty: [u64; 0] = [];
+        assert_eq!(to_int_indices(matches), empty);
+
+        let other = PrimitiveArray::from_nullable_vec(vec![
+            Some(1i32),
+            Some(2),
+            Some(3),
+            Some(4),
+            None,
+            Some(6),
+            Some(7),
+            Some(8),
+            Some(9),
+            None,
+            Some(10),
+            None,
+        ]);
+
+        let matches = arr
+            .compare_arrays(&other.to_array(), Operator::LessThanOrEqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [0u64, 1, 2, 3, 5, 6, 7, 8, 10]);
+
+        let matches = arr
+            .compare_arrays(&other.to_array(), Operator::LessThan)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [5u64, 6, 7, 8, 10]);
+
+        let matches = other
+            .compare_arrays(&arr.to_array(), Operator::GreaterThanOrEqualTo)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [0u64, 1, 2, 3, 5, 6, 7, 8, 10]);
+
+        let matches = other
+            .compare_arrays(&arr.to_array(), Operator::GreaterThan)
+            .unwrap()
+            .flatten_bool()
+            .unwrap();
+        assert_eq!(to_int_indices(matches), [5u64, 6, 7, 8, 10]);
+    }
+}