Skip to content

Commit 2cc04c4

Browse files
joseph-isaacsjoseph-isaacs
authored
perf(array): forward probe_scalar through pass-through encodings (#9905)
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 24bffe5 commit 2cc04c4

14 files changed

Lines changed: 355 additions & 63 deletions

File tree

‎encodings/alp/src/alp/ops.rs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,32 +3,41 @@
33

44
use vortex_array::ArrayView;
55
use vortex_array::ExecutionCtx;
6+
use vortex_array::ProbeState;
67
use vortex_array::scalar::Scalar;
78
use vortex_array::vtable::OperationsVTable;
89
use vortex_error::VortexExpect;
910
use vortex_error::VortexResult;
11+
use vortex_error::vortex_err;
1012

1113
use crate::ALP;
1214
use crate::ALPArrayExt;
13-
use crate::ALPArraySlotsExt;
1415
use crate::ALPFloat;
16+
use crate::ALPSlots;
1517
use crate::match_each_alp_float_ptype;
1618

1719
impl OperationsVTable<ALP> for ALP {
1820
type ProbeState = ();
1921

20-
fn scalar_at(
21-
array: ArrayView<'_, ALP>,
22+
fn probe_scalar(
23+
state: &mut ProbeState<'_, ALP>,
2224
index: usize,
2325
ctx: &mut ExecutionCtx,
2426
) -> VortexResult<Scalar> {
27+
if !state.is_valid(index, ctx)? {
28+
return Ok(Scalar::null(state.array().dtype().clone()));
29+
}
30+
let array = state.array();
2531
if let Some(patches) = array.patches()
2632
&& let Some(patch) = patches.get_patched(index)?
2733
{
2834
return patch.cast(array.dtype());
2935
}
3036

31-
let encoded_val = array.encoded().execute_scalar(index, ctx)?;
37+
let encoded_val = state
38+
.slot(ALPSlots::ENCODED)?
39+
.ok_or_else(|| vortex_err!("ALP encoded slot is missing"))?
40+
.execute_scalar(index, ctx)?;
3241

3342
Ok(match_each_alp_float_ptype!(array.dtype().as_ptype(), |T| {
3443
let encoded_val: <T as ALPFloat>::ALPInt =
@@ -39,4 +48,12 @@ impl OperationsVTable<ALP> for ALP {
3948
)
4049
}))
4150
}
51+
52+
fn scalar_at(
53+
array: ArrayView<'_, ALP>,
54+
index: usize,
55+
ctx: &mut ExecutionCtx,
56+
) -> VortexResult<Scalar> {
57+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
58+
}
4259
}

‎encodings/datetime-parts/src/ops.rs‎

Lines changed: 42 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,27 +3,30 @@
33

44
use vortex_array::ArrayView;
55
use vortex_array::ExecutionCtx;
6+
use vortex_array::ProbeState;
67
use vortex_array::dtype::DType;
78
use vortex_array::extension::datetime::Timestamp;
89
use vortex_array::scalar::Scalar;
910
use vortex_array::vtable::OperationsVTable;
1011
use vortex_error::VortexExpect;
1112
use vortex_error::VortexResult;
13+
use vortex_error::vortex_err;
1214
use vortex_error::vortex_panic;
1315

1416
use crate::DateTimeParts;
15-
use crate::array::DateTimePartsArraySlotsExt;
17+
use crate::DateTimePartsSlots;
1618
use crate::timestamp;
1719
use crate::timestamp::TimestampParts;
1820

1921
impl OperationsVTable<DateTimeParts> for DateTimeParts {
2022
type ProbeState = ();
2123

22-
fn scalar_at(
23-
array: ArrayView<'_, DateTimeParts>,
24+
fn probe_scalar(
25+
state: &mut ProbeState<'_, DateTimeParts>,
2426
index: usize,
2527
ctx: &mut ExecutionCtx,
2628
) -> VortexResult<Scalar> {
29+
let array = state.array();
2730
let DType::Extension(ext) = array.dtype().clone() else {
2831
vortex_panic!(
2932
"DateTimePartsArray must have extension dtype, found {}",
@@ -35,28 +38,19 @@ impl OperationsVTable<DateTimeParts> for DateTimeParts {
3538
vortex_panic!(Compute: "must decode TemporalMetadata from extension metadata");
3639
};
3740

38-
if !array.as_ref().is_valid(index, ctx)? {
41+
if !state.is_valid(index, ctx)? {
3942
return Ok(Scalar::null(DType::Extension(ext)));
4043
}
4144

42-
let days: i32 = array
43-
.days()
44-
.execute_scalar(index, ctx)?
45-
.as_primitive()
46-
.as_::<i32>()
47-
.vortex_expect("days fits in i32");
48-
let seconds: i32 = array
49-
.seconds()
50-
.execute_scalar(index, ctx)?
51-
.as_primitive()
52-
.as_::<i32>()
53-
.vortex_expect("seconds fits in i32");
54-
let subseconds: i32 = array
55-
.subseconds()
56-
.execute_scalar(index, ctx)?
57-
.as_primitive()
58-
.as_::<i32>()
59-
.vortex_expect("subseconds fits in i32");
45+
let days = part_at(state, DateTimePartsSlots::DAYS, "days", index, ctx)?;
46+
let seconds = part_at(state, DateTimePartsSlots::SECONDS, "seconds", index, ctx)?;
47+
let subseconds = part_at(
48+
state,
49+
DateTimePartsSlots::SUBSECONDS,
50+
"subseconds",
51+
index,
52+
ctx,
53+
)?;
6054

6155
let ts = timestamp::combine(
6256
TimestampParts {
@@ -72,4 +66,30 @@ impl OperationsVTable<DateTimeParts> for DateTimeParts {
7266
Scalar::primitive(ts, ext.storage_dtype().nullability()),
7367
))
7468
}
69+
70+
fn scalar_at(
71+
array: ArrayView<'_, DateTimeParts>,
72+
index: usize,
73+
ctx: &mut ExecutionCtx,
74+
) -> VortexResult<Scalar> {
75+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
76+
}
77+
}
78+
79+
/// Reads one timestamp part out of `slot`, through the probe so a repeated read keeps the
80+
/// child's preparation.
81+
fn part_at(
82+
state: &mut ProbeState<'_, DateTimeParts>,
83+
slot: usize,
84+
name: &'static str,
85+
index: usize,
86+
ctx: &mut ExecutionCtx,
87+
) -> VortexResult<i32> {
88+
Ok(state
89+
.slot(slot)?
90+
.ok_or_else(|| vortex_err!("DateTimeParts {name} slot is missing"))?
91+
.execute_scalar(index, ctx)?
92+
.as_primitive()
93+
.as_::<i32>()
94+
.vortex_expect("timestamp part fits in i32"))
7595
}

‎encodings/fastlanes/src/for_/tests.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,23 @@ fn drifting_u32(len: u32) -> PrimitiveArray {
146146
PrimitiveArray::from_iter((0..len).map(|i| (i / 1024) * 1_000_000 + i % 100))
147147
}
148148

149+
#[test]
150+
fn repeated_probe_across_sliced_chunks() -> VortexResult<()> {
151+
let mut ctx = SESSION.create_execution_ctx();
152+
let values = drifting_u32(3000);
153+
let encoded = FoR::encode_chunked(values.clone(), &mut ctx)?;
154+
let values = values.into_array();
155+
let sliced = encoded.into_array().slice(1000..2100)?;
156+
let mut probe = sliced.repeated_probe();
157+
for index in [1048, 24, 0, 1099, 23, 1047, 24] {
158+
assert_eq!(
159+
probe.execute_scalar(index, &mut ctx)?,
160+
values.execute_scalar(index + 1000, &mut ctx)?
161+
);
162+
}
163+
Ok(())
164+
}
165+
149166
#[rstest]
150167
#[case::empty(PrimitiveArray::from_iter(Vec::<u32>::new()))]
151168
#[case::one(PrimitiveArray::from_iter([7u32]))]

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

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,28 +3,37 @@
33

44
use vortex_array::ArrayView;
55
use vortex_array::ExecutionCtx;
6+
use vortex_array::ProbeState;
67
use vortex_array::match_each_integer_ptype;
78
use vortex_array::scalar::Scalar;
89
use vortex_array::vtable::OperationsVTable;
910
use vortex_error::VortexExpect;
1011
use vortex_error::VortexResult;
12+
use vortex_error::vortex_err;
1113

1214
use super::FoR;
1315
use crate::FL_CHUNK_SIZE;
1416
use crate::for_::array::FoRArrayExt;
15-
use crate::for_::array::FoRArraySlotsExt;
17+
use crate::for_::array::FoRSlots;
1618
impl OperationsVTable<FoR> for FoR {
1719
type ProbeState = ();
1820

19-
fn scalar_at(
20-
array: ArrayView<'_, FoR>,
21+
fn probe_scalar(
22+
state: &mut ProbeState<'_, FoR>,
2123
index: usize,
2224
ctx: &mut ExecutionCtx,
2325
) -> VortexResult<Scalar> {
24-
let encoded_pvalue = array.encoded().execute_scalar(index, ctx)?;
26+
let array = state.array();
27+
let encoded_pvalue = state
28+
.slot(FoRSlots::ENCODED)?
29+
.ok_or_else(|| vortex_err!("FoR encoded slot is missing"))?
30+
.execute_scalar(index, ctx)?;
2531
let encoded_pvalue = encoded_pvalue.as_primitive();
2632
let chunk = (usize::from(array.offset()) + index) / FL_CHUNK_SIZE;
27-
let reference = array.references().execute_scalar(chunk, ctx)?;
33+
let reference = state
34+
.slot(FoRSlots::REFERENCES)?
35+
.ok_or_else(|| vortex_err!("FoR references slot is missing"))?
36+
.execute_scalar(chunk, ctx)?;
2837
let reference = reference.as_primitive();
2938

3039
Ok(match_each_integer_ptype!(array.ptype(), |P| {
@@ -41,6 +50,14 @@ impl OperationsVTable<FoR> for FoR {
4150
.unwrap_or_else(|| Scalar::null(array.dtype().clone()))
4251
}))
4352
}
53+
54+
fn scalar_at(
55+
array: ArrayView<'_, FoR>,
56+
index: usize,
57+
ctx: &mut ExecutionCtx,
58+
) -> VortexResult<Scalar> {
59+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
60+
}
4461
}
4562

4663
#[cfg(test)]

‎encodings/fsst/src/ops.rs‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@
33

44
use vortex_array::ArrayView;
55
use vortex_array::ExecutionCtx;
6+
use vortex_array::IntoArray;
7+
use vortex_array::ProbeState;
8+
use vortex_array::RepeatedArrayProbe;
69
use vortex_array::arrays::varbin::varbin_scalar;
710
use vortex_array::scalar::Scalar;
811
use vortex_array::vtable::OperationsVTable;
@@ -13,18 +16,43 @@ use vortex_error::VortexResult;
1316
use crate::FSST;
1417
use crate::FSSTArrayExt;
1518

19+
/// The codes array is rebuilt from the codes buffer and the offsets slot on every read, so a
20+
/// repeated probe keeps one probe over it and the offsets child keeps its preparation.
21+
#[derive(Default)]
22+
pub struct FsstProbeState {
23+
codes: Option<RepeatedArrayProbe>,
24+
}
25+
1626
impl OperationsVTable<FSST> for FSST {
17-
type ProbeState = ();
27+
type ProbeState = FsstProbeState;
1828

19-
fn scalar_at(
20-
array: ArrayView<'_, FSST>,
29+
fn probe_scalar(
30+
state: &mut ProbeState<'_, FSST>,
2131
index: usize,
2232
ctx: &mut ExecutionCtx,
2333
) -> VortexResult<Scalar> {
24-
let compressed = array.codes().execute_scalar(index, ctx)?;
34+
if !state.is_valid(index, ctx)? {
35+
return Ok(Scalar::null(state.array().dtype().clone()));
36+
}
37+
let array = state.array();
38+
let compressed = match state.retained() {
39+
None => array.codes().into_array().execute_scalar(index, ctx)?,
40+
Some(retained) => retained
41+
.codes
42+
.get_or_insert_with(|| RepeatedArrayProbe::new(array.codes().into_array()))
43+
.execute_scalar(index, ctx)?,
44+
};
2545
let binary_datum = compressed.as_binary().value().vortex_expect("non-null");
2646

2747
let decoded_buffer = ByteBuffer::from(array.decompressor().decompress(binary_datum));
2848
Ok(varbin_scalar(decoded_buffer, array.dtype()))
2949
}
50+
51+
fn scalar_at(
52+
array: ArrayView<'_, FSST>,
53+
index: usize,
54+
ctx: &mut ExecutionCtx,
55+
) -> VortexResult<Scalar> {
56+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
57+
}
3058
}

‎encodings/zigzag/src/array.rs‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use vortex_array::EqMode;
1616
use vortex_array::ExecutionCtx;
1717
use vortex_array::ExecutionResult;
1818
use vortex_array::IntoArray;
19+
use vortex_array::ProbeState;
1920
use vortex_array::TypedArrayRef;
2021
use vortex_array::array_slots;
2122
use vortex_array::buffer::BufferHandle;
@@ -33,6 +34,7 @@ use vortex_error::VortexExpect;
3334
use vortex_error::VortexResult;
3435
use vortex_error::vortex_bail;
3536
use vortex_error::vortex_ensure;
37+
use vortex_error::vortex_err;
3638
use vortex_error::vortex_panic;
3739
use vortex_session::VortexSession;
3840
use vortex_session::registry::CachedId;
@@ -233,12 +235,16 @@ impl Default for ZigZagData {
233235
impl OperationsVTable<ZigZag> for ZigZag {
234236
type ProbeState = ();
235237

236-
fn scalar_at(
237-
array: ArrayView<'_, ZigZag>,
238+
fn probe_scalar(
239+
state: &mut ProbeState<'_, ZigZag>,
238240
index: usize,
239241
ctx: &mut ExecutionCtx,
240242
) -> VortexResult<Scalar> {
241-
let scalar = array.encoded().execute_scalar(index, ctx)?;
243+
let array = state.array();
244+
let scalar = state
245+
.slot(ZigZagSlots::ENCODED)?
246+
.ok_or_else(|| vortex_err!("ZigZag encoded slot is missing"))?
247+
.execute_scalar(index, ctx)?;
242248
if scalar.is_null() {
243249
return scalar.primitive_reinterpret_cast(ZigZagArrayExt::ptype(&array));
244250
}
@@ -255,6 +261,14 @@ impl OperationsVTable<ZigZag> for ZigZag {
255261
)
256262
}))
257263
}
264+
265+
fn scalar_at(
266+
array: ArrayView<'_, ZigZag>,
267+
index: usize,
268+
ctx: &mut ExecutionCtx,
269+
) -> VortexResult<Scalar> {
270+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
271+
}
258272
}
259273

260274
impl ValidityChild<ZigZag> for ZigZag {

‎vortex-array/src/arrays/chunked/vtable/operations.rs‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,24 +2,39 @@
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

44
use vortex_error::VortexResult;
5+
use vortex_error::vortex_err;
56

67
use crate::ExecutionCtx;
78
use crate::array::ArrayView;
89
use crate::array::OperationsVTable;
10+
use crate::array::ProbeState;
911
use crate::arrays::Chunked;
1012
use crate::arrays::chunked::ChunkedArrayExt;
13+
use crate::arrays::chunked::ChunkedSlots;
1114
use crate::scalar::Scalar;
1215

1316
impl OperationsVTable<Chunked> for Chunked {
1417
type ProbeState = ();
1518

19+
fn probe_scalar(
20+
state: &mut ProbeState<'_, Chunked>,
21+
index: usize,
22+
ctx: &mut ExecutionCtx,
23+
) -> VortexResult<Scalar> {
24+
let (chunk_index, chunk_offset) = state.array().find_chunk_idx(index)?;
25+
let slot = ChunkedSlots::CHUNKS_OFFSET + chunk_index;
26+
state
27+
.slot(slot)?
28+
.ok_or_else(|| vortex_err!("Chunked chunk slot {slot} is missing"))?
29+
.execute_scalar(chunk_offset, ctx)
30+
}
31+
1632
fn scalar_at(
1733
array: ArrayView<'_, Chunked>,
1834
index: usize,
1935
ctx: &mut ExecutionCtx,
2036
) -> VortexResult<Scalar> {
21-
let (chunk_index, chunk_offset) = array.find_chunk_idx(index)?;
22-
array.chunk(chunk_index).execute_scalar(chunk_offset, ctx)
37+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
2338
}
2439
}
2540

0 commit comments

Comments
 (0)