Skip to main content

vortex_array/arrays/decimal/compute/
between.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::BitBuffer;
5use vortex_error::VortexExpect;
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::IntoArray;
11use crate::array::ArrayView;
12use crate::arrays::BoolArray;
13use crate::arrays::Decimal;
14use crate::dtype::NativeDecimalType;
15use crate::dtype::Nullability;
16use crate::dtype::i256;
17use crate::match_each_decimal_value_type;
18use crate::scalar::Scalar;
19use crate::scalar_fn::fns::between::BetweenKernel;
20use crate::scalar_fn::fns::between::BetweenOptions;
21use crate::scalar_fn::fns::between::StrictComparison;
22
23impl BetweenKernel for Decimal {
24    fn between(
25        arr: ArrayView<'_, Decimal>,
26        lower: &ArrayRef,
27        upper: &ArrayRef,
28        options: &BetweenOptions,
29        ctx: &mut ExecutionCtx,
30    ) -> VortexResult<Option<ArrayRef>> {
31        // NOTE: We know that the precision and scale were already checked to be equal by the main
32        // `between` entrypoint function.
33
34        let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
35            return Ok(None);
36        };
37
38        // NOTE: we know that have checked before that the lower and upper bounds are not all null.
39        let nullability =
40            arr.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability();
41
42        match_each_decimal_value_type!(arr.values_type(), |D| {
43            between_unpack::<D>(arr, lower, upper, nullability, options, ctx)
44        })
45    }
46}
47
48fn between_unpack<T: NativeDecimalType>(
49    arr: ArrayView<'_, Decimal>,
50    lower: Scalar,
51    upper: Scalar,
52    nullability: Nullability,
53    options: &BetweenOptions,
54    ctx: &mut ExecutionCtx,
55) -> VortexResult<Option<ArrayRef>> {
56    let Some(lower_dv) = lower.as_decimal().decimal_value() else {
57        // Null lower bound — fall back to canonical path.
58        return Ok(None);
59    };
60    let Some(upper_dv) = upper.as_decimal().decimal_value() else {
61        // Null upper bound — fall back to canonical path.
62        return Ok(None);
63    };
64
65    // Try to cast the bound scalar to the array's storage type T.
66    //
67    // If the cast fails, the bound's value is outside [T::MIN, T::MAX].  For all signed
68    // NativeDecimalType implementations the minimum is negative and the maximum is positive, so
69    // we can determine the direction of overflow from the sign of the bound:
70    //   • non-negative and doesn't fit in T  ⟹  value > T::MAX
71    //   • negative and doesn't fit in T      ⟹  value < T::MIN
72    //
73    // From the direction we can answer the comparison immediately:
74    //   lower > T::MAX: no array value (≤ T::MAX) satisfies lower ≤ value  → all-false
75    //   lower < T::MIN: every array value (≥ T::MIN) satisfies lower ≤ value → no lower constraint
76    //   upper > T::MAX: every array value (≤ T::MAX) satisfies value ≤ upper → no upper constraint
77    //   upper < T::MIN: no array value (≥ T::MIN) satisfies value ≤ upper   → all-false
78    //
79    // Both the strict and non-strict forms lead to the same conclusion because the overflow is
80    // by at least one integer, so no boundary element can make the strict form differ.
81    let lower_value: Option<T> = match lower_dv.cast::<T>() {
82        Some(v) => Some(v),
83        None => {
84            if lower_dv.as_i256() >= i256::ZERO {
85                let validity = arr.validity()?.union_nullability(nullability);
86                return Ok(Some(
87                    BoolArray::new(BitBuffer::new_unset(arr.len()), validity).into_array(),
88                ));
89            }
90            None
91        }
92    };
93
94    let upper_value: Option<T> = match upper_dv.cast::<T>() {
95        Some(v) => Some(v),
96        None => {
97            if upper_dv.as_i256() < i256::ZERO {
98                let validity = arr.validity()?.union_nullability(nullability);
99                return Ok(Some(
100                    BoolArray::new(BitBuffer::new_unset(arr.len()), validity).into_array(),
101                ));
102            }
103            None
104        }
105    };
106
107    let lower_op = match options.lower_strict {
108        StrictComparison::Strict => |a, b| a < b,
109        StrictComparison::NonStrict => |a, b| a <= b,
110    };
111
112    let upper_op = match options.upper_strict {
113        StrictComparison::Strict => |a, b| a < b,
114        StrictComparison::NonStrict => |a, b| a <= b,
115    };
116
117    Ok(Some(between_impl::<T>(
118        arr,
119        lower_value,
120        upper_value,
121        nullability,
122        lower_op,
123        upper_op,
124        ctx,
125    )))
126}
127
128fn between_impl<T: NativeDecimalType>(
129    arr: ArrayView<'_, Decimal>,
130    lower: Option<T>,
131    upper: Option<T>,
132    nullability: Nullability,
133    lower_op: impl Fn(T, T) -> bool,
134    upper_op: impl Fn(T, T) -> bool,
135    ctx: &mut ExecutionCtx,
136) -> ArrayRef {
137    let buffer = arr.buffer::<T>();
138    BoolArray::new(
139        BitBuffer::collect_bool_multiversioned_in(
140            buffer.len(),
141            |idx| {
142                // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices
143                // `0..buffer.len()` only.
144                let value = unsafe { *buffer.get_unchecked(idx) };
145                lower.is_none_or(|l| lower_op(l, value)) & upper.is_none_or(|u| upper_op(value, u))
146            },
147            ctx.allocator().clone(),
148        ),
149        arr.validity()
150            .vortex_expect("validity should be derivable")
151            .union_nullability(nullability),
152    )
153    .into_array()
154}