Skip to content

Commit 1e689a1

Browse files
gatesngatesn
authored
Expose array buffer replacement (#8606)
## Rationale for this change I want to move array buffers between devices, but keep the same array tree. Signed-off-by: "Nicholas Gates" <nick@nickgates.com> Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 5ea73bd commit 1e689a1

53 files changed

Lines changed: 753 additions & 48 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎encodings/alp/src/alp/array.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,14 @@ impl VTable for ALP {
102102
None
103103
}
104104

105+
fn with_buffers(
106+
&self,
107+
array: ArrayView<'_, Self>,
108+
buffers: &[BufferHandle],
109+
) -> VortexResult<ArrayParts<Self>> {
110+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
111+
}
112+
105113
fn serialize(
106114
array: ArrayView<'_, Self>,
107115
_session: &VortexSession,

‎encodings/alp/src/alp_rd/array.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,14 @@ impl VTable for ALPRD {
129129
None
130130
}
131131

132+
fn with_buffers(
133+
&self,
134+
array: ArrayView<'_, Self>,
135+
buffers: &[BufferHandle],
136+
) -> VortexResult<ArrayParts<Self>> {
137+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
138+
}
139+
132140
fn serialize(
133141
array: ArrayView<'_, Self>,
134142
_session: &VortexSession,

‎encodings/bytebool/src/array.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,23 @@ impl VTable for ByteBool {
9494
}
9595
}
9696

97+
fn with_buffers(
98+
&self,
99+
array: ArrayView<'_, Self>,
100+
buffers: &[BufferHandle],
101+
) -> VortexResult<ArrayParts<Self>> {
102+
vortex_ensure!(
103+
buffers.len() == 1,
104+
"Expected 1 buffer, got {}",
105+
buffers.len()
106+
);
107+
let data = ByteBoolData::new(buffers[0].clone());
108+
Ok(
109+
ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
110+
.with_slots(array.slots().iter().cloned().collect()),
111+
)
112+
}
113+
97114
fn serialize(
98115
_array: ArrayView<'_, Self>,
99116
_session: &VortexSession,

‎encodings/datetime-parts/src/array.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,14 @@ impl VTable for DateTimeParts {
122122
vortex_panic!("DateTimePartsArray buffer_name index {idx} out of bounds")
123123
}
124124

125+
fn with_buffers(
126+
&self,
127+
array: ArrayView<'_, Self>,
128+
buffers: &[BufferHandle],
129+
) -> VortexResult<ArrayParts<Self>> {
130+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
131+
}
132+
125133
fn serialize(
126134
array: ArrayView<'_, Self>,
127135
_session: &VortexSession,

‎encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,14 @@ impl VTable for DecimalByteParts {
108108
vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds")
109109
}
110110

111+
fn with_buffers(
112+
&self,
113+
array: ArrayView<'_, Self>,
114+
buffers: &[BufferHandle],
115+
) -> VortexResult<ArrayParts<Self>> {
116+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
117+
}
118+
111119
fn serialize(
112120
array: ArrayView<'_, Self>,
113121
_session: &VortexSession,

‎encodings/experimental/onpair/src/array.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,24 @@ impl VTable for OnPair {
367367
}
368368
}
369369

370+
fn with_buffers(
371+
&self,
372+
array: ArrayView<'_, Self>,
373+
buffers: &[BufferHandle],
374+
) -> VortexResult<ArrayParts<Self>> {
375+
vortex_ensure!(
376+
buffers.len() == 1,
377+
"Expected 1 buffer, got {}",
378+
buffers.len()
379+
);
380+
let mut data = array.data().clone();
381+
data.dict_bytes = buffers[0].clone();
382+
Ok(
383+
ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
384+
.with_slots(array.slots().iter().cloned().collect()),
385+
)
386+
}
387+
370388
fn serialize(
371389
array: ArrayView<'_, Self>,
372390
_session: &VortexSession,

‎encodings/fastlanes/src/bitpacking/vtable/mod.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ use vortex_array::vtable::validity_to_child;
3535
use vortex_error::VortexExpect;
3636
use vortex_error::VortexResult;
3737
use vortex_error::vortex_bail;
38+
use vortex_error::vortex_ensure;
3839
use vortex_error::vortex_err;
3940
use vortex_error::vortex_panic;
4041
use vortex_session::VortexSession;
@@ -141,6 +142,24 @@ impl VTable for BitPacked {
141142
}
142143
}
143144

145+
fn with_buffers(
146+
&self,
147+
array: ArrayView<'_, Self>,
148+
buffers: &[BufferHandle],
149+
) -> VortexResult<ArrayParts<Self>> {
150+
vortex_ensure!(
151+
buffers.len() == 1,
152+
"Expected 1 buffer, got {}",
153+
buffers.len()
154+
);
155+
let mut data = array.data().clone();
156+
data.packed = buffers[0].clone();
157+
Ok(
158+
ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
159+
.with_slots(array.slots().iter().cloned().collect()),
160+
)
161+
}
162+
144163
fn serialize(
145164
array: ArrayView<'_, Self>,
146165
_session: &VortexSession,

‎encodings/fastlanes/src/delta/vtable/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,14 @@ impl VTable for Delta {
108108
None
109109
}
110110

111+
fn with_buffers(
112+
&self,
113+
array: ArrayView<'_, Self>,
114+
buffers: &[BufferHandle],
115+
) -> VortexResult<ArrayParts<Self>> {
116+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
117+
}
118+
111119
fn reduce_parent(
112120
array: ArrayView<'_, Self>,
113121
parent: &ArrayRef,

‎encodings/fastlanes/src/for/vtable/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,14 @@ impl VTable for FoR {
9898
None
9999
}
100100

101+
fn with_buffers(
102+
&self,
103+
array: ArrayView<'_, Self>,
104+
buffers: &[BufferHandle],
105+
) -> VortexResult<ArrayParts<Self>> {
106+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
107+
}
108+
101109
fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
102110
SLOT_NAMES[idx].to_string()
103111
}

‎encodings/fastlanes/src/rle/vtable/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,14 @@ impl VTable for RLE {
121121
None
122122
}
123123

124+
fn with_buffers(
125+
&self,
126+
array: ArrayView<'_, Self>,
127+
buffers: &[BufferHandle],
128+
) -> VortexResult<ArrayParts<Self>> {
129+
vortex_array::vtable::with_empty_buffers(self, array, buffers)
130+
}
131+
124132
fn reduce_parent(
125133
array: ArrayView<'_, Self>,
126134
parent: &ArrayRef,

0 commit comments

Comments
 (0)