|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! Block-streaming between kernel for [`BitPackedArray`] against constant bounds. |
| 5 | +//! |
| 6 | +//! Reuses the same single-block scratch buffer as the compare kernel and folds a |
| 7 | +//! `lower op_l v op_u upper` predicate per element, so the full primitive never |
| 8 | +//! materialises. |
| 9 | +
|
| 10 | +use vortex_array::ArrayRef; |
| 11 | +use vortex_array::ArrayView; |
| 12 | +use vortex_array::ExecutionCtx; |
| 13 | +use vortex_array::dtype::NativePType; |
| 14 | +use vortex_array::dtype::Nullability; |
| 15 | +use vortex_array::match_each_integer_ptype; |
| 16 | +use vortex_array::scalar_fn::fns::between::BetweenKernel; |
| 17 | +use vortex_array::scalar_fn::fns::between::BetweenOptions; |
| 18 | +use vortex_array::scalar_fn::fns::between::StrictComparison; |
| 19 | +use vortex_error::VortexExpect; |
| 20 | +use vortex_error::VortexResult; |
| 21 | + |
| 22 | +use crate::BitPacked; |
| 23 | +use crate::bitpacking::compute::stream_predicate::stream_predicate; |
| 24 | + |
| 25 | +impl BetweenKernel for BitPacked { |
| 26 | + fn between( |
| 27 | + array: ArrayView<'_, Self>, |
| 28 | + lower: &ArrayRef, |
| 29 | + upper: &ArrayRef, |
| 30 | + options: &BetweenOptions, |
| 31 | + ctx: &mut ExecutionCtx, |
| 32 | + ) -> VortexResult<Option<ArrayRef>> { |
| 33 | + // Only accelerate constant-bounds between; vary-by-row bounds fall through to the |
| 34 | + // default `compare + and` pipeline. |
| 35 | + let (Some(lower_const), Some(upper_const)) = (lower.as_constant(), upper.as_constant()) |
| 36 | + else { |
| 37 | + return Ok(None); |
| 38 | + }; |
| 39 | + let (Some(lower_prim), Some(upper_prim)) = ( |
| 40 | + lower_const.as_primitive_opt(), |
| 41 | + upper_const.as_primitive_opt(), |
| 42 | + ) else { |
| 43 | + return Ok(None); |
| 44 | + }; |
| 45 | + |
| 46 | + let nullability = |
| 47 | + array.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability(); |
| 48 | + let arr_ptype = array.dtype().as_ptype(); |
| 49 | + if lower_prim.ptype() != arr_ptype || upper_prim.ptype() != arr_ptype { |
| 50 | + return Ok(None); |
| 51 | + } |
| 52 | + |
| 53 | + let result = match_each_integer_ptype!(arr_ptype, |T| { |
| 54 | + let lo: T = lower_prim |
| 55 | + .typed_value::<T>() |
| 56 | + .vortex_expect("between precondition strips null lower"); |
| 57 | + let up: T = upper_prim |
| 58 | + .typed_value::<T>() |
| 59 | + .vortex_expect("between precondition strips null upper"); |
| 60 | + between_constant_typed::<T>(array, lo, up, options, nullability, ctx)? |
| 61 | + }); |
| 62 | + Ok(Some(result)) |
| 63 | + } |
| 64 | +} |
| 65 | + |
| 66 | +fn between_constant_typed<T>( |
| 67 | + array: ArrayView<'_, BitPacked>, |
| 68 | + lower: T, |
| 69 | + upper: T, |
| 70 | + options: &BetweenOptions, |
| 71 | + nullability: Nullability, |
| 72 | + ctx: &mut ExecutionCtx, |
| 73 | +) -> VortexResult<ArrayRef> |
| 74 | +where |
| 75 | + T: NativePType + Copy + crate::unpack_iter::BitPacked, |
| 76 | +{ |
| 77 | + // Branch on strictness once at the top so each call into `between_impl` monomorphises |
| 78 | + // a single tight predicate — same shape as `Primitive::between` in `vortex-array`. |
| 79 | + match (options.lower_strict, options.upper_strict) { |
| 80 | + (StrictComparison::Strict, StrictComparison::Strict) => between_impl( |
| 81 | + array, |
| 82 | + lower, |
| 83 | + NativePType::is_lt, |
| 84 | + upper, |
| 85 | + NativePType::is_lt, |
| 86 | + nullability, |
| 87 | + ctx, |
| 88 | + ), |
| 89 | + (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl( |
| 90 | + array, |
| 91 | + lower, |
| 92 | + NativePType::is_lt, |
| 93 | + upper, |
| 94 | + NativePType::is_le, |
| 95 | + nullability, |
| 96 | + ctx, |
| 97 | + ), |
| 98 | + (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl( |
| 99 | + array, |
| 100 | + lower, |
| 101 | + NativePType::is_le, |
| 102 | + upper, |
| 103 | + NativePType::is_lt, |
| 104 | + nullability, |
| 105 | + ctx, |
| 106 | + ), |
| 107 | + (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl( |
| 108 | + array, |
| 109 | + lower, |
| 110 | + NativePType::is_le, |
| 111 | + upper, |
| 112 | + NativePType::is_le, |
| 113 | + nullability, |
| 114 | + ctx, |
| 115 | + ), |
| 116 | + } |
| 117 | +} |
| 118 | + |
| 119 | +fn between_impl<T, Lo, Up>( |
| 120 | + array: ArrayView<'_, BitPacked>, |
| 121 | + lower: T, |
| 122 | + lower_fn: Lo, |
| 123 | + upper: T, |
| 124 | + upper_fn: Up, |
| 125 | + nullability: Nullability, |
| 126 | + ctx: &mut ExecutionCtx, |
| 127 | +) -> VortexResult<ArrayRef> |
| 128 | +where |
| 129 | + T: NativePType + Copy + crate::unpack_iter::BitPacked, |
| 130 | + Lo: Fn(T, T) -> bool, |
| 131 | + Up: Fn(T, T) -> bool, |
| 132 | +{ |
| 133 | + stream_predicate::<T, _>( |
| 134 | + array, |
| 135 | + nullability, |
| 136 | + |v| lower_fn(lower, v) && upper_fn(v, upper), |
| 137 | + ctx, |
| 138 | + ) |
| 139 | +} |
| 140 | + |
| 141 | +#[cfg(test)] |
| 142 | +mod tests { |
| 143 | + use rstest::rstest; |
| 144 | + use vortex_array::IntoArray; |
| 145 | + use vortex_array::LEGACY_SESSION; |
| 146 | + use vortex_array::VortexSessionExecute; |
| 147 | + use vortex_array::arrays::BoolArray; |
| 148 | + use vortex_array::arrays::ConstantArray; |
| 149 | + use vortex_array::arrays::PrimitiveArray; |
| 150 | + use vortex_array::assert_arrays_eq; |
| 151 | + use vortex_array::builtins::ArrayBuiltins; |
| 152 | + use vortex_array::scalar_fn::fns::between::BetweenOptions; |
| 153 | + use vortex_array::scalar_fn::fns::between::StrictComparison; |
| 154 | + use vortex_array::validity::Validity; |
| 155 | + use vortex_buffer::BufferMut; |
| 156 | + use vortex_error::VortexResult; |
| 157 | + |
| 158 | + use crate::BitPackedArrayExt; |
| 159 | + use crate::BitPackedData; |
| 160 | + |
| 161 | + fn opts(lower: StrictComparison, upper: StrictComparison) -> BetweenOptions { |
| 162 | + BetweenOptions { |
| 163 | + lower_strict: lower, |
| 164 | + upper_strict: upper, |
| 165 | + } |
| 166 | + } |
| 167 | + |
| 168 | + #[rstest] |
| 169 | + #[case(StrictComparison::NonStrict, StrictComparison::NonStrict)] |
| 170 | + #[case(StrictComparison::Strict, StrictComparison::NonStrict)] |
| 171 | + #[case(StrictComparison::NonStrict, StrictComparison::Strict)] |
| 172 | + #[case(StrictComparison::Strict, StrictComparison::Strict)] |
| 173 | + fn multi_chunk_against_primitive_baseline( |
| 174 | + #[case] lower_strict: StrictComparison, |
| 175 | + #[case] upper_strict: StrictComparison, |
| 176 | + ) -> VortexResult<()> { |
| 177 | + let mut ctx = LEGACY_SESSION.create_execution_ctx(); |
| 178 | + let values: BufferMut<u32> = (0..3000u32).map(|i| i % 257).collect(); |
| 179 | + let prim = PrimitiveArray::new(values.freeze(), Validity::NonNullable); |
| 180 | + let packed = BitPackedData::encode(&prim.clone().into_array(), 9, &mut ctx)?; |
| 181 | + |
| 182 | + let lower = ConstantArray::new(40u32, prim.len()).into_array(); |
| 183 | + let upper = ConstantArray::new(200u32, prim.len()).into_array(); |
| 184 | + let options = opts(lower_strict, upper_strict); |
| 185 | + |
| 186 | + let expected = prim |
| 187 | + .into_array() |
| 188 | + .between(lower.clone(), upper.clone(), options.clone())? |
| 189 | + .execute::<BoolArray>(&mut ctx)?; |
| 190 | + let actual = packed |
| 191 | + .into_array() |
| 192 | + .between(lower, upper, options)? |
| 193 | + .execute::<BoolArray>(&mut ctx)?; |
| 194 | + |
| 195 | + assert_arrays_eq!(actual, expected); |
| 196 | + Ok(()) |
| 197 | + } |
| 198 | + |
| 199 | + #[test] |
| 200 | + fn signed_with_patches_against_primitive_baseline() -> VortexResult<()> { |
| 201 | + let mut ctx = LEGACY_SESSION.create_execution_ctx(); |
| 202 | + let values: Vec<i32> = (0..1500) |
| 203 | + .map(|i| if i % 73 == 0 { 100_000 + i } else { i % 100 }) |
| 204 | + .collect(); |
| 205 | + let prim = PrimitiveArray::from_iter(values); |
| 206 | + let packed = BitPackedData::encode(&prim.clone().into_array(), 7, &mut ctx)?; |
| 207 | + assert!(packed.patches().is_some(), "test setup expects patches"); |
| 208 | + |
| 209 | + let lower = ConstantArray::new(20i32, prim.len()).into_array(); |
| 210 | + let upper = ConstantArray::new(80i32, prim.len()).into_array(); |
| 211 | + let options = opts(StrictComparison::NonStrict, StrictComparison::NonStrict); |
| 212 | + |
| 213 | + let expected = prim |
| 214 | + .into_array() |
| 215 | + .between(lower.clone(), upper.clone(), options.clone())? |
| 216 | + .execute::<BoolArray>(&mut ctx)?; |
| 217 | + let actual = packed |
| 218 | + .into_array() |
| 219 | + .between(lower, upper, options)? |
| 220 | + .execute::<BoolArray>(&mut ctx)?; |
| 221 | + |
| 222 | + assert_arrays_eq!(actual, expected); |
| 223 | + Ok(()) |
| 224 | + } |
| 225 | + |
| 226 | + #[test] |
| 227 | + fn nullable_propagates_validity() -> VortexResult<()> { |
| 228 | + let mut ctx = LEGACY_SESSION.create_execution_ctx(); |
| 229 | + let prim = |
| 230 | + PrimitiveArray::from_option_iter([Some(1u32), None, Some(3), Some(4), None, Some(6)]); |
| 231 | + let packed = BitPackedData::encode(&prim.clone().into_array(), 3, &mut ctx)?; |
| 232 | + |
| 233 | + let lower = ConstantArray::new(2u32, packed.len()).into_array(); |
| 234 | + let upper = ConstantArray::new(5u32, packed.len()).into_array(); |
| 235 | + let options = opts(StrictComparison::NonStrict, StrictComparison::NonStrict); |
| 236 | + |
| 237 | + let actual = packed |
| 238 | + .into_array() |
| 239 | + .between(lower.clone(), upper.clone(), options.clone())? |
| 240 | + .execute::<BoolArray>(&mut ctx)?; |
| 241 | + let expected = prim |
| 242 | + .into_array() |
| 243 | + .between(lower, upper, options)? |
| 244 | + .execute::<BoolArray>(&mut ctx)?; |
| 245 | + assert_arrays_eq!(actual, expected); |
| 246 | + Ok(()) |
| 247 | + } |
| 248 | +} |
0 commit comments