Skip to content

Commit 60847b0

Browse files
a10ya10y
authored
move BitPacked slice from kernel -> parent reduce (#7195)
This is free and cheap since we do not need to materialize any data to do it. --------- Signed-off-by: Andrew Duffy <andrew@a10y.dev>
1 parent 1f3206f commit 60847b0

File tree

6 files changed

+21
-48
lines changed

6 files changed

+21
-48
lines changed

‎encodings/fastlanes/public-api.lock‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,9 +138,9 @@ impl vortex_array::arrays::filter::kernel::FilterKernel for vortex_fastlanes::Bi
138138

139139
pub fn vortex_fastlanes::BitPacked::filter(array: &vortex_fastlanes::BitPackedArray, mask: &vortex_mask::Mask, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
140140

141-
impl vortex_array::arrays::slice::SliceKernel for vortex_fastlanes::BitPacked
141+
impl vortex_array::arrays::slice::SliceReduce for vortex_fastlanes::BitPacked
142142

143-
pub fn vortex_fastlanes::BitPacked::slice(array: &vortex_fastlanes::BitPackedArray, range: core::ops::range::Range<usize>, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
143+
pub fn vortex_fastlanes::BitPacked::slice(array: &vortex_fastlanes::BitPackedArray, range: core::ops::range::Range<usize>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
144144

145145
impl vortex_array::scalar_fn::fns::cast::kernel::CastReduce for vortex_fastlanes::BitPacked
146146

‎encodings/fastlanes/src/bitpacking/compute/slice.rs‎

Lines changed: 5 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,15 @@ use std::cmp::max;
55
use std::ops::Range;
66

77
use vortex_array::ArrayRef;
8-
use vortex_array::ExecutionCtx;
98
use vortex_array::IntoArray;
10-
use vortex_array::arrays::slice::SliceKernel;
9+
use vortex_array::arrays::slice::SliceReduce;
1110
use vortex_error::VortexResult;
1211

1312
use crate::BitPacked;
1413
use crate::BitPackedArray;
1514

16-
impl SliceKernel for BitPacked {
17-
fn slice(
18-
array: &BitPackedArray,
19-
range: Range<usize>,
20-
_ctx: &mut ExecutionCtx,
21-
) -> VortexResult<Option<ArrayRef>> {
15+
impl SliceReduce for BitPacked {
16+
fn slice(array: &BitPackedArray, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
2217
let offset_start = range.start + array.offset() as usize;
2318
let offset_stop = range.end + array.offset() as usize;
2419
let offset = offset_start % 1024;
@@ -51,35 +46,26 @@ impl SliceKernel for BitPacked {
5146

5247
#[cfg(test)]
5348
mod tests {
54-
use std::sync::LazyLock;
55-
5649
use vortex_array::DynArray;
5750
use vortex_array::IntoArray;
58-
use vortex_array::VortexSessionExecute;
5951
use vortex_array::arrays::PrimitiveArray;
6052
use vortex_array::arrays::SliceArray;
61-
use vortex_array::session::ArraySession;
6253
use vortex_error::VortexResult;
63-
use vortex_session::VortexSession;
6454

6555
use crate::BitPacked;
6656
use crate::bitpack_compress::bitpack_encode;
6757

68-
static SESSION: LazyLock<VortexSession> =
69-
LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
70-
7158
#[test]
72-
fn test_execute_parent_returns_bitpacked_slice() -> VortexResult<()> {
59+
fn test_reduce_parent_returns_bitpacked_slice() -> VortexResult<()> {
7360
let values = PrimitiveArray::from_iter(0u32..2048);
7461
let bitpacked = bitpack_encode(&values, 11, None)?;
7562

7663
let slice_array = SliceArray::new(bitpacked.clone().into_array(), 500..1500);
7764

78-
let mut ctx = SESSION.create_execution_ctx();
7965
let bitpacked_ref = bitpacked.into_array();
8066
let reduced = bitpacked_ref
8167
.vtable()
82-
.execute_parent(&bitpacked_ref, &slice_array.into_array(), 0, &mut ctx)?
68+
.reduce_parent(&bitpacked_ref, &slice_array.into_array(), 0)?
8369
.expect("expected slice kernel to execute");
8470

8571
assert!(reduced.is::<BitPacked>());

‎encodings/fastlanes/src/bitpacking/vtable/kernels.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,11 @@
33

44
use vortex_array::arrays::dict::TakeExecuteAdaptor;
55
use vortex_array::arrays::filter::FilterExecuteAdaptor;
6-
use vortex_array::arrays::slice::SliceExecuteAdaptor;
76
use vortex_array::kernel::ParentKernelSet;
87

98
use crate::BitPacked;
109

1110
pub(crate) const PARENT_KERNELS: ParentKernelSet<BitPacked> = ParentKernelSet::new(&[
1211
ParentKernelSet::lift(&FilterExecuteAdaptor(BitPacked)),
13-
ParentKernelSet::lift(&SliceExecuteAdaptor(BitPacked)),
1412
ParentKernelSet::lift(&TakeExecuteAdaptor(BitPacked)),
1513
]);

‎encodings/fastlanes/src/bitpacking/vtable/operations.rs‎

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,9 @@ impl OperationsVTable<BitPacked> for BitPacked {
3131
#[cfg(test)]
3232
mod test {
3333
use std::ops::Range;
34-
use std::sync::LazyLock;
3534

3635
use vortex_array::DynArray;
3736
use vortex_array::IntoArray;
38-
use vortex_array::VortexSessionExecute;
3937
use vortex_array::arrays::PrimitiveArray;
4038
use vortex_array::arrays::SliceArray;
4139
use vortex_array::assert_arrays_eq;
@@ -46,7 +44,6 @@ mod test {
4644
use vortex_array::dtype::PType;
4745
use vortex_array::patches::Patches;
4846
use vortex_array::scalar::Scalar;
49-
use vortex_array::session::ArraySession;
5047
use vortex_array::validity::Validity;
5148
use vortex_buffer::Alignment;
5249
use vortex_buffer::Buffer;
@@ -56,16 +53,12 @@ mod test {
5653
use crate::BitPacked;
5754
use crate::BitPackedArray;
5855

59-
static SESSION: LazyLock<vortex_session::VortexSession> =
60-
LazyLock::new(|| vortex_session::VortexSession::empty().with::<ArraySession>());
61-
62-
fn slice_via_kernel(array: &BitPackedArray, range: Range<usize>) -> BitPackedArray {
56+
fn slice_via_reduce(array: &BitPackedArray, range: Range<usize>) -> BitPackedArray {
6357
let array_ref = array.clone().into_array();
6458
let slice_array = SliceArray::new(array_ref.clone(), range);
65-
let mut ctx = SESSION.create_execution_ctx();
6659
let sliced = array_ref
6760
.vtable()
68-
.execute_parent(&array_ref, &slice_array.into_array(), 0, &mut ctx)
61+
.reduce_parent(&array_ref, &slice_array.into_array(), 0)
6962
.expect("execute_parent failed")
7063
.expect("expected slice kernel to execute");
7164
sliced.as_::<BitPacked>().clone()
@@ -78,7 +71,7 @@ mod test {
7871
6,
7972
)
8073
.unwrap();
81-
let sliced = slice_via_kernel(&arr, 1024..2048);
74+
let sliced = slice_via_reduce(&arr, 1024..2048);
8275
assert_nth_scalar!(sliced, 0, 1024u32 % 64);
8376
assert_nth_scalar!(sliced, 1023, 2047u32 % 64);
8477
assert_eq!(sliced.offset(), 0);
@@ -92,7 +85,7 @@ mod test {
9285
6,
9386
)
9487
.unwrap();
95-
let sliced = slice_via_kernel(&arr, 512..1434);
88+
let sliced = slice_via_reduce(&arr, 512..1434);
9689
assert_nth_scalar!(sliced, 0, 512u32 % 64);
9790
assert_nth_scalar!(sliced, 921, 1433u32 % 64);
9891
assert_eq!(sliced.offset(), 512);
@@ -132,12 +125,12 @@ mod test {
132125
6,
133126
)
134127
.unwrap();
135-
let sliced = slice_via_kernel(&arr, 512..1434);
128+
let sliced = slice_via_reduce(&arr, 512..1434);
136129
assert_nth_scalar!(sliced, 0, 512u32 % 64);
137130
assert_nth_scalar!(sliced, 921, 1433u32 % 64);
138131
assert_eq!(sliced.offset(), 512);
139132
assert_eq!(sliced.len(), 922);
140-
let doubly_sliced = slice_via_kernel(&sliced, 127..911);
133+
let doubly_sliced = slice_via_reduce(&sliced, 127..911);
141134
assert_nth_scalar!(doubly_sliced, 0, (512u32 + 127) % 64);
142135
assert_nth_scalar!(doubly_sliced, 783, (512u32 + 910) % 64);
143136
assert_eq!(doubly_sliced.offset(), 639);
@@ -155,7 +148,7 @@ mod test {
155148
assert_eq!(patch_indices.len(), 1);
156149

157150
// Slicing drops the empty patches array.
158-
let sliced_bp = slice_via_kernel(&array, 0..64);
151+
let sliced_bp = slice_via_reduce(&array, 0..64);
159152
assert!(sliced_bp.patches().is_none());
160153
}
161154

‎encodings/fastlanes/src/bitpacking/vtable/rules.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use vortex_array::arrays::slice::SliceReduceAdaptor;
45
use vortex_array::optimizer::rules::ParentRuleSet;
56
use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
67

78
use crate::BitPacked;
89

9-
pub(crate) const RULES: ParentRuleSet<BitPacked> =
10-
ParentRuleSet::new(&[ParentRuleSet::lift(&CastReduceAdaptor(BitPacked))]);
10+
pub(crate) const RULES: ParentRuleSet<BitPacked> = ParentRuleSet::new(&[
11+
ParentRuleSet::lift(&CastReduceAdaptor(BitPacked)),
12+
ParentRuleSet::lift(&SliceReduceAdaptor(BitPacked)),
13+
]);

‎vortex-cuda/src/kernel/encodings/bitpacked.rs‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -169,12 +169,10 @@ where
169169
mod tests {
170170
use futures::executor::block_on;
171171
use rstest::rstest;
172-
use vortex::array::ExecutionCtx;
173172
use vortex::array::IntoArray;
174173
use vortex::array::arrays::PrimitiveArray;
175174
use vortex::array::assert_arrays_eq;
176175
use vortex::array::dtype::NativePType;
177-
use vortex::array::session::ArraySession;
178176
use vortex::array::validity::Validity::NonNullable;
179177
use vortex::buffer::Buffer;
180178
use vortex::error::VortexExpect;
@@ -519,13 +517,8 @@ mod tests {
519517

520518
let bitpacked_array = BitPackedArray::encode(&primitive_array.into_array(), bit_width)
521519
.vortex_expect("operation should succeed in test");
522-
let slice_ref = bitpacked_array.clone().into_array().slice(67..3969)?;
523-
let bitpacked_ref = bitpacked_array.into_array();
524-
let mut exec_ctx = ExecutionCtx::new(VortexSession::empty().with::<ArraySession>());
525-
let sliced_array = bitpacked_ref
526-
.vtable()
527-
.execute_parent(&bitpacked_ref, &slice_ref, 0, &mut exec_ctx)?
528-
.expect("expected slice kernel to execute");
520+
let sliced_array = bitpacked_array.into_array().slice(67..3969)?;
521+
assert!(sliced_array.is::<BitPacked>());
529522
let cpu_result = sliced_array.to_canonical()?;
530523
let gpu_result = block_on(async {
531524
BitPackedExecutor

0 commit comments

Comments
 (0)