Skip to content

Commit 7477dcb

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

16 files changed

Lines changed: 480 additions & 88 deletions

File tree

‎encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs‎

Lines changed: 37 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use vortex_array::ArrayView;
1717
use vortex_array::EqMode;
1818
use vortex_array::ExecutionCtx;
1919
use vortex_array::ExecutionResult;
20+
use vortex_array::ProbeState;
2021
use vortex_array::TypedArrayRef;
2122
use vortex_array::array_slots;
2223
use vortex_array::arrays::Primitive;
@@ -357,32 +358,44 @@ impl<T: TypedArrayRef<DecimalByteParts>> DecimalBytePartsArrayExt for T {}
357358
impl OperationsVTable<DecimalByteParts> for DecimalByteParts {
358359
type ProbeState = ();
359360

360-
fn scalar_at(
361-
array: ArrayView<'_, DecimalByteParts>,
361+
fn probe_scalar(
362+
state: &mut ProbeState<'_, DecimalByteParts>,
362363
index: usize,
363364
ctx: &mut ExecutionCtx,
364365
) -> VortexResult<Scalar> {
365-
let scalar = array.msp().execute_scalar(index, ctx)?;
366+
let array = state.array();
367+
let scalar = state
368+
.slot(DecimalBytePartsSlots::MSP)?
369+
.vortex_expect("validated DecimalByteParts MSP slot")
370+
.execute_scalar(index, ctx)?;
371+
if scalar.is_null() {
372+
return Ok(Scalar::null(array.dtype().clone()));
373+
}
366374

367375
// Widen the MSP's signed value (i8/i16/i32/i64) to i64 for scalar reconstruction.
368376
// The array retains its original MSP storage type.
369377
let primitive_scalar = scalar.as_primitive();
370378
let msp = primitive_scalar.as_::<i64>().vortex_expect("non-null");
371379

372380
// Zero-extend each narrowed lower value to its 64-bit window.
373-
let lower_parts = array
374-
.lower_parts()
375-
.iter()
376-
.map(|part| {
377-
Ok(part
378-
.execute_scalar(index, ctx)?
379-
.as_primitive()
380-
.as_::<u64>()
381-
.vortex_expect("lower parts are non-nullable"))
382-
})
383-
.collect::<VortexResult<Vec<_>>>()?;
381+
let lower_part_count = array.slots().len() - DecimalBytePartsSlots::LOWER_PARTS_OFFSET;
382+
if lower_part_count > MAX_LOWER_PARTS {
383+
vortex_bail!(
384+
"at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}"
385+
);
386+
}
387+
let mut lower_parts = [0u64; MAX_LOWER_PARTS];
388+
for (part_idx, lower_part) in lower_parts.iter_mut().take(lower_part_count).enumerate() {
389+
*lower_part = state
390+
.slot(DecimalBytePartsSlots::LOWER_PARTS_OFFSET + part_idx)?
391+
.vortex_expect("validated DecimalByteParts lower-part slot")
392+
.execute_scalar(index, ctx)?
393+
.as_primitive()
394+
.as_::<u64>()
395+
.vortex_expect("lower parts are non-nullable");
396+
}
384397

385-
let value = match lower_parts.as_slice() {
398+
let value = match &lower_parts[..lower_part_count] {
386399
[] => DecimalValue::I64(msp),
387400
[first] => DecimalValue::I128(assemble_wide_decimal_value(msp, [*first])),
388401
[first, second] => {
@@ -393,12 +406,20 @@ impl OperationsVTable<DecimalByteParts> for DecimalByteParts {
393406
}
394407
_ => vortex_bail!(
395408
"at most {MAX_LOWER_PARTS} lower parts are supported, got {}",
396-
lower_parts.len()
409+
lower_part_count
397410
),
398411
};
399412

400413
Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value)))
401414
}
415+
416+
fn scalar_at(
417+
array: ArrayView<'_, DecimalByteParts>,
418+
index: usize,
419+
ctx: &mut ExecutionCtx,
420+
) -> VortexResult<Scalar> {
421+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
422+
}
402423
}
403424

404425
impl ValidityChild<DecimalByteParts> for DecimalByteParts {

‎vortex-array/benches/probe.rs‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ use vortex_array::VortexSessionExecute;
1515
use vortex_array::array_session;
1616
use vortex_array::arrays::BoolArray;
1717
use vortex_array::arrays::DictArray;
18+
use vortex_array::arrays::InterleaveArray;
19+
use vortex_array::arrays::ListArray;
1820
use vortex_array::arrays::PrimitiveArray;
1921
use vortex_array::arrays::StructArray;
2022
use vortex_array::arrays::VarBinViewArray;
@@ -72,6 +74,38 @@ fn nullable_struct() -> ArrayRef {
7274
.into_array()
7375
}
7476

77+
fn nullable_list() -> ArrayRef {
78+
let offsets = (0..=LEN)
79+
.map(|i| u32::try_from(i * 3).vortex_expect("list offsets fit u32"))
80+
.collect::<Buffer<u32>>()
81+
.into_array();
82+
ListArray::try_new(
83+
nullable_primitive(LEN * 3),
84+
offsets,
85+
Validity::from_iter((0..LEN).map(|i| i % 11 != 0)),
86+
)
87+
.vortex_expect("list fixture")
88+
.into_array()
89+
}
90+
91+
fn nullable_interleave() -> ArrayRef {
92+
let array_indices = (0..LEN)
93+
.map(|i| u8::try_from(i % 2).vortex_expect("branch index fits u8"))
94+
.collect::<Buffer<u8>>()
95+
.into_array();
96+
let row_indices = (0..LEN)
97+
.map(|i| u32::try_from((i * 7919) % LEN).vortex_expect("row index fits u32"))
98+
.collect::<Buffer<u32>>()
99+
.into_array();
100+
InterleaveArray::try_new(
101+
vec![nullable_primitive(LEN), nullable_primitive(LEN)],
102+
array_indices,
103+
row_indices,
104+
)
105+
.vortex_expect("interleave fixture")
106+
.into_array()
107+
}
108+
75109
fn indices() -> Vec<usize> {
76110
let mut seed = 42u64;
77111
(0..ACCESSES)
@@ -82,11 +116,21 @@ fn indices() -> Vec<usize> {
82116
.collect()
83117
}
84118

85-
const ARRAYS: &[&str] = &["primitive", "dict", "struct", "bool", "varbinview"];
119+
const ARRAYS: &[&str] = &[
120+
"primitive",
121+
"dict",
122+
"struct",
123+
"bool",
124+
"varbinview",
125+
"list",
126+
"interleave",
127+
];
86128

87129
fn array(name: &str) -> ArrayRef {
88130
match name {
89131
"primitive" => nullable_primitive(LEN),
132+
"list" => nullable_list(),
133+
"interleave" => nullable_interleave(),
90134
"dict" => dict_over_nullable_values(),
91135
"struct" => nullable_struct(),
92136
"bool" => BoolArray::from_iter((0..LEN).map(|i| (i % 7 != 0).then_some(i % 2 == 0)))

‎vortex-array/src/arrays/filter/vtable.rs‎

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use std::hash::Hasher;
66
use vortex_error::VortexExpect;
77
use vortex_error::VortexResult;
88
use vortex_error::vortex_bail;
9+
use vortex_error::vortex_err;
910
use vortex_error::vortex_ensure;
1011
use vortex_error::vortex_ensure_eq;
1112
use vortex_error::vortex_panic;
@@ -26,6 +27,7 @@ use crate::array::Array;
2627
use crate::array::ArrayId;
2728
use crate::array::ArrayView;
2829
use crate::array::OperationsVTable;
30+
use crate::array::ProbeState;
2931
use crate::array::VTable;
3032
use crate::array::ValidityVTable;
3133
use crate::array::with_empty_buffers;
@@ -200,13 +202,24 @@ impl VTable for Filter {
200202
impl OperationsVTable<Filter> for Filter {
201203
type ProbeState = ();
202204

205+
fn probe_scalar(
206+
state: &mut ProbeState<'_, Filter>,
207+
index: usize,
208+
ctx: &mut ExecutionCtx,
209+
) -> VortexResult<Scalar> {
210+
let rank_idx = state.array().mask.rank(index);
211+
state
212+
.slot(FilterSlots::CHILD)?
213+
.ok_or_else(|| vortex_err!("Filter child slot is missing"))?
214+
.execute_scalar(rank_idx, ctx)
215+
}
216+
203217
fn scalar_at(
204218
array: ArrayView<'_, Filter>,
205219
index: usize,
206220
ctx: &mut ExecutionCtx,
207221
) -> VortexResult<Scalar> {
208-
let rank_idx = array.mask.rank(index);
209-
array.child().execute_scalar(rank_idx, ctx)
222+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
210223
}
211224
}
212225

@@ -215,3 +228,34 @@ impl ValidityVTable<Filter> for Filter {
215228
array.child().validity()?.filter(&array.mask)
216229
}
217230
}
231+
232+
#[cfg(test)]
233+
mod tests {
234+
use vortex_error::VortexResult;
235+
use vortex_mask::Mask;
236+
237+
use crate::IntoArray;
238+
use crate::VortexSessionExecute;
239+
use crate::array_session;
240+
use crate::arrays::FilterArray;
241+
use crate::arrays::PrimitiveArray;
242+
use crate::compute::conformance::consistency::test_scalar_probe_consistency;
243+
244+
#[test]
245+
fn scalar_probe_with_filtered_nulls() -> VortexResult<()> {
246+
let child = PrimitiveArray::from_option_iter([
247+
Some(10i32),
248+
None,
249+
Some(30),
250+
None,
251+
Some(50),
252+
Some(60),
253+
])
254+
.into_array();
255+
let mask = Mask::from_iter([false, true, true, false, true, true]);
256+
let filtered = FilterArray::try_new(child, mask)?.into_array();
257+
258+
test_scalar_probe_consistency(&filtered, &mut array_session().create_execution_ctx());
259+
Ok(())
260+
}
261+
}

‎vortex-array/src/arrays/fixed_size_list/tests/basic.rs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,41 @@
44
use std::sync::Arc;
55

66
use vortex_buffer::buffer;
7+
use vortex_error::VortexResult;
78

9+
use super::common::create_basic_fsl;
10+
use super::common::create_nullable_fsl;
811
use crate::IntoArray;
912
use crate::VortexSessionExecute;
1013
use crate::array_session;
1114
use crate::arrays::FixedSizeListArray;
15+
use crate::arrays::PrimitiveArray;
1216
use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
1317
use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
18+
use crate::compute::conformance::consistency::test_scalar_probe_consistency;
1419
use crate::dtype::DType;
1520
use crate::dtype::Nullability;
1621
use crate::dtype::PType;
1722
use crate::scalar::Scalar;
1823
use crate::validity::Validity;
1924

25+
#[test]
26+
fn test_fixed_size_list_scalar_probe() -> VortexResult<()> {
27+
let empty_lists = FixedSizeListArray::new(
28+
PrimitiveArray::empty::<i32>(Nullability::NonNullable).into_array(),
29+
0,
30+
Validity::NonNullable,
31+
3,
32+
);
33+
let mut ctx = array_session().create_execution_ctx();
34+
for array in [create_basic_fsl(), create_nullable_fsl(), empty_lists] {
35+
let array = array.into_array();
36+
test_scalar_probe_consistency(&array, &mut ctx);
37+
test_scalar_probe_consistency(&array.slice(1..array.len())?, &mut ctx);
38+
}
39+
Ok(())
40+
}
41+
2042
#[test]
2143
fn test_basic_fixed_size_list() {
2244
let len = 4;

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

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,34 +2,54 @@
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::FixedSizeList;
10-
use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
12+
use crate::arrays::fixed_size_list::FixedSizeListSlots;
13+
use crate::dtype::DType;
1114
use crate::scalar::Scalar;
1215

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

16-
fn scalar_at(
17-
array: ArrayView<'_, FixedSizeList>,
19+
fn probe_scalar(
20+
state: &mut ProbeState<'_, FixedSizeList>,
1821
index: usize,
1922
ctx: &mut ExecutionCtx,
2023
) -> VortexResult<Scalar> {
21-
// By the preconditions we know that the list scalar is not null.
22-
let list = array.fixed_size_list_elements_at(index)?;
23-
let children_elements: Vec<Scalar> = (0..list.len())
24-
.map(|i| list.execute_scalar(i, ctx))
25-
.collect::<VortexResult<_>>()?;
24+
let array = state.array();
25+
if !state.is_valid(index, ctx)? {
26+
return Ok(Scalar::null(array.dtype().clone()));
27+
}
2628

27-
debug_assert_eq!(children_elements.len(), array.list_size() as usize);
29+
let DType::FixedSizeList(element_dtype, list_size, nullability) = array.dtype() else {
30+
unreachable!("FixedSizeList arrays have a fixed-size list dtype")
31+
};
32+
let start = *list_size as usize * index;
33+
let end = start + *list_size as usize;
34+
let mut elements = state
35+
.slot(FixedSizeListSlots::ELEMENTS)?
36+
.ok_or_else(|| vortex_err!("FixedSizeList elements slot is missing"))?;
37+
let children_elements = (start..end)
38+
.map(|i| elements.execute_scalar(i, ctx))
39+
.collect::<VortexResult<Vec<_>>>()?;
2840

2941
Ok(Scalar::fixed_size_list(
30-
list.dtype().clone(),
42+
element_dtype.clone(),
3143
children_elements,
32-
array.dtype().nullability(),
44+
*nullability,
3345
))
3446
}
47+
48+
fn scalar_at(
49+
array: ArrayView<'_, FixedSizeList>,
50+
index: usize,
51+
ctx: &mut ExecutionCtx,
52+
) -> VortexResult<Scalar> {
53+
Self::probe_scalar(&mut ProbeState::once(array), index, ctx)
54+
}
3555
}

0 commit comments

Comments
 (0)