Skip to content

Commit 058081d

Browse files
connortsui20connortsui20
committed
update all ListViewArray construction with is_zctl
Signed-off-by: Connor Tsui <[email protected]>
1 parent f2efda0 commit 058081d

File tree

21 files changed

+531
-743
lines changed

21 files changed

+531
-743
lines changed

‎encodings/sparse/src/canonical.rs‎

Lines changed: 47 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -492,8 +492,8 @@ mod test {
492492

493493
use rstest::rstest;
494494
use vortex_array::arrays::{
495-
BoolArray, DecimalArray, FixedSizeListArray, ListArray, ListViewArray, ListViewShape,
496-
PrimitiveArray, StructArray, VarBinArray, VarBinViewArray,
495+
BoolArray, DecimalArray, FixedSizeListArray, ListArray, ListViewArray, PrimitiveArray,
496+
StructArray, VarBinArray, VarBinViewArray,
497497
};
498498
use vortex_array::arrow::IntoArrowArray as _;
499499
use vortex_array::validity::Validity;
@@ -930,14 +930,15 @@ mod test {
930930
// List 3: [2] at offset 3, size 1
931931
let offsets = buffer![0u32, 1, 2, 3].into_array();
932932
let sizes = buffer![1u32, 1, 1, 1].into_array();
933-
let lists = ListViewArray::try_new(
934-
elements,
935-
offsets,
936-
sizes,
937-
Validity::AllValid,
938-
ListViewShape::as_zero_copy_to_list(),
939-
)
940-
.unwrap()
933+
let lists = unsafe {
934+
ListViewArray::new_unchecked(
935+
elements,
936+
offsets,
937+
sizes,
938+
Validity::AllValid,
939+
true, // Is zero-copy to list.
940+
)
941+
}
941942
.into_array();
942943

943944
let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array();
@@ -983,14 +984,15 @@ mod test {
983984
let elements = buffer![1i32, 2, 1, 2, 1, 2, 1, 2].into_array();
984985
let offsets = buffer![0u32, 1, 2, 3, 4, 5, 6, 7].into_array();
985986
let sizes = buffer![1u32, 1, 1, 1, 1, 1, 1, 1].into_array();
986-
let lists = ListViewArray::try_new(
987-
elements,
988-
offsets,
989-
sizes,
990-
Validity::AllValid,
991-
ListViewShape::as_zero_copy_to_list(),
992-
)
993-
.unwrap()
987+
let lists = unsafe {
988+
ListViewArray::new_unchecked(
989+
elements,
990+
offsets,
991+
sizes,
992+
Validity::AllValid,
993+
true, // Is zero-copy to list.
994+
)
995+
}
994996
.into_array();
995997

996998
// Slice to get lists 2..6, which are: [1], [2], [1], [2]
@@ -1033,14 +1035,15 @@ mod test {
10331035
let elements = buffer![1i32, 2, 1, 2].into_array();
10341036
let offsets = buffer![0u32, 1, 2, 3].into_array();
10351037
let sizes = buffer![1u32, 1, 1, 1].into_array();
1036-
let lists = ListViewArray::try_new(
1037-
elements,
1038-
offsets,
1039-
sizes,
1040-
Validity::AllValid,
1041-
ListViewShape::as_zero_copy_to_list(),
1042-
)
1043-
.unwrap()
1038+
let lists = unsafe {
1039+
ListViewArray::new_unchecked(
1040+
elements,
1041+
offsets,
1042+
sizes,
1043+
Validity::AllValid,
1044+
true, // Is zero-copy to list.
1045+
)
1046+
}
10441047
.into_array();
10451048

10461049
let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array();
@@ -1401,14 +1404,15 @@ mod test {
14011404
let offsets = buffer![0u32, 3, 5].into_array();
14021405
let sizes = buffer![3u32, 2, 4].into_array();
14031406

1404-
let list_view = ListViewArray::try_new(
1405-
elements.clone(),
1406-
offsets,
1407-
sizes,
1408-
Validity::AllValid,
1409-
ListViewShape::as_zero_copy_to_list(),
1410-
)
1411-
.unwrap();
1407+
let list_view = unsafe {
1408+
ListViewArray::new_unchecked(
1409+
elements.clone(),
1410+
offsets,
1411+
sizes,
1412+
Validity::AllValid,
1413+
true, // Is zero-copy to list.
1414+
)
1415+
};
14121416

14131417
let list_dtype = list_view.dtype().clone();
14141418

@@ -1483,14 +1487,15 @@ mod test {
14831487
let offsets = buffer![0u32, 2, 5, 6, 8].into_array();
14841488
let sizes = buffer![2u32, 3, 1, 2, 2].into_array();
14851489

1486-
let full_listview = ListViewArray::try_new(
1487-
elements,
1488-
offsets,
1489-
sizes,
1490-
Validity::AllValid,
1491-
ListViewShape::as_zero_copy_to_list(),
1492-
)
1493-
.unwrap()
1490+
let full_listview = unsafe {
1491+
ListViewArray::new_unchecked(
1492+
elements,
1493+
offsets,
1494+
sizes,
1495+
Validity::AllValid,
1496+
true, // Is zero-copy to list.
1497+
)
1498+
}
14941499
.into_array();
14951500

14961501
// Slice to get lists 1, 2, 3 (indices 1..4)

‎fuzz/src/array/mask.rs‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -57,14 +57,18 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult<A
5757
}
5858
Canonical::List(array) => {
5959
let new_validity = apply_mask_to_validity(array.validity(), mask);
60-
ListViewArray::try_new(
61-
array.elements().clone(),
62-
array.offsets().clone(),
63-
array.sizes().clone(),
64-
new_validity,
65-
array.shape(),
66-
)
67-
.vortex_unwrap()
60+
61+
// SAFETY: Since we are only masking the validity and everything else comes from an
62+
// already valid `ListViewArray`, all of the invariants are still upheld.
63+
unsafe {
64+
ListViewArray::new_unchecked(
65+
array.elements().clone(),
66+
array.offsets().clone(),
67+
array.sizes().clone(),
68+
new_validity,
69+
array.is_zero_copy_to_list(),
70+
)
71+
}
6872
.into_array()
6973
}
7074
Canonical::FixedSizeList(array) => {
@@ -131,8 +135,8 @@ fn apply_mask_to_validity(validity: &Validity, mask: &Mask) -> Validity {
131135
#[cfg(test)]
132136
mod tests {
133137
use vortex_array::arrays::{
134-
BoolArray, DecimalArray, FixedSizeListArray, ListViewArray, ListViewShape, NullArray,
135-
PrimitiveArray, StructArray, VarBinViewArray,
138+
BoolArray, DecimalArray, FixedSizeListArray, ListViewArray, NullArray, PrimitiveArray,
139+
StructArray, VarBinViewArray,
136140
};
137141
use vortex_array::{Array, IntoArray};
138142
use vortex_dtype::{DecimalDType, FieldNames, Nullability};
@@ -244,14 +248,15 @@ mod tests {
244248
let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]).into_array();
245249
let offsets = PrimitiveArray::from_iter([0i32, 2, 4]).into_array();
246250
let sizes = PrimitiveArray::from_iter([2i32, 2, 2]).into_array();
247-
let array = ListViewArray::try_new(
248-
elements,
249-
offsets,
250-
sizes,
251-
Nullability::NonNullable.into(),
252-
ListViewShape::as_zero_copy_to_list(),
253-
)
254-
.unwrap();
251+
let array = unsafe {
252+
ListViewArray::new_unchecked(
253+
elements,
254+
offsets,
255+
sizes,
256+
Nullability::NonNullable.into(),
257+
true, // Is zero-copy to list.
258+
)
259+
};
255260

256261
let mask = Mask::from_iter([false, true, false]);
257262

‎fuzz/src/array/slice.rs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,19 @@ pub fn slice_canonical_array(
7474
// Since the list view elements can be stored out of order, we cannot slice it.
7575
let elements = list_array.elements().clone();
7676

77-
ListViewArray::try_new(elements, offsets, sizes, validity, list_array.shape())
78-
.map(|a| a.into_array())
77+
// SAFETY: If the array was already zero-copyable to list, slicing the offsets and sizes
78+
// only causes there to be leading and trailing garbage data, which is still
79+
// zero-copyable to a `ListArray`.
80+
Ok(unsafe {
81+
ListViewArray::new_unchecked(
82+
elements,
83+
offsets,
84+
sizes,
85+
validity,
86+
list_array.is_zero_copy_to_list(),
87+
)
88+
}
89+
.into_array())
7990
}
8091
DType::FixedSizeList(..) => {
8192
let fsl_array = array.to_fixed_size_list();

‎vortex-array/src/arrays/listview/array.rs‎

Lines changed: 30 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -118,14 +118,8 @@ impl ListViewArray {
118118
///
119119
/// Panics if the provided components do not satisfy the invariants documented
120120
/// in [`ListViewArray::new_unchecked`].
121-
pub fn new(
122-
elements: ArrayRef,
123-
offsets: ArrayRef,
124-
sizes: ArrayRef,
125-
validity: Validity,
126-
is_zctl: bool,
127-
) -> Self {
128-
Self::try_new(elements, offsets, sizes, validity, is_zctl)
121+
pub fn new(elements: ArrayRef, offsets: ArrayRef, sizes: ArrayRef, validity: Validity) -> Self {
122+
Self::try_new(elements, offsets, sizes, validity)
129123
.vortex_expect("`ListViewArray` construction failed")
130124
}
131125

@@ -140,16 +134,27 @@ impl ListViewArray {
140134
offsets: ArrayRef,
141135
sizes: ArrayRef,
142136
validity: Validity,
143-
is_zctl: bool,
144137
) -> VortexResult<Self> {
145-
Self::validate(&elements, &offsets, &sizes, &validity, is_zctl)?;
146-
147-
// SAFETY: validate ensures all invariants are met.
148-
Ok(unsafe { Self::new_unchecked(elements, offsets, sizes, validity, is_zctl) })
138+
Self::validate(&elements, &offsets, &sizes, &validity, false)?;
139+
140+
Ok(Self {
141+
dtype: DType::List(Arc::new(elements.dtype().clone()), validity.nullability()),
142+
elements,
143+
offsets,
144+
sizes,
145+
validity,
146+
is_zero_copy_to_list: false,
147+
stats_set: Default::default(),
148+
})
149149
}
150150

151151
/// Creates a new [`ListViewArray`] without validation.
152152
///
153+
/// Note that this constructor is slightly different from [`new()`] and [`try_new()`] in that it
154+
/// also takes an `is_zctl` flag. This is an optimization flag that allows conversion to a
155+
/// [`ListArray`] more efficient. The caller must ensure that the flag is correct when passing
156+
/// it to this unsafe constructor.
157+
///
153158
/// # Safety
154159
///
155160
/// The caller must ensure all of the following invariants are satisfied:
@@ -170,17 +175,18 @@ impl ListViewArray {
170175
is_zctl: bool,
171176
) -> Self {
172177
if cfg!(debug_assertions) {
173-
Self::new(elements, offsets, sizes, validity, is_zctl)
174-
} else {
175-
Self {
176-
dtype: DType::List(Arc::new(elements.dtype().clone()), validity.nullability()),
177-
elements,
178-
offsets,
179-
sizes,
180-
validity,
181-
is_zero_copy_to_list: is_zctl,
182-
stats_set: Default::default(),
183-
}
178+
Self::validate(&elements, &offsets, &sizes, &validity, is_zctl)
179+
.vortex_expect("Failed to crate `ListViewArray`");
180+
}
181+
182+
Self {
183+
dtype: DType::List(Arc::new(elements.dtype().clone()), validity.nullability()),
184+
elements,
185+
offsets,
186+
sizes,
187+
validity,
188+
is_zero_copy_to_list: is_zctl,
189+
stats_set: Default::default(),
184190
}
185191
}
186192

‎vortex-array/src/arrays/listview/compute/mask.rs‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,18 @@ use crate::{ArrayRef, IntoArray, register_kernel};
1111

1212
impl MaskKernel for ListViewVTable {
1313
fn mask(&self, array: &ListViewArray, mask: &Mask) -> VortexResult<ArrayRef> {
14-
ListViewArray::try_new(
15-
array.elements().clone(),
16-
array.offsets().clone(),
17-
array.sizes().clone(),
18-
array.validity().mask(mask),
19-
array.is_zero_copy_to_list(),
20-
)
21-
.map(|a| a.into_array())
14+
// SAFETY: Since we are only masking the validity and everything else comes from an already
15+
// valid `ListViewArray`, all of the invariants are still upheld.
16+
Ok(unsafe {
17+
ListViewArray::new_unchecked(
18+
array.elements().clone(),
19+
array.offsets().clone(),
20+
array.sizes().clone(),
21+
array.validity().mask(mask),
22+
array.is_zero_copy_to_list(),
23+
)
24+
}
25+
.into_array())
2226
}
2327
}
2428

0 commit comments

Comments
 (0)