between(array, lower, lower_op, upper, upper_op) compute function (#2391)

Author: joseph-isaacs

Cargo.lock

@@ -1,14 +1,19 @@
 mod compare;
 
+use std::fmt::Debug;
+
+use vortex_array::arrays::ConstantArray;
 use vortex_array::compute::{
-    filter, scalar_at, slice, take, CompareFn, FilterFn, ScalarAtFn, SliceFn, TakeFn,
+    between, filter, scalar_at, slice, take, BetweenFn, BetweenOptions, CompareFn, FilterFn,
+    ScalarAtFn, SliceFn, StrictComparison, TakeFn,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::ComputeVTable;
 use vortex_array::{Array, IntoArray};
+use vortex_dtype::NativePType;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
-use vortex_scalar::Scalar;
+use vortex_scalar::{Scalar, ScalarType};
 
 use crate::{match_each_alp_float_ptype, ALPArray, ALPEncoding, ALPFloat};
 
@@ -17,6 +22,10 @@ impl ComputeVTable for ALPEncoding {
         Some(self)
     }
 
+    fn between_fn(&self) -> Option<&dyn BetweenFn<Array>> {
+        Some(self)
+    }
+
     fn filter_fn(&self) -> Option<&dyn FilterFn<Array>> {
         Some(self)
     }
@@ -107,3 +116,156 @@ impl FilterFn<ALPArray> for ALPEncoding {
         )
     }
 }
+
+impl BetweenFn<ALPArray> for ALPEncoding {
+    fn between(
+        &self,
+        array: &ALPArray,
+        lower: &Array,
+        upper: &Array,
+        options: &BetweenOptions,
+    ) -> VortexResult<Option<Array>> {
+        let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
+            return Ok(None);
+        };
+
+        if array.patches().is_some() {
+            return Ok(None);
+        }
+
+        match_each_alp_float_ptype!(array.ptype(), |$F| {
+            between_impl::<$F>(array, $F::try_from(lower)?, $F::try_from(upper)?, options)
+        })
+        .map(Some)
+    }
+}
+
+fn between_impl<T: NativePType + ALPFloat>(
+    array: &ALPArray,
+    lower: T,
+    upper: T,
+    options: &BetweenOptions,
+) -> VortexResult<Array>
+where
+    Scalar: From<T::ALPInt>,
+    <T as ALPFloat>::ALPInt: ScalarType + Debug,
+{
+    let exponents = array.exponents();
+
+    // There are always compared
+    // the below bound is `value {< | <=} x`, either value encodes into the ALPInt domain
+    // in which case we can leave the comparison unchanged `enc(value) {< | <=} x` or it doesn't
+    // and we encode into value below enc_below(value) < value < x, in which case the comparison
+    // becomes enc(value) < x. See `alp_scalar_compare` for more details.
+    // note that if the value doesn't encode than value != x, so must use strict comparison.
+    let (lower_enc, lower_strict) = T::encode_single(lower, exponents)
+        .map(|x| (x, options.lower_strict))
+        .unwrap_or_else(|| (T::encode_below(lower, exponents), StrictComparison::Strict));
+
+    // the upper value `x { < | <= } value` similarly encodes or `x < value < enc_above(value())`
+    let (upper_enc, upper_strict) = T::encode_single(upper, exponents)
+        .map(|x| (x, options.upper_strict))
+        .unwrap_or_else(|| (T::encode_above(upper, exponents), StrictComparison::Strict));
+
+    let options = BetweenOptions {
+        lower_strict,
+        upper_strict,
+    };
+
+    between(
+        array.encoded(),
+        ConstantArray::new(lower_enc, array.len()),
+        ConstantArray::new(upper_enc, array.len()),
+        &options,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use itertools::Itertools;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::compute::{BetweenOptions, StrictComparison};
+    use vortex_array::IntoArrayVariant;
+
+    use crate::alp::compute::between_impl;
+    use crate::ALPArray;
+
+    fn between_test(arr: &ALPArray, lower: f32, upper: f32, options: &BetweenOptions) -> bool {
+        let res = between_impl(arr, lower, upper, options)
+            .unwrap()
+            .into_bool()
+            .unwrap()
+            .boolean_buffer()
+            .iter()
+            .collect_vec();
+        assert_eq!(res.len(), 1);
+
+        res[0]
+    }
+
+    #[test]
+    fn comparison_range() {
+        let value = 0.0605_f32;
+        let array = PrimitiveArray::from_iter([value; 1]);
+        let encoded = crate::alp::compress::alp_encode(&array).unwrap();
+        assert!(encoded.patches().is_none());
+        assert_eq!(
+            encoded
+                .encoded()
+                .into_primitive()
+                .unwrap()
+                .as_slice::<i32>(),
+            vec![605; 1]
+        );
+
+        assert!(between_test(
+            &encoded,
+            0.0605_f32,
+            0.0605,
+            &BetweenOptions {
+                lower_strict: StrictComparison::NonStrict,
+                upper_strict: StrictComparison::NonStrict,
+            },
+        ));
+
+        assert!(!between_test(
+            &encoded,
+            0.0605_f32,
+            0.0605,
+            &BetweenOptions {
+                lower_strict: StrictComparison::Strict,
+                upper_strict: StrictComparison::NonStrict,
+            },
+        ));
+
+        assert!(!between_test(
+            &encoded,
+            0.0605_f32,
+            0.0605,
+            &BetweenOptions {
+                lower_strict: StrictComparison::NonStrict,
+                upper_strict: StrictComparison::Strict,
+            },
+        ));
+
+        assert!(between_test(
+            &encoded,
+            0.060499_f32,
+            0.06051,
+            &BetweenOptions {
+                lower_strict: StrictComparison::NonStrict,
+                upper_strict: StrictComparison::NonStrict,
+            },
+        ));
+
+        assert!(between_test(
+            &encoded,
+            0.06_f32,
+            0.06051,
+            &BetweenOptions {
+                lower_strict: StrictComparison::NonStrict,
+                upper_strict: StrictComparison::Strict,
+            },
+        ))
+    }
+}

encodings/alp/src/alp/compute/mod.rs

@@ -41,6 +41,8 @@ impl ComputeVTable for DateTimePartsEncoding {
     fn compare_fn(&self) -> Option<&dyn CompareFn<Array>> {
         Some(self)
     }
+
+    // TODO(joe): implement `between_fn` this is used at lot.
 }
 
 impl SliceFn<DateTimePartsArray> for DateTimePartsEncoding {

encodings/datetime-parts/src/compute/mod.rs

@@ -37,6 +37,7 @@ divan = { workspace = true }
 rand = { workspace = true }
 vortex-array = { workspace = true, features = ["test-harness"] }
 vortex-fastlanes = { path = ".", features = ["test-harness"] }
+vortex-alp = { workspace = true }
 
 [features]
 test-harness = ["dep:rand"]
@@ -53,3 +54,9 @@ harness = false
 name = "canonicalize_bench"
 harness = false
 required-features = ["test-harness"]
+
+[[bench]]
+name = "compute_between"
+harness = false
+required-features = ["test-harness"]
+