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