Skip to content

Commit ee44dd6

Browse files
committed
fastlanes: streaming compare + between kernels for BitPacked
Adds `CompareKernel` and `BetweenKernel` for `BitPacked` that walk the encoded array one 1024-element FastLanes block at a time through a single reused scratch buffer, splice any `Patches` into the block in place via a sorted-index cursor, then fold a `Fn(T) -> bool` predicate over the block and write the bits directly into the output bit buffer. The materialised primitive never appears. The inner predicate-fold matches the canonical `BitBuffer::collect_bool` shape — pack 64 bools into a `u64` in a tight loop — which rustc auto-vectorises into the same `pcmpeq` + `psllq` (vector shift to bit position) + `por` (OR into accumulator) pattern that `arrow-ord::apply_op` lowers to. Verified via `objdump` on the bench binary (344 monomorphised `stream_predicate` variants emit those SIMD instructions in the inner loop). Smallest possible diff: only adds the two kernels and a private helper shared between them, no benches, no public-API expansion beyond the two trait impls. encodings/fastlanes/public-api.lock | 8 + encodings/fastlanes/src/bitpacking/compute/between.rs | 248 + encodings/fastlanes/src/bitpacking/compute/compare.rs | 187 + encodings/fastlanes/src/bitpacking/compute/mod.rs | 3 + encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs | 211 + encodings/fastlanes/src/bitpacking/vtable/kernels.rs | 4 + 6 files changed, 661 insertions(+) Checks: - cargo nextest run -p vortex-fastlanes (278 passed) - cargo clippy -p vortex-fastlanes --all-targets --all-features - cargo +nightly fmt -p vortex-fastlanes --check - ./scripts/public-api.sh (only adds the two new trait impls) Signed-off-by: Claude <noreply@anthropic.com>
1 parent 7b47788 commit ee44dd6

6 files changed

Lines changed: 661 additions & 0 deletions

File tree

encodings/fastlanes/public-api.lock

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,14 @@ impl vortex_array::arrays::slice::SliceReduce for vortex_fastlanes::BitPacked
190190

191191
pub fn vortex_fastlanes::BitPacked::slice(vortex_array::array::view::ArrayView<'_, Self>, core::ops::range::Range<usize>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
192192

193+
impl vortex_array::scalar_fn::fns::between::kernel::BetweenKernel for vortex_fastlanes::BitPacked
194+
195+
pub fn vortex_fastlanes::BitPacked::between(vortex_array::array::view::ArrayView<'_, Self>, &vortex_array::array::erased::ArrayRef, &vortex_array::array::erased::ArrayRef, &vortex_array::scalar_fn::fns::between::BetweenOptions, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
196+
197+
impl vortex_array::scalar_fn::fns::binary::compare::CompareKernel for vortex_fastlanes::BitPacked
198+
199+
pub fn vortex_fastlanes::BitPacked::compare(vortex_array::array::view::ArrayView<'_, Self>, &vortex_array::array::erased::ArrayRef, vortex_array::scalar_fn::fns::operators::CompareOperator, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
200+
193201
impl vortex_array::scalar_fn::fns::cast::kernel::CastKernel for vortex_fastlanes::BitPacked
194202

195203
pub fn vortex_fastlanes::BitPacked::cast(vortex_array::array::view::ArrayView<'_, Self>, &vortex_array::dtype::DType, &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
Lines changed: 248 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,248 @@
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

Comments
 (0)