Introduce MaybeCompareFn trait to allow for partial compare specializations

encodings/alp/src/compute.rs

@@ -1,8 +1,13 @@
-use vortex::compute::unary::{scalar_at_unchecked, ScalarAtFn};
-use vortex::compute::{slice, take, ArrayCompute, SliceFn, TakeFn};
+use vortex::array::{BoolArray, ConstantArray};
+use vortex::compute::unary::{scalar_at, scalar_at_unchecked, ScalarAtFn};
+use vortex::compute::{
+    compare, filter, slice, take, ArrayCompute, FilterFn, MaybeCompareFn, Operator, SliceFn, TakeFn,
+};
+use vortex::stats::{ArrayStatistics, Stat};
+use vortex::validity::Validity;
 use vortex::{Array, ArrayDType, IntoArray};
-use vortex_error::VortexResult;
-use vortex_scalar::Scalar;
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_scalar::{PValue, Scalar};
 
 use crate::{match_each_alp_float_ptype, ALPArray, ALPFloat};
 
@@ -18,6 +23,14 @@ impl ArrayCompute for ALPArray {
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
+
+    fn compare(&self, other: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        MaybeCompareFn::maybe_compare(self, other, operator)
+    }
+
+    fn filter(&self) -> Option<&dyn FilterFn> {
+        Some(self)
+    }
 }
 
 impl ScalarAtFn for ALPArray {
@@ -67,3 +80,146 @@ impl SliceFn for ALPArray {
         .into_array())
     }
 }
+
+impl FilterFn for ALPArray {
+    fn filter(&self, predicate: &Array) -> VortexResult<Array> {
+        Ok(Self::try_new(
+            filter(&self.encoded(), predicate)?,
+            self.exponents(),
+            self.patches().map(|p| filter(&p, predicate)).transpose()?,
+        )?
+        .into_array())
+    }
+}
+
+impl MaybeCompareFn for ALPArray {
+    fn maybe_compare(&self, array: &Array, operator: Operator) -> Option<VortexResult<Array>> {
+        if ConstantArray::try_from(array).is_ok()
+            || array
+                .statistics()
+                .get_as::<bool>(Stat::IsConstant)
+                .unwrap_or_default()
+        {
+            let rhs = scalar_at(array, 0).vortex_expect("should be scalar");
+            let pvalue = rhs
+                .value()
+                .as_pvalue()
+                .vortex_expect("Expected primitive value");
+
+            match pvalue {
+                Some(PValue::F32(f)) => Some(alp_scalar_compare(self, f, operator)),
+                Some(PValue::F64(f)) => Some(alp_scalar_compare(self, f, operator)),
+                None => Some(Ok(BoolArray::from_vec(
+                    vec![false; self.len()],
+                    Validity::AllValid,
+                )
+                .into_array())),
+                _ => unreachable!(),
+            }
+        } else {
+            None
+        }
+    }
+}
+
+fn alp_scalar_compare<F: ALPFloat + Into<Scalar>>(
+    alp: &ALPArray,
+    value: F,
+    operator: Operator,
+) -> VortexResult<Array>
+where
+    F::ALPInt: Into<Scalar>,
+{
+    let encoded = F::encode_single(value, alp.exponents());
+    match encoded {
+        Ok(encoded) => {
+            let s = ConstantArray::new(encoded, alp.len());
+            compare(&alp.encoded(), s.array(), operator)
+        }
+        Err(exception) => {
+            if let Some(patches) = alp.patches().as_ref() {
+                let s = ConstantArray::new(exception, alp.len());
+                compare(patches, s.array(), operator)
+            } else {
+                Ok(BoolArray::from_vec(vec![false; alp.len()], Validity::AllValid).into_array())
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex::array::PrimitiveArray;
+    use vortex::IntoArrayVariant;
+    use vortex_dtype::{DType, Nullability, PType};
+
+    use super::*;
+    use crate::alp_encode;
+
+    #[test]
+    fn basic_comparison_test() {
+        let array = PrimitiveArray::from(vec![1.234f32; 1025]);
+        let encoded = alp_encode(&array).unwrap();
+        assert!(encoded.patches().is_none());
+        assert_eq!(
+            encoded.encoded().as_primitive().maybe_null_slice::<i32>(),
+            vec![1234; 1025]
+        );
+
+        let r = alp_scalar_compare(&encoded, 1.3_f32, Operator::Eq)
+            .unwrap()
+            .into_bool()
+            .unwrap();
+
+        for v in r.boolean_buffer().iter() {
+            assert!(!v);
+        }
+
+        let r = alp_scalar_compare(&encoded, 1.234f32, Operator::Eq)
+            .unwrap()
+            .into_bool()
+            .unwrap();
+
+        for v in r.boolean_buffer().iter() {
+            assert!(v);
+        }
+    }
+
+    #[test]
+    fn compare_with_patches() {
+        let array =
+            PrimitiveArray::from(vec![1.234f32, 1.5, 19.0, std::f32::consts::E, 1_000_000.9]);
+        let encoded = alp_encode(&array).unwrap();
+        assert!(encoded.patches().is_some());
+
+        let r = alp_scalar_compare(&encoded, 1_000_000.9_f32, Operator::Eq)
+            .unwrap()
+            .into_bool()
+            .unwrap();
+
+        let buffer = r.boolean_buffer();
+        assert!(buffer.value(buffer.len() - 1));
+    }
+
+    #[test]
+    fn compare_to_null() {
+        let array = PrimitiveArray::from(vec![1.234f32; 1025]);
+        let encoded = alp_encode(&array).unwrap();
+
+        let other = ConstantArray::new(
+            Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)),
+            array.len(),
+        );
+
+        let r = encoded
+            .maybe_compare(other.array(), Operator::Eq)
+            .unwrap()
+            .unwrap()
+            .into_bool()
+            .unwrap();
+
+        for v in r.boolean_buffer().iter() {
+            assert!(!v);
+        }
+    }
+}

encodings/bytebool/src/compute/mod.rs

@@ -1,22 +1,15 @@
-use std::ops::{BitAnd, BitOr, BitXor, Not};
-
-use arrow_buffer::BooleanBuffer;
 use num_traits::AsPrimitive;
 use vortex::compute::unary::{FillForwardFn, ScalarAtFn};
-use vortex::compute::{ArrayCompute, CompareFn, Operator, SliceFn, TakeFn};
+use vortex::compute::{ArrayCompute, SliceFn, TakeFn};
 use vortex::validity::{ArrayValidity, Validity};
-use vortex::{Array, ArrayDType, IntoArray, IntoArrayVariant};
+use vortex::{Array, ArrayDType, IntoArray};
 use vortex_dtype::{match_each_integer_ptype, Nullability};
 use vortex_error::{vortex_err, VortexResult};
 use vortex_scalar::Scalar;
 
 use super::ByteBoolArray;
 
 impl ArrayCompute for ByteBoolArray {
-    fn compare(&self) -> Option<&dyn CompareFn> {
-        Some(self)
-    }
-
     fn fill_forward(&self) -> Option<&dyn FillForwardFn> {
         None
     }
@@ -99,29 +92,6 @@ impl TakeFn for ByteBoolArray {
     }
 }
 
-impl CompareFn for ByteBoolArray {
-    fn compare(&self, other: &Array, op: Operator) -> VortexResult<Array> {
-        let canonical = other.clone().into_bool()?;
-        let lhs = BooleanBuffer::from(self.maybe_null_slice());
-        let rhs = canonical.boolean_buffer();
-
-        let result_buf = match op {
-            Operator::Eq => lhs.bitxor(&rhs).not(),
-            Operator::NotEq => lhs.bitxor(&rhs),
-
-            Operator::Gt => lhs.bitand(&rhs.not()),
-            Operator::Gte => lhs.bitor(&rhs.not()),
-            Operator::Lt => lhs.not().bitand(&rhs),
-            Operator::Lte => lhs.not().bitor(&rhs),
-        };
-
-        let validity = self.validity().and(canonical.validity())?;
-
-        ByteBoolArray::try_from_vec(Vec::from_iter(result_buf.iter()), validity)
-            .map(ByteBoolArray::into_array)
-    }
-}
-
 impl FillForwardFn for ByteBoolArray {
     fn fill_forward(&self) -> VortexResult<Array> {
         let validity = self.logical_validity();
@@ -165,7 +135,7 @@ impl FillForwardFn for ByteBoolArray {
 #[cfg(test)]
 mod tests {
     use vortex::compute::unary::{scalar_at, scalar_at_unchecked};
-    use vortex::compute::{compare, slice};
+    use vortex::compute::{compare, slice, Operator};
     use vortex::AsArray as _;
     use vortex_scalar::ScalarValue;
 

encodings/datetime-parts/src/array.rs

@@ -1,11 +1,14 @@
 use std::fmt::Debug;
 
 use serde::{Deserialize, Serialize};
+use vortex::array::StructArray;
 use vortex::stats::{ArrayStatisticsCompute, StatsSet};
 use vortex::validity::{ArrayValidity, LogicalValidity};
 use vortex::variants::{ArrayVariants, ExtensionArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
-use vortex::{impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoCanonical};
+use vortex::{
+    impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoArray, IntoCanonical,
+};
 use vortex_dtype::DType;
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 
@@ -89,7 +92,19 @@ impl ArrayVariants for DateTimePartsArray {
     }
 }
 
-impl ExtensionArrayTrait for DateTimePartsArray {}
+impl ExtensionArrayTrait for DateTimePartsArray {
+    fn storage_array(&self) -> Array {
+        // FIXME(ngates): this needs to be a tuple array so we can implement Compare
+        StructArray::try_new(
+            vec!["days".into(), "seconds".into(), "subseconds".into()].into(),
+            [self.days(), self.seconds(), self.subsecond()].into(),
+            self.len(),
+            self.logical_validity().into_validity(),
+        )
+        .vortex_expect("Failed to create struct array")
+        .into_array()
+    }
+}
 
 impl IntoCanonical for DateTimePartsArray {
     fn into_canonical(self) -> VortexResult<Canonical> {

encodings/fastlanes/src/for/compute.rs

@@ -3,8 +3,8 @@ use std::ops::{AddAssign, Shl, Shr};
 use num_traits::{WrappingAdd, WrappingSub};
 use vortex::compute::unary::{scalar_at_unchecked, ScalarAtFn};
 use vortex::compute::{
-    search_sorted, slice, take, ArrayCompute, SearchResult, SearchSortedFn, SearchSortedSide,
-    SliceFn, TakeFn,
+    filter, search_sorted, slice, take, ArrayCompute, FilterFn, SearchResult, SearchSortedFn,
+    SearchSortedSide, SliceFn, TakeFn,
 };
 use vortex::{Array, ArrayDType, IntoArray};
 use vortex_dtype::{match_each_integer_ptype, NativePType};
@@ -29,6 +29,10 @@ impl ArrayCompute for FoRArray {
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
+
+    fn filter(&self) -> Option<&dyn FilterFn> {
+        Some(self)
+    }
 }
 
 impl TakeFn for FoRArray {
@@ -42,6 +46,17 @@ impl TakeFn for FoRArray {
     }
 }
 
+impl FilterFn for FoRArray {
+    fn filter(&self, predicate: &Array) -> VortexResult<Array> {
+        Self::try_new(
+            filter(&self.encoded(), predicate)?,
+            self.reference().clone(),
+            self.shift(),
+        )
+        .map(|a| a.into_array())
+    }
+}
+
 impl ScalarAtFn for FoRArray {
     fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
         Ok(self.scalar_at_unchecked(index))