vortex_array/compute/
between.rs

1use vortex_dtype::{DType, Nullability};
2use vortex_error::{VortexExpect, VortexResult};
3
4use crate::compute::{BinaryOperator, Operator, binary_boolean, compare};
5use crate::{Array, ArrayRef, Canonical, Encoding, IntoArray};
6
7pub trait BetweenFn<A> {
8    fn between(
9        &self,
10        arr: A,
11        lower: &dyn Array,
12        upper: &dyn Array,
13        options: &BetweenOptions,
14    ) -> VortexResult<Option<ArrayRef>>;
15}
16
17#[derive(Debug, Clone, PartialEq, Eq, Hash)]
18pub struct BetweenOptions {
19    pub lower_strict: StrictComparison,
20    pub upper_strict: StrictComparison,
21}
22
23#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
24pub enum StrictComparison {
25    Strict,
26    NonStrict,
27}
28
29impl StrictComparison {
30    pub const fn to_operator(&self) -> Operator {
31        match self {
32            StrictComparison::Strict => Operator::Lt,
33            StrictComparison::NonStrict => Operator::Lte,
34        }
35    }
36}
37
38impl<E: Encoding> BetweenFn<&dyn Array> for E
39where
40    E: for<'a> BetweenFn<&'a E::Array>,
41{
42    fn between(
43        &self,
44        arr: &dyn Array,
45        lower: &dyn Array,
46        upper: &dyn Array,
47        options: &BetweenOptions,
48    ) -> VortexResult<Option<ArrayRef>> {
49        let array_ref = arr
50            .as_any()
51            .downcast_ref::<E::Array>()
52            .vortex_expect("Failed to downcast array");
53        BetweenFn::between(self, array_ref, lower, upper, options)
54    }
55}
56
57/// Compute between (a <= x <= b), this can be implemented using compare and boolean andn but this
58/// will likely have a lower runtime.
59///
60/// This semantics is equivalent to:
61/// ```
62/// use vortex_array::{Array, ArrayRef};
63/// use vortex_array::compute::{binary_boolean, compare, BetweenOptions, BinaryOperator, Operator};///
64/// use vortex_error::VortexResult;
65///
66/// fn between(
67///    arr: &dyn Array,
68///    lower: &dyn Array,
69///    upper: &dyn Array,
70///    options: &BetweenOptions
71/// ) -> VortexResult<ArrayRef> {
72///     binary_boolean(
73///         &compare(lower, arr, options.lower_strict.to_operator())?,
74///         &compare(arr, upper,  options.upper_strict.to_operator())?,
75///         BinaryOperator::And
76///     )
77/// }
78///  ```
79///
80/// The BetweenOptions { lower: StrictComparison, upper: StrictComparison } defines if the
81/// value is < (strict) or <= (non-strict).
82///
83pub fn between(
84    arr: &dyn Array,
85    lower: &dyn Array,
86    upper: &dyn Array,
87    options: &BetweenOptions,
88) -> VortexResult<ArrayRef> {
89    debug_assert!(arr.dtype().eq_ignore_nullability(lower.dtype()));
90    debug_assert!(arr.dtype().eq_ignore_nullability(upper.dtype()));
91    debug_assert_eq!(arr.len(), lower.len());
92    debug_assert_eq!(arr.len(), upper.len());
93
94    let result = between_impl(arr, lower, upper, options)?;
95
96    debug_assert_eq!(result.len(), arr.len());
97    debug_assert_eq!(
98        result.dtype(),
99        &DType::Bool(
100            arr.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability()
101        )
102    );
103
104    Ok(result)
105}
106
107fn between_impl(
108    arr: &dyn Array,
109    lower: &dyn Array,
110    upper: &dyn Array,
111    options: &BetweenOptions,
112) -> VortexResult<ArrayRef> {
113    if lower.as_constant().is_some_and(|v| v.is_null())
114        || upper.as_constant().is_some_and(|v| v.is_null())
115    {
116        return Ok(
117            Canonical::empty(&arr.dtype().with_nullability(Nullability::Nullable)).into_array(),
118        );
119    }
120
121    if let Some(result) = arr
122        .vtable()
123        .between_fn()
124        .and_then(|f| f.between(arr, lower, upper, options).transpose())
125        .transpose()?
126    {
127        return Ok(result);
128    }
129
130    // TODO(joe): should we try to canonicalize the array and try between
131    binary_boolean(
132        &compare(lower, arr, options.lower_strict.to_operator())?,
133        &compare(arr, upper, options.upper_strict.to_operator())?,
134        BinaryOperator::And,
135    )
136}
vortex_array/compute/between.rs

vortex_array/compute/
between.rs
