1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
use vortex_buffer::ByteBuffer;
use vortex_error::VortexExpect;
use crate::pipeline::N;
use crate::pipeline::bits::{BitView, BitViewMut};
use crate::pipeline::types::{Element, VType};
use crate::pipeline::vec::Selection;
pub struct View<'a> {
/// The physical type of the vector, which defines how the elements are stored.
pub(super) vtype: VType,
/// A pointer to the allocated elements buffer.
/// Alignment is at least the size of the element type.
/// The capacity of the elements buffer is N * `size_of::<T>()` where T is the element type.
pub(super) elements: *const u8,
/// The validity mask for the vector, indicating which elements in the buffer are valid.
/// This value can be `None` if the expected DType is `NonNullable`.
// TODO: support validity
#[allow(dead_code)]
pub(super) validity: Option<BitView<'a>>,
// Indicates where the selected elements are positioned within the vector.
pub(super) selection: Selection,
/// Additional buffers of data used by the vector, such as string data.
#[allow(dead_code)]
pub(super) data: Vec<ByteBuffer>,
/// Marker defining the lifetime of the contents of the vector.
pub(super) _marker: std::marker::PhantomData<&'a ()>,
}
impl<'a> View<'a> {
#[inline(always)]
pub fn selection(&self) -> Selection {
self.selection
}
pub fn as_array<T>(&self) -> &'a [T; N]
where
T: Element,
{
debug_assert_eq!(self.vtype, T::vtype(), "Invalid type for canonical view");
// SAFETY: We assume that the elements are of type T and that the view is valid.
unsafe { &*(self.elements.cast::<T>() as *const [T; N]) }
}
/// Re-interpret cast the vector into a new type where the element has the same width.
#[inline(always)]
pub fn reinterpret_as<E: Element>(&mut self) {
assert_eq!(
self.vtype.byte_width(),
size_of::<E>(),
"Cannot reinterpret {} as {}",
self.vtype,
E::vtype()
);
self.vtype = E::vtype();
}
}
pub struct ViewMut<'a> {
/// The physical type of the vector, which defines how the elements are stored.
pub(super) vtype: VType,
/// A pointer to the allocated elements buffer.
/// Alignment is at least the size of the element type.
/// The capacity of the elements buffer is N * `size_of::<T>()` where T is the element type.
// TODO(ngates): it would be nice to guarantee _wider_ alignment, ideally 128 bytes, so that
// we can use aligned load/store instructions for wide SIMD lanes.
pub(super) elements: *mut u8,
/// The validity mask for the vector, indicating which elements in the buffer are valid.
/// This value can be `None` if the expected DType is `NonNullable`.
pub(super) validity: Option<BitViewMut<'a>>,
/// Additional buffers of data used by the vector, such as string data.
// TODO(ngates): ideally these buffers are compressed somehow? E.g. using FSST?
#[allow(dead_code)]
pub(super) data: Vec<ByteBuffer>,
/// The position of the selected values of this buffer.
/// One of:
/// * All - all N values are selected.
/// * Prefix - the first n values are selected where i is the true count of the kernel mask.
/// * Mask - the values are in the positions indicated by the kernel mask.
pub(super) selection: Selection,
/// Marker defining the lifetime of the contents of the vector.
pub(super) _marker: std::marker::PhantomData<&'a mut ()>,
}
impl<'a> ViewMut<'a> {
pub fn new<E: Element>(elements: &'a mut [E], validity: Option<BitViewMut<'a>>) -> Self {
assert_eq!(elements.len(), N);
Self {
vtype: E::vtype(),
elements: elements.as_mut_ptr().cast(),
validity,
data: vec![],
selection: Selection::Prefix,
_marker: Default::default(),
}
}
/// Re-interpret cast the vector into a new type where the element has the same width.
#[inline(always)]
pub fn reinterpret_as<E: Element>(&mut self) {
assert_eq!(
self.vtype.byte_width(),
size_of::<E>(),
"Cannot reinterpret {} as {}",
self.vtype,
E::vtype()
);
self.vtype = E::vtype();
}
/// Returns an immutable array of the elements in the vector.
#[inline(always)]
pub fn as_array<E: Element>(&self) -> &'a [E; N] {
debug_assert_eq!(self.vtype, E::vtype(), "Invalid type for canonical view");
unsafe { &*(self.elements.cast::<E>() as *const [E; N]) }
}
/// Returns a mutable array of the elements in the vector, allowing for modification.
#[inline(always)]
pub fn as_array_mut<E: Element>(&mut self) -> &'a mut [E; N] {
debug_assert_eq!(self.vtype, E::vtype(), "Invalid type for canonical view");
unsafe { &mut *(self.elements.cast::<E>() as *mut [E; N]) }
}
/// Access the validity mask of the vector.
///
/// ## Panics
///
/// Panics if the vector does not support validity, i.e. if the DType was non-nullable when
/// it was created.
pub fn validity(&mut self) -> &mut BitViewMut<'a> {
self.validity
.as_mut()
.vortex_expect("Vector does not support validity")
}
pub fn add_buffer(&mut self, buffer: ByteBuffer) {
self.data.push(buffer);
}
#[inline(always)]
pub fn selection(&self) -> Selection {
self.selection
}
pub fn set_selection(&mut self, selection: Selection) {
self.selection = selection;
}
/// Flatten the view by bringing the selected elements of the mask to the beginning of
pub fn flatten<E: Element>(&mut self, selection: &BitView<'_>) {
assert_eq!(
self.vtype,
E::vtype(),
"ViewMut::flatten_mask: type mismatch"
);
if matches!(self.selection, Selection::Prefix) {
// Nothing to do, all elements are already selected.
return;
}
match selection.true_count() {
0 | N => {
// If the mask has no true bits or all true bits, we are already flattened.
}
n if n > 3 * N / 4 => {
// High density: use iter_zeros to compact by removing gaps
let slice = self.as_array_mut::<E>();
let mut write_idx = 0;
let mut read_idx = 0;
selection.iter_zeros(|zero_idx| {
// Copy elements from read_idx to zero_idx (exclusive) to write_idx
let count = zero_idx - read_idx;
unsafe {
// SAFETY: We assume that the elements are of type E and that the view is valid.
// Using memmove for potentially overlapping regions
std::ptr::copy(
slice.as_ptr().add(read_idx),
slice.as_mut_ptr().add(write_idx),
count,
);
write_idx += count;
}
read_idx = zero_idx + 1;
});
// Copy any remaining elements after the last zero
unsafe {
std::ptr::copy(
slice.as_ptr().add(read_idx),
slice.as_mut_ptr().add(write_idx),
N - read_idx,
);
}
}
_ => {
let mut offset = 0;
let slice = self.as_array_mut::<E>();
selection.iter_ones(|idx| {
unsafe {
// SAFETY: We assume that the elements are of type E and that the view is valid.
let value = *slice.get_unchecked(idx);
// TODO(joe): use ptr increment (not offset).
*slice.get_unchecked_mut(offset) = value;
offset += 1;
}
});
}
}
self.selection = Selection::Prefix
}
}