vortex_array/compute/
between.rs

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

vortex_array/compute/
between.rs
