-
Notifications
You must be signed in to change notification settings - Fork 203
Expand file tree
/
Copy pathmod.rs
More file actions
347 lines (313 loc) · 12.6 KB
/
Copy pathmod.rs
File metadata and controls
347 lines (313 loc) · 12.6 KB
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
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
//! This module contains the VTable definitions for a Vortex encoding.
//!
//! A Vortex array encoding is implemented by a small static vtable type plus an associated
//! `TypedArrayData` value stored in each array instance. The vtable owns behavior such as
//! validation, serialization, execution, child traversal, scalar access, and validity access.
//!
//! The public [`ArrayRef`] API performs common precondition checks before calling
//! into these traits. Implementations should focus on encoding-specific work and uphold the
//! documented postconditions.
mod operations;
mod validity;
use std::fmt::Debug;
use std::fmt::Display;
use std::fmt::Formatter;
use std::hash::Hasher;
pub use operations::*;
pub use validity::*;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_panic;
use vortex_session::VortexSession;
use crate::Array;
use crate::ArrayRef;
use crate::ArrayView;
use crate::Canonical;
use crate::EqMode;
use crate::ExecutionResult;
use crate::IntoArray;
pub use crate::array::plugin::*;
use crate::arrays::ConstantArray;
use crate::arrays::constant::Constant;
use crate::buffer::BufferHandle;
use crate::builders::ArrayBuilder;
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::executor::ExecutionCtx;
use crate::hash::ArrayEq;
use crate::hash::ArrayHash;
use crate::patches::Patches;
use crate::scalar::ScalarValue;
use crate::serde::ArrayChildren;
use crate::validity::Validity;
/// The array [`VTable`] encapsulates logic for an Array type within Vortex.
///
/// The logic is split across several "VTable" traits to enable easier code organization than
/// simply lumping everything into a single trait.
///
/// From this [`VTable`] trait, we derive implementations for the sealed `DynArrayData` trait and the
/// public [`ArrayPlugin`] registry trait.
///
/// The functions defined in these vtable traits will typically document their pre- and
/// post-conditions. The pre-conditions are validated inside the `DynArrayData` and [`ArrayRef`]
/// implementations so do not need to be checked in the vtable implementations (for example, index
/// out of bounds). Post-conditions are validated after invocation of the vtable function and will
/// panic if violated.
pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug {
/// Per-array data owned by this encoding, excluding child arrays.
///
/// Child arrays belong in [`ArrayParts::slots`](crate::ArrayParts::slots) so traversal,
/// serialization, and layout writers can discover them generically.
type TypedArrayData: 'static + Send + Sync + Clone + Debug + Display + ArrayHash + ArrayEq;
/// Scalar and element-wise operation hooks for this encoding.
type OperationsVTable: OperationsVTable<Self>;
/// Validity hook for nullable instances of this encoding.
type ValidityVTable: ValidityVTable<Self>;
/// Returns the ID of the array.
fn id(&self) -> ArrayId;
/// Validates that externally supplied logical metadata matches the array data.
///
/// This is called by [`Array::try_from_parts`](crate::Array::try_from_parts) before the array
/// is published. Implementations should check dtype, length, slot count, child dtypes/lengths,
/// metadata bounds, and any buffer shape invariants that unsafe accessors depend on.
fn validate(
&self,
data: &Self::TypedArrayData,
dtype: &DType,
len: usize,
slots: &[Option<ArrayRef>],
) -> VortexResult<()>;
/// Returns the number of top-level buffers in the array.
fn nbuffers(array: ArrayView<'_, Self>) -> usize;
/// Returns the buffer at the given index.
///
/// # Panics
/// Panics if `idx >= nbuffers(array)`.
fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle;
/// Returns the name of the buffer at the given index, or `None` if unnamed.
fn buffer_name(array: ArrayView<'_, Self>, idx: usize) -> Option<String>;
/// Rebuild this array with replacement top-level buffers.
///
/// This is for physical rewrites that preserve `dtype`, `len`, child slots, buffer count, and
/// buffer lengths. The caller checks the generic invariants before dispatching here;
/// implementations should interpret the replacement buffers for their encoding-specific
/// in-memory representation.
fn with_buffers(
&self,
array: ArrayView<'_, Self>,
buffers: &[BufferHandle],
) -> VortexResult<ArrayParts<Self>>;
/// Serialize encoding metadata into a byte buffer for IPC or file storage.
///
/// Return `None` if the array cannot be serialized by this encoding. Buffers and children are
/// serialized separately through [`buffer`](Self::buffer), [`nbuffers`](Self::nbuffers), and
/// child traversal.
fn serialize(
array: ArrayView<'_, Self>,
session: &VortexSession,
) -> VortexResult<Option<Vec<u8>>>;
/// Deserialize an array from serialized metadata, buffers, and children.
///
/// The returned [`ArrayParts`] are still validated by the generic adapter.
/// Deserializers should use the provided `session` to resolve plugin-owned metadata instead of
/// relying on global state.
fn deserialize(
&self,
dtype: &DType,
len: usize,
metadata: &[u8],
buffers: &[BufferHandle],
children: &dyn ArrayChildren,
session: &VortexSession,
) -> VortexResult<ArrayParts<Self>>;
/// Writes the array's logical values into a canonical builder.
///
/// The default implementation executes the full array to [`Canonical`] and appends that result.
/// Encodings may override this to avoid materializing an intermediate canonical array.
fn append_to_builder(
array: ArrayView<'_, Self>,
builder: &mut dyn ArrayBuilder,
ctx: &mut ExecutionCtx,
) -> VortexResult<()> {
let canonical = array
.array()
.clone()
.execute::<Canonical>(ctx)?
.into_array();
canonical.append_to_builder(builder, ctx)
}
/// Returns the name of the slot at the given index.
///
/// # Panics
/// Panics if `idx >= slots(array).len()`.
fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String;
/// Execute this array by returning an [`ExecutionResult`].
///
/// Execution is **iterative**, not recursive. Instead of recursively executing children,
/// implementations should return [`ExecutionResult::execute_slot`] to request that the
/// scheduler execute a slot first, or [`ExecutionResult::done`] when the encoding can
/// produce a result directly.
///
/// For good examples of this pattern, see:
/// - [`Dict::execute`](crate::arrays::dict::vtable::Dict::execute) — demonstrates
/// requiring children via `require_child!` and producing a result once they are canonical.
/// - `BitPacked::execute` (in `vortex-fastlanes`) — demonstrates requiring patches and
/// validity via `require_patches!`/`require_validity!`.
///
/// Array execution is designed such that repeated execution of an array will eventually
/// converge to a canonical representation. Implementations of this function should therefore
/// ensure they make progress towards that goal.
///
/// The returned array (in `Done`) must be logically equivalent to the input array. In other
/// words, the recursively canonicalized forms of both arrays must be equal.
///
/// Debug builds will panic if the returned array is of the wrong type, wrong length, or
/// incorrectly contains null values.
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult>;
/// Attempt to reduce the array to a simpler representation without changing logical values.
///
/// Reductions are opportunistic and may return `Ok(None)` when no cheaper representation is
/// known.
fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
_ = array;
Ok(None)
}
/// Attempt to reduce `parent` after this array appears as one of its children.
///
/// This is used by lazy arrays to let child execution unlock parent simplifications.
fn reduce_parent(
array: ArrayView<'_, Self>,
parent: &ArrayRef,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>> {
_ = (array, parent, child_idx);
Ok(None)
}
}
/// Alias for migration — downstream code can start using `ArrayVTable`.
pub use VTable as ArrayVTable;
use crate::array::ArrayId;
use crate::array::ArrayParts;
/// Empty array metadata struct for encodings with no per-array metadata.
#[derive(Clone, Debug, Default)]
pub struct EmptyArrayData;
impl ArrayEq for EmptyArrayData {
fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool {
true
}
}
impl ArrayHash for EmptyArrayData {
fn array_hash<H: Hasher>(&self, _state: &mut H, _accuracy: EqMode) {}
}
impl Display for EmptyArrayData {
fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
Ok(())
}
}
/// Rebuild an array that has no top-level buffers.
#[inline]
pub fn with_empty_buffers<V: VTable>(
vtable: &V,
array: ArrayView<'_, V>,
buffers: &[BufferHandle],
) -> VortexResult<ArrayParts<V>> {
vortex_ensure!(
buffers.is_empty(),
"Array {} expects 0 buffers, got {}",
array.encoding_id(),
buffers.len()
);
Ok(ArrayParts::new(
vtable.clone(),
array.dtype().clone(),
array.len(),
array.data().clone(),
)
.with_slots(array.slots().iter().cloned().collect()))
}
/// Reject buffer replacement for encodings whose exposed buffers are not runtime backing buffers.
#[inline]
pub fn unsupported_buffer_replacement<V: VTable>(
array: ArrayView<'_, V>,
_buffers: &[BufferHandle],
) -> VortexResult<ArrayParts<V>> {
vortex_bail!(
"Array {} does not support in-memory buffer replacement",
array.encoding_id()
)
}
/// Placeholder type used to indicate when a particular vtable is not supported by the encoding.
pub struct NotSupported;
/// Returns the validity as a child array if it produces one.
#[inline]
pub fn validity_to_child(validity: &Validity, len: usize) -> Option<ArrayRef> {
match validity {
Validity::NonNullable | Validity::AllValid => None,
Validity::AllInvalid => Some(ConstantArray::new(false, len).into_array()),
Validity::Array(array) => Some(array.clone()),
}
}
/// Reconstruct a [`Validity`] from an optional child array and nullability.
///
/// This is the inverse of [`validity_to_child`].
#[inline]
pub fn child_to_validity(child: Option<&ArrayRef>, nullability: Nullability) -> Validity {
match child {
Some(arr) => {
// Detect constant bool arrays created by validity_to_child.
// Use direct ScalarValue matching to avoid expensive scalar conversion.
if let Some(c) = arr.as_opt::<Constant>()
&& let Some(ScalarValue::Bool(val)) = c.scalar().value()
{
return if *val {
Validity::AllValid
} else {
Validity::AllInvalid
};
}
Validity::Array(arr.clone())
}
None => Validity::from(nullability),
}
}
/// Returns 1 if validity produces a child, 0 otherwise.
#[inline]
pub fn validity_nchildren(validity: &Validity) -> usize {
match validity {
Validity::NonNullable | Validity::AllValid => 0,
Validity::AllInvalid | Validity::Array(_) => 1,
}
}
/// Returns the number of children produced by patches.
#[inline]
pub fn patches_nchildren(patches: &Patches) -> usize {
2 + patches.chunk_offsets().is_some() as usize
}
/// Returns the child at the given index within a patches component.
#[inline]
pub fn patches_child(patches: &Patches, idx: usize) -> ArrayRef {
match idx {
0 => patches.indices().clone(),
1 => patches.values().clone(),
2 => patches
.chunk_offsets()
.as_ref()
.vortex_expect("patch_chunk_offsets child out of bounds")
.clone(),
_ => vortex_panic!("patches child index {idx} out of bounds"),
}
}
/// Returns the name of the child at the given index within a patches component.
#[inline]
pub fn patches_child_name(idx: usize) -> &'static str {
match idx {
0 => "patch_indices",
1 => "patch_values",
2 => "patch_chunk_offsets",
_ => vortex_panic!("patches child name index {idx} out of bounds"),
}
}