Skip to content

Commit a371bd1

Browse files
connortsui20connortsui20
committed
Move primitive storage behind fixed-width abstraction
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 2ea51cb commit a371bd1

4 files changed

Lines changed: 130 additions & 30 deletions

File tree

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

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Shared physical storage for canonical fixed-width arrays.
5+
6+
use std::ops::Deref;
7+
8+
use vortex_error::VortexResult;
9+
use vortex_error::vortex_ensure;
10+
use vortex_error::vortex_panic;
11+
12+
use crate::ArrayParts;
13+
use crate::array::ArrayView;
14+
use crate::array::VTable;
15+
use crate::buffer::BufferHandle;
16+
17+
/// A single buffer containing contiguous fixed-width values.
18+
#[derive(Clone, Debug)]
19+
pub(crate) struct FixedWidthData {
20+
buffer: BufferHandle,
21+
}
22+
23+
impl FixedWidthData {
24+
pub(crate) fn new(buffer: BufferHandle) -> Self {
25+
Self { buffer }
26+
}
27+
28+
pub(crate) fn buffer_handle(&self) -> &BufferHandle {
29+
&self.buffer
30+
}
31+
32+
pub(crate) fn into_buffer_handle(self) -> BufferHandle {
33+
self.buffer
34+
}
35+
}
36+
37+
impl Deref for FixedWidthData {
38+
type Target = BufferHandle;
39+
40+
fn deref(&self) -> &Self::Target {
41+
&self.buffer
42+
}
43+
}
44+
45+
pub(crate) trait FixedWidth {
46+
fn fixed_width(&self) -> &FixedWidthData;
47+
48+
fn fixed_width_mut(&mut self) -> &mut FixedWidthData;
49+
}
50+
51+
pub(crate) fn nbuffers<V>(_array: ArrayView<'_, V>) -> usize
52+
where
53+
V: VTable,
54+
V::TypedArrayData: FixedWidth,
55+
{
56+
1
57+
}
58+
59+
pub(crate) fn buffer<V>(array: ArrayView<'_, V>, idx: usize) -> BufferHandle
60+
where
61+
V: VTable,
62+
V::TypedArrayData: FixedWidth,
63+
{
64+
match idx {
65+
0 => array.data().fixed_width().buffer_handle().clone(),
66+
_ => vortex_panic!("fixed-width buffer index {idx} out of bounds"),
67+
}
68+
}
69+
70+
pub(crate) fn buffer_name<V>(_array: ArrayView<'_, V>, idx: usize) -> Option<String>
71+
where
72+
V: VTable,
73+
V::TypedArrayData: FixedWidth,
74+
{
75+
match idx {
76+
0 => Some("values".to_string()),
77+
_ => None,
78+
}
79+
}
80+
81+
pub(crate) fn with_buffers<V>(
82+
vtable: &V,
83+
array: ArrayView<'_, V>,
84+
buffers: &[BufferHandle],
85+
) -> VortexResult<ArrayParts<V>>
86+
where
87+
V: VTable,
88+
V::TypedArrayData: FixedWidth,
89+
{
90+
vortex_ensure!(
91+
buffers.len() == 1,
92+
"Expected 1 buffer, got {}",
93+
buffers.len()
94+
);
95+
let mut data = array.data().clone();
96+
data.fixed_width_mut().buffer = buffers[0].clone();
97+
Ok(
98+
ArrayParts::new(vtable.clone(), array.dtype().clone(), array.len(), data)
99+
.with_slots(array.slots().iter().cloned().collect()),
100+
)
101+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ pub mod filter;
6060
pub use filter::Filter;
6161
pub use filter::FilterArray;
6262

63+
pub(crate) mod fixed_width;
64+
6365
pub mod fixed_size_list;
6466
pub use fixed_size_list::FixedSizeList;
6567
pub use fixed_size_list::FixedSizeListArray;

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

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ use crate::array::TypedArrayRef;
2424
use crate::arrays::BoolArray;
2525
use crate::arrays::Primitive;
2626
use crate::arrays::PrimitiveArray;
27+
use crate::arrays::fixed_width::FixedWidth;
28+
use crate::arrays::fixed_width::FixedWidthData;
2729
use crate::dtype::DType;
2830
use crate::dtype::NativePType;
2931
use crate::dtype::Nullability;
@@ -88,7 +90,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"];
8890
#[derive(Clone, Debug)]
8991
pub struct PrimitiveData {
9092
pub(super) ptype: PType,
91-
pub(super) buffer: BufferHandle,
93+
pub(super) buffer: FixedWidthData,
9294
}
9395

9496
impl Display for PrimitiveData {
@@ -97,6 +99,16 @@ impl Display for PrimitiveData {
9799
}
98100
}
99101

102+
impl FixedWidth for PrimitiveData {
103+
fn fixed_width(&self) -> &FixedWidthData {
104+
&self.buffer
105+
}
106+
107+
fn fixed_width_mut(&mut self) -> &mut FixedWidthData {
108+
&mut self.buffer
109+
}
110+
}
111+
100112
pub struct PrimitiveDataParts {
101113
pub ptype: PType,
102114
pub buffer: BufferHandle,
@@ -256,7 +268,7 @@ impl PrimitiveData {
256268
) -> Self {
257269
Self {
258270
ptype,
259-
buffer: handle,
271+
buffer: FixedWidthData::new(handle),
260272
}
261273
}
262274

@@ -307,7 +319,7 @@ impl PrimitiveData {
307319

308320
Self {
309321
ptype: T::PTYPE,
310-
buffer: BufferHandle::new_host(buffer.into_byte_buffer()),
322+
buffer: FixedWidthData::new(BufferHandle::new_host(buffer.into_byte_buffer())),
311323
}
312324
}
313325

@@ -475,7 +487,7 @@ impl Array<Primitive> {
475487
let data = self.into_data();
476488
PrimitiveDataParts {
477489
ptype,
478-
buffer: data.buffer,
490+
buffer: data.buffer.into_buffer_handle(),
479491
validity,
480492
}
481493
}
@@ -528,7 +540,7 @@ impl PrimitiveData {
528540
pub fn from_buffer_handle(handle: BufferHandle, ptype: PType, _validity: Validity) -> Self {
529541
Self {
530542
ptype,
531-
buffer: handle,
543+
buffer: FixedWidthData::new(handle),
532544
}
533545
}
534546

@@ -592,7 +604,7 @@ impl PrimitiveData {
592604
self.ptype()
593605
)
594606
}
595-
Buffer::from_byte_buffer(self.buffer.into_host_sync())
607+
Buffer::from_byte_buffer(self.buffer.into_buffer_handle().into_host_sync())
596608
}
597609

598610
/// Extract a mutable buffer from the PrimitiveData. Attempts to do this with zero-copy
@@ -614,7 +626,8 @@ impl PrimitiveData {
614626
self.ptype()
615627
)
616628
}
617-
let buffer = Buffer::<T>::from_byte_buffer(self.buffer.into_host_sync());
629+
let buffer =
630+
Buffer::<T>::from_byte_buffer(self.buffer.into_buffer_handle().into_host_sync());
618631
buffer.try_into_mut()
619632
}
620633
}

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

Lines changed: 7 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
use vortex_error::VortexResult;
55
use vortex_error::vortex_bail;
66
use vortex_error::vortex_ensure;
7-
use vortex_error::vortex_panic;
87

98
use crate::ArrayParts;
109
use crate::ArrayRef;
@@ -13,6 +12,7 @@ use crate::ExecutionResult;
1312
use crate::array::Array;
1413
use crate::array::ArrayView;
1514
use crate::array::VTable;
15+
use crate::arrays::fixed_width;
1616
use crate::arrays::primitive::PrimitiveData;
1717
use crate::buffer::BufferHandle;
1818
use crate::builders::ArrayBuilder;
@@ -69,40 +69,24 @@ impl VTable for Primitive {
6969
*ID
7070
}
7171

72-
fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
73-
1
72+
fn nbuffers(array: ArrayView<'_, Self>) -> usize {
73+
fixed_width::nbuffers(array)
7474
}
7575

7676
fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
77-
match idx {
78-
0 => array.buffer_handle().clone(),
79-
_ => vortex_panic!("PrimitiveArray buffer index {idx} out of bounds"),
80-
}
77+
fixed_width::buffer(array, idx)
8178
}
8279

83-
fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
84-
match idx {
85-
0 => Some("values".to_string()),
86-
_ => None,
87-
}
80+
fn buffer_name(array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
81+
fixed_width::buffer_name(array, idx)
8882
}
8983

9084
fn with_buffers(
9185
&self,
9286
array: ArrayView<'_, Self>,
9387
buffers: &[BufferHandle],
9488
) -> VortexResult<ArrayParts<Self>> {
95-
vortex_ensure!(
96-
buffers.len() == 1,
97-
"Expected 1 buffer, got {}",
98-
buffers.len()
99-
);
100-
let mut data = array.data().clone();
101-
data.buffer = buffers[0].clone();
102-
Ok(
103-
ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
104-
.with_slots(array.slots().iter().cloned().collect()),
105-
)
89+
fixed_width::with_buffers(self, array, buffers)
10690
}
10791

10892
fn serialize(

0 commit comments

Comments
 (0)