Skip to content

Commit 5c92e4e

Browse files
gatesngatesn
authored
refactor(array): propagate allocators through builders (#9670)
## Summary Require array builder call sites to choose a buffer allocator explicitly while keeping the old APIs as deprecated compatibility shims. ## Changes - Add allocator-taking `_in` APIs and deprecate the allocator-free builder constructors. - Take `&BufferAllocatorRef` at API boundaries; clone only when a buffer or builder retains it. - Propagate the allocator through nested, dictionary, and offset-based string builders. - Use `ctx.allocator()` in execution paths; use the static allocator only where no execution context exists. - Add custom allocator coverage. --------- Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 7915813 commit 5c92e4e

96 files changed

Lines changed: 2190 additions & 585 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/fastlanes/benches/canonicalize_bench.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -95,11 +95,13 @@ fn canonical_into_non_nullable(
9595
bencher
9696
.with_inputs(|| {
9797
let chunked = ChunkedArray::from_iter(chunks.clone()).into_array();
98-
let primitive_builder = PrimitiveBuilder::<i32>::with_capacity(
98+
let ctx = SESSION.create_execution_ctx();
99+
let primitive_builder = PrimitiveBuilder::<i32>::with_capacity_in(
99100
chunked.dtype().nullability(),
100101
chunk_len * chunk_count,
102+
ctx.allocator(),
101103
);
102-
(chunked, primitive_builder, SESSION.create_execution_ctx())
104+
(chunked, primitive_builder, ctx)
103105
})
104106
.bench_refs(|(chunked, primitive_builder, ctx)| {
105107
chunked
@@ -174,11 +176,13 @@ fn canonical_into_nullable(
174176
bencher
175177
.with_inputs(|| {
176178
let chunked = ChunkedArray::from_iter(chunks.clone()).into_array();
177-
let primitive_builder = PrimitiveBuilder::<i32>::with_capacity(
179+
let ctx = SESSION.create_execution_ctx();
180+
let primitive_builder = PrimitiveBuilder::<i32>::with_capacity_in(
178181
chunked.dtype().nullability(),
179182
chunk_len * chunk_count,
183+
ctx.allocator(),
180184
);
181-
(chunked, primitive_builder, SESSION.create_execution_ctx())
185+
(chunked, primitive_builder, ctx)
182186
})
183187
.bench_refs(|(chunked, primitive_builder, ctx)| {
184188
chunked

‎encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -512,15 +512,21 @@ mod test {
512512
let chunked = ChunkedArray::from_iter(chunks).into_array();
513513

514514
let into_ca = chunked.clone().execute::<PrimitiveArray>(&mut ctx)?;
515-
let mut primitive_builder =
516-
PrimitiveBuilder::<i32>::with_capacity(chunked.dtype().nullability(), 10 * 100);
515+
let mut primitive_builder = PrimitiveBuilder::<i32>::with_capacity_in(
516+
chunked.dtype().nullability(),
517+
10 * 100,
518+
vortex_buffer::BufferAllocatorRef::static_ref(),
519+
);
517520
chunked.append_to_builder(&mut primitive_builder, &mut ctx)?;
518521
let ca_into = primitive_builder.finish();
519522

520523
assert_arrays_eq!(into_ca, ca_into, &mut ctx);
521524

522-
let mut primitive_builder =
523-
PrimitiveBuilder::<i32>::with_capacity(chunked.dtype().nullability(), 10 * 100);
525+
let mut primitive_builder = PrimitiveBuilder::<i32>::with_capacity_in(
526+
chunked.dtype().nullability(),
527+
10 * 100,
528+
vortex_buffer::BufferAllocatorRef::static_ref(),
529+
);
524530
chunked.append_to_builder(&mut primitive_builder, &mut ctx)?;
525531
let ca_into = primitive_builder.finish();
526532

‎encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,11 @@ pub fn unpack_primitive_array<T: BitPackedUnpack>(
4040
array: ArrayView<'_, BitPacked>,
4141
ctx: &mut ExecutionCtx,
4242
) -> VortexResult<PrimitiveArray> {
43-
let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len());
43+
let mut builder = PrimitiveBuilder::with_capacity_in(
44+
array.dtype().nullability(),
45+
array.len(),
46+
ctx.allocator(),
47+
);
4448
unpack_into_primitive_builder::<T>(array, &mut builder, ctx)?;
4549
assert_eq!(builder.len(), array.len());
4650
Ok(builder.finish_into_primitive())
@@ -294,22 +298,24 @@ mod tests {
294298
Validity::NonNullable,
295299
)?;
296300

297-
let mut listview_builder = ListViewBuilder::<u64, u32>::with_capacity(
301+
let mut listview_builder = ListViewBuilder::<u64, u32>::with_capacity_in(
298302
Arc::clone(&element_dtype),
299303
Nullability::NonNullable,
300304
0,
301305
0,
306+
ctx.allocator(),
302307
);
303308
list.clone()
304309
.into_array()
305310
.append_to_builder(&mut listview_builder, &mut ctx)?;
306311
assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
307312

308-
let mut list_builder = ListBuilder::<u64>::with_capacity(
313+
let mut list_builder = ListBuilder::<u64>::with_capacity_in(
309314
Arc::clone(&element_dtype),
310315
Nullability::NonNullable,
311316
0,
312317
0,
318+
ctx.allocator(),
313319
);
314320
list.clone()
315321
.into_array()
@@ -323,8 +329,13 @@ mod tests {
323329
Buffer::from_iter(std::iter::repeat_n(64u32, 48)).into_array(),
324330
Validity::NonNullable,
325331
)?;
326-
let mut list_builder =
327-
ListBuilder::<u64>::with_capacity(element_dtype, Nullability::NonNullable, 0, 0);
332+
let mut list_builder = ListBuilder::<u64>::with_capacity_in(
333+
element_dtype,
334+
Nullability::NonNullable,
335+
0,
336+
0,
337+
vortex_buffer::BufferAllocatorRef::static_ref(),
338+
);
328339
listview
329340
.into_array()
330341
.append_to_builder(&mut list_builder, &mut ctx)?;
@@ -452,7 +463,10 @@ mod tests {
452463
let empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::<u32>::new());
453464
let bitpacked = encode(&empty, 0);
454465

455-
let mut builder = PrimitiveBuilder::<u32>::new(Nullability::NonNullable);
466+
let mut builder = PrimitiveBuilder::<u32>::new_in(
467+
Nullability::NonNullable,
468+
vortex_buffer::BufferAllocatorRef::static_ref(),
469+
);
456470
unpack_map_into_builder(
457471
bitpacked.as_view(),
458472
&mut builder,
@@ -481,7 +495,11 @@ mod tests {
481495
let bitpacked = encode(&array, 3);
482496

483497
// Unpack into a new builder.
484-
let mut builder = PrimitiveBuilder::<u32>::with_capacity(Nullability::Nullable, 5);
498+
let mut builder = PrimitiveBuilder::<u32>::with_capacity_in(
499+
Nullability::Nullable,
500+
5,
501+
vortex_buffer::BufferAllocatorRef::static_ref(),
502+
);
485503
unpack_map_into_builder(
486504
bitpacked.as_view(),
487505
&mut builder,
@@ -519,7 +537,11 @@ mod tests {
519537
);
520538

521539
// Unpack into a new builder.
522-
let mut builder = PrimitiveBuilder::<u32>::with_capacity(Nullability::NonNullable, 100);
540+
let mut builder = PrimitiveBuilder::<u32>::with_capacity_in(
541+
Nullability::NonNullable,
542+
100,
543+
vortex_buffer::BufferAllocatorRef::static_ref(),
544+
);
523545
unpack_map_into_builder(
524546
bitpacked.as_view(),
525547
&mut builder,

‎encodings/fastlanes/src/bitpacking/compute/cast.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,11 @@ impl CastKernel for BitPacked {
100100
.cast_nullability(tgt_nullability, array.len(), ctx)?;
101101

102102
let result = match_each_integer_ptype!(tgt, |T| {
103-
let mut builder = PrimitiveBuilder::<T>::with_capacity(tgt_nullability, array.len());
103+
let mut builder = PrimitiveBuilder::<T>::with_capacity_in(
104+
tgt_nullability,
105+
array.len(),
106+
ctx.allocator(),
107+
);
104108
match_each_integer_ptype!(src, |F| {
105109
unpack_map_into_builder::<F, T, _>(array, &mut builder, ctx, |v: F| v.as_())?;
106110
});

‎encodings/fastlanes/src/for/array/for_decompress.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,10 @@ pub(crate) fn fused_decompress<
109109
&mut scratch,
110110
)?;
111111

112-
let mut builder = PrimitiveBuilder::<T>::with_capacity(
112+
let mut builder = PrimitiveBuilder::<T>::with_capacity_in(
113113
for_.reference_scalar().dtype().nullability(),
114114
bp.len(),
115+
ctx.allocator(),
115116
);
116117
let mut uninit_range = builder.uninit_range(bp.len());
117118
unsafe {

‎encodings/fsst/benches/chunked_dict_fsst_builder.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use vortex_array::Canonical;
99
use vortex_array::IntoArray;
1010
use vortex_array::VortexSessionExecute;
1111
use vortex_array::arrays::ChunkedArray;
12-
use vortex_array::builders::builder_with_capacity;
12+
use vortex_array::builders::builder_with_capacity_in;
1313
use vortex_array::dtype::NativePType;
1414
use vortex_error::VortexExpect;
1515
use vortex_fsst::test_utils::gen_dict_fsst_test_data;
@@ -57,7 +57,8 @@ fn chunked_dict_fsst_canonical_into(
5757
bencher
5858
.with_inputs(|| (&chunk, SESSION.create_execution_ctx()))
5959
.bench_refs(|(chunk, ctx)| {
60-
let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count);
60+
let mut builder =
61+
builder_with_capacity_in(chunk.dtype(), len * chunk_count, ctx.allocator());
6162
chunk
6263
.append_to_builder(builder.as_mut(), ctx)
6364
.vortex_expect("append failed");

‎encodings/fsst/benches/fsst_compress.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,10 @@ fn chunked_canonicalize_into(
152152
bencher
153153
.with_inputs(|| (&array, SESSION.create_execution_ctx()))
154154
.bench_refs(|(array, ctx)| {
155-
let mut builder = VarBinViewBuilder::with_capacity(
155+
let mut builder = VarBinViewBuilder::with_capacity_in(
156156
DType::Binary(Nullability::NonNullable),
157157
array.len(),
158+
ctx.allocator().clone(),
158159
);
159160
array.append_to_builder(&mut builder, ctx).unwrap();
160161
builder.finish()

‎encodings/fsst/src/canonical.rs‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -220,8 +220,11 @@ mod tests {
220220
let mut ctx = SESSION.create_execution_ctx();
221221
let (chunked_arr, data) = make_data_chunked();
222222

223-
let mut builder =
224-
VarBinViewBuilder::with_capacity(chunked_arr.dtype().clone(), chunked_arr.len());
223+
let mut builder = VarBinViewBuilder::with_capacity_in(
224+
chunked_arr.dtype().clone(),
225+
chunked_arr.len(),
226+
ctx.allocator().clone(),
227+
);
225228
chunked_arr
226229
.clone()
227230
.into_array()
@@ -249,8 +252,11 @@ mod tests {
249252
};
250253

251254
{
252-
let mut builder =
253-
VarBinBuilder::<i32>::with_capacity(chunked_arr.dtype().clone(), data.len());
255+
let mut builder = VarBinBuilder::<i32>::with_capacity_in(
256+
chunked_arr.dtype().clone(),
257+
data.len(),
258+
ctx.allocator(),
259+
);
254260
chunked_arr
255261
.into_array()
256262
.append_to_builder(&mut builder, &mut ctx)?;
@@ -267,7 +273,11 @@ mod tests {
267273
#[test]
268274
fn test_append_after_in_progress_buffer() -> VortexResult<()> {
269275
let dtype = DType::Binary(Nullability::NonNullable);
270-
let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 2);
276+
let mut builder = VarBinViewBuilder::with_capacity_in(
277+
dtype.clone(),
278+
2,
279+
vortex_buffer::BufferAllocatorRef::statically_allocated(),
280+
);
271281
builder.append_value(b"long enough!!!");
272282

273283
let varbin = VarBinArray::from_iter(

‎encodings/fsst/src/compress.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use vortex_array::dtype::DType;
2727
use vortex_array::dtype::IntegerPType;
2828
use vortex_array::dtype::OffsetBuilderPType;
2929
use vortex_array::match_each_integer_ptype;
30+
use vortex_buffer::BufferAllocatorRef;
3031
use vortex_buffer::BufferMut;
3132
use vortex_error::VortexExpect;
3233
use vortex_error::VortexResult;
@@ -206,6 +207,7 @@ where
206207
DType::Binary(strings.dtype().nullability()),
207208
strings.len(),
208209
compressor,
210+
ctx.allocator(),
209211
);
210212
let views = strings.views();
211213
let buffers = strings.data_buffers();
@@ -243,6 +245,7 @@ where
243245
DType::Binary(strings.dtype().nullability()),
244246
strings.len(),
245247
compressor,
248+
ctx.allocator(),
246249
);
247250
let bytes = strings.bytes().as_slice();
248251
match_each_integer_ptype!(offsets.ptype(), |I| {
@@ -288,11 +291,16 @@ struct FsstSink<'c, O: OffsetBuilderPType + 'static> {
288291
}
289292

290293
impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> {
291-
fn with_capacity(dtype: DType, len: usize, compressor: &'c Compressor) -> Self {
294+
fn with_capacity(
295+
dtype: DType,
296+
len: usize,
297+
compressor: &'c Compressor,
298+
allocator: &BufferAllocatorRef,
299+
) -> Self {
292300
Self {
293301
buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN),
294-
builder: VarBinBuilder::<O>::with_capacity(dtype, len),
295-
uncompressed_lengths: BufferMut::with_capacity(len),
302+
builder: VarBinBuilder::<O>::with_capacity_in(dtype, len, allocator),
303+
uncompressed_lengths: BufferMut::with_capacity_in(len, allocator.clone()),
296304
compressor,
297305
}
298306
}

‎encodings/fsst/src/kernel.rs‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,11 @@ mod tests {
6161
});
6262

6363
fn build_test_fsst_array() -> ArrayRef {
64-
let mut builder =
65-
VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::NonNullable), 10);
64+
let mut builder = VarBinBuilder::<i32>::with_capacity_in(
65+
DType::Utf8(Nullability::NonNullable),
66+
10,
67+
vortex_buffer::BufferAllocatorRef::static_ref(),
68+
);
6669
builder.append_value(b"hello world");
6770
builder.append_value(b"foo bar baz");
6871
builder.append_value(b"testing fsst compression");
@@ -132,8 +135,11 @@ mod tests {
132135
// Test case with special characters and nulls
133136
// Values: ["", "", "", "", "", "", "", "", "", "", "", ",", "A<<<<<<<", "", "", "", "", null, null, null, null, null, null]
134137
// Mask: only the last element is selected (true at index 22)
135-
let mut builder =
136-
VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::Nullable), 23);
138+
let mut builder = VarBinBuilder::<i32>::with_capacity_in(
139+
DType::Utf8(Nullability::Nullable),
140+
23,
141+
vortex_buffer::BufferAllocatorRef::static_ref(),
142+
);
137143
// 11 empty strings
138144
for _ in 0..11 {
139145
builder.append_value(b"");
@@ -174,8 +180,11 @@ mod tests {
174180

175181
#[test]
176182
fn filter_only_null() -> VortexResult<()> {
177-
let mut builder =
178-
VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::Nullable), 3);
183+
let mut builder = VarBinBuilder::<i32>::with_capacity_in(
184+
DType::Utf8(Nullability::Nullable),
185+
3,
186+
vortex_buffer::BufferAllocatorRef::static_ref(),
187+
);
179188
builder.push_null();
180189
builder.append_value(b"A");
181190
builder.push_null();
@@ -216,8 +225,11 @@ mod tests {
216225

217226
#[test]
218227
fn test_fsst_byte_length() -> VortexResult<()> {
219-
let mut builder =
220-
VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::NonNullable), 3);
228+
let mut builder = VarBinBuilder::<i32>::with_capacity_in(
229+
DType::Utf8(Nullability::NonNullable),
230+
3,
231+
vortex_buffer::BufferAllocatorRef::static_ref(),
232+
);
221233
builder.append_value(b"hello");
222234
builder.append_value(b"world!!");
223235
builder.append_value("Пуховички"); // 9 characters, 18 bytes

0 commit comments

Comments
 (0)