vortex_array/arrays/decimal/compute/
between.rs1use 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 let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
36 return Ok(None);
37 };
38
39 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 return Ok(None);
60 };
61 let Some(upper_dv) = upper.as_decimal().decimal_value() else {
62 return Ok(None);
64 };
65
66 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 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}