Skip to content

Commit e9ad244

Browse files
jiengupjiengup
committed
feat(array): migrate simple encodings to scalar probes
Signed-off-by: Gunther Xing <jiengup@gmail.com>
1 parent ed119bf commit e9ad244

9 files changed

Lines changed: 129 additions & 22 deletions

File tree

‎encodings/bytebool/src/array.rs‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use vortex_array::EqMode;
1818
use vortex_array::ExecutionCtx;
1919
use vortex_array::ExecutionResult;
2020
use vortex_array::IntoArray;
21+
use vortex_array::ProbeState;
2122
use vortex_array::TypedArrayRef;
2223
use vortex_array::array_slots;
2324
use vortex_array::arrays::BoolArray;
@@ -308,16 +309,28 @@ impl ValidityVTable<ByteBool> for ByteBool {
308309
impl OperationsVTable<ByteBool> for ByteBool {
309310
type ProbeState = ();
310311

311-
fn scalar_at(
312-
array: ArrayView<'_, ByteBool>,
312+
fn probe_scalar(
313+
state: &mut ProbeState<'_, ByteBool>,
313314
index: usize,
314-
_ctx: &mut ExecutionCtx,
315+
ctx: &mut ExecutionCtx,
315316
) -> VortexResult<Scalar> {
317+
let array = state.array();
318+
if !state.is_valid(index, ctx)? {
319+
return Ok(Scalar::null(array.dtype().clone()));
320+
}
316321
Ok(Scalar::bool(
317322
array.buffer.as_host()[index] == 1,
318323
array.dtype().nullability(),
319324
))
320325
}
326+
327+
fn scalar_at(
328+
array: ArrayView<'_, ByteBool>,
329+
index: usize,
330+
ctx: &mut ExecutionCtx,
331+
) -> VortexResult<Scalar> {
332+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
333+
}
321334
}
322335

323336
#[cfg(test)]

‎encodings/sequence/src/array.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use vortex_array::ArrayView;
2020
use vortex_array::EqMode;
2121
use vortex_array::ExecutionCtx;
2222
use vortex_array::ExecutionResult;
23+
use vortex_array::ProbeState;
2324
use vortex_array::buffer::BufferHandle;
2425
use vortex_array::dtype::DType;
2526
use vortex_array::dtype::NativePType;
@@ -434,16 +435,25 @@ impl VTable for Sequence {
434435
impl OperationsVTable<Sequence> for Sequence {
435436
type ProbeState = ();
436437

437-
fn scalar_at(
438-
array: ArrayView<'_, Sequence>,
438+
fn probe_scalar(
439+
state: &mut ProbeState<'_, Sequence>,
439440
index: usize,
440441
_ctx: &mut ExecutionCtx,
441442
) -> VortexResult<Scalar> {
443+
let array = state.array();
442444
Scalar::try_new(
443445
array.dtype().clone(),
444446
Some(ScalarValue::Primitive(array.index_value(index))),
445447
)
446448
}
449+
450+
fn scalar_at(
451+
array: ArrayView<'_, Sequence>,
452+
index: usize,
453+
ctx: &mut ExecutionCtx,
454+
) -> VortexResult<Scalar> {
455+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
456+
}
447457
}
448458

449459
impl ValidityVTable<Sequence> for Sequence {

‎vortex-array/benches/probe.rs‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,14 @@ use vortex_array::ArrayRef;
1313
use vortex_array::IntoArray;
1414
use vortex_array::VortexSessionExecute;
1515
use vortex_array::array_session;
16+
use vortex_array::arrays::BoolArray;
1617
use vortex_array::arrays::DictArray;
1718
use vortex_array::arrays::PrimitiveArray;
1819
use vortex_array::arrays::StructArray;
20+
use vortex_array::arrays::VarBinViewArray;
21+
use vortex_array::dtype::DType;
1922
use vortex_array::dtype::FieldNames;
23+
use vortex_array::dtype::Nullability;
2024
use vortex_array::validity::Validity;
2125
use vortex_buffer::Buffer;
2226
use vortex_error::VortexExpect;
@@ -78,13 +82,27 @@ fn indices() -> Vec<usize> {
7882
.collect()
7983
}
8084

81-
const ARRAYS: &[&str] = &["primitive", "dict", "struct"];
85+
const ARRAYS: &[&str] = &["primitive", "dict", "struct", "bool", "varbinview"];
8286

8387
fn array(name: &str) -> ArrayRef {
8488
match name {
8589
"primitive" => nullable_primitive(LEN),
8690
"dict" => dict_over_nullable_values(),
8791
"struct" => nullable_struct(),
92+
"bool" => BoolArray::from_iter((0..LEN).map(|i| (i % 7 != 0).then_some(i % 2 == 0)))
93+
.into_array(),
94+
"varbinview" => VarBinViewArray::from_iter(
95+
[
96+
Some("short"),
97+
None,
98+
Some("this is a much longer string to force outlining"),
99+
]
100+
.into_iter()
101+
.cycle()
102+
.take(LEN),
103+
DType::Utf8(Nullability::Nullable),
104+
)
105+
.into_array(),
88106
_ => unreachable!("unknown fixture {name}"),
89107
}
90108
}

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,23 +6,36 @@ use vortex_error::VortexResult;
66
use crate::ExecutionCtx;
77
use crate::array::ArrayView;
88
use crate::array::OperationsVTable;
9+
use crate::array::ProbeState;
910
use crate::arrays::Bool;
1011
use crate::arrays::bool::BoolArrayExt;
1112
use crate::scalar::Scalar;
1213

1314
impl OperationsVTable<Bool> for Bool {
1415
type ProbeState = ();
1516

16-
fn scalar_at(
17-
array: ArrayView<'_, Bool>,
17+
fn probe_scalar(
18+
state: &mut ProbeState<'_, Bool>,
1819
index: usize,
19-
_ctx: &mut ExecutionCtx,
20+
ctx: &mut ExecutionCtx,
2021
) -> VortexResult<Scalar> {
22+
let array = state.array();
23+
if !state.is_valid(index, ctx)? {
24+
return Ok(Scalar::null(array.dtype().clone()));
25+
}
2126
Ok(Scalar::bool(
2227
array.bit_buffer_view().value(index),
2328
array.dtype().nullability(),
2429
))
2530
}
31+
32+
fn scalar_at(
33+
array: ArrayView<'_, Bool>,
34+
index: usize,
35+
ctx: &mut ExecutionCtx,
36+
) -> VortexResult<Scalar> {
37+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
38+
}
2639
}
2740

2841
#[cfg(test)]

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

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,27 @@ use vortex_error::VortexResult;
66
use crate::ExecutionCtx;
77
use crate::array::ArrayView;
88
use crate::array::OperationsVTable;
9+
use crate::array::ProbeState;
910
use crate::arrays::Constant;
1011
use crate::scalar::Scalar;
1112

1213
impl OperationsVTable<Constant> for Constant {
1314
type ProbeState = ();
1415

15-
fn scalar_at(
16-
array: ArrayView<'_, Constant>,
16+
fn probe_scalar(
17+
state: &mut ProbeState<'_, Constant>,
1718
_index: usize,
1819
_ctx: &mut ExecutionCtx,
1920
) -> VortexResult<Scalar> {
20-
Ok(array.scalar.clone())
21+
Ok(state.array().scalar.clone())
22+
}
23+
24+
fn scalar_at(
25+
array: ArrayView<'_, Constant>,
26+
index: usize,
27+
ctx: &mut ExecutionCtx,
28+
) -> VortexResult<Scalar> {
29+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
2130
}
2231
}
2332

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use vortex_error::VortexResult;
66
use crate::ExecutionCtx;
77
use crate::array::ArrayView;
88
use crate::array::OperationsVTable;
9+
use crate::array::ProbeState;
910
use crate::arrays::Decimal;
1011
use crate::match_each_decimal_value_type;
1112
use crate::scalar::DecimalValue;
@@ -14,11 +15,15 @@ use crate::scalar::Scalar;
1415
impl OperationsVTable<Decimal> for Decimal {
1516
type ProbeState = ();
1617

17-
fn scalar_at(
18-
array: ArrayView<'_, Decimal>,
18+
fn probe_scalar(
19+
state: &mut ProbeState<'_, Decimal>,
1920
index: usize,
20-
_ctx: &mut ExecutionCtx,
21+
ctx: &mut ExecutionCtx,
2122
) -> VortexResult<Scalar> {
23+
let array = state.array();
24+
if !state.is_valid(index, ctx)? {
25+
return Ok(Scalar::null(array.dtype().clone()));
26+
}
2227
Ok(match_each_decimal_value_type!(array.values_type(), |D| {
2328
Scalar::decimal(
2429
DecimalValue::from(array.buffer::<D>()[index]),
@@ -27,6 +32,14 @@ impl OperationsVTable<Decimal> for Decimal {
2732
)
2833
}))
2934
}
35+
36+
fn scalar_at(
37+
array: ArrayView<'_, Decimal>,
38+
index: usize,
39+
ctx: &mut ExecutionCtx,
40+
) -> VortexResult<Scalar> {
41+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
42+
}
3043
}
3144

3245
#[cfg(test)]

‎vortex-array/src/arrays/null/mod.rs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use crate::array::ArrayParts;
1919
use crate::array::ArrayView;
2020
use crate::array::EmptyArrayData;
2121
use crate::array::OperationsVTable;
22+
use crate::array::ProbeState;
2223
use crate::array::VTable;
2324
use crate::array::ValidityVTable;
2425
use crate::array::with_empty_buffers;
@@ -186,13 +187,21 @@ impl Array<Null> {
186187
impl OperationsVTable<Null> for Null {
187188
type ProbeState = ();
188189

189-
fn scalar_at(
190-
_array: ArrayView<'_, Null>,
190+
fn probe_scalar(
191+
_state: &mut ProbeState<'_, Null>,
191192
_index: usize,
192193
_ctx: &mut ExecutionCtx,
193194
) -> VortexResult<Scalar> {
194195
Ok(Scalar::null(DType::Null))
195196
}
197+
198+
fn scalar_at(
199+
array: ArrayView<'_, Null>,
200+
index: usize,
201+
ctx: &mut ExecutionCtx,
202+
) -> VortexResult<Scalar> {
203+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
204+
}
196205
}
197206

198207
impl ValidityVTable<Null> for Null {

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,33 @@ use vortex_error::VortexResult;
66
use crate::ExecutionCtx;
77
use crate::array::ArrayView;
88
use crate::array::OperationsVTable;
9+
use crate::array::ProbeState;
910
use crate::arrays::VarBinView;
1011
use crate::arrays::varbin::varbin_scalar_unchecked;
1112
use crate::scalar::Scalar;
1213

1314
impl OperationsVTable<VarBinView> for VarBinView {
1415
type ProbeState = ();
1516

16-
fn scalar_at(
17-
array: ArrayView<'_, VarBinView>,
17+
fn probe_scalar(
18+
state: &mut ProbeState<'_, VarBinView>,
1819
index: usize,
19-
_ctx: &mut ExecutionCtx,
20+
ctx: &mut ExecutionCtx,
2021
) -> VortexResult<Scalar> {
22+
let array = state.array();
23+
if !state.is_valid(index, ctx)? {
24+
return Ok(Scalar::null(array.dtype().clone()));
25+
}
2126
// SAFETY: `VarBinViewArray` has a Utf8 or Binary dtype, and it checks that all Utf8 data
2227
// is valid UTF-8.
2328
Ok(unsafe { varbin_scalar_unchecked(array.bytes_at(index), array.dtype()) })
2429
}
30+
31+
fn scalar_at(
32+
array: ArrayView<'_, VarBinView>,
33+
index: usize,
34+
ctx: &mut ExecutionCtx,
35+
) -> VortexResult<Scalar> {
36+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
37+
}
2538
}

‎vortex-python/src/arrays/py/vtable.rs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use vortex::array::EqMode;
1515
use vortex::array::ExecutionCtx;
1616
use vortex::array::ExecutionResult;
1717
use vortex::array::OperationsVTable;
18+
use vortex::array::ProbeState;
1819
use vortex::array::VTable;
1920
use vortex::array::ValidityVTable;
2021
use vortex::array::buffer::BufferHandle;
@@ -124,13 +125,21 @@ impl VTable for PythonVTable {
124125
impl OperationsVTable<PythonVTable> for PythonVTable {
125126
type ProbeState = ();
126127

127-
fn scalar_at(
128-
_array: ArrayView<'_, PythonVTable>,
128+
fn probe_scalar(
129+
_state: &mut ProbeState<'_, PythonVTable>,
129130
_index: usize,
130131
_ctx: &mut ExecutionCtx,
131132
) -> VortexResult<Scalar> {
132133
vortex_bail!("PythonArray scalar_at is not supported");
133134
}
135+
136+
fn scalar_at(
137+
array: ArrayView<'_, PythonVTable>,
138+
index: usize,
139+
ctx: &mut ExecutionCtx,
140+
) -> VortexResult<Scalar> {
141+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
142+
}
134143
}
135144

136145
impl ValidityVTable<PythonVTable> for PythonVTable {

0 commit comments

Comments
 (0)