Skip to content

Commit e8e4a97

Browse files
a10ya10y
committed
update VTable::build to return ArrayRef
This lets us return something other than the original array encoding at read time. Currently we'll want this so that BitPacked::build returns a LazyPatched, but this is applicable for pretty much any back-compat preserving encoding rewrites. Signed-off-by: Andrew Duffy <andrew@a10y.dev>
1 parent 44fbb7f commit e8e4a97

File tree

45 files changed

+309
-184
lines changed

Some content is hidden

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

45 files changed

+309
-184
lines changed

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ impl VTable for ALP {
161161
metadata: &Self::Metadata,
162162
_buffers: &[BufferHandle],
163163
children: &dyn ArrayChildren,
164-
) -> VortexResult<ALPArray> {
164+
) -> VortexResult<ArrayRef> {
165165
let encoded_ptype = match &dtype {
166166
DType::Primitive(PType::F32, n) => DType::Primitive(PType::I32, *n),
167167
DType::Primitive(PType::F64, n) => DType::Primitive(PType::I64, *n),
@@ -183,14 +183,15 @@ impl VTable for ALP {
183183
})
184184
.transpose()?;
185185

186-
ALPArray::try_new(
186+
Ok(ALPArray::try_new(
187187
encoded,
188188
Exponents {
189189
e: u8::try_from(metadata.exp_e)?,
190190
f: u8::try_from(metadata.exp_f)?,
191191
},
192192
patches,
193-
)
193+
)?
194+
.into_array())
194195
}
195196

196197
fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ impl VTable for ALPRD {
168168
metadata: &Self::Metadata,
169169
_buffers: &[BufferHandle],
170170
children: &dyn ArrayChildren,
171-
) -> VortexResult<ALPRDArray> {
171+
) -> VortexResult<ArrayRef> {
172172
if children.len() < 2 {
173173
vortex_bail!(
174174
"Expected at least 2 children for ALPRD encoding, found {}",
@@ -216,7 +216,7 @@ impl VTable for ALPRD {
216216
})
217217
.transpose()?;
218218

219-
ALPRDArray::try_new(
219+
Ok(ALPRDArray::try_new(
220220
dtype.clone(),
221221
left_parts,
222222
left_parts_dictionary,
@@ -228,7 +228,8 @@ impl VTable for ALPRD {
228228
)
229229
})?,
230230
left_parts_patches,
231-
)
231+
)?
232+
.into_array())
232233
}
233234

234235
fn slots(array: &ALPRDArray) -> &[Option<ArrayRef>] {

‎encodings/bytebool/src/array.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ impl VTable for ByteBool {
126126
_metadata: &Self::Metadata,
127127
buffers: &[BufferHandle],
128128
children: &dyn ArrayChildren,
129-
) -> VortexResult<ByteBoolArray> {
129+
) -> VortexResult<ArrayRef> {
130130
let validity = if children.is_empty() {
131131
Validity::from(dtype.nullability())
132132
} else if children.len() == 1 {
@@ -141,7 +141,7 @@ impl VTable for ByteBool {
141141
}
142142
let buffer = buffers[0].clone();
143143

144-
Ok(ByteBoolArray::new(buffer, validity))
144+
Ok(ByteBoolArray::new(buffer, validity).into_array())
145145
}
146146

147147
fn slots(array: &ByteBoolArray) -> &[Option<ArrayRef>] {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ impl VTable for DateTimeParts {
164164
metadata: &Self::Metadata,
165165
_buffers: &[BufferHandle],
166166
children: &dyn ArrayChildren,
167-
) -> VortexResult<DateTimePartsArray> {
167+
) -> VortexResult<ArrayRef> {
168168
if children.len() != 3 {
169169
vortex_bail!(
170170
"Expected 3 children for datetime-parts encoding, found {}",
@@ -188,7 +188,7 @@ impl VTable for DateTimeParts {
188188
len,
189189
)?;
190190

191-
DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)
191+
Ok(DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)?.into_array())
192192
}
193193

194194
fn slots(array: &DateTimePartsArray) -> &[Option<ArrayRef>] {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ impl VTable for DecimalByteParts {
142142
metadata: &Self::Metadata,
143143
_buffers: &[BufferHandle],
144144
children: &dyn ArrayChildren,
145-
) -> VortexResult<DecimalBytePartsArray> {
145+
) -> VortexResult<ArrayRef> {
146146
let Some(decimal_dtype) = dtype.as_decimal_opt() else {
147147
vortex_bail!("decoding decimal but given non decimal dtype {}", dtype)
148148
};
@@ -156,7 +156,7 @@ impl VTable for DecimalByteParts {
156156
"lower_part_count > 0 not currently supported"
157157
);
158158

159-
DecimalBytePartsArray::try_new(msp, *decimal_dtype)
159+
Ok(DecimalBytePartsArray::try_new(msp, *decimal_dtype)?.into_array())
160160
}
161161

162162
fn slots(array: &DecimalBytePartsArray) -> &[Option<ArrayRef>] {

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

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use vortex_array::IntoArray;
1515
use vortex_array::Precision;
1616
use vortex_array::ProstMetadata;
1717
use vortex_array::SerializeMetadata;
18+
use vortex_array::arrays::lazy_patched::LazyPatchedArray;
1819
use vortex_array::buffer::BufferHandle;
1920
use vortex_array::builders::ArrayBuilder;
2021
use vortex_array::dtype::DType;
@@ -50,6 +51,7 @@ use crate::bitpacking::array::SLOT_NAMES;
5051
use crate::bitpacking::array::VALIDITY_SLOT;
5152
use crate::bitpacking::vtable::kernels::PARENT_KERNELS;
5253
use crate::bitpacking::vtable::rules::RULES;
54+
5355
mod kernels;
5456
mod operations;
5557
mod rules;
@@ -208,7 +210,7 @@ impl VTable for BitPacked {
208210
metadata: &Self::Metadata,
209211
buffers: &[BufferHandle],
210212
children: &dyn ArrayChildren,
211-
) -> VortexResult<BitPackedArray> {
213+
) -> VortexResult<ArrayRef> {
212214
if buffers.len() != 1 {
213215
vortex_bail!("Expected 1 buffer, got {}", buffers.len());
214216
}
@@ -238,25 +240,11 @@ impl VTable for BitPacked {
238240

239241
let validity = load_validity(validity_idx)?;
240242

241-
let patches = metadata
242-
.patches
243-
.map(|p| {
244-
let indices = children.get(0, &p.indices_dtype()?, p.len()?)?;
245-
let values = children.get(1, dtype, p.len()?)?;
246-
let chunk_offsets = p
247-
.chunk_offsets_dtype()?
248-
.map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize))
249-
.transpose()?;
250-
251-
Patches::new(len, p.offset()?, indices, values, chunk_offsets)
252-
})
253-
.transpose()?;
254-
255-
BitPackedArray::try_new(
243+
let bitpacked = BitPackedArray::try_new(
256244
packed,
257245
PType::try_from(dtype)?,
258246
validity,
259-
patches,
247+
None,
260248
u8::try_from(metadata.bit_width).map_err(|_| {
261249
vortex_err!(
262250
"BitPackedMetadata bit_width {} does not fit in u8",
@@ -270,7 +258,24 @@ impl VTable for BitPacked {
270258
metadata.offset
271259
)
272260
})?,
273-
)
261+
)?
262+
.into_array();
263+
264+
match metadata.patches {
265+
Some(p) => {
266+
let indices = children.get(0, &p.indices_dtype()?, p.len()?)?;
267+
let values = children.get(1, dtype, p.len()?)?;
268+
let chunk_offsets = p
269+
.chunk_offsets_dtype()?
270+
.map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize))
271+
.transpose()?;
272+
273+
let patches = Patches::new(len, p.offset()?, indices, values, chunk_offsets)?;
274+
275+
Ok(LazyPatchedArray::try_new(bitpacked, patches)?.into_array())
276+
}
277+
None => Ok(bitpacked),
278+
}
274279
}
275280

276281
fn append_to_builder(

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ impl VTable for Delta {
161161
metadata: &Self::Metadata,
162162
_buffers: &[BufferHandle],
163163
children: &dyn ArrayChildren,
164-
) -> VortexResult<DeltaArray> {
164+
) -> VortexResult<ArrayRef> {
165165
assert_eq!(children.len(), 2);
166166
let ptype = PType::try_from(dtype)?;
167167
let lanes = match_each_unsigned_integer_ptype!(ptype, |T| { <T as FastLanes>::LANES });
@@ -176,7 +176,7 @@ impl VTable for Delta {
176176
let bases = children.get(0, dtype, bases_len)?;
177177
let deltas = children.get(1, dtype, deltas_len)?;
178178

179-
DeltaArray::try_new(bases, deltas, metadata.0.offset as usize, len)
179+
Ok(DeltaArray::try_new(bases, deltas, metadata.0.offset as usize, len)?.into_array())
180180
}
181181

182182
fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ impl VTable for FoR {
139139
metadata: &Self::Metadata,
140140
_buffers: &[BufferHandle],
141141
children: &dyn ArrayChildren,
142-
) -> VortexResult<FoRArray> {
142+
) -> VortexResult<ArrayRef> {
143143
if children.len() != 1 {
144144
vortex_bail!(
145145
"Expected 1 child for FoR encoding, found {}",
@@ -149,7 +149,7 @@ impl VTable for FoR {
149149

150150
let encoded = children.get(0, dtype, len)?;
151151

152-
FoRArray::try_new(encoded, metadata.clone())
152+
Ok(FoRArray::try_new(encoded, metadata.clone())?.into_array())
153153
}
154154

155155
fn reduce_parent(

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ impl VTable for RLE {
174174
metadata: &Self::Metadata,
175175
_buffers: &[BufferHandle],
176176
children: &dyn ArrayChildren,
177-
) -> VortexResult<RLEArray> {
177+
) -> VortexResult<ArrayRef> {
178178
let metadata = &metadata.0;
179179
let values = children.get(
180180
0,
@@ -197,13 +197,14 @@ impl VTable for RLE {
197197
usize::try_from(metadata.values_idx_offsets_len)?,
198198
)?;
199199

200-
RLEArray::try_new(
200+
Ok(RLEArray::try_new(
201201
values,
202202
indices,
203203
values_idx_offsets,
204204
metadata.offset as usize,
205205
len,
206-
)
206+
)?
207+
.into_array())
207208
}
208209

209210
fn execute_parent(

‎encodings/fsst/src/array.rs‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ impl VTable for FSST {
199199
metadata: &Self::Metadata,
200200
buffers: &[BufferHandle],
201201
children: &dyn ArrayChildren,
202-
) -> VortexResult<FSSTArray> {
202+
) -> VortexResult<ArrayRef> {
203203
let symbols = Buffer::<Symbol>::from_byte_buffer(buffers[0].clone().try_to_host_sync()?);
204204
let symbol_lengths = Buffer::<u8>::from_byte_buffer(buffers[1].clone().try_to_host_sync()?);
205205

@@ -227,13 +227,14 @@ impl VTable for FSST {
227227
len,
228228
)?;
229229

230-
return FSSTArray::try_new(
230+
return Ok(FSSTArray::try_new(
231231
dtype.clone(),
232232
symbols,
233233
symbol_lengths,
234234
codes,
235235
uncompressed_lengths,
236-
);
236+
)?
237+
.into_array());
237238
}
238239

239240
// Check for the current deserialization path.
@@ -274,13 +275,14 @@ impl VTable for FSST {
274275
codes_validity,
275276
)?;
276277

277-
return FSSTArray::try_new(
278+
return Ok(FSSTArray::try_new(
278279
dtype.clone(),
279280
symbols,
280281
symbol_lengths,
281282
codes,
282283
uncompressed_lengths,
283-
);
284+
)?
285+
.into_array());
284286
}
285287

286288
vortex_bail!(

0 commit comments

Comments
 (0)