Skip to content

Commit 94de103

Browse files
gatesngatesn
committed
Move execute-parent kernels into sessions
Signed-off-by: "Nicholas Gates" <nick@nickgates.com>
1 parent 48c33e8 commit 94de103

105 files changed

Lines changed: 883 additions & 750 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎encodings/alp/src/alp/array.rs‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@ use vortex_session::registry::CachedId;
4848
use crate::ALPFloat;
4949
use crate::alp::Exponents;
5050
use crate::alp::decompress::execute_decompress;
51-
use crate::alp::rules::PARENT_KERNELS;
5251
use crate::alp::rules::RULES;
5352

5453
/// A [`ALP`]-encoded Vortex array.
@@ -188,15 +187,6 @@ impl VTable for ALP {
188187
) -> VortexResult<Option<ArrayRef>> {
189188
RULES.evaluate(array, parent, child_idx)
190189
}
191-
192-
fn execute_parent(
193-
array: ArrayView<'_, Self>,
194-
parent: &ArrayRef,
195-
child_idx: usize,
196-
ctx: &mut ExecutionCtx,
197-
) -> VortexResult<Option<ArrayRef>> {
198-
PARENT_KERNELS.execute(array, parent, child_idx, ctx)
199-
}
200190
}
201191

202192
#[array_slots(ALP)]

‎encodings/alp/src/alp/mod.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,9 +61,14 @@ use vortex_array::dtype::NativePType;
6161
use vortex_array::scalar::PValue;
6262
use vortex_buffer::Buffer;
6363
use vortex_buffer::BufferMut;
64+
use vortex_session::VortexSession;
6465

6566
const SAMPLE_SIZE: usize = 32;
6667

68+
pub(crate) fn initialize(session: &VortexSession) {
69+
rules::initialize(session);
70+
}
71+
6772
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
6873
pub struct Exponents {
6974
pub e: u8,

‎encodings/alp/src/alp/rules.rs‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,35 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use vortex_array::ArrayVTable;
5+
use vortex_array::arrays::Dict;
6+
use vortex_array::arrays::Filter;
7+
use vortex_array::arrays::Slice;
48
use vortex_array::arrays::dict::TakeExecuteAdaptor;
59
use vortex_array::arrays::filter::FilterExecuteAdaptor;
610
use vortex_array::arrays::slice::SliceExecuteAdaptor;
7-
use vortex_array::kernel::ParentKernelSet;
11+
use vortex_array::optimizer::kernels::ArrayKernelsExt;
812
use vortex_array::optimizer::rules::ParentRuleSet;
13+
use vortex_array::scalar_fn::ScalarFnVTable;
914
use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor;
15+
use vortex_array::scalar_fn::fns::binary::Binary;
1016
use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor;
1117
use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor;
18+
use vortex_array::scalar_fn::fns::mask::Mask;
1219
use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor;
1320
use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor;
21+
use vortex_session::VortexSession;
1422

1523
use crate::ALP;
1624

17-
pub(super) const PARENT_KERNELS: ParentKernelSet<ALP> = ParentKernelSet::new(&[
18-
ParentKernelSet::lift(&CompareExecuteAdaptor(ALP)),
19-
ParentKernelSet::lift(&FilterExecuteAdaptor(ALP)),
20-
ParentKernelSet::lift(&MaskExecuteAdaptor(ALP)),
21-
ParentKernelSet::lift(&SliceExecuteAdaptor(ALP)),
22-
ParentKernelSet::lift(&TakeExecuteAdaptor(ALP)),
23-
]);
25+
pub(super) fn initialize(session: &VortexSession) {
26+
let kernels = session.kernels();
27+
kernels.register_execute_parent_kernel(Binary.id(), ALP, CompareExecuteAdaptor(ALP));
28+
kernels.register_execute_parent_kernel(Filter.id(), ALP, FilterExecuteAdaptor(ALP));
29+
kernels.register_execute_parent_kernel(Mask.id(), ALP, MaskExecuteAdaptor(ALP));
30+
kernels.register_execute_parent_kernel(Slice.id(), ALP, SliceExecuteAdaptor(ALP));
31+
kernels.register_execute_parent_kernel(Dict.id(), ALP, TakeExecuteAdaptor(ALP));
32+
}
2433

2534
pub(super) const RULES: ParentRuleSet<ALP> = ParentRuleSet::new(&[
2635
ParentRuleSet::lift(&BetweenReduceAdaptor(ALP)),

‎encodings/alp/src/alp_rd/array.rs‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ use vortex_error::vortex_panic;
5353
use vortex_session::VortexSession;
5454
use vortex_session::registry::CachedId;
5555

56-
use crate::alp_rd::kernel::PARENT_KERNELS;
5756
use crate::alp_rd::rules::RULES;
5857
use crate::alp_rd_decode;
5958

@@ -307,15 +306,6 @@ impl VTable for ALPRD {
307306
) -> VortexResult<Option<ArrayRef>> {
308307
RULES.evaluate(array, parent, child_idx)
309308
}
310-
311-
fn execute_parent(
312-
array: ArrayView<'_, Self>,
313-
parent: &ArrayRef,
314-
child_idx: usize,
315-
ctx: &mut ExecutionCtx,
316-
) -> VortexResult<Option<ArrayRef>> {
317-
PARENT_KERNELS.execute(array, parent, child_idx, ctx)
318-
}
319309
}
320310

321311
/// The left (most significant) parts of the real-double encoded values.

‎encodings/alp/src/alp_rd/kernel.rs‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,21 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use vortex_array::ArrayVTable;
5+
use vortex_array::arrays::Dict;
6+
use vortex_array::arrays::Filter;
7+
use vortex_array::arrays::Slice;
48
use vortex_array::arrays::dict::TakeExecuteAdaptor;
59
use vortex_array::arrays::filter::FilterExecuteAdaptor;
610
use vortex_array::arrays::slice::SliceExecuteAdaptor;
7-
use vortex_array::kernel::ParentKernelSet;
11+
use vortex_array::optimizer::kernels::ArrayKernelsExt;
12+
use vortex_session::VortexSession;
813

914
use crate::alp_rd::ALPRD;
1015

11-
pub(crate) static PARENT_KERNELS: ParentKernelSet<ALPRD> = ParentKernelSet::new(&[
12-
ParentKernelSet::lift(&SliceExecuteAdaptor(ALPRD)),
13-
ParentKernelSet::lift(&FilterExecuteAdaptor(ALPRD)),
14-
ParentKernelSet::lift(&TakeExecuteAdaptor(ALPRD)),
15-
]);
16+
pub(crate) fn initialize(session: &VortexSession) {
17+
let kernels = session.kernels();
18+
kernels.register_execute_parent_kernel(Slice.id(), ALPRD, SliceExecuteAdaptor(ALPRD));
19+
kernels.register_execute_parent_kernel(Filter.id(), ALPRD, FilterExecuteAdaptor(ALPRD));
20+
kernels.register_execute_parent_kernel(Dict.id(), ALPRD, TakeExecuteAdaptor(ALPRD));
21+
}

‎encodings/alp/src/alp_rd/mod.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ use vortex_buffer::BufferMut;
3636
use vortex_error::VortexExpect;
3737
use vortex_error::VortexResult;
3838
use vortex_error::vortex_panic;
39+
use vortex_session::VortexSession;
3940
use vortex_utils::aliases::hash_map::HashMap;
4041

4142
use crate::match_each_alp_float_ptype;
@@ -55,6 +56,10 @@ const CUT_LIMIT: usize = 16;
5556

5657
const MAX_DICT_SIZE: u8 = 8;
5758

59+
pub(crate) fn initialize(session: &VortexSession) {
60+
kernel::initialize(session);
61+
}
62+
5863
mod private {
5964
pub trait Sealed {}
6065

‎encodings/alp/src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ pub fn initialize(session: &VortexSession) {
3939
session.arrays().register(ALP);
4040
}
4141
session.arrays().register(ALPRD);
42+
alp::initialize(session);
43+
alp_rd::initialize(session);
4244

4345
// Register the ALP-specific NaN count aggregate kernel.
4446
session.aggregate_fns().register_aggregate_kernel(

‎encodings/bytebool/src/array.rs‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,6 @@ use vortex_error::vortex_panic;
3939
use vortex_session::VortexSession;
4040
use vortex_session::registry::CachedId;
4141

42-
use crate::kernel::PARENT_KERNELS;
43-
4442
/// A [`ByteBool`]-encoded Vortex array.
4543
pub type ByteBoolArray = Array<ByteBool>;
4644

@@ -157,15 +155,6 @@ impl VTable for ByteBool {
157155
BoolArray::new(boolean_buffer, validity).into_array(),
158156
))
159157
}
160-
161-
fn execute_parent(
162-
array: ArrayView<'_, Self>,
163-
parent: &ArrayRef,
164-
child_idx: usize,
165-
ctx: &mut ExecutionCtx,
166-
) -> VortexResult<Option<ArrayRef>> {
167-
PARENT_KERNELS.execute(array, parent, child_idx, ctx)
168-
}
169158
}
170159

171160
/// The validity bitmap indicating which elements are non-null.

‎encodings/bytebool/src/kernel.rs‎

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

4+
use vortex_array::ArrayVTable;
5+
use vortex_array::arrays::Dict;
46
use vortex_array::arrays::dict::TakeExecuteAdaptor;
5-
use vortex_array::kernel::ParentKernelSet;
7+
use vortex_array::optimizer::kernels::ArrayKernelsExt;
8+
use vortex_array::scalar_fn::ScalarFnVTable;
9+
use vortex_array::scalar_fn::fns::cast::Cast;
610
use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor;
11+
use vortex_session::VortexSession;
712

813
use crate::ByteBool;
914

10-
pub(crate) const PARENT_KERNELS: ParentKernelSet<ByteBool> = ParentKernelSet::new(&[
11-
ParentKernelSet::lift(&CastExecuteAdaptor(ByteBool)),
12-
ParentKernelSet::lift(&TakeExecuteAdaptor(ByteBool)),
13-
]);
15+
pub(crate) fn initialize(session: &VortexSession) {
16+
let kernels = session.kernels();
17+
kernels.register_execute_parent_kernel(Cast.id(), ByteBool, CastExecuteAdaptor(ByteBool));
18+
kernels.register_execute_parent_kernel(Dict.id(), ByteBool, TakeExecuteAdaptor(ByteBool));
19+
}

‎encodings/bytebool/src/lib.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,17 @@
4141
//! [spec]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#bit-boolean
4242
4343
pub use array::*;
44+
use vortex_array::session::ArraySessionExt;
45+
use vortex_session::VortexSession;
4446

4547
mod array;
4648
mod compute;
4749
mod kernel;
4850
mod rules;
4951
mod slice;
52+
53+
/// Initialize bytebool encoding in the given session.
54+
pub fn initialize(session: &VortexSession) {
55+
session.arrays().register(ByteBool);
56+
kernel::initialize(session);
57+
}

0 commit comments

Comments
 (0)