Skip to content

Commit 68bd7c9

Browse files
committed
Move primitive storage behind fixed-width abstraction
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 5b688a9 commit 68bd7c9

7 files changed

Lines changed: 249 additions & 124 deletions

File tree

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Shared physical storage helpers for canonical fixed-width arrays.
5+
//!
6+
//! Primitive and decimal arrays have distinct VTables and logical behavior, but use the same
7+
//! one-buffer-plus-validity physical representation.
8+
9+
use std::fmt::Display;
10+
use std::fmt::Formatter;
11+
use std::hash::Hash;
12+
use std::hash::Hasher;
13+
14+
use vortex_error::VortexResult;
15+
use vortex_error::vortex_bail;
16+
use vortex_error::vortex_ensure;
17+
use vortex_error::vortex_panic;
18+
19+
use crate::ArrayParts;
20+
use crate::ArrayRef;
21+
use crate::EqMode;
22+
use crate::array::ArrayView;
23+
use crate::array::VTable;
24+
use crate::buffer::BufferHandle;
25+
use crate::dtype::DecimalType;
26+
use crate::dtype::Nullability;
27+
use crate::dtype::PType;
28+
use crate::hash::ArrayEq;
29+
use crate::hash::ArrayHash;
30+
use crate::serde::ArrayChildren;
31+
use crate::validity::Validity;
32+
33+
/// The element-level physical interpretation of a canonical fixed-width values buffer.
34+
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
35+
pub enum FixedWidthType {
36+
/// A native primitive value.
37+
Primitive(PType),
38+
/// An integer storage type used for logical decimal values.
39+
Decimal(DecimalType),
40+
}
41+
42+
impl FixedWidthType {
43+
/// Returns the number of bytes occupied by one value.
44+
pub fn byte_width(self) -> usize {
45+
match self {
46+
Self::Primitive(ptype) => ptype.byte_width(),
47+
Self::Decimal(decimal_type) => decimal_type.byte_width(),
48+
}
49+
}
50+
}
51+
52+
/// The shared physical data stored by the primitive and decimal VTables.
53+
#[derive(Clone, Debug)]
54+
pub struct FixedWidthData {
55+
pub(crate) physical_type: FixedWidthType,
56+
pub(crate) buffer: BufferHandle,
57+
pub(crate) len: usize,
58+
}
59+
60+
impl Display for FixedWidthData {
61+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
62+
match self.physical_type {
63+
FixedWidthType::Primitive(ptype) => write!(f, "physical_type: {ptype}"),
64+
FixedWidthType::Decimal(values_type) => write!(f, "physical_type: {values_type}"),
65+
}
66+
}
67+
}
68+
69+
impl ArrayHash for FixedWidthData {
70+
fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
71+
self.buffer.array_hash(state, accuracy);
72+
self.physical_type.hash(state);
73+
self.len.hash(state);
74+
}
75+
}
76+
77+
impl ArrayEq for FixedWidthData {
78+
fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
79+
self.buffer.array_eq(&other.buffer, accuracy)
80+
&& self.physical_type == other.physical_type
81+
&& self.len == other.len
82+
}
83+
}
84+
85+
pub(crate) fn nbuffers<V>(_array: ArrayView<'_, V>) -> usize
86+
where
87+
V: VTable<TypedArrayData = FixedWidthData>,
88+
{
89+
1
90+
}
91+
92+
pub(crate) fn buffer<V>(array: ArrayView<'_, V>, idx: usize) -> BufferHandle
93+
where
94+
V: VTable<TypedArrayData = FixedWidthData>,
95+
{
96+
match idx {
97+
0 => array.data().buffer_handle().clone(),
98+
_ => vortex_panic!("Fixed-width buffer index {idx} out of bounds"),
99+
}
100+
}
101+
102+
pub(crate) fn buffer_name<V>(_array: ArrayView<'_, V>, idx: usize) -> Option<String>
103+
where
104+
V: VTable<TypedArrayData = FixedWidthData>,
105+
{
106+
match idx {
107+
0 => Some("values".to_string()),
108+
_ => None,
109+
}
110+
}
111+
112+
pub(crate) fn with_buffers<V>(
113+
vtable: &V,
114+
array: ArrayView<'_, V>,
115+
buffers: &[BufferHandle],
116+
) -> VortexResult<ArrayParts<V>>
117+
where
118+
V: VTable<TypedArrayData = FixedWidthData>,
119+
{
120+
vortex_ensure!(
121+
buffers.len() == 1,
122+
"Expected 1 buffer, got {}",
123+
buffers.len()
124+
);
125+
let mut data = array.data().clone();
126+
data.buffer = buffers[0].clone();
127+
Ok(
128+
ArrayParts::new(vtable.clone(), array.dtype().clone(), array.len(), data)
129+
.with_slots(array.slots().iter().cloned().collect()),
130+
)
131+
}
132+
133+
pub(crate) fn validate_layout(
134+
data: &FixedWidthData,
135+
len: usize,
136+
slots: &[Option<ArrayRef>],
137+
nullability: Nullability,
138+
name: &str,
139+
) -> VortexResult<()> {
140+
vortex_ensure!(slots.len() == 1, "{name} expects one validity slot");
141+
vortex_ensure!(
142+
data.len() == len,
143+
InvalidArgument: "{name} length {} does not match outer length {len}",
144+
data.len(),
145+
);
146+
let validity = crate::array::child_to_validity(slots[0].as_ref(), nullability);
147+
if let Some(validity_len) = validity.maybe_len() {
148+
vortex_ensure!(
149+
validity_len == len,
150+
InvalidArgument: "{name} validity len {validity_len} does not match outer length {len}",
151+
);
152+
}
153+
Ok(())
154+
}
155+
156+
pub(crate) fn deserialize_validity(
157+
dtype_nullability: Nullability,
158+
len: usize,
159+
children: &dyn ArrayChildren,
160+
) -> VortexResult<Validity> {
161+
if children.is_empty() {
162+
Ok(Validity::from(dtype_nullability))
163+
} else if children.len() == 1 {
164+
Ok(Validity::Array(children.get(0, &Validity::DTYPE, len)?))
165+
} else {
166+
vortex_bail!("Expected 0 or 1 child, got {}", children.len())
167+
}
168+
}
169+
170+
pub(crate) fn validity<V>(array: ArrayView<'_, V>) -> VortexResult<Validity>
171+
where
172+
V: VTable<TypedArrayData = FixedWidthData>,
173+
{
174+
Ok(crate::array::child_to_validity(
175+
array.slots()[0].as_ref(),
176+
array.dtype().nullability(),
177+
))
178+
}
179+
180+
pub(crate) fn slot_name(idx: usize) -> String {
181+
match idx {
182+
0 => "validity".to_string(),
183+
_ => vortex_panic!("Fixed-width slot index {idx} out of bounds"),
184+
}
185+
}

vortex-array/src/arrays/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ pub mod fixed_size_list;
6464
pub use fixed_size_list::FixedSizeList;
6565
pub use fixed_size_list::FixedSizeListArray;
6666

67+
pub mod fixed_width;
68+
6769
pub mod interleave;
6870
pub use interleave::Interleave;
6971
pub use interleave::InterleaveArray;

vortex-array/src/arrays/primitive/array/mod.rs

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4-
use std::fmt::Display;
5-
use std::fmt::Formatter;
64
use std::iter::repeat;
75

86
use smallvec::smallvec;
@@ -24,6 +22,8 @@ use crate::array::TypedArrayRef;
2422
use crate::arrays::BoolArray;
2523
use crate::arrays::Primitive;
2624
use crate::arrays::PrimitiveArray;
25+
use crate::arrays::fixed_width::FixedWidthData;
26+
use crate::arrays::fixed_width::FixedWidthType;
2727
use crate::dtype::DType;
2828
use crate::dtype::NativePType;
2929
use crate::dtype::Nullability;
@@ -50,8 +50,6 @@ use crate::builtins::ArrayBuiltins;
5050

5151
/// The validity bitmap indicating which elements are non-null.
5252
pub(super) const VALIDITY_SLOT: usize = 0;
53-
pub(super) const NUM_SLOTS: usize = 1;
54-
pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"];
5553

5654
/// A primitive array that stores [native types][crate::dtype::NativePType] in a contiguous buffer
5755
/// of memory, along with an optional validity child.
@@ -85,17 +83,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"];
8583
/// # Ok(())
8684
/// # }
8785
/// ```
88-
#[derive(Clone, Debug)]
89-
pub struct PrimitiveData {
90-
pub(super) ptype: PType,
91-
pub(super) buffer: BufferHandle,
92-
}
93-
94-
impl Display for PrimitiveData {
95-
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
96-
write!(f, "ptype: {}", self.ptype)
97-
}
98-
}
86+
pub type PrimitiveData = FixedWidthData;
9987

10088
pub struct PrimitiveDataParts {
10189
pub ptype: PType,
@@ -237,9 +225,9 @@ pub trait PrimitiveArrayExt: TypedArrayRef<Primitive> {
237225
impl<T: TypedArrayRef<Primitive>> PrimitiveArrayExt for T {}
238226

239227
// TODO(connor): There are a lot of places where we could be using `new_unchecked` in the codebase.
240-
impl PrimitiveData {
228+
impl FixedWidthData {
241229
/// Build the slots vector for this array.
242-
pub(super) fn make_slots(validity: &Validity, len: usize) -> ArraySlots {
230+
pub(crate) fn make_slots(validity: &Validity, len: usize) -> ArraySlots {
243231
smallvec![validity_to_child(validity, len)]
244232
}
245233

@@ -255,7 +243,8 @@ impl PrimitiveData {
255243
_validity: Validity,
256244
) -> Self {
257245
Self {
258-
ptype,
246+
physical_type: FixedWidthType::Primitive(ptype),
247+
len: handle.len() / ptype.byte_width(),
259248
buffer: handle,
260249
}
261250
}
@@ -306,7 +295,8 @@ impl PrimitiveData {
306295
.vortex_expect("[Debug Assertion]: Invalid `PrimitiveArray` parameters");
307296

308297
Self {
309-
ptype: T::PTYPE,
298+
physical_type: FixedWidthType::Primitive(T::PTYPE),
299+
len: buffer.len(),
310300
buffer: BufferHandle::new_host(buffer.into_byte_buffer()),
311301
}
312302
}
@@ -506,18 +496,27 @@ impl Array<Primitive> {
506496
}
507497
}
508498

509-
impl PrimitiveData {
499+
impl FixedWidthData {
510500
pub fn len(&self) -> usize {
511-
self.buffer.len() / self.ptype.byte_width()
501+
self.len
512502
}
513503

514504
/// Returns `true` if the array is empty.
515505
pub fn is_empty(&self) -> bool {
516-
self.buffer.is_empty()
506+
self.len == 0
517507
}
518508

519509
pub fn ptype(&self) -> PType {
520-
self.ptype
510+
match self.physical_type {
511+
FixedWidthType::Primitive(ptype) => ptype,
512+
FixedWidthType::Decimal(_) => {
513+
vortex_panic!("decimal fixed-width data does not have a primitive type")
514+
}
515+
}
516+
}
517+
518+
pub fn physical_type(&self) -> FixedWidthType {
519+
self.physical_type
521520
}
522521

523522
/// Get access to the buffer handle backing the array.
@@ -527,7 +526,8 @@ impl PrimitiveData {
527526

528527
pub fn from_buffer_handle(handle: BufferHandle, ptype: PType, _validity: Validity) -> Self {
529528
Self {
530-
ptype,
529+
physical_type: FixedWidthType::Primitive(ptype),
530+
len: handle.len() / ptype.byte_width(),
531531
buffer: handle,
532532
}
533533
}

vortex-array/src/arrays/primitive/compute/zip.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use crate::IntoArray;
1515
use crate::array::ArrayView;
1616
use crate::arrays::Primitive;
1717
use crate::arrays::PrimitiveArray;
18+
use crate::dtype::DType;
1819
use crate::dtype::NativePType;
1920
use crate::match_each_native_ptype;
2021
use crate::scalar_fn::fns::zip::ZipKernel;
@@ -33,6 +34,9 @@ impl ZipKernel for Primitive {
3334
mask: &ArrayRef,
3435
ctx: &mut ExecutionCtx,
3536
) -> VortexResult<Option<ArrayRef>> {
37+
if !matches!(if_true.dtype(), DType::Primitive(..)) {
38+
return Ok(None);
39+
}
3640
let Some(if_false) = if_false.as_opt::<Primitive>() else {
3741
return Ok(None);
3842
};

vortex-array/src/arrays/primitive/mod.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ pub use array::chunk_range;
99
pub use array::patch_chunk;
1010
pub use vtable::PrimitiveArray;
1111

12+
pub use crate::arrays::fixed_width::FixedWidthData;
13+
pub use crate::arrays::fixed_width::FixedWidthType;
14+
1215
pub(crate) mod compute;
1316

1417
mod vtable;

0 commit comments

Comments
 (0)