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::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 let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else {
35 return Ok(None);
36 };
37
38 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 return Ok(None);
59 };
60 let Some(upper_dv) = upper.as_decimal().decimal_value() else {
61 return Ok(None);
63 };
64
65 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 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}