Skip to content

Commit a92a971

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 f6375f8 commit a92a971

File tree

42 files changed

+186
-135
lines changed

Some content is hidden

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

42 files changed

+186
-135
lines changed

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

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

188-
ALPArray::try_new(
188+
Ok(ALPArray::try_new(
189189
encoded,
190190
Exponents {
191191
e: u8::try_from(metadata.exp_e)?,
192192
f: u8::try_from(metadata.exp_f)?,
193193
},
194194
patches,
195-
)
195+
)?
196+
.into_array())
196197
}
197198

198199
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ impl VTable for ALPRD {
199199
metadata: &Self::Metadata,
200200
_buffers: &[BufferHandle],
201201
children: &dyn ArrayChildren,
202-
) -> VortexResult<ALPRDArray> {
202+
) -> VortexResult<ArrayRef> {
203203
if children.len() < 2 {
204204
vortex_bail!(
205205
"Expected at least 2 children for ALPRD encoding, found {}",
@@ -247,7 +247,7 @@ impl VTable for ALPRD {
247247
})
248248
.transpose()?;
249249

250-
ALPRDArray::try_new(
250+
Ok(ALPRDArray::try_new(
251251
dtype.clone(),
252252
left_parts,
253253
left_parts_dictionary,
@@ -259,7 +259,8 @@ impl VTable for ALPRD {
259259
)
260260
})?,
261261
left_parts_patches,
262-
)
262+
)?
263+
.into_array())
263264
}
264265

265266
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {

‎encodings/bytebool/src/array.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ impl VTable for ByteBool {
146146
_metadata: &Self::Metadata,
147147
buffers: &[BufferHandle],
148148
children: &dyn ArrayChildren,
149-
) -> VortexResult<ByteBoolArray> {
149+
) -> VortexResult<ArrayRef> {
150150
let validity = if children.is_empty() {
151151
Validity::from(dtype.nullability())
152152
} else if children.len() == 1 {
@@ -161,7 +161,7 @@ impl VTable for ByteBool {
161161
}
162162
let buffer = buffers[0].clone();
163163

164-
Ok(ByteBoolArray::new(buffer, validity))
164+
Ok(ByteBoolArray::new(buffer, validity).into_array())
165165
}
166166

167167
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ impl VTable for DateTimeParts {
185185
metadata: &Self::Metadata,
186186
_buffers: &[BufferHandle],
187187
children: &dyn ArrayChildren,
188-
) -> VortexResult<DateTimePartsArray> {
188+
) -> VortexResult<ArrayRef> {
189189
if children.len() != 3 {
190190
vortex_bail!(
191191
"Expected 3 children for datetime-parts encoding, found {}",
@@ -209,7 +209,7 @@ impl VTable for DateTimeParts {
209209
len,
210210
)?;
211211

212-
DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)
212+
Ok(DateTimePartsArray::try_new(dtype.clone(), days, seconds, subseconds)?.into_array())
213213
}
214214

215215
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ impl VTable for DecimalByteParts {
160160
metadata: &Self::Metadata,
161161
_buffers: &[BufferHandle],
162162
children: &dyn ArrayChildren,
163-
) -> VortexResult<DecimalBytePartsArray> {
163+
) -> VortexResult<ArrayRef> {
164164
let Some(decimal_dtype) = dtype.as_decimal_opt() else {
165165
vortex_bail!("decoding decimal but given non decimal dtype {}", dtype)
166166
};
@@ -174,7 +174,7 @@ impl VTable for DecimalByteParts {
174174
"lower_part_count > 0 not currently supported"
175175
);
176176

177-
DecimalBytePartsArray::try_new(msp, *decimal_dtype)
177+
Ok(DecimalBytePartsArray::try_new(msp, *decimal_dtype)?.into_array())
178178
}
179179

180180
fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {

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

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use vortex_array::IntoArray;
1414
use vortex_array::Precision;
1515
use vortex_array::ProstMetadata;
1616
use vortex_array::SerializeMetadata;
17+
use vortex_array::arrays::lazy_patched::LazyPatchedArray;
1718
use vortex_array::buffer::BufferHandle;
1819
use vortex_array::builders::ArrayBuilder;
1920
use vortex_array::dtype::DType;
@@ -46,6 +47,7 @@ use crate::bitpack_decompress::unpack_array;
4647
use crate::bitpack_decompress::unpack_into_primitive_builder;
4748
use crate::bitpacking::vtable::kernels::PARENT_KERNELS;
4849
use crate::bitpacking::vtable::rules::RULES;
50+
4951
mod kernels;
5052
mod operations;
5153
mod rules;
@@ -277,7 +279,7 @@ impl VTable for BitPacked {
277279
metadata: &Self::Metadata,
278280
buffers: &[BufferHandle],
279281
children: &dyn ArrayChildren,
280-
) -> VortexResult<BitPackedArray> {
282+
) -> VortexResult<ArrayRef> {
281283
if buffers.len() != 1 {
282284
vortex_bail!("Expected 1 buffer, got {}", buffers.len());
283285
}
@@ -307,25 +309,11 @@ impl VTable for BitPacked {
307309

308310
let validity = load_validity(validity_idx)?;
309311

310-
let patches = metadata
311-
.patches
312-
.map(|p| {
313-
let indices = children.get(0, &p.indices_dtype()?, p.len()?)?;
314-
let values = children.get(1, dtype, p.len()?)?;
315-
let chunk_offsets = p
316-
.chunk_offsets_dtype()?
317-
.map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize))
318-
.transpose()?;
319-
320-
Patches::new(len, p.offset()?, indices, values, chunk_offsets)
321-
})
322-
.transpose()?;
323-
324-
BitPackedArray::try_new(
312+
let bitpacked = BitPackedArray::try_new(
325313
packed,
326314
PType::try_from(dtype)?,
327315
validity,
328-
patches,
316+
None,
329317
u8::try_from(metadata.bit_width).map_err(|_| {
330318
vortex_err!(
331319
"BitPackedMetadata bit_width {} does not fit in u8",
@@ -339,7 +327,24 @@ impl VTable for BitPacked {
339327
metadata.offset
340328
)
341329
})?,
342-
)
330+
)?
331+
.into_array();
332+
333+
match metadata.patches {
334+
Some(p) => {
335+
let indices = children.get(0, &p.indices_dtype()?, p.len()?)?;
336+
let values = children.get(1, dtype, p.len()?)?;
337+
let chunk_offsets = p
338+
.chunk_offsets_dtype()?
339+
.map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize))
340+
.transpose()?;
341+
342+
let patches = Patches::new(len, p.offset()?, indices, values, chunk_offsets)?;
343+
344+
Ok(LazyPatchedArray::try_new(bitpacked, patches)?.into_array())
345+
}
346+
None => Ok(bitpacked),
347+
}
343348
}
344349

345350
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
@@ -176,7 +176,7 @@ impl VTable for Delta {
176176
metadata: &Self::Metadata,
177177
_buffers: &[BufferHandle],
178178
children: &dyn ArrayChildren,
179-
) -> VortexResult<DeltaArray> {
179+
) -> VortexResult<ArrayRef> {
180180
assert_eq!(children.len(), 2);
181181
let ptype = PType::try_from(dtype)?;
182182
let lanes = match_each_unsigned_integer_ptype!(ptype, |T| { <T as FastLanes>::LANES });
@@ -191,7 +191,7 @@ impl VTable for Delta {
191191
let bases = children.get(0, dtype, bases_len)?;
192192
let deltas = children.get(1, dtype, deltas_len)?;
193193

194-
DeltaArray::try_new(bases, deltas, metadata.0.offset as usize, len)
194+
Ok(DeltaArray::try_new(bases, deltas, metadata.0.offset as usize, len)?.into_array())
195195
}
196196

197197
fn execute(array: Arc<Self::Array>, 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
@@ -150,7 +150,7 @@ impl VTable for FoR {
150150
metadata: &Self::Metadata,
151151
_buffers: &[BufferHandle],
152152
children: &dyn ArrayChildren,
153-
) -> VortexResult<FoRArray> {
153+
) -> VortexResult<ArrayRef> {
154154
if children.len() != 1 {
155155
vortex_bail!(
156156
"Expected 1 child for FoR encoding, found {}",
@@ -160,7 +160,7 @@ impl VTable for FoR {
160160

161161
let encoded = children.get(0, dtype, len)?;
162162

163-
FoRArray::try_new(encoded, metadata.clone())
163+
Ok(FoRArray::try_new(encoded, metadata.clone())?.into_array())
164164
}
165165

166166
fn reduce_parent(

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ impl VTable for RLE {
195195
metadata: &Self::Metadata,
196196
_buffers: &[BufferHandle],
197197
children: &dyn ArrayChildren,
198-
) -> VortexResult<RLEArray> {
198+
) -> VortexResult<ArrayRef> {
199199
let metadata = &metadata.0;
200200
let values = children.get(
201201
0,
@@ -218,13 +218,14 @@ impl VTable for RLE {
218218
usize::try_from(metadata.values_idx_offsets_len)?,
219219
)?;
220220

221-
RLEArray::try_new(
221+
Ok(RLEArray::try_new(
222222
values,
223223
indices,
224224
values_idx_offsets,
225225
metadata.offset as usize,
226226
len,
227-
)
227+
)?
228+
.into_array())
228229
}
229230

230231
fn execute_parent(

‎encodings/fsst/src/array.rs‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,7 @@ impl VTable for FSST {
222222
metadata: &Self::Metadata,
223223
buffers: &[BufferHandle],
224224
children: &dyn ArrayChildren,
225-
) -> VortexResult<FSSTArray> {
225+
) -> VortexResult<ArrayRef> {
226226
let symbols = Buffer::<Symbol>::from_byte_buffer(buffers[0].clone().try_to_host_sync()?);
227227
let symbol_lengths = Buffer::<u8>::from_byte_buffer(buffers[1].clone().try_to_host_sync()?);
228228

@@ -250,13 +250,14 @@ impl VTable for FSST {
250250
len,
251251
)?;
252252

253-
return FSSTArray::try_new(
253+
return Ok(FSSTArray::try_new(
254254
dtype.clone(),
255255
symbols,
256256
symbol_lengths,
257257
codes,
258258
uncompressed_lengths,
259-
);
259+
)?
260+
.into_array());
260261
}
261262

262263
// Check for the current deserialization path.
@@ -297,13 +298,14 @@ impl VTable for FSST {
297298
codes_validity,
298299
)?;
299300

300-
return FSSTArray::try_new(
301+
return Ok(FSSTArray::try_new(
301302
dtype.clone(),
302303
symbols,
303304
symbol_lengths,
304305
codes,
305306
uncompressed_lengths,
306-
);
307+
)?
308+
.into_array());
307309
}
308310

309311
vortex_bail!(

0 commit comments

Comments
 (0)