Skip to content

Commit 00a0ccc

Browse files
joseph-isaacsjoseph-isaacs
authored
perf(array): reuse probe state in RLE, RunEnd and PCO (#9844)
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 27a2f10 commit 00a0ccc

13 files changed

Lines changed: 654 additions & 40 deletions

File tree

‎encodings/fastlanes/src/rle/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ pub use array::RLESlots;
99

1010
mod compute;
1111
mod kernel;
12+
mod probe;
1213

1314
mod vtable;
1415
pub use vtable::RLE;

‎encodings/fastlanes/src/rle/probe.rs‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! RLE scalar access through the indices, offsets, and values slots.
5+
6+
use vortex_array::ExecutionCtx;
7+
use vortex_array::ProbeState;
8+
use vortex_array::scalar::Scalar;
9+
use vortex_error::VortexResult;
10+
use vortex_error::vortex_err;
11+
12+
use crate::FL_CHUNK_SIZE;
13+
use crate::RLE;
14+
use crate::rle::RLEArrayExt;
15+
use crate::rle::RLESlots;
16+
17+
/// The slice's base value offset, retained across repeated lookups.
18+
#[derive(Default)]
19+
pub struct RleProbeState {
20+
base: Option<usize>,
21+
}
22+
23+
pub(crate) fn scalar_at(
24+
state: &mut ProbeState<'_, RLE>,
25+
index: usize,
26+
ctx: &mut ExecutionCtx,
27+
) -> VortexResult<Scalar> {
28+
let array = state.array();
29+
let logical_index = array.offset() + index;
30+
// Hold the retained state while reading children: the base offset is cached across rows.
31+
let (mut retained, mut children) = state.split();
32+
let code = children
33+
.slot(RLESlots::INDICES)?
34+
.ok_or_else(|| vortex_err!("RLE indices slot is missing"))?
35+
.execute_scalar(logical_index, ctx)?;
36+
let Some(code) = code.as_primitive().as_::<usize>() else {
37+
return Ok(Scalar::null(array.dtype().clone()));
38+
};
39+
40+
let chunk = logical_index / FL_CHUNK_SIZE;
41+
let offset = if chunk == 0 {
42+
0
43+
} else {
44+
let base = match retained.as_deref().and_then(|state| state.base) {
45+
Some(base) => base,
46+
None => {
47+
let value = read_offset(
48+
children
49+
.slot(RLESlots::VALUES_IDX_OFFSETS)?
50+
.ok_or_else(|| vortex_err!("RLE offsets slot is missing"))?
51+
.execute_scalar(0, ctx)?,
52+
)?;
53+
if let Some(state) = retained.as_mut() {
54+
state.base = Some(value);
55+
}
56+
value
57+
}
58+
};
59+
read_offset(
60+
children
61+
.slot(RLESlots::VALUES_IDX_OFFSETS)?
62+
.ok_or_else(|| vortex_err!("RLE offsets slot is missing"))?
63+
.execute_scalar(chunk, ctx)?,
64+
)?
65+
.checked_sub(base)
66+
.ok_or_else(|| vortex_err!("RLE offsets precede the slice base"))?
67+
};
68+
let value_index = offset
69+
.checked_add(code)
70+
.ok_or_else(|| vortex_err!("RLE value index overflow"))?;
71+
let scalar = children
72+
.slot(RLESlots::VALUES)?
73+
.ok_or_else(|| vortex_err!("RLE values slot is missing"))?
74+
.execute_scalar(value_index, ctx)?;
75+
Scalar::try_new(array.dtype().clone(), scalar.into_value())
76+
}
77+
78+
fn read_offset(scalar: Scalar) -> VortexResult<usize> {
79+
scalar
80+
.as_primitive()
81+
.as_::<usize>()
82+
.ok_or_else(|| vortex_err!("RLE offset must be a non-null usize"))
83+
}
84+
85+
#[cfg(test)]
86+
mod tests;

‎encodings/fastlanes/src/rle/probe/tests.rs‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use rstest::rstest;
5+
use vortex_array::ArrayProbe;
6+
use vortex_array::ArrayRef;
7+
use vortex_array::IntoArray;
8+
use vortex_array::RepeatedArrayProbe;
9+
use vortex_array::VortexSessionExecute;
10+
use vortex_array::arrays::PrimitiveArray;
11+
use vortex_array::assert_arrays_eq;
12+
use vortex_array::builders::builder_with_capacity_in;
13+
use vortex_array::builtins::ArrayBuiltins;
14+
use vortex_array::scalar_fn::fns::operators::Operator;
15+
use vortex_array::validity::Validity;
16+
use vortex_error::VortexResult;
17+
18+
use crate::RLE;
19+
use crate::RLEData;
20+
21+
/// A one-off or retained probe over `array`, the retained one living in `retained`.
22+
fn probe_for<'a>(
23+
array: &'a ArrayRef,
24+
retained: &'a mut Option<RepeatedArrayProbe>,
25+
repeated: bool,
26+
) -> ArrayProbe<'a> {
27+
if repeated {
28+
retained.insert(array.repeated_probe()).as_probe()
29+
} else {
30+
array.probe()
31+
}
32+
}
33+
34+
#[rstest]
35+
fn random_access_across_chunks_and_nulls(
36+
#[values(false, true)] repeated: bool,
37+
#[values(false, true)] sliced: bool,
38+
) -> VortexResult<()> {
39+
let mut ctx = crate::test::SESSION.create_execution_ctx();
40+
let input =
41+
PrimitiveArray::from_option_iter((0..8192u32).map(|i| (i % 11 != 0).then_some(i / 16)));
42+
let encoded = RLEData::encode(input.as_view(), &mut ctx)?.into_array();
43+
let range = if sliced { 1777..7333 } else { 0..8192 };
44+
let source = if sliced {
45+
encoded
46+
.slice(range.clone())?
47+
.execute::<ArrayRef>(&mut ctx)?
48+
} else {
49+
encoded
50+
};
51+
assert!(source.is::<RLE>());
52+
let input = input.slice(range)?;
53+
let indices = [0u32, 1, 1023, 1024, 2048, 2047, 4097, 11, 33, 17, 0];
54+
let mut actual = builder_with_capacity_in(source.dtype(), indices.len(), ctx.allocator());
55+
let mut retained = None;
56+
let mut probe = probe_for(&source, &mut retained, repeated);
57+
for index in indices {
58+
actual.append_scalar(&probe.execute_scalar(index as usize, &mut ctx)?)?;
59+
}
60+
assert_arrays_eq!(
61+
actual.finish(),
62+
input.take(PrimitiveArray::from_iter(indices).into_array())?,
63+
&mut ctx
64+
);
65+
assert!(probe.execute_scalar(source.len(), &mut ctx).is_err());
66+
Ok(())
67+
}
68+
69+
#[rstest]
70+
fn lazy_validity_does_not_evaluate_unrequested_rows(
71+
#[values(false, true)] repeated: bool,
72+
) -> VortexResult<()> {
73+
let mut ctx = crate::test::SESSION.create_execution_ctx();
74+
let numerators = PrimitiveArray::from_iter(vec![1u32; 1024]).into_array();
75+
let denominators =
76+
PrimitiveArray::from_iter((0..1024).map(|i| u32::from(i != 1023))).into_array();
77+
let validity = numerators
78+
.binary(denominators, Operator::Div)?
79+
.binary(numerators, Operator::Eq)?;
80+
let array = RLE::try_new(
81+
PrimitiveArray::from_iter([42u32]).into_array(),
82+
PrimitiveArray::new(vec![0u16; 1024], Validity::Array(validity)).into_array(),
83+
PrimitiveArray::from_iter([0u64]).into_array(),
84+
0,
85+
1024,
86+
)?
87+
.into_array();
88+
let expected = array.execute_scalar(0, &mut ctx)?;
89+
let mut retained = None;
90+
let mut probe = probe_for(&array, &mut retained, repeated);
91+
assert_eq!(probe.execute_scalar(0, &mut ctx)?, expected);
92+
assert!(probe.execute_scalar(1023, &mut ctx).is_err());
93+
Ok(())
94+
}

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

Lines changed: 14 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,42 +3,32 @@
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;
8-
use vortex_error::VortexExpect;
99
use vortex_error::VortexResult;
1010

1111
use super::RLE;
12-
use crate::FL_CHUNK_SIZE;
13-
use crate::rle::RLEArrayExt;
14-
use crate::rle::RLEArraySlotsExt;
12+
use crate::rle::probe;
13+
use crate::rle::probe::RleProbeState;
1514

1615
impl OperationsVTable<RLE> for RLE {
17-
type ProbeState = ();
16+
type ProbeState = RleProbeState;
17+
18+
fn probe_scalar(
19+
state: &mut ProbeState<'_, RLE>,
20+
index: usize,
21+
ctx: &mut ExecutionCtx,
22+
) -> VortexResult<Scalar> {
23+
probe::scalar_at(state, index, ctx)
24+
}
1825

1926
fn scalar_at(
2027
array: ArrayView<'_, RLE>,
2128
index: usize,
2229
ctx: &mut ExecutionCtx,
2330
) -> VortexResult<Scalar> {
24-
let offset_in_chunk = array.offset();
25-
let chunk_relative_idx = array
26-
.indices()
27-
.execute_scalar(offset_in_chunk + index, ctx)?;
28-
29-
let chunk_relative_idx = chunk_relative_idx
30-
.as_primitive()
31-
.as_::<usize>()
32-
.vortex_expect("Index must not be null");
33-
34-
let chunk_id = (offset_in_chunk + index) / FL_CHUNK_SIZE;
35-
let value_idx_offset = array.values_idx_offset(chunk_id, ctx);
36-
37-
let scalar = array
38-
.values()
39-
.execute_scalar(value_idx_offset + chunk_relative_idx, ctx)?;
40-
41-
Scalar::try_new(array.dtype().clone(), scalar.into_value())
31+
probe::scalar_at(&mut ProbeState::once(array), index, ctx)
4232
}
4333
}
4434

@@ -54,9 +44,9 @@ mod tests {
5444
use vortex_array::validity::Validity;
5545
use vortex_buffer::Buffer;
5646
use vortex_buffer::buffer;
47+
use vortex_error::VortexExpect;
5748
use vortex_session::VortexSession;
5849

59-
use super::*;
6050
use crate::RLE;
6151
use crate::RLEArray;
6252
use crate::RLEData;

‎encodings/pco/benches/scalar.rs‎

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

4-
//! Scalar reads out of a PCO array. Cases are `(access_count, nullable, scattered)`; clustered
5-
//! indices stay inside one page so a retained decode can be reused, scattered ones cross pages.
4+
//! Scalar reads out of a PCO array through one repeated probe. Cases are
5+
//! `(access_count, nullable, scattered)`; clustered indices stay inside one page so the retained
6+
//! decode is reused, scattered ones cross pages.
67
78
use std::sync::LazyLock;
89

@@ -69,11 +70,17 @@ fn scalar_access(bencher: Bencher, (count, nullable, scattered): (usize, bool, b
6970
let array = pco(nullable);
7071
let indices = indices(count, scattered);
7172
bencher
72-
.with_inputs(|| (SESSION.create_execution_ctx(), Vec::with_capacity(count)))
73-
.bench_refs(|(ctx, scalars)| {
73+
.with_inputs(|| {
74+
(
75+
SESSION.create_execution_ctx(),
76+
array.repeated_probe(),
77+
Vec::with_capacity(count),
78+
)
79+
})
80+
.bench_refs(|(ctx, probe, scalars)| {
7481
for &index in &indices {
7582
scalars.push(
76-
array
83+
probe
7784
.execute_scalar(index, ctx)
7885
.vortex_expect("scalar access"),
7986
);

‎encodings/pco/src/array.rs‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ use vortex_array::EqMode;
2929
use vortex_array::ExecutionCtx;
3030
use vortex_array::ExecutionResult;
3131
use vortex_array::IntoArray;
32+
use vortex_array::ProbeState;
3233
use vortex_array::TypedArrayRef;
3334
use vortex_array::array_slots;
3435
use vortex_array::arrays::Primitive;
@@ -778,19 +779,23 @@ impl ValidityVTable<Pco> for Pco {
778779
}
779780

780781
impl OperationsVTable<Pco> for Pco {
781-
type ProbeState = ();
782+
type ProbeState = crate::probe::PcoProbeState;
783+
784+
fn probe_scalar(
785+
state: &mut ProbeState<'_, Pco>,
786+
index: usize,
787+
ctx: &mut ExecutionCtx,
788+
) -> VortexResult<Scalar> {
789+
let array = state.array();
790+
crate::probe::scalar_at(array, index, state.retained(), ctx)
791+
}
782792

783793
fn scalar_at(
784794
array: ArrayView<'_, Pco>,
785795
index: usize,
786796
ctx: &mut ExecutionCtx,
787797
) -> VortexResult<Scalar> {
788-
let unsliced_validity = array.unsliced_validity();
789-
array
790-
._slice(index, index + 1)
791-
.decompress(&unsliced_validity, ctx)?
792-
.into_array()
793-
.execute_scalar(0, ctx)
798+
crate::probe::scalar_at(array, index, None, ctx)
794799
}
795800
}
796801

‎encodings/pco/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
2121
mod array;
2222
mod compute;
23+
mod probe;
2324
mod rules;
2425
mod slice;
2526

0 commit comments

Comments
 (0)