diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index 30ee0756960..a6dcfb5e300 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -95,11 +95,13 @@ fn canonical_into_non_nullable( bencher .with_inputs(|| { let chunked = ChunkedArray::from_iter(chunks.clone()).into_array(); + let ctx = SESSION.create_execution_ctx(); let primitive_builder = PrimitiveBuilder::::with_capacity( chunked.dtype().nullability(), chunk_len * chunk_count, + ctx.allocator(), ); - (chunked, primitive_builder, SESSION.create_execution_ctx()) + (chunked, primitive_builder, ctx) }) .bench_refs(|(chunked, primitive_builder, ctx)| { chunked @@ -174,11 +176,13 @@ fn canonical_into_nullable( bencher .with_inputs(|| { let chunked = ChunkedArray::from_iter(chunks.clone()).into_array(); + let ctx = SESSION.create_execution_ctx(); let primitive_builder = PrimitiveBuilder::::with_capacity( chunked.dtype().nullability(), chunk_len * chunk_count, + ctx.allocator(), ); - (chunked, primitive_builder, SESSION.create_execution_ctx()) + (chunked, primitive_builder, ctx) }) .bench_refs(|(chunked, primitive_builder, ctx)| { chunked diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index a393db6ecc8..36b1e070ddd 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -512,15 +512,21 @@ mod test { let chunked = ChunkedArray::from_iter(chunks).into_array(); let into_ca = chunked.clone().execute::(&mut ctx)?; - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); + let mut primitive_builder = PrimitiveBuilder::::with_capacity( + chunked.dtype().nullability(), + 10 * 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); assert_arrays_eq!(into_ca, ca_into, &mut ctx); - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); + let mut primitive_builder = PrimitiveBuilder::::with_capacity( + chunked.dtype().nullability(), + 10 * 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index d192e50d04f..49c7b6de572 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -40,7 +40,8 @@ pub fn unpack_primitive_array( array: ArrayView<'_, BitPacked>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len()); + let mut builder = + PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len(), ctx.allocator()); unpack_into_primitive_builder::(array, &mut builder, ctx)?; assert_eq!(builder.len(), array.len()); Ok(builder.finish_into_primitive()) @@ -299,6 +300,7 @@ mod tests { Nullability::NonNullable, 0, 0, + ctx.allocator(), ); list.clone() .into_array() @@ -310,6 +312,7 @@ mod tests { Nullability::NonNullable, 0, 0, + ctx.allocator(), ); list.clone() .into_array() @@ -323,8 +326,13 @@ mod tests { Buffer::from_iter(std::iter::repeat_n(64u32, 48)).into_array(), Validity::NonNullable, )?; - let mut list_builder = - ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + let mut list_builder = ListBuilder::::with_capacity( + element_dtype, + Nullability::NonNullable, + 0, + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); listview .into_array() .append_to_builder(&mut list_builder, &mut ctx)?; @@ -452,7 +460,10 @@ mod tests { let empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::::new()); let bitpacked = encode(&empty, 0); - let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); + let mut builder = PrimitiveBuilder::::new( + Nullability::NonNullable, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); unpack_map_into_builder( bitpacked.as_view(), &mut builder, @@ -481,7 +492,11 @@ mod tests { let bitpacked = encode(&array, 3); // Unpack into a new builder. - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 5, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); unpack_map_into_builder( bitpacked.as_view(), &mut builder, @@ -519,7 +534,11 @@ mod tests { ); // Unpack into a new builder. - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 100); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::NonNullable, + 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); unpack_map_into_builder( bitpacked.as_view(), &mut builder, diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index f7cfc56acf0..39817b39537 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -100,7 +100,8 @@ impl CastKernel for BitPacked { .cast_nullability(tgt_nullability, array.len(), ctx)?; let result = match_each_integer_ptype!(tgt, |T| { - let mut builder = PrimitiveBuilder::::with_capacity(tgt_nullability, array.len()); + let mut builder = + PrimitiveBuilder::::with_capacity(tgt_nullability, array.len(), ctx.allocator()); match_each_integer_ptype!(src, |F| { unpack_map_into_builder::(array, &mut builder, ctx, |v: F| v.as_())?; }); diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index 6e429abc419..1e57a6a094a 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -112,6 +112,7 @@ pub(crate) fn fused_decompress< let mut builder = PrimitiveBuilder::::with_capacity( for_.reference_scalar().dtype().nullability(), bp.len(), + ctx.allocator(), ); let mut uninit_range = builder.uninit_range(bp.len()); unsafe { diff --git a/encodings/fsst/benches/chunked_dict_fsst_builder.rs b/encodings/fsst/benches/chunked_dict_fsst_builder.rs index a59ae9113de..e0fa49de37e 100644 --- a/encodings/fsst/benches/chunked_dict_fsst_builder.rs +++ b/encodings/fsst/benches/chunked_dict_fsst_builder.rs @@ -57,7 +57,8 @@ fn chunked_dict_fsst_canonical_into( bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 6a0c5f310a2..a3bfd29b957 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -155,6 +155,7 @@ fn chunked_canonicalize_into( let mut builder = VarBinViewBuilder::with_capacity( DType::Binary(Nullability::NonNullable), array.len(), + ctx.allocator().clone(), ); array.append_to_builder(&mut builder, ctx).unwrap(); builder.finish() diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 28942648aa4..4c3925539ac 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -220,8 +220,11 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let (chunked_arr, data) = make_data_chunked(); - let mut builder = - VarBinViewBuilder::with_capacity(chunked_arr.dtype().clone(), chunked_arr.len()); + let mut builder = VarBinViewBuilder::with_capacity( + chunked_arr.dtype().clone(), + chunked_arr.len(), + ctx.allocator().clone(), + ); chunked_arr .clone() .into_array() @@ -249,8 +252,11 @@ mod tests { }; { - let mut builder = - VarBinBuilder::::with_capacity(chunked_arr.dtype().clone(), data.len()); + let mut builder = VarBinBuilder::::with_capacity( + chunked_arr.dtype().clone(), + data.len(), + ctx.allocator(), + ); chunked_arr .into_array() .append_to_builder(&mut builder, &mut ctx)?; @@ -267,7 +273,11 @@ mod tests { #[test] fn test_append_after_in_progress_buffer() -> VortexResult<()> { let dtype = DType::Binary(Nullability::NonNullable); - let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 2); + let mut builder = VarBinViewBuilder::with_capacity( + dtype.clone(), + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value(b"long enough!!!"); let varbin = VarBinArray::from_iter( diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index 9df8f6989dd..32f8bbab8fa 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -27,6 +27,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::IntegerPType; use vortex_array::dtype::OffsetBuilderPType; use vortex_array::match_each_integer_ptype; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -206,6 +207,7 @@ where DType::Binary(strings.dtype().nullability()), strings.len(), compressor, + ctx.allocator(), ); let views = strings.views(); let buffers = strings.data_buffers(); @@ -243,6 +245,7 @@ where DType::Binary(strings.dtype().nullability()), strings.len(), compressor, + ctx.allocator(), ); let bytes = strings.bytes().as_slice(); match_each_integer_ptype!(offsets.ptype(), |I| { @@ -288,11 +291,16 @@ struct FsstSink<'c, O: OffsetBuilderPType + 'static> { } impl<'c, O: OffsetBuilderPType + 'static> FsstSink<'c, O> { - fn with_capacity(dtype: DType, len: usize, compressor: &'c Compressor) -> Self { + fn with_capacity( + dtype: DType, + len: usize, + compressor: &'c Compressor, + allocator: &BufferAllocatorRef, + ) -> Self { Self { buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN), - builder: VarBinBuilder::::with_capacity(dtype, len), - uncompressed_lengths: BufferMut::with_capacity(len), + builder: VarBinBuilder::::with_capacity(dtype, len, allocator), + uncompressed_lengths: BufferMut::with_capacity_in(len, allocator.clone()), compressor, } } diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 1ae34dfc800..99b97f346d6 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -61,8 +61,11 @@ mod tests { }); fn build_test_fsst_array() -> ArrayRef { - let mut builder = - VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 10); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::NonNullable), + 10, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"hello world"); builder.append_value(b"foo bar baz"); builder.append_value(b"testing fsst compression"); @@ -132,8 +135,11 @@ mod tests { // Test case with special characters and nulls // Values: ["", "", "", "", "", "", "", "", "", "", "", ",", "A<<<<<<<", "", "", "", "", null, null, null, null, null, null] // Mask: only the last element is selected (true at index 22) - let mut builder = - VarBinBuilder::::with_capacity(DType::Utf8(Nullability::Nullable), 23); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::Nullable), + 23, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); // 11 empty strings for _ in 0..11 { builder.append_value(b""); @@ -174,8 +180,11 @@ mod tests { #[test] fn filter_only_null() -> VortexResult<()> { - let mut builder = - VarBinBuilder::::with_capacity(DType::Utf8(Nullability::Nullable), 3); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::Nullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.push_null(); builder.append_value(b"A"); builder.push_null(); @@ -216,8 +225,11 @@ mod tests { #[test] fn test_fsst_byte_length() -> VortexResult<()> { - let mut builder = - VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 3); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"hello"); builder.append_value(b"world!!"); builder.append_value("Пуховички"); // 9 characters, 18 bytes diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 2b908a8997f..28bd65c8e60 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -30,8 +30,11 @@ static SESSION: LazyLock = LazyLock::new(|| { /// this function is VERY slow on miri, so we only want to run it once pub(crate) fn build_fsst_array(ctx: &mut ExecutionCtx) -> ArrayRef { - let mut input_array = - VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 3); + let mut input_array = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); input_array.append_value(b"The Greeks never said that the limit could not be overstepped"); input_array.append_value( b"They said it existed and that whoever dared to exceed it was mercilessly struck down", @@ -161,7 +164,11 @@ fn fsst_compress_offsets_overflow_i32() { println!("building large VarBinArray"); let string = vec![b'a'; STRING_LEN]; - let mut builder = VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), N); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::NonNullable), + N, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..N { builder.append_value(&string); } diff --git a/encodings/onpair/src/tests.rs b/encodings/onpair/src/tests.rs index ff541f020a0..b4af5269764 100644 --- a/encodings/onpair/src/tests.rs +++ b/encodings/onpair/src/tests.rs @@ -66,7 +66,11 @@ fn test_direct_offset_builder() -> vortex_error::VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let input = sample_input(); let encoded = compress_onpair(input.as_ref(), &mut ctx)?; - let mut builder = VarBinBuilder::::with_capacity(input.dtype().clone(), input.len()); + let mut builder = VarBinBuilder::::with_capacity( + input.dtype().clone(), + input.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); encoded .into_array() .append_to_builder(&mut builder, &mut ctx)?; diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 9a358138bac..621550ad564 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -577,7 +577,11 @@ mod tests { Some("c"), ]) .into_array(); - let mut builder = VarBinBuilder::::with_capacity(arr.dtype().clone(), arr.len()); + let mut builder = VarBinBuilder::::with_capacity( + arr.dtype().clone(), + arr.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); arr.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); assert_arrays_eq!(arr.into_array(), expected, &mut ctx); diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 727c094cd84..727435e5fa0 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -426,10 +426,11 @@ fn execute_sparse_lists_inner( nullability, total_canonical_values, len, + ctx.allocator(), ); // The fill's elements become an array once, up front. Every gap then appends that same array, // so the fill's elements are stored once for the whole result however many gaps reference them. - let fill_elements = list_scalar_elements_array(fill_value.as_list()); + let fill_elements = list_scalar_elements_array(fill_value.as_list(), ctx.allocator()); // One mask for the whole patch array rather than a validity lookup per patch. let patch_validity = patch_values @@ -476,9 +477,12 @@ fn execute_sparse_lists_inner( } /// Materializes a list scalar's elements into an array, or `None` if the scalar is null. -fn list_scalar_elements_array(list: ListScalar) -> Option { +fn list_scalar_elements_array( + list: ListScalar, + allocator: &vortex_buffer::BufferAllocatorRef, +) -> Option { list.elements().map(|elements| { - let mut builder = builder_with_capacity(list.element_dtype(), elements.len()); + let mut builder = builder_with_capacity(list.element_dtype(), elements.len(), allocator); for element in elements { builder .append_scalar(&element) @@ -560,12 +564,13 @@ fn execute_sparse_fixed_size_list_inner( list_size, nullability, array_len, + ctx.allocator(), ); // The fill's elements become an array once, up front, so that a gap does not rebuild them. // They are tiled per row rather than shared - a fixed-size list holds its elements back to // back - unless they are all the same scalar, in which case the tile stays constant-encoded // and the tiling costs nothing. - let fill_elements = fixed_size_list_fill_tile(fill_value.as_list(), list_size); + let fill_elements = fixed_size_list_fill_tile(fill_value.as_list(), list_size, ctx.allocator()); // One mask for the whole patch array rather than a validity lookup per patch. let patch_validity = values @@ -623,13 +628,18 @@ fn execute_sparse_fixed_size_list_inner( /// /// Elements that are all the same scalar stay a constant array, so tiling them over a gap costs /// nothing however many rows it covers. -fn fixed_size_list_fill_tile(fill: ListScalar, list_size: u32) -> Option { +fn fixed_size_list_fill_tile( + fill: ListScalar, + list_size: u32, + allocator: &vortex_buffer::BufferAllocatorRef, +) -> Option { let elements = fill.elements()?; Some(match elements.iter().all_equal_value() { Ok(uniform) => ConstantArray::new(uniform.clone(), list_size as usize).into_array(), Err(_) => { - let mut builder = builder_with_capacity(fill.element_dtype(), elements.len()); + let mut builder = + builder_with_capacity(fill.element_dtype(), elements.len(), allocator); for element in &elements { builder .append_scalar(element) @@ -775,7 +785,8 @@ fn execute_sparse_decimal( len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut builder = DecimalBuilder::with_capacity::(len, decimal_dtype, nullability); + let mut builder = + DecimalBuilder::with_capacity::(len, decimal_dtype, nullability, ctx.allocator()); match fill_value.decimal_value() { Some(fill_value) => { let fill_value = fill_value @@ -2189,11 +2200,18 @@ mod test { for candidate in [array.clone(), array.slice(1..9)?] { let expected = candidate.clone().execute::(&mut ctx)?; - let mut view_builder = VarBinViewBuilder::with_capacity(candidate.dtype().clone(), 4); + let mut view_builder = VarBinViewBuilder::with_capacity( + candidate.dtype().clone(), + 4, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); candidate.append_to_builder(&mut view_builder, &mut ctx)?; assert_arrays_eq!(view_builder.finish_into_varbinview(), expected, &mut ctx); - let mut varbin_builder = VarBinBuilder::::new(candidate.dtype().clone()); + let mut varbin_builder = VarBinBuilder::::new( + candidate.dtype().clone(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); candidate.append_to_builder(&mut varbin_builder, &mut ctx)?; assert_arrays_eq!(varbin_builder.finish_into_varbin(), expected, &mut ctx); } diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index e18106f4ca4..1a653db4d4d 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -957,7 +957,11 @@ mod test { Some("last"), ]) .into_array(); - let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + let mut builder = VarBinBuilder::::with_capacity( + array.dtype().clone(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); } @@ -974,7 +978,11 @@ mod test { ) .unwrap(); let expected = VarBinViewArray::from_iter_str(["fill", "second"]).into_array(); - let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + let mut builder = VarBinBuilder::::with_capacity( + array.dtype().clone(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(&mut builder, &mut ctx).unwrap(); assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx); } diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 81dda5940fd..9a161e7e474 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1770,7 +1770,10 @@ mod tests { #[test] fn test_append_to_varbin_copies_the_stored_values() -> VortexResult<()> { let slice = decompressed_slice(make_interleaved(&[b"hello", b"world"]), 0, 2, 0, 2); - let mut builder = VarBinBuilder::::new(DType::Utf8(NonNullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(NonNullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); append_slice_to_varbin(&slice, &Mask::new_true(2), &mut builder)?; let appended = builder.finish_into_varbin(); @@ -1790,7 +1793,10 @@ mod tests { buffer.extend_from_slice(&1u32.to_le_bytes()); let slice = decompressed_slice(ByteBuffer::copy_from(buffer.as_slice()), 0, 2, 0, 2); - let mut builder = VarBinBuilder::::new(DType::Utf8(NonNullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(NonNullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); assert!(append_slice_to_varbin(&slice, &Mask::new_true(2), &mut builder).is_err()); // The builder rejected the values before committing any of them. assert_eq!(builder.finish_into_varbin().len(), 0); diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index f65e5f6fe98..548b61b98b9 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -233,8 +233,11 @@ fn test_zstd_append_to_offset_builder() { .unwrap() .slice(1..4) .unwrap(); - let mut builder = - VarBinBuilder::::with_capacity(compressed.dtype().clone(), compressed.len()); + let mut builder = VarBinBuilder::::with_capacity( + compressed.dtype().clone(), + compressed.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); compressed .append_to_builder(&mut builder, &mut ctx) .unwrap(); @@ -261,7 +264,11 @@ fn test_zstd_append_to_view_builder_keeps_only_the_sliced_bytes() -> VortexResul // Seeded with a value of its own so the pushed buffers land after an in-progress buffer, and // appended to twice so the second push has to rebase past the first. - let mut builder = VarBinViewBuilder::with_capacity(compressed.dtype().clone(), 9); + let mut builder = VarBinViewBuilder::with_capacity( + compressed.dtype().clone(), + 9, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value(&values[0]); compressed.append_to_builder(&mut builder, &mut ctx)?; compressed.append_to_builder(&mut builder, &mut ctx)?; @@ -373,8 +380,11 @@ fn test_zstd_rejects_corrupt_frame_metadata( assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); - let mut builder = - VarBinBuilder::::with_capacity(compressed.dtype().clone(), compressed.len()); + let mut builder = VarBinBuilder::::with_capacity( + compressed.dtype().clone(), + compressed.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); assert!( compressed .append_to_builder(&mut builder, &mut ctx) @@ -414,9 +424,17 @@ fn test_zstd_rejects_a_frame_ending_in_a_dangling_length_prefix() -> VortexResul )?; assert!(Zstd::decompress(&compressed, &mut ctx).is_err()); - let mut varbin = VarBinBuilder::::with_capacity(dtype.clone(), 2); + let mut varbin = VarBinBuilder::::with_capacity( + dtype.clone(), + 2, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); assert!(compressed.append_to_builder(&mut varbin, &mut ctx).is_err()); - let mut views = VarBinViewBuilder::with_capacity(dtype, 2); + let mut views = VarBinViewBuilder::with_capacity( + dtype, + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); assert!(compressed.append_to_builder(&mut views, &mut ctx).is_err()); Ok(()) } diff --git a/fuzz/src/array/filter.rs b/fuzz/src/array/filter.rs index 147b9c166ac..245edfe1690 100644 --- a/fuzz/src/array/filter.rs +++ b/fuzz/src/array/filter.rs @@ -123,8 +123,11 @@ pub fn filter_canonical_array( .map(|a| a.into_array()) } DType::Map(..) => { - let mut builder = - builder_with_capacity(array.dtype(), filter.iter().filter(|b| **b).count()); + let mut builder = builder_with_capacity( + array.dtype(), + filter.iter().filter(|b| **b).count(), + ctx.allocator(), + ); for (idx, keep) in filter.iter().enumerate() { if *keep { builder.append_scalar(&array.execute_scalar(idx, ctx)?)?; diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index de8547c2306..aeab21fad2b 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -142,7 +142,7 @@ pub fn mask_canonical_array( } Canonical::Map(array) => { let result_dtype = array.dtype().as_nullable(); - let mut builder = builder_with_capacity(&result_dtype, array.len()); + let mut builder = builder_with_capacity(&result_dtype, array.len(), ctx.allocator()); for idx in 0..array.len() { if mask.value(idx) { builder.append_scalar(&array.execute_scalar(idx, ctx)?.cast(&result_dtype)?)?; diff --git a/fuzz/src/array/slice.rs b/fuzz/src/array/slice.rs index 86228a4e538..715987d46b6 100644 --- a/fuzz/src/array/slice.rs +++ b/fuzz/src/array/slice.rs @@ -127,7 +127,7 @@ pub fn slice_canonical_array( .map(|a| a.into_array()) } DType::Map(..) => { - let mut builder = builder_with_capacity(array.dtype(), stop - start); + let mut builder = builder_with_capacity(array.dtype(), stop - start, ctx.allocator()); for idx in start..stop { builder.append_scalar(&array.execute_scalar(idx, ctx)?)?; } diff --git a/fuzz/src/array/take.rs b/fuzz/src/array/take.rs index dacd13fbbb3..3867d00238c 100644 --- a/fuzz/src/array/take.rs +++ b/fuzz/src/array/take.rs @@ -118,6 +118,7 @@ pub fn take_canonical_array( let mut builder = builder_with_capacity( &array.dtype().union_nullability(nullable), indices_slice_non_opt.len(), + ctx.allocator(), ); for idx in indices { if let Some(idx) = idx { @@ -150,7 +151,7 @@ pub fn take_canonical_array( } DType::Map(..) => { let result_dtype = array.dtype().union_nullability(nullable); - let mut builder = builder_with_capacity(&result_dtype, indices.len()); + let mut builder = builder_with_capacity(&result_dtype, indices.len(), ctx.allocator()); for idx in indices { if let Some(idx) = idx { builder diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index 098b3ec3007..99eae6dd0da 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -49,7 +49,8 @@ fn chunked_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usi bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -64,7 +65,8 @@ fn chunked_opt_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -91,6 +93,7 @@ fn chunked_varbinview_canonical_into(bencher: Bencher, (len, chunk_count): (usiz let mut builder = VarBinViewBuilder::with_capacity( DType::Utf8(chunk.dtype().nullability()), len * chunk_count, + ctx.allocator().clone(), ); chunk .append_to_builder(&mut builder, ctx) @@ -118,6 +121,7 @@ fn chunked_varbinview_opt_canonical_into(bencher: Bencher, (len, chunk_count): ( let mut builder = VarBinViewBuilder::with_capacity( DType::Utf8(chunk.dtype().nullability()), len * chunk_count, + ctx.allocator().clone(), ); chunk .append_to_builder(&mut builder, ctx) @@ -149,8 +153,11 @@ fn chunked_varbin_to_varbinview_builder(bencher: Bencher, (len, chunk_count): (u bencher .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = - VarBinViewBuilder::with_capacity(chunk.dtype().clone(), len * chunk_count); + let mut builder = VarBinViewBuilder::with_capacity( + chunk.dtype().clone(), + len * chunk_count, + ctx.allocator().clone(), + ); chunk .append_to_builder(&mut builder, ctx) .vortex_expect("append failed"); @@ -165,8 +172,11 @@ fn chunked_varbin_opt_to_varbinview_builder(bencher: Bencher, (len, chunk_count) bencher .with_inputs(|| (&chunks, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = - VarBinViewBuilder::with_capacity(chunk.dtype().clone(), len * chunk_count); + let mut builder = VarBinViewBuilder::with_capacity( + chunk.dtype().clone(), + len * chunk_count, + ctx.allocator().clone(), + ); chunk .append_to_builder(&mut builder, ctx) .vortex_expect("append failed"); @@ -190,7 +200,8 @@ fn chunked_constant_i32_append_to_builder(bencher: Bencher, (len, chunk_count): bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -214,7 +225,8 @@ fn chunked_constant_utf8_append_to_builder( bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); @@ -282,7 +294,11 @@ fn make_varbin_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRe (0..chunk_count) .map(|_| { - let mut builder = VarBinBuilder::::with_capacity(dtype.clone(), len); + let mut builder = VarBinBuilder::::with_capacity( + dtype.clone(), + len, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); (0..len).for_each(|_| { if nullable && rng.random_bool(0.2) { builder.push_null() @@ -305,7 +321,11 @@ fn make_string_chunks(nullable: bool, len: usize, chunk_count: usize) -> ArrayRe (0..chunk_count) .map(|_| { - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(nullable.into()), len); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(nullable.into()), + len, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); (0..len).for_each(|_| { if nullable && rng.random_bool(0.2) { builder.append_null() diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index ffaf6069dee..8d9684c659d 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -43,7 +43,8 @@ fn chunked_dict_primitive_canonical_into( bencher .with_inputs(|| (&chunk, SESSION.create_execution_ctx())) .bench_refs(|(chunk, ctx)| { - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = + builder_with_capacity(chunk.dtype(), len * chunk_count, ctx.allocator()); chunk .append_to_builder(builder.as_mut(), ctx) .vortex_expect("append failed"); diff --git a/vortex-array/benches/listview_builder_extend.rs b/vortex-array/benches/listview_builder_extend.rs index 9f77a8f7013..cb65b5aa0c1 100644 --- a/vortex-array/benches/listview_builder_extend.rs +++ b/vortex-array/benches/listview_builder_extend.rs @@ -73,6 +73,7 @@ fn extend_from_array_zctl(bencher: Bencher, (num_lists, list_size): (usize, usiz NonNullable, num_lists * list_size, num_lists, + ctx.allocator(), ); source.append_to_builder(&mut builder, &mut ctx).unwrap(); divan::black_box(builder.finish_into_listview()) @@ -96,6 +97,7 @@ fn extend_from_array_non_zctl_overlapping( Nullable, num_lists * list_size, num_lists, + ctx.allocator(), ); source.append_to_builder(&mut builder, &mut ctx).unwrap(); divan::black_box(builder.finish_into_listview()) diff --git a/vortex-array/benches/varbinview_compact.rs b/vortex-array/benches/varbinview_compact.rs index 884cbd4eccf..b9fc8c432c4 100644 --- a/vortex-array/benches/varbinview_compact.rs +++ b/vortex-array/benches/varbinview_compact.rs @@ -93,7 +93,11 @@ fn compact_sliced_impl(bencher: Bencher, (output_size, utilization_pct): (usize, /// Creates a base VarBinViewArray with mix of inlined and outlined strings. fn build_varbinview_fixture(len: usize) -> VarBinViewArray { - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), len); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::NonNullable), + len, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); let mut rng = StdRng::seed_from_u64(42); for _ in 0..len { diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index a4b4e8fdd4e..fd8514919f7 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/vortex-array/src/aggregate_fn/accumulator_grouped.rs @@ -361,7 +361,7 @@ impl GroupedAccumulator { self.options.clone(), self.dtype.clone(), )?; - let mut states = builder_with_capacity(&self.partial_dtype, grouped.len()); + let mut states = builder_with_capacity(&self.partial_dtype, grouped.len(), ctx.allocator()); let group_ranges = grouped.group_ranges(ctx)?; let group_validity = grouped.group_validity(ctx)?; diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index ee0acd4462b..9eea6a4103f 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -477,8 +477,12 @@ mod tests { false, )?; let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); - let mut builder = - MapBuilder::::with_capacity(map_dtype, Nullability::Nullable, rows.len()); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + Nullability::Nullable, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for row in rows { let scalar = match row { diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index efe0825d4d6..76673092e5f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -202,7 +202,11 @@ mod tests { let sum_dtype = Sum .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) .expect("sum partial dtype"); - let mut builder = builder_with_capacity(&sum_dtype, ranges.len()); + let mut builder = builder_with_capacity( + &sum_dtype, + ranges.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for (i, &(offset, size)) in ranges.iter().enumerate() { if group_valid[i] { let slice = elements.slice(offset..offset + size)?; diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index bc6dfcbf28a..df1428a754d 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -386,7 +386,11 @@ mod tests { /// the top level - recursively canonicalize it before measuring. fn materialized_uncompressed_size_in_bytes(array: &ArrayRef) -> u64 { let mut ctx = array_session().create_execution_ctx(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .append_to_builder(builder.as_mut(), &mut ctx) .vortex_expect("appended"); diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs index e77d196ba94..eefb7e47493 100644 --- a/vortex-array/src/arrays/arbitrary.rs +++ b/vortex-array/src/arrays/arbitrary.rs @@ -144,7 +144,11 @@ fn random_array_chunk( d @ DType::Decimal(decimal, n) => { let elem_len = chunk_len.unwrap_or(u.int_in_range(0..=20)?); match_each_decimal_value_type!(DecimalType::smallest_decimal_value_type(decimal), |D| { - let mut builder = DecimalBuilder::new::(*decimal, *n); + let mut builder = DecimalBuilder::new::( + *decimal, + *n, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _i in 0..elem_len { let random_decimal = random_scalar(u, d)?; builder.append_scalar(&random_decimal).vortex_expect( @@ -213,7 +217,12 @@ fn random_map( let key_dtype = map_dtype.key_dtype(); let value_dtype = map_dtype.value_dtype(); let dtype = DType::Map(map_dtype.clone(), nullability); - let mut builder = MapBuilder::::with_capacity(map_dtype, nullability, array_length); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + nullability, + array_length, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..array_length { if nullability == Nullability::Nullable && u.arbitrary::()? { @@ -250,8 +259,13 @@ fn random_fixed_size_list( ) -> Result { let array_length = chunk_len.unwrap_or(u.int_in_range(0..=20)?); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(elem_dtype), list_size, null, array_length); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(elem_dtype), + list_size, + null, + array_length, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..array_length { if null == Nullability::Nullable && u.arbitrary::()? { @@ -304,8 +318,13 @@ fn random_list_with_offset_type( null: Nullability, array_length: usize, ) -> Result { - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(elem_dtype), null, array_length, 10); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(elem_dtype), + null, + array_length, + 10, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..array_length { if null == Nullability::Nullable && u.arbitrary::()? { diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index f9f57c48bb2..0edcad87d09 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -46,11 +47,11 @@ enum ChunkFlattener { } impl ChunkFlattener { - fn new(dtype: &DType, capacity: usize) -> Self { + fn new(dtype: &DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { if dtype.is_nested() { Self::Chunks(Vec::new()) } else { - Self::Builder(builder_with_capacity(dtype, capacity)) + Self::Builder(builder_with_capacity(dtype, capacity, allocator)) } } @@ -102,7 +103,7 @@ fn take_chunked_via_sort( let chunk_offsets = array.chunk_offset_values(); let nchunks = array.nchunks(); - let mut flattener = ChunkFlattener::new(array.dtype(), pairs.len()); + let mut flattener = ChunkFlattener::new(array.dtype(), pairs.len(), ctx.allocator()); let mut final_take = BufferMut::::zeroed(n); let mut cursor = 0usize; let mut dedup_idx = 0u64; @@ -227,7 +228,8 @@ fn take_chunked( buckets[chunk_idx].push((local_index, original_position)); } - let mut flattener = ChunkFlattener::new(array.dtype(), indices_mask.true_count()); + let mut flattener = + ChunkFlattener::new(array.dtype(), indices_mask.true_count(), ctx.allocator()); let mut final_take = (!monotonic || indices.dtype().is_nullable()).then(|| BufferMut::::zeroed(n)); let mut grouped_position = 0u64; @@ -504,7 +506,7 @@ fn take_piecewise_chunked( // the flattened result in output order. let mut bases = vec![0usize; nchunks]; let mut running = 0usize; - let mut flattener = ChunkFlattener::new(array.dtype(), output_len); + let mut flattener = ChunkFlattener::new(array.dtype(), output_len, ctx.allocator()); for (chunk_idx, plan) in plans.into_iter().enumerate() { if plan.starts.is_empty() { continue; diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index d5a993cf04d..ea2fb4b8cb3 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -54,7 +54,8 @@ fn builder_kernel_path_canonicalizes_primitive_chunks() { let dtype = array.dtype().clone(); let len = array.len(); - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); // Clone the array into the builder path — the test also holds `array` so refcount > 1 on // entry, which previously caused `take_slot_unchecked` to silently keep slots populated. let mut builder = execute_into_builder(array.clone(), builder, &mut ctx).unwrap(); @@ -93,7 +94,8 @@ fn builder_kernel_nested_chunked_of_chunked() { let dtype = outer.dtype().clone(); let len = outer.len(); - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); let mut builder = execute_into_builder(outer, builder, &mut ctx).unwrap(); let output = builder.finish(); @@ -120,13 +122,15 @@ fn builder_kernel_path_repeated_shared_chunked_dict_execution() { .into_array(); let first = { - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); let mut builder = execute_into_builder(array.clone(), builder, &mut ctx).unwrap(); builder.finish() }; let second = { - let builder = builder_with_capacity(&dtype, len); + let builder = + builder_with_capacity(&dtype, len, vortex_buffer::BufferAllocatorRef::static_ref()); let mut builder = execute_into_builder(array, builder, &mut ctx).unwrap(); builder.finish() }; @@ -177,7 +181,11 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { let keep_alive = outer.clone(); let expected = { - let mut builder = builder_with_capacity(outer.dtype(), outer.len()); + let mut builder = builder_with_capacity( + outer.dtype(), + outer.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); inner_1 .append_to_builder(builder.as_mut(), &mut ctx) .unwrap(); diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index df6e5bd3739..f8683aa4985 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -24,7 +24,7 @@ use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::variant::VariantArraySlotsExt; -use crate::builders::builder_with_capacity_in; +use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -69,8 +69,7 @@ pub(super) fn _canonicalize( } DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?), _ => { - let mut builder = - builder_with_capacity_in(ctx.allocator().clone(), array.dtype(), array.len()); + let mut builder = builder_with_capacity(array.dtype(), array.len(), ctx.allocator()); array.array().append_to_builder(builder.as_mut(), ctx)?; builder.finish_into_canonical(ctx) } diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index a668831e224..6e40160ae79 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -127,7 +128,11 @@ pub(crate) fn constant_canonicalize( array.len(), )) } - DType::List(..) => Canonical::List(constant_canonical_list_array(scalar, array.len())), + DType::List(..) => Canonical::List(constant_canonical_list_array( + scalar, + array.len(), + ctx.allocator(), + )), DType::Map(map_dtype, nullability) => { let entries_scalar = Scalar::try_new( DType::List(Arc::new(map_dtype.entries_dtype()), *nullability), @@ -135,7 +140,7 @@ pub(crate) fn constant_canonicalize( )?; Canonical::Map(MapArray::try_new( map_dtype.clone(), - constant_canonical_list_array(&entries_scalar, array.len()), + constant_canonical_list_array(&entries_scalar, array.len(), ctx.allocator()), )?) } DType::FixedSizeList(element_dtype, list_size, _) => { @@ -147,6 +152,7 @@ pub(crate) fn constant_canonicalize( *list_size, value.dtype().nullability(), array.len(), + ctx.allocator(), )) } DType::Struct(struct_dtype, _) => { @@ -248,7 +254,11 @@ fn constant_canonical_byte_view( /// /// We basically just project the list scalar value into list view components. If the caller wants /// a fully decompressed and non-overlapping array, they can rebuild the array. -fn constant_canonical_list_array(scalar: &Scalar, len: usize) -> ListViewArray { +fn constant_canonical_list_array( + scalar: &Scalar, + len: usize, + allocator: &BufferAllocatorRef, +) -> ListViewArray { let list = scalar.as_list(); // Since "canonicalize" only applies to the top level array, we can simply have 1 scalar in our @@ -260,6 +270,7 @@ fn constant_canonical_list_array(scalar: &Scalar, len: usize) -> ListViewArray { .as_list_element_opt() .vortex_expect("list scalar somehow did not have a list DType"), list.len(), + allocator, ); for scalar in &elements { builder @@ -302,13 +313,14 @@ fn constant_canonical_fixed_size_list_array( list_size: u32, list_nullability: Nullability, len: usize, + allocator: &BufferAllocatorRef, ) -> FixedSizeListArray { match values { None => { // Even though the scalar is null, we still have to allocate the correct amount of space // for the given `DType`. let elements_len = list_size as usize * len; - let mut element_builder = builder_with_capacity(element_dtype, elements_len); + let mut element_builder = builder_with_capacity(element_dtype, elements_len, allocator); element_builder.append_defaults(elements_len); let elements = element_builder.finish(); @@ -319,7 +331,8 @@ fn constant_canonical_fixed_size_list_array( } } Some(values) => { - let mut elements_builder = builder_with_capacity(element_dtype, len * values.len()); + let mut elements_builder = + builder_with_capacity(element_dtype, len * values.len(), allocator); for _ in 0..len { for v in &values { diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 7e950bfe5b5..f5c9c2148fb 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -326,7 +326,8 @@ fn append_repeated_list_run( return Ok(()); }; - let mut elements_builder = builder_with_capacity(scalar.element_dtype(), elements.len()); + let mut elements_builder = + builder_with_capacity(scalar.element_dtype(), elements.len(), ctx.allocator()); for element in &elements { elements_builder.append_scalar(element)?; } @@ -363,7 +364,8 @@ fn append_constant_fixed_size_list_run( let tile = match elements.iter().all_equal_value() { Ok(uniform) => ConstantArray::new(uniform.clone(), elements.len()).into_array(), Err(_) => { - let mut tile_builder = builder_with_capacity(builder.element_dtype(), elements.len()); + let mut tile_builder = + builder_with_capacity(builder.element_dtype(), elements.len(), ctx.allocator()); for element in &elements { tile_builder.append_scalar(element)?; } @@ -449,7 +451,11 @@ mod tests { let mut ctx = crate::array_session().create_execution_ctx(); let expected = constant_canonicalize(array.as_view(), &mut ctx)?.into_array(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -585,7 +591,11 @@ mod tests { ); let array = ConstantArray::new(scalar, 1_000); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -613,8 +623,13 @@ mod tests { ); let array = ConstantArray::new(scalar, 4).into_array(); - let mut builder = - ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + element_dtype, + Nullability::NonNullable, + 0, + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(&mut builder, &mut ctx)?; assert_arrays_eq!(&builder.finish(), &array, &mut ctx); @@ -671,7 +686,11 @@ mod tests { ); let array = ConstantArray::new(scalar, 1_000); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -721,7 +740,11 @@ mod tests { ); let array = ConstantArray::new(scalar, 1_000); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -752,7 +775,11 @@ mod tests { let scalar = Scalar::extension::(TimeUnit::Days, Scalar::from(Some(42i32))); let array = ConstantArray::new(scalar, 1_000); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 24832003052..e2e66ffee3c 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -449,8 +449,9 @@ mod test { let values = VarBinViewArray::from_iter_str(["zero", "one", "two"]); let dict = DictArray::try_new(buffer![2u8, 0, 2, 1].into_array(), values.into_array())?; let expected = VarBinViewArray::from_iter_str(["two", "zero", "two", "one"]); - let mut builder = VarBinBuilder::::with_capacity(dict.dtype().clone(), dict.len()); let mut ctx = array_session().create_execution_ctx(); + let mut builder = + VarBinBuilder::::with_capacity(dict.dtype().clone(), dict.len(), ctx.allocator()); dict.into_array() .append_to_builder(&mut builder, &mut ctx)?; @@ -469,6 +470,7 @@ mod test { let mut builder = builder_with_capacity( &DType::Primitive(PType::U64, NonNullable), len * chunk_count, + ctx.allocator(), ); array.append_to_builder( builder.as_mut(), diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index 92f80f45908..913fd261662 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -370,7 +370,7 @@ mod tests { .into_array(), )?; - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable), ctx.allocator()); dict.append_to_builder(&mut builder, &mut ctx)?; let expected = VarBinViewArray::from_iter( @@ -394,7 +394,11 @@ mod tests { .into_array(), )?; - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullable), 8); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullable), + 8, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value(LONG); dict.append_to_builder(&mut builder, &mut ctx)?; @@ -430,7 +434,11 @@ mod tests { values.into_array(), )?; - let mut builder = VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullable), 8); + let mut builder = VarBinViewBuilder::with_buffer_deduplication( + DType::Utf8(Nullable), + 8, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); first.append_to_builder(&mut builder, &mut ctx)?; second.append_to_builder(&mut builder, &mut ctx)?; assert_eq!(builder.completed_block_count(), 1); diff --git a/vortex-array/src/arrays/fixed_size_list/compute/take.rs b/vortex-array/src/arrays/fixed_size_list/compute/take.rs index b6834a38408..6706dc5d77f 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/compute/take.rs @@ -77,7 +77,7 @@ fn take_empty_fsl( "FixedSizeList take output length overflow: {new_len} lists of size {list_size}" ) })?; - let new_elements = default_elements(array, elements_len); + let new_elements = default_elements(array, elements_len, ctx.allocator()); let new_validity = if new_len == 0 { array.validity()?.take(indices)? } else { @@ -337,8 +337,12 @@ fn bounds_check_valid_indices( Ok(()) } -fn default_elements(array: ArrayView<'_, FixedSizeList>, len: usize) -> ArrayRef { - let mut builder = builder_with_capacity(array.elements().dtype(), len); +fn default_elements( + array: ArrayView<'_, FixedSizeList>, + len: usize, + allocator: &vortex_buffer::BufferAllocatorRef, +) -> ArrayRef { + let mut builder = builder_with_capacity(array.elements().dtype(), len, allocator); builder.append_defaults(len); builder.finish() } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs index 40fd2bc066a..0848cf61514 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -413,6 +413,7 @@ fn test_fsl_of_list() { Nullability::NonNullable, 12, 6, + vortex_buffer::BufferAllocatorRef::static_ref(), ); // Add 6 lists (2 FSL * 3 lists each). @@ -494,6 +495,7 @@ fn test_fsl_of_nullable_list() { Nullability::Nullable, 8, 4, + vortex_buffer::BufferAllocatorRef::static_ref(), ); // Add 4 lists (2 FSL * 2 lists each). diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 6fcd650bb0e..0e33b180160 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -314,6 +314,7 @@ fn test_take_nullable_arrays_fsl_specific( list_size, Nullability::Nullable, array_values.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), ); for value in array_values { diff --git a/vortex-array/src/arrays/list/test_harness.rs b/vortex-array/src/arrays/list/test_harness.rs index 78c3dc97a72..e835dc46a70 100644 --- a/vortex-array/src/arrays/list/test_harness.rs +++ b/vortex-array/src/arrays/list/test_harness.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use itertools::Itertools; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use crate::arrays::ListArray; @@ -31,6 +32,7 @@ impl ListArray { crate::dtype::Nullability::NonNullable, 2 * iter.size_hint().0, iter.size_hint().0, + BufferAllocatorRef::static_ref(), ); for v in iter { @@ -58,6 +60,7 @@ impl ListArray { crate::dtype::Nullability::Nullable, 2 * iter.size_hint().0, iter.size_hint().0, + BufferAllocatorRef::static_ref(), ); for v in iter { diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 906f4e03588..92af98abaf4 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -386,8 +386,13 @@ fn test_list_filter_variable_sizes() { #[test] fn test_offset_to_0() { - let mut builder = - ListBuilder::::with_capacity(Arc::new(I32.into()), Nullability::NonNullable, 10, 5); + let mut builder = ListBuilder::::with_capacity( + Arc::new(I32.into()), + Nullability::NonNullable, + 10, + 5, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder .append_value( Scalar::list( diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index e0401460f65..67ed55743bd 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -51,7 +52,7 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes // Create `sizes` array by computing differences between consecutive offsets. // We use the same `DType` for the sizes as the `offsets` array to ensure compatibility. let sizes = match_each_integer_ptype!(list_offsets.ptype(), |O| { - build_sizes_from_offsets::(&list_offsets)? + build_sizes_from_offsets::(&list_offsets, ctx.allocator())? }); // We need to slice the `offsets` to remove the last element (`ListArray` has `n + 1` offsets). @@ -75,12 +76,16 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes /// Builds a sizes array by computing differences between consecutive offsets. /// /// `offsets` **must** be the `n + 1` sorted offsets of a non-empty `ListArray` of `n` rows. -fn build_sizes_from_offsets(offsets: &PrimitiveArray) -> VortexResult { +fn build_sizes_from_offsets( + offsets: &PrimitiveArray, + allocator: &BufferAllocatorRef, +) -> VortexResult { let offsets_slice = offsets.as_slice::(); debug_assert!(offsets_slice.is_sorted()); let len = offsets_slice.len() - 1; - let mut sizes_builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len); + let mut sizes_builder = + PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len, allocator); // Create `UninitRange` for direct memory access. let mut sizes_range = sizes_builder.uninit_range(len); @@ -149,7 +154,7 @@ unsafe fn build_list_offsets_from_list_view( ) -> ArrayRef { let len = list_view.len(); let mut offsets_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len + 1); + PrimitiveBuilder::::with_capacity(Nullability::NonNullable, len + 1, ctx.allocator()); // Create uninit range for direct memory access. let mut offsets_range = offsets_builder.uninit_range(len + 1); diff --git a/vortex-array/src/arrays/map/tests.rs b/vortex-array/src/arrays/map/tests.rs index 2b5570a8d16..497fefcbb9f 100644 --- a/vortex-array/src/arrays/map/tests.rs +++ b/vortex-array/src/arrays/map/tests.rs @@ -80,7 +80,12 @@ fn map_array_from_rows( ) -> VortexResult { let rows = rows.into_iter().collect::>(); let dtype = DType::Map(map_dtype.clone(), nullability); - let mut builder = MapBuilder::::with_capacity(map_dtype, nullability, rows.len()); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + nullability, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for row in rows { let scalar = match row { @@ -222,7 +227,12 @@ fn scalar_access_preserves_variable_entry_counts_and_utf8_pairs() -> VortexResul false, )?; let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); - let mut builder = MapBuilder::::with_capacity(map_dtype, Nullability::Nullable, 4); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + Nullability::Nullable, + 4, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let rows = [ Some(vec![ ( @@ -648,6 +658,7 @@ fn builder_appends_existing_map_arrays() -> VortexResult<()> { source.map_dtype().clone(), source.dtype().nullability(), 0, + vortex_buffer::BufferAllocatorRef::static_ref(), ); let mut ctx = array_session().create_execution_ctx(); builder.append_map_array(source.as_view(), &mut ctx)?; diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index e2473c19749..159c8e50ee4 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -457,7 +457,11 @@ mod tests { .unwrap() .into_array(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); let result = builder.finish(); @@ -492,7 +496,11 @@ mod tests { .slice(3..1024) .unwrap(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); let result = builder.finish(); @@ -527,7 +535,11 @@ mod tests { .unwrap() .into_array(); - let mut builder = builder_with_capacity(array.dtype(), array.len()); + let mut builder = builder_with_capacity( + array.dtype(), + array.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_to_builder(builder.as_mut(), &mut ctx).unwrap(); let result = builder.finish(); diff --git a/vortex-array/src/arrays/varbin/array.rs b/vortex-array/src/arrays/varbin/array.rs index 7cabbefa0bd..ec95df096dc 100644 --- a/vortex-array/src/arrays/varbin/array.rs +++ b/vortex-array/src/arrays/varbin/array.rs @@ -6,6 +6,7 @@ use std::fmt::Formatter; use num_traits::AsPrimitive; use vortex_array::arrays::PrimitiveArray; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -383,7 +384,11 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinBuilder::::with_capacity(dtype, iter.size_hint().0); + let mut builder = VarBinBuilder::::with_capacity( + dtype, + iter.size_hint().0, + BufferAllocatorRef::static_ref(), + ); for v in iter { builder.append(v.as_ref().map(|o| o.as_ref())); } @@ -395,7 +400,11 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinBuilder::::with_capacity(dtype, iter.size_hint().0); + let mut builder = VarBinBuilder::::with_capacity( + dtype, + iter.size_hint().0, + BufferAllocatorRef::static_ref(), + ); for v in iter { builder.append_value(v); } @@ -407,7 +416,8 @@ impl Array { O: OffsetBuilderPType, T: AsRef<[u8]>, { - let mut builder = VarBinBuilder::::with_capacity(dtype, vec.len()); + let mut builder = + VarBinBuilder::::with_capacity(dtype, vec.len(), BufferAllocatorRef::static_ref()); for v in vec { builder.append_value(v.as_ref()); } diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 5990ecee371..99caa0f3819 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -5,7 +5,9 @@ use std::any::Any; use std::mem::MaybeUninit; use num_traits::AsPrimitive; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -63,29 +65,34 @@ pub struct VarBinBuilder { impl VarBinBuilder { /// Creates an empty builder for `dtype`. - pub fn new(dtype: DType) -> Self { - Self::with_capacity(dtype, 0) + pub fn new(dtype: DType, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity(dtype, 0, allocator) } /// Creates a builder for `dtype` with room for `capacity` values. - pub fn with_capacity(dtype: DType, capacity: usize) -> Self { + pub fn with_capacity(dtype: DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { assert!( matches!(dtype, DType::Utf8(_) | DType::Binary(_)), "VarBinBuilder dtype must be Utf8 or Binary, got {dtype}" ); - let mut offsets = BufferMut::with_capacity(capacity + 1); + let mut offsets = BufferMut::with_capacity_in(capacity + 1, allocator.clone()); offsets.push(O::zero()); Self { dtype, offsets, - data: BufferMut::empty(), - validity: BitBufferMut::with_capacity(capacity), + data: BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + validity: BitBufferMut::with_capacity_in(capacity, allocator.clone()), } } /// Creates a builder for `dtype` with room for `capacity` values totalling `bytes` bytes. - pub fn with_capacity_bytes(dtype: DType, capacity: usize, bytes: usize) -> Self { - let mut builder = Self::with_capacity(dtype, capacity); + pub fn with_capacity_bytes( + dtype: DType, + capacity: usize, + bytes: usize, + allocator: &BufferAllocatorRef, + ) -> Self { + let mut builder = Self::with_capacity(dtype, capacity, allocator); builder.reserve_data(bytes); builder } @@ -388,14 +395,19 @@ impl VarBinBuilder { "The offset count must be one more than the validity length" ); - let mut fresh_offsets = BufferMut::with_capacity(1); + let allocator = self.offsets.allocator().clone(); + let mut fresh_offsets = BufferMut::with_capacity_in(1, allocator.clone()); fresh_offsets.push(O::zero()); let offsets = PrimitiveArray::new( std::mem::replace(&mut self.offsets, fresh_offsets).freeze(), Validity::NonNullable, ); - let data = std::mem::replace(&mut self.data, BufferMut::empty()); - let nulls = std::mem::replace(&mut self.validity, BitBufferMut::empty()).freeze(); + let data = std::mem::replace( + &mut self.data, + BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + ); + let nulls = + std::mem::replace(&mut self.validity, BitBufferMut::empty_in(allocator)).freeze(); let validity = Validity::from_bit_buffer(nulls, self.dtype.nullability()); @@ -724,7 +736,11 @@ mod tests { #[test] fn test_builder() { - let mut builder = VarBinBuilder::::with_capacity(DType::Utf8(Nullable), 0); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullable), + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append(Some(b"hello")); builder.append(None); builder.append(Some(b"world")); @@ -774,7 +790,10 @@ mod tests { #[test] fn append_n_values_offset_overflow_returns_error() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); // The limit is checked before anything is reserved, so the huge byte total is free. let result = builder.append_n_values(b"hello", i32::MAX as usize / 5 + 1); @@ -787,7 +806,10 @@ mod tests { #[test] fn append_values_rejects_a_short_offset_count() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let result = builder.append_values(b"ab", [1usize, 2].into_iter(), &Mask::new_true(3)); @@ -798,7 +820,10 @@ mod tests { #[test] fn append_values_rejects_non_monotonic_offsets() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let result = builder.append_values(b"ab", [2usize, 1].into_iter(), &Mask::new_true(2)); @@ -809,7 +834,10 @@ mod tests { #[test] fn append_decoded_writes_into_the_builder_storage() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); // SAFETY: the closure initializes exactly the 6 bytes it reports. unsafe { @@ -837,7 +865,10 @@ mod tests { /// should not pay for a full decompression first. #[test] fn append_decoded_rejects_an_offset_overflow_without_decoding() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let num_bytes = i32::MAX as usize + 1; let mut decoded = false; @@ -859,7 +890,10 @@ mod tests { /// column wider than the offset type errors like the bulk appends do. #[test] fn append_scalar_repeated_rejects_an_offset_overflow() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let result = builder .append_scalar_repeated(&Scalar::utf8("hello", Nullable), i32::MAX as usize / 5 + 1); @@ -872,7 +906,10 @@ mod tests { #[test] fn append_decoded_rejects_a_short_decode() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); // SAFETY: the closure initializes the 3 bytes it reports (none, and it reports 3 — but the // length mismatch is rejected before anything is published). @@ -895,7 +932,10 @@ mod tests { #[case::all_valid(Mask::new_true(2))] #[case::some_valid(Mask::from_iter([true, false, true]))] fn append_valid_slices_rejects_a_byte_count_mismatch(#[case] validity: Mask) { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let values = [b"foo".as_slice(), b"quux".as_slice()]; let mut next = 0; @@ -913,7 +953,10 @@ mod tests { #[test] #[should_panic(expected = "The offset count must be one more than the validity length")] fn finish_rejects_mismatched_validity() { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.validity.append_true(); drop(builder.finish_into_varbin()); } @@ -973,10 +1016,10 @@ mod tests { /// only matched the signed pair would send an unsigned builder down a downcast that assumes /// `VarBinViewBuilder` and panic. #[rstest] - #[case::u32(VarBinBuilder::::new(DType::Utf8(Nullable)))] - #[case::u64(VarBinBuilder::::new(DType::Utf8(Nullable)))] - #[case::i32(VarBinBuilder::::new(DType::Utf8(Nullable)))] - #[case::i64(VarBinBuilder::::new(DType::Utf8(Nullable)))] + #[case::u32(VarBinBuilder::::new(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + #[case::u64(VarBinBuilder::::new(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + #[case::i32(VarBinBuilder::::new(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] + #[case::i64(VarBinBuilder::::new(DType::Utf8(Nullable), vortex_buffer::BufferAllocatorRef::static_ref()))] fn append_to_every_offset_width(#[case] mut builder: impl ArrayBuilder) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let long = "a string that is far too long to be inlined in a view"; @@ -997,7 +1040,11 @@ mod tests { #[test] fn offsets_have_is_sorted_stat() -> VortexResult<()> { - let mut builder = VarBinBuilder::::with_capacity(DType::Utf8(Nullable), 0); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullable), + 0, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"aaa"); builder.push_null(); builder.append_value(b"bbb"); @@ -1013,7 +1060,10 @@ mod tests { #[test] fn empty_builder_offsets_have_is_sorted_stat() -> VortexResult<()> { - let mut builder = VarBinBuilder::::new(DType::Utf8(Nullable)); + let mut builder = VarBinBuilder::::new( + DType::Utf8(Nullable), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let array = builder.finish_into_varbin(); let is_sorted = array @@ -1031,11 +1081,19 @@ mod tests { f: impl FnOnce(&mut dyn ArrayBuilder) -> VortexResult<()>, ) -> VortexResult { if large_offsets { - let mut builder = VarBinBuilder::::with_capacity(dtype, 8); + let mut builder = VarBinBuilder::::with_capacity( + dtype, + 8, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); f(&mut builder)?; Ok(builder.finish_into_varbin()) } else { - let mut builder = VarBinBuilder::::with_capacity(dtype, 8); + let mut builder = VarBinBuilder::::with_capacity( + dtype, + 8, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); f(&mut builder)?; Ok(builder.finish_into_varbin()) } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 1fcd313fa46..9653ed44ca6 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -303,8 +303,11 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = - VarBinBuilder::::with_capacity(DType::Utf8(Nullability::NonNullable), 3); + let mut builder = VarBinBuilder::::with_capacity( + DType::Utf8(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"abc"); builder.append_value(b"xyz"); builder.append_value(b"abc"); @@ -325,8 +328,11 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant_binary() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = - VarBinBuilder::::with_capacity(DType::Binary(Nullability::NonNullable), 3); + let mut builder = VarBinBuilder::::with_capacity( + DType::Binary(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(b"abc"); builder.append_value(b"xyz"); builder.append_value(b"abc"); diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index e6ff4b9734c..bdf92bb99c2 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -3,6 +3,7 @@ use itertools::Itertools; use num_traits::AsPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -75,6 +76,7 @@ fn filter_select_var_bin_by_slice( .execute_mask(values.as_ref().len(), ctx) .vortex_expect("Failed to compute validity mask"), selection_count, + ctx.allocator(), ) }) }) @@ -87,13 +89,14 @@ fn filter_select_var_bin_by_slice_primitive_offset( mask_slices: &[(usize, usize)], logical_validity: Mask, selection_count: usize, + allocator: &BufferAllocatorRef, ) -> VortexResult where O: IntegerPType, B: OffsetBuilderPType, usize: AsPrimitive, { - let mut builder = VarBinBuilder::::with_capacity(dtype, selection_count); + let mut builder = VarBinBuilder::::with_capacity(dtype, selection_count, allocator); match logical_validity.bit_buffer() { AllOr::All => { for &(start, end) in mask_slices { @@ -180,6 +183,7 @@ fn filter_select_var_bin_by_index( .execute_mask(values.as_ref().len(), ctx) .vortex_expect("Failed to compute validity mask"), selection_count, + ctx.allocator(), ) }) }) @@ -192,6 +196,7 @@ fn filter_select_var_bin_by_index_primitive_offset VortexResult { let value_at = |idx: usize| -> VortexResult<&[u8]> { let start = offsets[idx] @@ -203,7 +208,7 @@ fn filter_select_var_bin_by_index_primitive_offset::with_capacity(dtype, selection_count); + let mut builder = VarBinBuilder::::with_capacity(dtype, selection_count, allocator); match mask.bit_buffer() { AllOr::All => { for idx in mask_indices.iter().copied() { diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index 971d4397381..2b4954b825b 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -62,7 +62,11 @@ mod tests { #[case(DType::Utf8(Nullability::Nullable))] #[case(DType::Binary(Nullability::Nullable))] fn test_canonical_varbin_sliced(#[case] dtype: DType) { - let mut varbin = VarBinBuilder::::with_capacity(dtype.clone(), 10); + let mut varbin = VarBinBuilder::::with_capacity( + dtype.clone(), + 10, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); varbin.push_null(); varbin.push_null(); // inlined value @@ -128,7 +132,11 @@ mod tests { let second = VarBinArray::from_iter([Some(longer), Some("tiny"), None], dtype.clone()); let view = VarBinViewArray::from_iter([Some(long), None], dtype.clone()); - let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 8); + let mut builder = VarBinViewBuilder::with_capacity( + dtype.clone(), + 8, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); first .as_array() .clone() @@ -167,7 +175,12 @@ mod tests { // Every value inlines, so nothing references the value bytes at all. let array = VarBinArray::from_iter_nonnull(["short", "tiny", "small"], dtype.clone()); - let mut builder = VarBinViewBuilder::with_compaction(dtype, 4, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + dtype, + 4, + 1.0, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); array .as_array() .clone() diff --git a/vortex-array/src/arrays/varbinview/array.rs b/vortex-array/src/arrays/varbinview/array.rs index 70f6012db80..e0e722cc62e 100644 --- a/vortex-array/src/arrays/varbinview/array.rs +++ b/vortex-array/src/arrays/varbinview/array.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use smallvec::smallvec; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -575,7 +576,11 @@ impl VarBinViewData { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity(dtype, iter.size_hint().0); + let mut builder = VarBinViewBuilder::with_capacity( + dtype, + iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), + ); for item in iter { match item { @@ -592,6 +597,7 @@ impl VarBinViewData { let mut builder = VarBinViewBuilder::with_capacity( DType::Utf8(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -608,6 +614,7 @@ impl VarBinViewData { let mut builder = VarBinViewBuilder::with_capacity( DType::Utf8(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -625,6 +632,7 @@ impl VarBinViewData { let mut builder = VarBinViewBuilder::with_capacity( DType::Binary(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -641,6 +649,7 @@ impl VarBinViewData { let mut builder = VarBinViewBuilder::with_capacity( DType::Binary(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for item in iter { @@ -693,7 +702,11 @@ impl Array { dtype: DType, ) -> Self { let iter = iter.into_iter(); - let mut builder = VarBinViewBuilder::with_capacity(dtype, iter.size_hint().0); + let mut builder = VarBinViewBuilder::with_capacity( + dtype, + iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), + ); for value in iter { match value { Some(value) => builder.append_value(value), @@ -708,6 +721,7 @@ impl Array { let mut builder = VarBinViewBuilder::with_capacity( DType::Utf8(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { builder.append_value(value.as_ref()); @@ -722,6 +736,7 @@ impl Array { let mut builder = VarBinViewBuilder::with_capacity( DType::Utf8(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { match value { @@ -737,6 +752,7 @@ impl Array { let mut builder = VarBinViewBuilder::with_capacity( DType::Binary(Nullability::NonNullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { builder.append_value(value.as_ref()); @@ -751,6 +767,7 @@ impl Array { let mut builder = VarBinViewBuilder::with_capacity( DType::Binary(Nullability::Nullable), iter.size_hint().0, + BufferAllocatorRef::statically_allocated(), ); for value in iter { match value { diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 75ed5d813db..c11e52eacd4 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -146,6 +146,7 @@ impl VarBinViewArray { self.dtype().clone(), self.len(), buffer_utilization_threshold, + ctx.allocator().clone(), ); builder.append_varbinview_array(self, ctx)?; Ok(builder.finish_into_varbinview()) diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index 8966a327f9d..2711a7e4fc6 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-array/src/arrays/varbinview/compute/zip.rs @@ -52,7 +52,7 @@ impl ZipKernel for VarBinView { buffers.extend_from_iter(if_false.data_buffers().iter().map(|b| b.as_host().clone())); let mut views_builder = BufferMut::::with_capacity(len); - let mut validity_builder = LazyBitBufferBuilder::new(len); + let mut validity_builder = LazyBitBufferBuilder::new(len, ctx.allocator().clone()); let true_validity = if_true.varbinview_validity().execute_mask(len, ctx)?; let false_validity = if_false.varbinview_validity().execute_mask(len, ctx)?; diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 3587f6fc3ba..a3092809040 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::mem; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -28,14 +29,19 @@ pub struct BoolBuilder { } impl BoolBuilder { - pub fn new(nullability: Nullability) -> Self { - Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY) + pub fn new(nullability: Nullability, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY, allocator) } - pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self { + /// Creates a builder with the given capacity and allocator. + pub fn with_capacity( + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { Self { - inner: BitBufferMut::with_capacity(capacity), - nulls: LazyBitBufferBuilder::new(capacity), + inner: BitBufferMut::with_capacity_in(capacity, allocator.clone()), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), dtype: DType::Bool(nullability), } } @@ -61,8 +67,10 @@ impl BoolBuilder { "Null count and value count should match when calling BoolBuilder::finish." ); + let allocator = self.inner.allocator().clone(); + let inner = mem::replace(&mut self.inner, BitBufferMut::empty_in(allocator)).freeze(); BoolArray::new( - mem::take(&mut self.inner).freeze(), + inner, self.nulls.finish_with_nullability(self.dtype.nullability()), ) } @@ -140,6 +148,7 @@ mod tests { use rand::RngExt; use rand::SeedableRng; use rand::prelude::StdRng; + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use crate::ArrayRef; @@ -181,7 +190,11 @@ mod tests { let chunk = make_opt_bool_chunks(len, chunk_count); let mut ctx = array_session().create_execution_ctx(); - let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); + let mut builder = builder_with_capacity( + chunk.dtype(), + len * chunk_count, + BufferAllocatorRef::static_ref(), + ); chunk .clone() .append_to_builder(builder.as_mut(), &mut ctx)?; @@ -201,7 +214,8 @@ mod tests { #[test] fn test_append_scalar() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); + let mut builder = + BoolBuilder::with_capacity(Nullability::Nullable, 10, BufferAllocatorRef::static_ref()); // Test appending true value. let true_scalar = Scalar::bool(true, Nullability::Nullable); @@ -220,7 +234,11 @@ mod tests { assert_arrays_eq!(&array, &expected, &mut ctx); // Test wrong dtype error. - let mut builder = BoolBuilder::with_capacity(Nullability::NonNullable, 10); + let mut builder = BoolBuilder::with_capacity( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index d1ff5f2954e..4a19d51699a 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -42,12 +43,13 @@ pub struct ChildBuilder { impl ChildBuilder { /// Creates a new `ChildBuilder` whose scalar builder is pre-allocated for `capacity` values. - pub fn with_capacity(dtype: &DType, capacity: usize) -> Self { + /// Creates a child builder with the provided allocator. + pub fn with_capacity(dtype: &DType, capacity: usize, allocator: &BufferAllocatorRef) -> Self { Self { dtype: dtype.clone(), chunks: Vec::new(), chunks_len: 0, - pending: builder_with_capacity(dtype, capacity), + pending: builder_with_capacity(dtype, capacity, allocator), } } @@ -155,6 +157,7 @@ impl ChildBuilder { #[cfg(test)] mod tests { use rstest::rstest; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -194,7 +197,8 @@ mod tests { #[test] fn test_appended_arrays_are_kept_as_chunks() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; @@ -214,7 +218,8 @@ mod tests { #[test] fn test_short_arrays_are_kept_as_chunks_too() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.append_array(&constant(1, 1), &mut ctx)?; builder.append_array(&constant(2, 1), &mut ctx)?; @@ -231,7 +236,8 @@ mod tests { #[test] fn test_scalars_interleaved_with_chunks_keep_their_order() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.append_scalar(&1i32.into())?; builder.append_array(&constant(2, CHUNK_LEN), &mut ctx)?; @@ -257,7 +263,8 @@ mod tests { #[test] fn test_single_chunk_is_not_wrapped() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.append_array(&constant(7, CHUNK_LEN), &mut ctx)?; @@ -270,7 +277,8 @@ mod tests { #[test] fn test_empty_arrays_never_become_chunks() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); let empty = constant(1, CHUNK_LEN).slice(0..0)?; builder.append_array(&empty, &mut ctx)?; @@ -289,7 +297,8 @@ mod tests { #[test] fn test_empty_child_finishes_without_chunks() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.append_array(&constant(1, CHUNK_LEN).slice(0..0)?, &mut ctx)?; @@ -307,7 +316,8 @@ mod tests { #[case::non_empty(CHUNK_LEN)] fn test_appending_a_mismatched_dtype_is_rejected(#[case] len: usize) { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); let wrong_dtype = ConstantArray::new(1i64, len).into_array(); assert!(builder.append_array(&wrong_dtype, &mut ctx).is_err()); @@ -318,7 +328,7 @@ mod tests { fn test_zeros_and_nulls_around_chunks_keep_their_order() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(I32, Nullable); - let mut builder = ChildBuilder::with_capacity(&dtype, 0); + let mut builder = ChildBuilder::with_capacity(&dtype, 0, BufferAllocatorRef::static_ref()); builder.append_array(&nullable_constant(1, CHUNK_LEN), &mut ctx)?; builder.append_nulls(2); @@ -344,7 +354,8 @@ mod tests { #[test] fn test_finish_resets_the_builder() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ChildBuilder::with_capacity(&DType::from(I32), 0); + let mut builder = + ChildBuilder::with_capacity(&DType::from(I32), 0, BufferAllocatorRef::static_ref()); builder.append_array(&constant(1, CHUNK_LEN), &mut ctx)?; builder.append_scalar(&2i32.into())?; diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index 027b7e62c75..a54253d650f 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -3,6 +3,7 @@ use std::any::Any; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -94,8 +95,12 @@ macro_rules! delegate_fn { impl DecimalBuilder { /// Creates a new `DecimalBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(decimal: DecimalDType, nullability: Nullability) -> Self { - Self::with_capacity::(DEFAULT_BUILDER_CAPACITY, decimal, nullability) + pub fn new( + decimal: DecimalDType, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity::(DEFAULT_BUILDER_CAPACITY, decimal, nullability, allocator) } /// Creates a new `DecimalBuilder` with the given `capacity`. @@ -103,13 +108,17 @@ impl DecimalBuilder { capacity: usize, decimal: DecimalDType, nullability: Nullability, + allocator: &BufferAllocatorRef, ) -> Self { Self { dtype: DType::Decimal(decimal, nullability), values: match_each_decimal_value_type!(T::DECIMAL_TYPE, |D| { - DecimalBuffer::from(BufferMut::::with_capacity(capacity)) + DecimalBuffer::from(BufferMut::::with_capacity_in( + capacity, + allocator.clone(), + )) }), - nulls: LazyBitBufferBuilder::new(capacity), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), } } @@ -153,7 +162,7 @@ impl DecimalBuilder { let decimal_dtype = *self.decimal_dtype(); - delegate_fn!(std::mem::take(&mut self.values), |T, values| { + delegate_fn!(self.values.take(), |T, values| { DecimalArray::new::(values.freeze(), decimal_dtype, validity) }) } @@ -228,6 +237,13 @@ impl ArrayBuilder for DecimalBuilder { } impl DecimalBuffer { + fn take(&mut self) -> Self { + delegate_fn!(self, |T, buffer| { + let allocator = buffer.allocator(); + DecimalBuffer::from(std::mem::replace(buffer, allocator.with_capacity(0))) + }) + } + fn push(&mut self, value: V) { delegate_fn!(self, |T, buffer| { buffer.push( @@ -291,14 +307,10 @@ impl_from_buffer!(i64, I64); impl_from_buffer!(i128, I128); impl_from_buffer!(i256, I256); -impl Default for DecimalBuffer { - fn default() -> Self { - Self::I8(BufferMut::::empty()) - } -} - #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; + use crate::VortexSessionExecute; use crate::array_session; use crate::assert_arrays_eq; @@ -311,13 +323,21 @@ mod tests { fn test_mixed_extend() { let values = 42i8; - let mut i8s = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut i8s = DecimalBuilder::new::( + DecimalDType::new(2, 1), + false.into(), + BufferAllocatorRef::static_ref(), + ); for v in 0..values { i8s.append_value(v); } let i8s = i8s.finish(); - let mut i128s = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut i128s = DecimalBuilder::new::( + DecimalDType::new(2, 1), + false.into(), + BufferAllocatorRef::static_ref(), + ); i8s.append_to_builder(&mut i128s, &mut array_session().create_execution_ctx()) .unwrap(); let i128s = i128s.finish(); @@ -339,7 +359,11 @@ mod tests { use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. - let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); + let mut builder = DecimalBuilder::new::( + DecimalDType::new(10, 2), + true.into(), + BufferAllocatorRef::static_ref(), + ); builder.append_value(1234i64); builder.append_value(5678i64); builder.append_null(); @@ -352,7 +376,11 @@ mod tests { assert_arrays_eq!(&array, &expected, &mut ctx); // Test by taking a scalar from the array and appending it to a new builder. - let mut builder2 = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); + let mut builder2 = DecimalBuilder::new::( + DecimalDType::new(10, 2), + true.into(), + BufferAllocatorRef::static_ref(), + ); for i in 0..array.len() { let scalar = array .execute_scalar(i, &mut array_session().create_execution_ctx()) @@ -364,7 +392,11 @@ mod tests { assert_arrays_eq!(&array2, &array, &mut ctx); // Test wrong dtype error. - let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), false.into()); + let mut builder = DecimalBuilder::new::( + DecimalDType::new(10, 2), + false.into(), + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(true); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 48444a63000..e7ff28a3723 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -8,7 +8,9 @@ use std::sync::Arc; use num_traits::AsPrimitive; use vortex_array::ExecutionCtx; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -50,29 +52,44 @@ pub struct BytesDictBuilder { dtype: DType, max_dict_bytes: usize, max_dict_len: usize, + allocator: BufferAllocatorRef, } -pub fn bytes_dict_builder(dtype: DType, constraints: &DictConstraints) -> Box { +pub fn bytes_dict_builder( + dtype: DType, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, +) -> Box { match constraints.max_len as u64 { - max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - _ => Box::new(BytesDictBuilder::::new(dtype, constraints)), + max if max <= u8::MAX as u64 => { + Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)) + } + max if max <= u16::MAX as u64 => { + Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)) + } + max if max <= u32::MAX as u64 => { + Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)) + } + _ => Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)), } } impl BytesDictBuilder { - pub fn new(dtype: DType, constraints: &DictConstraints) -> Self { + pub fn new(dtype: DType, constraints: &DictConstraints, allocator: BufferAllocatorRef) -> Self { Self { lookup: Some(HashTable::new()), - views: BufferMut::::empty(), + views: BufferMut::::empty_aligned_in( + Alignment::of::(), + allocator.clone(), + ), null_code: OnceCell::new(), - values: BufferMut::empty(), - values_nulls: BitBufferMut::empty(), + values: BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + values_nulls: BitBufferMut::empty_in(allocator.clone()), hasher: DefaultHashBuilder::default(), dtype, max_dict_bytes: constraints.max_bytes.min(u32::MAX as usize), max_dict_len: constraints.max_len, + allocator, } } @@ -172,7 +189,7 @@ impl BytesDictBuilder { F: FnMut(usize) -> &'a [u8], { let mut local_lookup = self.lookup.take().vortex_expect("Must have a lookup dict"); - let mut codes: BufferMut = BufferMut::with_capacity(len); + let mut codes = BufferMut::::with_capacity_in(len, self.allocator.clone()); match validity_mask.bit_buffer() { AllOr::All => { @@ -290,9 +307,21 @@ impl DictEncoder for BytesDictBuilder { lookup.clear(); } self.null_code = OnceCell::new(); - let views = mem::take(&mut self.views).freeze(); - let buffer = mem::take(&mut self.values).freeze(); - let value_nulls = mem::take(&mut self.values_nulls).freeze(); + let views = mem::replace( + &mut self.views, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + let buffer = mem::replace( + &mut self.values, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + let value_nulls = mem::replace( + &mut self.values_nulls, + BitBufferMut::empty_in(self.allocator.clone()), + ) + .freeze(); // SAFETY: we build the views explicitly and the bytes should be checked before feeding // to the encoder. @@ -318,6 +347,7 @@ mod test { use std::sync::LazyLock; use vortex_buffer::Buffer; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -443,7 +473,7 @@ mod test { fn reset_clears_dict() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let first = VarBinViewArray::from_iter_str(["one", "two"]).into_array(); - let mut encoder = dict_encoder(&first, &UNCONSTRAINED); + let mut encoder = dict_encoder(&first, &UNCONSTRAINED, ctx.allocator().clone()); assert_arrays_eq!( encoder.encode(&first, &mut ctx)?, @@ -465,8 +495,11 @@ mod test { #[test] fn max_dict_bytes_cannot_exceed_the_view_offset_range() { - let builder = - BytesDictBuilder::::new(DType::Utf8(Nullability::NonNullable), &UNCONSTRAINED); + let builder = BytesDictBuilder::::new( + DType::Utf8(Nullability::NonNullable), + &UNCONSTRAINED, + BufferAllocatorRef::statically_allocated(), + ); assert_eq!(builder.max_dict_bytes, u32::MAX as usize); } } diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 94834f83cc4..53e43ff79d7 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -3,6 +3,7 @@ use bytes::bytes_dict_builder; use primitive::primitive_dict_builder; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; @@ -44,15 +45,19 @@ pub trait DictEncoder: Send { fn codes_ptype(&self) -> PType; } -pub fn dict_encoder(array: &ArrayRef, constraints: &DictConstraints) -> Box { +pub fn dict_encoder( + array: &ArrayRef, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, +) -> Box { let dict_builder: Box = if let Some(pa) = array.as_opt::() { match_each_native_ptype!(pa.ptype(), |P| { - primitive_dict_builder::

(pa.dtype().nullability(), constraints) + primitive_dict_builder::

(pa.dtype().nullability(), constraints, allocator) }) } else if let Some(vbv) = array.as_opt::() { - bytes_dict_builder(vbv.dtype().clone(), constraints) + bytes_dict_builder(vbv.dtype().clone(), constraints, allocator) } else if let Some(vb) = array.as_opt::() { - bytes_dict_builder(vb.dtype().clone(), constraints) + bytes_dict_builder(vb.dtype().clone(), constraints, allocator) } else { vortex_panic!("Can only encode primitive or varbin/view arrays") }; @@ -67,7 +72,7 @@ pub fn dict_encode_with_constraints( constraints: &DictConstraints, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut encoder = dict_encoder(array, constraints); + let mut encoder = dict_encoder(array, constraints, ctx.allocator().clone()); let codes = encoder.encode(array, ctx)?.narrow(ctx)?; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index 4c8d8ae6856..cb6d12efe90 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -6,7 +6,9 @@ use std::hash::Hash; use std::mem; use rustc_hash::FxBuildHasher; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -31,6 +33,7 @@ use crate::validity::Validity; pub fn primitive_dict_builder( nullability: Nullability, constraints: &DictConstraints, + allocator: BufferAllocatorRef, ) -> Box where NativeValue: Hash + Eq, @@ -44,20 +47,25 @@ where width => vortex_panic!("invalid bit_width: {width}"), }); match max_possible_len { - max if max <= u8::MAX as u64 => { - Box::new(PrimitiveDictBuilder::::new(nullability, constraints)) - } + max if max <= u8::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( + nullability, + constraints, + allocator, + )), max if max <= u16::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( nullability, constraints, + allocator, )), max if max <= u32::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( nullability, constraints, + allocator, )), _ => Box::new(PrimitiveDictBuilder::::new( nullability, constraints, + allocator, )), } } @@ -68,15 +76,20 @@ where NativeValue: Hash + Eq, Code: UnsignedPType, { - pub fn new(nullability: Nullability, constraints: &DictConstraints) -> Self { + pub fn new( + nullability: Nullability, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, + ) -> Self { let max_dict_len = constraints .max_len .min(constraints.max_bytes / T::PTYPE.byte_width()); Self { lookup: HashMap::with_hasher(FxBuildHasher), null_code: OnceCell::new(), - values: BufferMut::::empty(), - values_nulls: BitBufferMut::empty(), + values: BufferMut::::empty_aligned_in(Alignment::of::(), allocator.clone()), + values_nulls: BitBufferMut::empty_in(allocator.clone()), + allocator, nullability, max_dict_len, } @@ -135,6 +148,7 @@ pub struct PrimitiveDictBuilder { values_nulls: BitBufferMut, nullability: Nullability, max_dict_len: usize, + allocator: BufferAllocatorRef, } impl DictEncoder for PrimitiveDictBuilder @@ -144,7 +158,7 @@ where Code: UnsignedPType, { fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let mut codes = BufferMut::::with_capacity(array.len()); + let mut codes = BufferMut::::with_capacity_in(array.len(), self.allocator.clone()); let prim = array.clone().execute::(ctx)?; match prim.validity()?.execute_mask(array.len(), ctx)? { @@ -185,11 +199,16 @@ where fn reset(&mut self) -> ArrayRef { self.lookup.clear(); self.null_code = OnceCell::new(); - PrimitiveArray::new( - mem::take(&mut self.values), - Validity::from_bit_buffer(mem::take(&mut self.values_nulls).freeze(), self.nullability), + let values = mem::replace( + &mut self.values, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ); + let nulls = mem::replace( + &mut self.values_nulls, + BitBufferMut::empty_in(self.allocator.clone()), ) - .into_array() + .freeze(); + PrimitiveArray::new(values, Validity::from_bit_buffer(nulls, self.nullability)).into_array() } fn codes_ptype(&self) -> PType { @@ -256,7 +275,7 @@ mod test { fn reset_clears_dict() { let mut ctx = SESSION.create_execution_ctx(); let first = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); - let mut encoder = dict_encoder(&first, &UNCONSTRAINED); + let mut encoder = dict_encoder(&first, &UNCONSTRAINED, ctx.allocator().clone()); assert_arrays_eq!( encoder.encode(&first, &mut ctx).unwrap(), diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index c1a91f202d0..6b99431af6e 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -3,6 +3,7 @@ use std::any::Any; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -28,14 +29,18 @@ pub struct ExtensionBuilder { impl ExtensionBuilder { /// Creates a new `ExtensionBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(ext_dtype: ExtDTypeRef) -> Self { - Self::with_capacity(ext_dtype, DEFAULT_BUILDER_CAPACITY) + pub fn new(ext_dtype: ExtDTypeRef, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity(ext_dtype, DEFAULT_BUILDER_CAPACITY, allocator) } /// Creates a new `ExtensionBuilder` with the given `capacity`. - pub fn with_capacity(ext_dtype: ExtDTypeRef, capacity: usize) -> Self { + pub fn with_capacity( + ext_dtype: ExtDTypeRef, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { Self { - storage: ChildBuilder::with_capacity(ext_dtype.storage_dtype(), capacity), + storage: ChildBuilder::with_capacity(ext_dtype.storage_dtype(), capacity, allocator), dtype: DType::Extension(ext_dtype), } } @@ -140,7 +145,8 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); - let mut builder = ExtensionBuilder::new(ext_dtype.clone()); + let mut builder = + ExtensionBuilder::new(ext_dtype.clone(), BufferAllocatorRef::static_ref()); // Test appending a valid extension value. let storage1 = Scalar::from(Some(42i32)); @@ -170,7 +176,7 @@ mod tests { assert_eq!(array.len(), 3); // Test wrong dtype error. - let mut builder = ExtensionBuilder::new(ext_dtype); + let mut builder = ExtensionBuilder::new(ext_dtype, BufferAllocatorRef::static_ref()); let wrong_scalar = Scalar::from(true); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 98acd2c468e..022107b37bc 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -44,12 +45,18 @@ pub struct FixedSizeListBuilder { impl FixedSizeListBuilder { /// Creates a new `FixedSizeListBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(element_dtype: Arc, list_size: u32, nullability: Nullability) -> Self { + pub fn new( + element_dtype: Arc, + list_size: u32, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { Self::with_capacity( element_dtype, list_size, nullability, DEFAULT_BUILDER_CAPACITY, + allocator, ) } @@ -59,12 +66,14 @@ impl FixedSizeListBuilder { list_size: u32, nullability: Nullability, capacity: usize, + allocator: &BufferAllocatorRef, ) -> Self { let elements_capacity = capacity * list_size as usize; - let elements_builder = ChildBuilder::with_capacity(&element_dtype, elements_capacity); + let elements_builder = + ChildBuilder::with_capacity(&element_dtype, elements_capacity, allocator); let fsl_dtype = DType::FixedSizeList(element_dtype, list_size, nullability); - let nulls = ValidityBuilder::new(capacity); + let nulls = ValidityBuilder::new(capacity, allocator); Self { dtype: fsl_dtype, @@ -320,6 +329,7 @@ impl ArrayBuilder for FixedSizeListBuilder { mod tests { use std::sync::Arc; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -341,8 +351,13 @@ mod tests { #[test] fn test_empty() { - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::new(I32.into()), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::new(I32.into()), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); let fsl = builder.finish(); assert_eq!(fsl.len(), 0); @@ -351,8 +366,13 @@ mod tests { #[test] fn test_values() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -388,8 +408,13 @@ mod tests { #[test] fn test_degenerate_size_zero_non_nullable() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 0, NonNullable, 10000000); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 0, + NonNullable, + 10000000, + BufferAllocatorRef::static_ref(), + ); // Append multiple "empty" lists. for _ in 0..100 { @@ -414,8 +439,13 @@ mod tests { fn test_degenerate_size_zero_nullable() { // Use nullable elements since we'll be appending nulls let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 0, Nullable, 10000000); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 0, + Nullable, + 10000000, + BufferAllocatorRef::static_ref(), + ); // Mix of null and non-null empty lists. for i in 0..100 { @@ -443,8 +473,13 @@ mod tests { fn test_capacity_growth() { let dtype: Arc = Arc::new(I32.into()); // Start with capacity 0. - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 2, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Add more items than initial capacity. for i in 0..5 { @@ -472,7 +507,13 @@ mod tests { fn test_large_size_zero_capacity_empty_result() { let dtype: Arc = Arc::new(I32.into()); // Large list size but zero capacity and no appends. - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 100000000, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 100000000, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); let fsl = builder.finish(); assert_eq!(fsl.len(), 0); @@ -487,7 +528,13 @@ mod tests { fn test_nullable_lists_non_nullable_elements() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(DType::Primitive(I32, NonNullable)); - let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -538,8 +585,13 @@ mod tests { #[test] fn test_non_nullable_lists_nullable_elements() { let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -582,7 +634,13 @@ mod tests { #[test] fn test_append_zeros() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_zeros(5); @@ -609,7 +667,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Elements must be nullable if we're going to append null lists let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); assert_eq!(builder.dtype().nullability(), Nullable); builder.append_nulls(3); @@ -638,7 +702,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Elements must be nullable if we're going to append null lists let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); assert_eq!(builder.dtype().nullability(), Nullable); builder @@ -665,7 +735,13 @@ mod tests { #[test] fn test_append_zeros_degenerate() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 0, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 0, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); assert_eq!(builder.len(), 0); builder.append_zeros(1000); @@ -683,8 +759,13 @@ mod tests { #[test] fn test_invalid_size_error() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Try to append a list with wrong size. let result = builder.append_value( @@ -718,7 +799,13 @@ mod tests { 3, ); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); let source_array = source.into_array(); source_array @@ -799,7 +886,13 @@ mod tests { 2, ); - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 0, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 0, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); source1 .into_array() @@ -868,8 +961,13 @@ mod tests { 0, ); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 3, + NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Add some initial data. builder @@ -898,7 +996,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Use nullable elements since we'll be appending nulls let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); - let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 0); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 2, + Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); // Mix of operations. builder @@ -985,7 +1089,13 @@ mod tests { fn test_append_scalar() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); - let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 2, + Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid fixed-size list. let list_scalar1 = @@ -1045,7 +1155,13 @@ mod tests { ); // Test wrong dtype error. - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 2, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } @@ -1054,8 +1170,13 @@ mod tests { fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); let mut ctx = array_session().create_execution_ctx(); - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 3, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); @@ -1105,8 +1226,13 @@ mod tests { ); // Test dtype mismatch error. - let mut builder = - FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity( + Arc::clone(&dtype), + 3, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); assert!( builder @@ -1115,7 +1241,13 @@ mod tests { ); // Test length mismatch error. - let mut builder = FixedSizeListBuilder::with_capacity(dtype, 3, NonNullable, 10); + let mut builder = FixedSizeListBuilder::with_capacity( + dtype, + 3, + NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_len_arr = buffer![1i32, 2].into_array(); assert!( builder diff --git a/vortex-array/src/builders/lazy_null_builder.rs b/vortex-array/src/builders/lazy_null_builder.rs index 8a9f62d1a03..d46d731cf25 100644 --- a/vortex-array/src/builders/lazy_null_builder.rs +++ b/vortex-array/src/builders/lazy_null_builder.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -18,16 +19,18 @@ pub struct LazyBitBufferBuilder { inner: Option, len: usize, capacity: usize, + allocator: BufferAllocatorRef, } impl LazyBitBufferBuilder { /// Creates a new empty builder. /// `capacity` is the number of bits in the null buffer. - pub fn new(capacity: usize) -> Self { + pub fn new(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { inner: None, len: 0, capacity, + allocator, } } @@ -148,7 +151,8 @@ impl LazyBitBufferBuilder { #[inline(never)] fn materialize(&mut self) { if self.inner.is_none() { - let mut bit_mut = BitBufferMut::with_capacity(self.len.max(self.capacity)); + let mut bit_mut = + BitBufferMut::with_capacity_in(self.len.max(self.capacity), self.allocator.clone()); bit_mut.append_n(true, self.len); self.inner = Some(bit_mut); } diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index ca74b7b9328..bffd8ca7f79 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::sync::Arc; use num_traits::AsPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -56,7 +57,11 @@ pub struct ListBuilder { impl ListBuilder { /// Creates a new `ListBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(value_dtype: Arc, nullability: Nullability) -> Self { + pub fn new( + value_dtype: Arc, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { Self::with_capacity( value_dtype, nullability, @@ -64,6 +69,7 @@ impl ListBuilder { // elements builder since we cannot know this ahead of time. DEFAULT_BUILDER_CAPACITY * 2, DEFAULT_BUILDER_CAPACITY, + allocator, ) } @@ -79,9 +85,12 @@ impl ListBuilder { nullability: Nullability, elements_capacity: usize, capacity: usize, + allocator: &BufferAllocatorRef, ) -> Self { - let elements_builder = ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity); - let mut offsets_builder = PrimitiveBuilder::::with_capacity(NonNullable, capacity + 1); + let elements_builder = + ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity, allocator); + let mut offsets_builder = + PrimitiveBuilder::::with_capacity(NonNullable, capacity + 1, allocator); // The first offset is always 0 and represents an empty list. offsets_builder.append_zero(); @@ -89,7 +98,7 @@ impl ListBuilder { Self { elements_builder, offsets_builder, - nulls: ValidityBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity, allocator), dtype: DType::List(value_dtype, nullability), } } @@ -395,6 +404,7 @@ mod tests { use Nullability::NonNullable; use Nullability::Nullable; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -423,8 +433,13 @@ mod tests { #[test] fn test_empty() { - let mut builder = - ListBuilder::::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + Arc::new(I32.into()), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); let list = builder.finish(); assert_eq!(list.len(), 0); @@ -433,7 +448,13 @@ mod tests { #[test] fn test_values() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -470,7 +491,13 @@ mod tests { #[test] fn test_append_empty_list() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); assert!( builder @@ -482,7 +509,13 @@ mod tests { #[test] fn test_nullable_values() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -532,7 +565,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ListBuilder::::with_capacity(Arc::new(I32.into()), Nullable, 18, 9); + let mut builder = ListBuilder::::with_capacity( + Arc::new(I32.into()), + Nullable, + 18, + 9, + BufferAllocatorRef::static_ref(), + ); list.append_to_builder(&mut builder, &mut ctx).unwrap(); list.append_to_builder(&mut builder, &mut ctx).unwrap(); list.slice(0..0) @@ -604,34 +643,62 @@ mod tests { // `builder_with_capacity` produces a `ListViewBuilder` for `DType::List`; appending the // `List`-encoded array must dispatch into it instead of bailing. - let mut listview_builder = builder_with_capacity(list.dtype(), list.len()); + let mut listview_builder = + builder_with_capacity(list.dtype(), list.len(), BufferAllocatorRef::static_ref()); list.append_to_builder(listview_builder.as_mut(), &mut ctx)?; assert_arrays_eq!(listview_builder.finish(), list, &mut ctx); // A `ListViewBuilder` with non-`u64` (including signed) offset and size types must work // for both source encodings. - let mut lv_u64_u32 = - ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + let mut lv_u64_u32 = ListViewBuilder::::with_capacity( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); list.append_to_builder(&mut lv_u64_u32, &mut ctx)?; assert_arrays_eq!(lv_u64_u32.finish(), list, &mut ctx); - let mut lv_i64_i32 = - ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + let mut lv_i64_i32 = ListViewBuilder::::with_capacity( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); list.append_to_builder(&mut lv_i64_i32, &mut ctx)?; assert_arrays_eq!(lv_i64_i32.finish(), list, &mut ctx); - let mut lv_u32_u32 = - ListViewBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + let mut lv_u32_u32 = ListViewBuilder::::with_capacity( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); listview.append_to_builder(&mut lv_u32_u32, &mut ctx)?; assert_arrays_eq!(lv_u32_u32.finish(), list, &mut ctx); // Both source encodings appended into `ListBuilder`s with non-`u64` (including signed) // offset types. - let mut list_builder = ListBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + let mut list_builder = ListBuilder::::with_capacity( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); list.append_to_builder(&mut list_builder, &mut ctx)?; assert_arrays_eq!(list_builder.finish(), list, &mut ctx); - let mut list_builder_i32 = ListBuilder::::with_capacity(elem_dtype(), Nullable, 8, 4); + let mut list_builder_i32 = ListBuilder::::with_capacity( + elem_dtype(), + Nullable, + 8, + 4, + BufferAllocatorRef::static_ref(), + ); listview.append_to_builder(&mut list_builder_i32, &mut ctx)?; assert_arrays_eq!(list_builder_i32.finish(), list, &mut ctx); @@ -654,13 +721,25 @@ mod tests { )?; // Appending twice checks growth from a non-empty builder and offset rebasing. - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_list_array(source.as_view(), &mut ctx)?; builder.append_list_array(source.as_view(), &mut ctx)?; assert_arrays_eq!(builder.finish(), expected, &mut ctx); let source_listview = source.into_array().execute::(&mut ctx)?; - let mut builder = ListBuilder::::with_capacity(dtype, Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + dtype, + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_listview_array(source_listview.as_view(), &mut ctx)?; builder.append_listview_array(source_listview.as_view(), &mut ctx)?; assert_arrays_eq!(builder.finish(), expected, &mut ctx); @@ -705,7 +784,13 @@ mod tests { .slice(1..2)? .execute::(&mut ctx)?; - let mut builder = ListBuilder::::with_capacity(dtype, Nullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + dtype, + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_listview_array(overlapping.as_view(), &mut ctx)?; builder.append_listview_array(sliced.as_view(), &mut ctx)?; @@ -786,7 +871,13 @@ mod tests { #[test] fn test_append_scalar() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 20, 10); + let mut builder = ListBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid list. let list_scalar1 = @@ -857,7 +948,13 @@ mod tests { ); // Test wrong dtype error. - let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); + let mut builder = ListBuilder::::with_capacity( + dtype, + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } @@ -866,8 +963,13 @@ mod tests { fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); let mut ctx = array_session().create_execution_ctx(); - let mut builder = - ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); + let mut builder = ListBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); @@ -916,7 +1018,13 @@ mod tests { ); // Test dtype mismatch error. - let mut builder = ListBuilder::::with_capacity(dtype, NonNullable, 20, 10); + let mut builder = ListBuilder::::with_capacity( + dtype, + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); assert!( builder diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 2efc2d45d74..12f42aff794 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use num_traits::ToPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -80,7 +81,11 @@ pub struct ListViewBuilder { impl ListViewBuilder { /// Creates a new `ListViewBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(element_dtype: Arc, nullability: Nullability) -> Self { + pub fn new( + element_dtype: Arc, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { Self::with_capacity( element_dtype, nullability, @@ -88,6 +93,7 @@ impl ListViewBuilder { // elements builder since we cannot know this ahead of time. DEFAULT_BUILDER_CAPACITY * 2, DEFAULT_BUILDER_CAPACITY, + allocator, ) } @@ -103,15 +109,17 @@ impl ListViewBuilder { nullability: Nullability, elements_capacity: usize, capacity: usize, + allocator: &BufferAllocatorRef, ) -> Self { - let elements_builder = ChildBuilder::with_capacity(&element_dtype, elements_capacity); + let elements_builder = + ChildBuilder::with_capacity(&element_dtype, elements_capacity, allocator); let offsets_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); + PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity, allocator); let sizes_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); + PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity, allocator); - let nulls = ValidityBuilder::new(capacity); + let nulls = ValidityBuilder::new(capacity, allocator); Self { dtype: DType::List(element_dtype, nullability), @@ -635,6 +643,7 @@ fn extend_converted_sizes( mod tests { use std::sync::Arc; + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -660,8 +669,13 @@ mod tests { #[test] fn test_empty() { - let mut builder = - ListViewBuilder::::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::new(I32.into()), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); let listview = builder.finish(); assert_eq!(listview.len(), 0); @@ -671,8 +685,13 @@ mod tests { fn test_basic_append_and_nulls() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Append a regular list. builder @@ -736,8 +755,13 @@ mod tests { let mut ctx = array_session().create_execution_ctx(); // Test u64 offsets with u32 sizes. let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value( @@ -780,8 +804,13 @@ mod tests { // Test i64 offsets with i32 sizes. let dtype2: Arc = Arc::new(I32.into()); - let mut builder2 = - ListViewBuilder::::with_capacity(Arc::clone(&dtype2), NonNullable, 0, 0); + let mut builder2 = ListViewBuilder::::with_capacity( + Arc::clone(&dtype2), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); for i in 0..5 { builder2 @@ -808,8 +837,13 @@ mod tests { fn test_builder_trait_methods() { let mut ctx = array_session().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Test append_zeros (creates empty lists). builder.append_zeros(2); @@ -869,8 +903,13 @@ mod tests { ) .unwrap(); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Add initial data. builder @@ -946,8 +985,13 @@ mod tests { (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect(); let source = ListArray::from_iter_opt_slow::(lists.clone(), Arc::clone(&dtype))?; - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_list_array(source.as_view(), &mut ctx)?; // Append a second time to check growth from a non-empty builder and offset rebasing. builder.append_list_array(source.as_view(), &mut ctx)?; @@ -979,8 +1023,13 @@ mod tests { ); let constant = ConstantArray::new(fill, ROWS).into_array(); - let mut builder = - ListViewBuilder::::with_capacity(element_dtype, NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + element_dtype, + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); constant.append_to_builder(&mut builder, &mut ctx)?; let listview = builder.finish_into_listview(); @@ -1012,8 +1061,13 @@ mod tests { .execute::(&mut ctx)?; let middle = source.slice(1..4)?.execute::(&mut ctx)?; - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder.append_listview_array(middle.as_view(), &mut ctx)?; // A second append has to rebase onto the elements already in the builder. builder.append_listview_array(middle.as_view(), &mut ctx)?; @@ -1057,8 +1111,13 @@ mod tests { }; assert!(!source.is_zero_copy_to_list()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + Nullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_listview_array(source.as_view(), &mut ctx) .unwrap(); @@ -1092,8 +1151,13 @@ mod tests { #[test] fn test_error_append_null_to_non_nullable() { let dtype: Arc = Arc::new(I32.into()); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); // Create a null list with nullable type (since Scalar::null requires nullable type). let null_scalar = Scalar::null(DType::List(dtype, Nullable)); @@ -1114,8 +1178,13 @@ mod tests { fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); let mut ctx = array_session().create_execution_ctx(); - let mut builder = - ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); + let mut builder = ListViewBuilder::::with_capacity( + Arc::clone(&dtype), + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); // Append a primitive array as a single list entry. let arr1 = buffer![1i32, 2, 3].into_array(); @@ -1171,7 +1240,13 @@ mod tests { ); // Test dtype mismatch error. - let mut builder = ListViewBuilder::::with_capacity(dtype, NonNullable, 20, 10); + let mut builder = ListViewBuilder::::with_capacity( + dtype, + NonNullable, + 20, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_dtype_arr = buffer![1i64, 2, 3].into_array(); assert!( builder diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs index 6714cf8f881..740748da780 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -39,17 +40,27 @@ pub struct MapBuilder { impl MapBuilder { /// Creates a map builder with the default capacity. - pub fn new(map_dtype: MapDType, nullability: Nullability) -> Self { - Self::with_capacity(map_dtype, nullability, DEFAULT_BUILDER_CAPACITY) + pub fn new( + map_dtype: MapDType, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity(map_dtype, nullability, DEFAULT_BUILDER_CAPACITY, allocator) } /// Creates a map builder with space for `capacity` map rows. - pub fn with_capacity(map_dtype: MapDType, nullability: Nullability, capacity: usize) -> Self { + pub fn with_capacity( + map_dtype: MapDType, + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { let entries_builder = ListViewBuilder::with_capacity( Arc::new(map_dtype.entries_dtype()), nullability, capacity.saturating_mul(2), capacity, + allocator, ); let dtype = DType::Map(map_dtype.clone(), nullability); Self { diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index ea8ee4d74a0..d36c6cae8dd 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -17,10 +17,11 @@ //! ``` //! use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; //! use vortex_array::dtype::{DType, Nullability}; +//! use vortex_array::memory::BufferAllocatorRef; //! use vortex_array::{VortexSessionExecute, array_session}; //! //! // Create a new builder for string data. -//! let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); +//! let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4, BufferAllocatorRef::static_ref()); //! //! builder.append_scalar(&"a".into()).unwrap(); //! builder.append_scalar(&"b".into()).unwrap(); @@ -368,10 +369,11 @@ macro_rules! __match_each_map_builder_size { /// ``` /// use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; /// use vortex_array::dtype::{DType, Nullability}; +/// use vortex_array::memory::BufferAllocatorRef; /// use vortex_array::{VortexSessionExecute, array_session}; /// /// // Create a new builder for string data. -/// let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); +/// let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4, BufferAllocatorRef::static_ref()); /// /// builder.append_scalar(&"a".into()).unwrap(); /// builder.append_scalar(&"b".into()).unwrap(); @@ -386,13 +388,19 @@ macro_rules! __match_each_map_builder_size { /// assert_eq!(strings.execute_scalar(2, &mut ctx).unwrap(), "c".into()); /// assert_eq!(strings.execute_scalar(3, &mut ctx).unwrap(), "d".into()); /// ``` -pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { +pub fn builder_with_capacity( + dtype: &DType, + capacity: usize, + allocator: &BufferAllocatorRef, +) -> Box { match dtype { DType::Null => Box::new(NullBuilder::new()), - DType::Bool(n) => Box::new(BoolBuilder::with_capacity(*n, capacity)), + DType::Bool(n) => Box::new(BoolBuilder::with_capacity(*n, capacity, allocator)), DType::Primitive(ptype, n) => { match_each_native_ptype!(ptype, |P| { - Box::new(PrimitiveBuilder::

::with_capacity(*n, capacity)) + Box::new(PrimitiveBuilder::

::with_capacity( + *n, capacity, allocator, + )) }) } DType::Decimal(decimal_type, n) => { @@ -403,25 +411,33 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box Box::new(VarBinViewBuilder::with_capacity(DType::Utf8(*n), capacity)), + DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity( + DType::Utf8(*n), + capacity, + allocator.clone(), + )), DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity( DType::Binary(*n), capacity, + allocator.clone(), )), DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity( Arc::clone(dtype), *n, 2 * capacity, // Arbitrarily choose 2 times the `offsets` capacity here. capacity, + allocator, )), DType::Map(map_dtype, nullability) => Box::new(MapBuilder::::with_capacity( map_dtype.clone(), *nullability, capacity, + allocator, )), DType::FixedSizeList(elem_dtype, list_size, null) => { Box::new(FixedSizeListBuilder::with_capacity( @@ -429,30 +445,23 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box Box::new(StructBuilder::with_capacity( struct_dtype.clone(), *n, capacity, + allocator, )), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => { unimplemented!() } - DType::Extension(ext_dtype) => { - Box::new(ExtensionBuilder::with_capacity(ext_dtype.clone(), capacity)) - } + DType::Extension(ext_dtype) => Box::new(ExtensionBuilder::with_capacity( + ext_dtype.clone(), + capacity, + allocator, + )), } } - -/// Construct a new canonical builder for the given [`DType`] using a host -/// [`vortex_buffer::BufferAllocator`]. -pub fn builder_with_capacity_in( - allocator: BufferAllocatorRef, - dtype: &DType, - capacity: usize, -) -> Box { - let _allocator = allocator; - builder_with_capacity(dtype, capacity) -} diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index aca2db36286..daa603426e8 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::mem::MaybeUninit; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -32,15 +33,19 @@ pub struct PrimitiveBuilder { impl PrimitiveBuilder { /// Creates a new `PrimitiveBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(nullability: Nullability) -> Self { - Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY) + pub fn new(nullability: Nullability, allocator: &BufferAllocatorRef) -> Self { + Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY, allocator) } /// Creates a new `PrimitiveBuilder` with the given `capacity`. - pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self { + pub fn with_capacity( + nullability: Nullability, + capacity: usize, + allocator: &BufferAllocatorRef, + ) -> Self { Self { - values: BufferMut::with_capacity(capacity), - nulls: LazyBitBufferBuilder::new(capacity), + values: BufferMut::with_capacity_in(capacity, allocator.clone()), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), dtype: DType::Primitive(T::PTYPE, nullability), } } @@ -83,10 +88,11 @@ impl PrimitiveBuilder { /// use std::mem::MaybeUninit; /// use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder}; /// use vortex_array::dtype::Nullability; + /// use vortex_array::memory::BufferAllocatorRef; /// /// // Create a new builder. /// let mut builder: PrimitiveBuilder = - /// PrimitiveBuilder::with_capacity(Nullability::NonNullable, 5); + /// PrimitiveBuilder::with_capacity(Nullability::NonNullable, 5, BufferAllocatorRef::static_ref()); /// /// // Populate the values. /// let mut uninit_range = builder.uninit_range(5); @@ -120,7 +126,9 @@ impl PrimitiveBuilder { .nulls .finish_with_nullability(self.dtype().nullability()); - PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity) + let allocator = self.values.allocator().clone(); + let values = std::mem::replace(&mut self.values, allocator.with_capacity(0)).freeze(); + PrimitiveArray::new(values, validity) } /// Extends the primitive array with an iterator. @@ -371,7 +379,11 @@ mod tests { #[test] fn test_multiple_uninit_ranges_correct_offsets() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); // First range. let mut range1 = builder.uninit_range(3); @@ -414,7 +426,11 @@ mod tests { /// This test ensures the new API works correctly. #[test] fn test_append_mask_on_uninit_range() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 5, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(3); // Create a mask for 3 values. @@ -464,7 +480,11 @@ mod tests { expected = "Tried to append a mask to an `UninitRange` that was beyond the allowed range" )] fn test_append_mask_wrong_length_panics() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(5); // Try to append a mask with wrong length (3 instead of 5). @@ -482,7 +502,11 @@ mod tests { #[test] fn test_copy_from_slice_with_offsets() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(6); // Copy to different offsets. @@ -510,7 +534,11 @@ mod tests { /// modify individual bits with relative indexing. #[test] fn test_set_bit_relative_indexing() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // First add some values to the builder. builder.append_value(100); @@ -584,7 +612,10 @@ mod tests { #[test] #[should_panic(expected = "cannot create an uninit range of length 0")] fn test_zero_length_uninit_range_panics() { - let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); + let mut builder = PrimitiveBuilder::::new( + Nullability::NonNullable, + BufferAllocatorRef::static_ref(), + ); let _range = builder.uninit_range(0); } @@ -592,7 +623,11 @@ mod tests { #[test] #[should_panic(expected = "uninit_range of len 261 exceeds builder with length 0 and capacity")] fn test_uninit_range_exceeds_capacity_panics() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::NonNullable, + 5, + BufferAllocatorRef::static_ref(), + ); let _range = builder.uninit_range(261); } @@ -603,7 +638,11 @@ mod tests { #[cfg(debug_assertions)] #[should_panic(expected = "tried to copy a slice into a `UninitRange` past its boundary")] fn test_copy_from_slice_out_of_bounds() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(3); // Try to copy 3 elements starting at offset 1 (would need 4 slots total). @@ -615,7 +654,11 @@ mod tests { /// This test demonstrates proper usage of the unsafe `finish` method. #[test] fn test_finish_unsafe_contract() { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 5, + BufferAllocatorRef::static_ref(), + ); let mut range = builder.uninit_range(3); // Set validity mask. @@ -643,7 +686,11 @@ mod tests { use crate::dtype::DType; use crate::scalar::Scalar; - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 10, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid primitive value. let scalar1 = Scalar::primitive(42i32, Nullability::Nullable); @@ -694,7 +741,11 @@ mod tests { ); // Test wrong dtype error. - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::NonNullable, + 10, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(true); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 66a2c66358d..7966952d7f0 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -4,6 +4,7 @@ use std::any::Any; use itertools::Itertools; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -35,8 +36,17 @@ pub struct StructBuilder { impl StructBuilder { /// Creates a new `StructBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`]. - pub fn new(struct_dtype: StructFields, nullability: Nullability) -> Self { - Self::with_capacity(struct_dtype, nullability, DEFAULT_BUILDER_CAPACITY) + pub fn new( + struct_dtype: StructFields, + nullability: Nullability, + allocator: &BufferAllocatorRef, + ) -> Self { + Self::with_capacity( + struct_dtype, + nullability, + DEFAULT_BUILDER_CAPACITY, + allocator, + ) } /// Creates a new `StructBuilder` with the given `capacity`. @@ -44,15 +54,16 @@ impl StructBuilder { struct_dtype: StructFields, nullability: Nullability, capacity: usize, + allocator: &BufferAllocatorRef, ) -> Self { let builders = struct_dtype .fields() - .map(|dt| ChildBuilder::with_capacity(&dt, capacity)) + .map(|dt| ChildBuilder::with_capacity(&dt, capacity, allocator)) .collect(); Self { builders, - nulls: ValidityBuilder::new(capacity), + nulls: ValidityBuilder::new(capacity, allocator), dtype: DType::Struct(struct_dtype, nullability), } } @@ -200,6 +211,8 @@ impl ArrayBuilder for StructBuilder { #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -220,7 +233,12 @@ mod tests { fn test_struct_builder() { let sdt = StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]); let dtype = DType::Struct(sdt.clone(), Nullability::NonNullable); - let mut builder = StructBuilder::with_capacity(sdt, Nullability::NonNullable, 0); + let mut builder = StructBuilder::with_capacity( + sdt, + Nullability::NonNullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value(Scalar::struct_(dtype.clone(), vec![1.into(), 2.into()]).as_struct()) @@ -235,7 +253,12 @@ mod tests { fn test_append_nullable_struct() { let sdt = StructFields::new(["a", "b"].into(), vec![I32.into(), I32.into()]); let dtype = DType::Struct(sdt.clone(), Nullability::Nullable); - let mut builder = StructBuilder::with_capacity(sdt, Nullability::Nullable, 0); + let mut builder = StructBuilder::with_capacity( + sdt, + Nullability::Nullable, + 0, + BufferAllocatorRef::static_ref(), + ); builder .append_value(Scalar::struct_(dtype.clone(), vec![1.into(), 2.into()]).as_struct()) @@ -271,7 +294,11 @@ mod tests { DType::Struct(fields, _) => fields.clone(), _ => panic!("Expected struct dtype"), }; - let mut builder = StructBuilder::new(struct_fields, Nullability::Nullable); + let mut builder = StructBuilder::new( + struct_fields, + Nullability::Nullable, + BufferAllocatorRef::static_ref(), + ); // Test appending a valid struct value. let struct_scalar1 = Scalar::struct_( @@ -326,7 +353,11 @@ mod tests { DType::Struct(fields, _) => fields.clone(), _ => panic!("Expected struct dtype"), }; - let mut builder = StructBuilder::new(struct_fields, Nullability::NonNullable); + let mut builder = StructBuilder::new( + struct_fields, + Nullability::NonNullable, + BufferAllocatorRef::static_ref(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 914cd04e1b1..4e941ae0f0e 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -1,10 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::alloc::Layout; +use std::ptr::NonNull; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use allocator_api2::alloc::AllocError; +use allocator_api2::alloc::Allocator; +use allocator_api2::alloc::Global; use rstest::rstest; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -53,6 +61,56 @@ use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; use crate::validity::Validity; +#[derive(Debug)] +struct CountingAllocator { + allocations: Arc, +} + +// SAFETY: this forwards memory operations to Global and only counts allocations. +unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.allocations.fetch_add(1, Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } + } +} + +#[test] +fn builder_reuses_its_allocator_after_finish() { + let allocations = Arc::new(AtomicUsize::new(0)); + let allocator = BufferAllocatorRef::new(CountingAllocator { + allocations: Arc::clone(&allocations), + }); + let dtype = DType::Struct( + StructFields::from_iter([ + ( + "number", + DType::Primitive(PType::I32, Nullability::Nullable), + ), + ("text", DType::Utf8(Nullability::Nullable)), + ]), + Nullability::Nullable, + ); + let mut builder = builder_with_capacity(&dtype, 1, &allocator); + + builder.append_null(); + drop(builder.finish()); + let first = allocations.load(Ordering::Relaxed); + + builder.append_null(); + drop(builder.finish()); + let second = allocations.load(Ordering::Relaxed); + + assert!( + second - first >= 5, + "all nested buffers must reuse the allocator" + ); +} + /// Test that `append_zeros` produces the same result as manually appending `Scalar::default_value`. /// /// This test verifies that the implementation of `append_zeros` correctly matches the behavior @@ -107,12 +165,14 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { let num_elements = 5; // Builder 1: Use append_zeros. - let mut builder_zeros = builder_with_capacity(&dtype, num_elements); + let mut builder_zeros = + builder_with_capacity(&dtype, num_elements, BufferAllocatorRef::static_ref()); builder_zeros.append_zeros(num_elements); let array_zeros = builder_zeros.finish(); // Builder 2: Manually append default values. - let mut builder_manual = builder_with_capacity(&dtype, num_elements); + let mut builder_manual = + builder_with_capacity(&dtype, num_elements, BufferAllocatorRef::static_ref()); let default_scalar = Scalar::zero_value(&dtype); for _ in 0..num_elements { builder_manual.append_scalar(&default_scalar).unwrap(); @@ -207,7 +267,7 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { )] #[should_panic(expected = "non-nullable")] fn test_append_nulls_panics_on_non_nullable(#[case] dtype: DType, #[case] count: usize) { - let mut builder = builder_with_capacity(&dtype, count); + let mut builder = builder_with_capacity(&dtype, count, BufferAllocatorRef::static_ref()); builder.append_nulls(count); } @@ -220,7 +280,7 @@ fn test_append_nulls_panics_on_non_nullable(#[case] dtype: DType, #[case] count: #[case::nullable_utf8(DType::Utf8(Nullability::Nullable), true)] #[case::non_nullable_utf8(DType::Utf8(Nullability::NonNullable), false)] fn test_append_defaults_behavior(#[case] dtype: DType, #[case] should_be_null: bool) { - let mut builder = builder_with_capacity(&dtype, 3); + let mut builder = builder_with_capacity(&dtype, 3, BufferAllocatorRef::static_ref()); builder.append_defaults(3); let array = builder.finish(); @@ -259,8 +319,8 @@ where F: FnMut(&mut dyn ArrayBuilder), { // Create two identical builders. - let mut builder1 = builder_with_capacity(dtype, 10); - let mut builder2 = builder_with_capacity(dtype, 10); + let mut builder1 = builder_with_capacity(dtype, 10, BufferAllocatorRef::static_ref()); + let mut builder2 = builder_with_capacity(dtype, 10, BufferAllocatorRef::static_ref()); // Fill both builders with the same data. fill_builder(builder1.as_mut()); @@ -554,7 +614,8 @@ fn test_to_canonical_f32() { ))] fn test_append_scalar_comprehensive(#[case] dtype: DType) { let num_elements = 3; - let mut builder = builder_with_capacity(&dtype, num_elements * 2); + let mut builder = + builder_with_capacity(&dtype, num_elements * 2, BufferAllocatorRef::static_ref()); // Create test scalars based on the dtype. let scalars = create_test_scalars_for_dtype(&dtype, num_elements); @@ -731,7 +792,7 @@ fn assert_scalars_equal(actual: &Scalar, expected: &Scalar, dtype: &DType, index #[case::utf8(DType::Utf8(Nullability::Nullable))] #[case::binary(DType::Binary(Nullability::Nullable))] fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { - let mut builder = builder_with_capacity(&dtype, 6); + let mut builder = builder_with_capacity(&dtype, 6, BufferAllocatorRef::static_ref()); // Create a pattern of valid, null, valid, null, valid. let test_scalars = create_test_scalars_for_dtype(&dtype, 3); @@ -809,7 +870,11 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { #[test] fn test_append_scalar_wrong_dtype_rejection() { // Test bool builder rejecting i32 scalar. - let mut bool_builder = builder_with_capacity(&DType::Bool(Nullability::NonNullable), 1); + let mut bool_builder = builder_with_capacity( + &DType::Bool(Nullability::NonNullable), + 1, + BufferAllocatorRef::static_ref(), + ); let i32_scalar = Scalar::from(42i32); assert!( bool_builder.append_scalar(&i32_scalar).is_err(), @@ -817,8 +882,11 @@ fn test_append_scalar_wrong_dtype_rejection() { ); // Test i32 builder rejecting string scalar. - let mut i32_builder = - builder_with_capacity(&DType::Primitive(PType::I32, Nullability::NonNullable), 1); + let mut i32_builder = builder_with_capacity( + &DType::Primitive(PType::I32, Nullability::NonNullable), + 1, + BufferAllocatorRef::static_ref(), + ); let string_scalar = Scalar::utf8("test", Nullability::NonNullable); assert!( i32_builder.append_scalar(&string_scalar).is_err(), @@ -826,7 +894,11 @@ fn test_append_scalar_wrong_dtype_rejection() { ); // Test string builder rejecting binary scalar. - let mut string_builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 1); + let mut string_builder = builder_with_capacity( + &DType::Utf8(Nullability::NonNullable), + 1, + BufferAllocatorRef::static_ref(), + ); let binary_scalar = Scalar::binary(vec![0u8, 1, 2], Nullability::NonNullable); assert!( string_builder.append_scalar(&binary_scalar).is_err(), @@ -839,7 +911,7 @@ fn test_append_scalar_wrong_dtype_rejection() { #[test] fn test_append_scalar_repeated_same_instance() { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let mut builder = builder_with_capacity(&dtype, 5); + let mut builder = builder_with_capacity(&dtype, 5, BufferAllocatorRef::static_ref()); let scalar = Scalar::primitive(42i32, Nullability::NonNullable); @@ -915,7 +987,7 @@ fn test_children_are_not_canonicalized( ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); array.append_to_builder(builder.as_mut(), &mut ctx)?; array.append_to_builder(builder.as_mut(), &mut ctx)?; let built = builder.finish(); @@ -947,7 +1019,7 @@ fn test_children_are_chunked_on_the_boundaries_they_are_appended_on() -> VortexR let elements = ConstantArray::new(1i32, 2).into_array(); let array = FixedSizeListArray::new(elements, 2, Validity::NonNullable, 1).into_array(); - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); for _ in 0..CHUNK_LEN { array.append_to_builder(builder.as_mut(), &mut ctx)?; } @@ -965,7 +1037,7 @@ fn test_children_are_chunked_on_the_boundaries_they_are_appended_on() -> VortexR ); // The same values appended as scalars land in a single canonical child. - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); let scalar = array.execute_scalar(0, &mut ctx)?; for _ in 0..CHUNK_LEN { builder.append_scalar(&scalar)?; @@ -994,7 +1066,7 @@ fn test_struct_builder_interleaves_arrays_and_scalars() -> VortexResult<()> { vec![Scalar::primitive(1i32, Nullability::NonNullable)], ); - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); builder.append_scalar(&scalar)?; array.append_to_builder(builder.as_mut(), &mut ctx)?; builder.append_scalar(&scalar)?; @@ -1030,7 +1102,7 @@ fn test_appended_validity_is_not_materialized() -> VortexResult<()> { )? .into_array(); - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); array.append_to_builder(builder.as_mut(), &mut ctx)?; array.append_to_builder(builder.as_mut(), &mut ctx)?; let built = builder.finish(); @@ -1107,6 +1179,7 @@ fn test_list_offsets_are_rebased_across_element_chunks( Nullability::NonNullable, ), 0, + BufferAllocatorRef::static_ref(), ); lists.append_to_builder(builder.as_mut(), &mut ctx)?; lists.append_to_builder(builder.as_mut(), &mut ctx)?; @@ -1130,8 +1203,13 @@ fn test_list_builder_offsets_are_rebased_across_element_chunks() -> VortexResult let mut ctx = array_session().create_execution_ctx(); let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); - let mut builder = - ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + let mut builder = ListBuilder::::with_capacity( + element_dtype, + Nullability::NonNullable, + 0, + 0, + BufferAllocatorRef::static_ref(), + ); for value in 0..3i32 { builder .append_array_as_list(&ConstantArray::new(value, CHUNK_LEN).into_array(), &mut ctx)?; @@ -1181,7 +1259,7 @@ fn test_validity_survives_chunked_children( ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); array.append_to_builder(builder.as_mut(), &mut ctx)?; builder.append_nulls(1); array.append_to_builder(builder.as_mut(), &mut ctx)?; @@ -1189,7 +1267,7 @@ fn test_validity_survives_chunked_children( assert!(child_of(&built).is::()); - let mut null = builder_with_capacity(array.dtype(), 1); + let mut null = builder_with_capacity(array.dtype(), 1, BufferAllocatorRef::static_ref()); null.append_nulls(1); let expected = ChunkedArray::try_new( vec![array.clone(), null.finish(), array], @@ -1206,7 +1284,7 @@ fn test_chunked_children_canonicalize_recursively() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = StructArray::try_from_iter([("a", constant_i32())])?.into_array(); - let mut builder = builder_with_capacity(array.dtype(), 0); + let mut builder = builder_with_capacity(array.dtype(), 0, BufferAllocatorRef::static_ref()); array.append_to_builder(builder.as_mut(), &mut ctx)?; array.append_to_builder(builder.as_mut(), &mut ctx)?; let built = builder.finish(); diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs index d5417407b2f..522b3b750d7 100644 --- a/vortex-array/src/builders/validity.rs +++ b/vortex-array/src/builders/validity.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::vortex_panic; @@ -37,11 +38,12 @@ pub(crate) struct ValidityBuilder { impl ValidityBuilder { /// Creates a new `ValidityBuilder` whose null buffer is pre-allocated for `capacity` bits. - pub fn new(capacity: usize) -> Self { + /// Creates a validity builder with the provided allocator. + pub fn new(capacity: usize, allocator: &BufferAllocatorRef) -> Self { Self { runs: Vec::new(), runs_len: 0, - pending: LazyBitBufferBuilder::new(capacity), + pending: LazyBitBufferBuilder::new(capacity, allocator.clone()), } } @@ -132,6 +134,7 @@ impl ValidityBuilder { #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -154,7 +157,7 @@ mod tests { #[test] fn test_whole_array_validity_is_kept_as_a_run() { - let mut builder = ValidityBuilder::new(0); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); builder.append_validity(array_backed(RUN_LEN), RUN_LEN); builder.append_validity(array_backed(RUN_LEN), RUN_LEN); @@ -172,7 +175,7 @@ mod tests { /// Uniform runs collapse instead of becoming a bool array. #[test] fn test_all_valid_runs_stay_lazy() { - let mut builder = ValidityBuilder::new(0); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); builder.append_validity(Validity::AllValid, RUN_LEN); builder.append_validity(Validity::AllValid, RUN_LEN); @@ -187,7 +190,7 @@ mod tests { /// at a time does not pay a bool array for validity it never had. #[test] fn test_short_validity_is_kept_as_a_run_too() { - let mut builder = ValidityBuilder::new(0); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); for _ in 0..RUN_LEN { builder.append_validity(Validity::AllInvalid, 1); @@ -204,7 +207,7 @@ mod tests { #[test] fn test_bits_and_runs_keep_their_order() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut builder = ValidityBuilder::new(0); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); builder.append_n_nulls(1); builder.append_validity(Validity::AllValid, RUN_LEN); @@ -228,7 +231,7 @@ mod tests { #[test] fn test_non_nullable_finishes_non_nullable() { - let mut builder = ValidityBuilder::new(0); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); builder.append_validity(Validity::NonNullable, RUN_LEN); builder.append_n_non_nulls(1); @@ -241,7 +244,7 @@ mod tests { #[test] fn test_finish_resets_the_builder() { - let mut builder = ValidityBuilder::new(0); + let mut builder = ValidityBuilder::new(0, BufferAllocatorRef::static_ref()); builder.append_validity(Validity::AllInvalid, RUN_LEN); assert_eq!(builder.finish_with_nullability(Nullable).maybe_len(), None); diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index cde20d43019..4ea67c2c5a0 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -9,6 +9,7 @@ use itertools::Itertools; use num_traits::AsPrimitive; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -46,30 +47,49 @@ pub struct VarBinViewBuilder { in_progress: Option, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, + allocator: BufferAllocatorRef, } impl VarBinViewBuilder { - pub fn with_capacity(dtype: DType, capacity: usize) -> Self { - Self::new(dtype, capacity, Default::default(), Default::default(), 0.0) + pub fn with_capacity(dtype: DType, capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::new( + dtype, + capacity, + Default::default(), + Default::default(), + 0.0, + allocator, + ) } - pub fn with_buffer_deduplication(dtype: DType, capacity: usize) -> Self { + pub fn with_buffer_deduplication( + dtype: DType, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { Self::new( dtype, capacity, CompletedBuffers::Deduplicated(Default::default()), Default::default(), 0.0, + allocator, ) } - pub fn with_compaction(dtype: DType, capacity: usize, compaction_threshold: f64) -> Self { + pub fn with_compaction( + dtype: DType, + capacity: usize, + compaction_threshold: f64, + allocator: BufferAllocatorRef, + ) -> Self { Self::new( dtype, capacity, Default::default(), Default::default(), compaction_threshold, + allocator, ) } @@ -79,23 +99,26 @@ impl VarBinViewBuilder { completed: CompletedBuffers, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, + allocator: BufferAllocatorRef, ) -> Self { assert!( matches!(dtype, DType::Utf8(_) | DType::Binary(_)), "VarBinViewBuilder DType must be Utf8 or Binary." ); Self { - views_builder: BufferMut::with_capacity_preferred_aligned( + views_builder: BufferMut::with_capacity_preferred_aligned_in( capacity, Alignment::of::(), None, + allocator.clone(), ), - nulls: LazyBitBufferBuilder::new(capacity), + nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()), completed, in_progress: None, dtype, growth_strategy, compaction_threshold, + allocator, } } @@ -153,10 +176,11 @@ impl VarBinViewBuilder { fn init_in_progress(&mut self, min_len: usize) { let next_buffer_size = self.growth_strategy.next_size() as usize; let to_reserve = next_buffer_size.max(min_len); - self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned( + self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned_in( to_reserve, Alignment::of::(), None, + self.allocator.clone(), )); } @@ -616,7 +640,7 @@ impl VarBinViewBuilder { referenced: usize, ) { let buf_index = self.completed.len(); - let mut compact = ByteBufferMut::with_capacity(referenced); + let mut compact = ByteBufferMut::with_capacity_in(referenced, self.allocator.clone()); self.views_builder.reserve(count); let data = bytes.as_slice(); @@ -667,14 +691,15 @@ impl VarBinViewBuilder { let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); + let views = std::mem::replace( + &mut self.views_builder, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + // SAFETY: the builder methods check safety at each step. unsafe { - VarBinViewArray::new_unchecked( - std::mem::take(&mut self.views_builder).freeze(), - buffers.finish(), - self.dtype.clone(), - validity, - ) + VarBinViewArray::new_unchecked(views, buffers.finish(), self.dtype.clone(), validity) } } @@ -1326,6 +1351,7 @@ impl RewritingViewAdjustment { #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -1348,7 +1374,11 @@ mod tests { #[test] fn test_append_buffer_with_lengths() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); // Stages an in-progress buffer the bulk append has to flush first. builder.append_value(LONG); @@ -1378,7 +1408,11 @@ mod tests { #[test] fn test_append_buffer_with_offsets() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); let heap = ByteBuffer::copy_from(format!("..{LONG}tiny!!")); let long_len = u32::try_from(LONG.len()).unwrap(); @@ -1411,8 +1445,12 @@ mod tests { #[test] fn test_append_buffer_with_lengths_compacts_underutilized_heap() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 4, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 4, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); let heap = ByteBuffer::copy_from([b"short".as_slice(), LONG.as_bytes(), b"tiny"].concat()); let lengths = [5u32, u32::try_from(LONG.len()).unwrap(), 4]; @@ -1439,8 +1477,12 @@ mod tests { #[test] fn test_append_buffer_with_lengths_compaction_skips_null_bytes() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 2, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 2, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); let heap = ByteBuffer::copy_from([LONG.as_bytes(), LONG.as_bytes()].concat()); let lengths = [u32::try_from(LONG.len()).unwrap(); 2]; @@ -1462,8 +1504,11 @@ mod tests { /// re-pushed buffer back to its existing index instead of holding it twice. #[test] fn test_push_buffers_deduplicates() { - let mut builder = - VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 8); + let mut builder = VarBinViewBuilder::with_buffer_deduplication( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); let first = ByteBuffer::copy_from(LONG); let second = ByteBuffer::copy_from("another value far too long to inline"); @@ -1492,7 +1537,11 @@ mod tests { .map(|buffer| buffer.as_host().clone()) .collect::>(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); // Stages an in-progress buffer that the gather has to flush ahead of its own buffers. builder.append_value(LONG); @@ -1539,7 +1588,11 @@ mod tests { 0, ); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 8, + BufferAllocatorRef::statically_allocated(), + ); let views = patch_values.views(); builder.append_views_scattered( buffers, @@ -1577,8 +1630,12 @@ mod tests { .map(|buffer| buffer.as_host().clone()) .collect::>(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 8, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); let codes: [usize; 5] = [2, 0, usize::MAX, 2, 0]; builder.append_views_gathered( buffers, @@ -1615,8 +1672,12 @@ mod tests { let dictionary = >::from_iter([Some(LONG), Some(TAIL)]); let heap = dictionary.data_buffers()[0].as_host().clone(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 8, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); builder.append_views_gathered( [heap.clone()], dictionary.views(), @@ -1646,8 +1707,12 @@ mod tests { >::from_iter([Some(LONG), Some(SHORTER)]); let heap = dictionary.data_buffers()[0].as_host().clone(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 0.5); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 8, + 0.5, + BufferAllocatorRef::statically_allocated(), + ); builder.append_views_gathered( [heap.clone()], dictionary.views(), @@ -1676,8 +1741,12 @@ mod tests { .map(|buffer| buffer.as_host().clone()) .collect::>(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 8, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); builder.append_views_gathered(buffers, dictionary.views(), &Mask::new_false(3), |_| { usize::MAX }); @@ -1696,8 +1765,12 @@ mod tests { #[test] fn test_append_buffer_with_lengths_drops_fully_inlined_heap() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = - VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 4, 1.0); + let mut builder = VarBinViewBuilder::with_compaction( + DType::Utf8(Nullability::Nullable), + 4, + 1.0, + BufferAllocatorRef::statically_allocated(), + ); let heap = ByteBuffer::copy_from(b"shorttinysmall".as_slice()); builder.append_buffer_with_lengths(heap, &[5u32, 4, 5], &Mask::new_true(3)); @@ -1719,7 +1792,11 @@ mod tests { #[test] fn test_utf8_builder() { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Hello"); builder.append_null(); @@ -1748,13 +1825,20 @@ mod tests { fn test_utf8_builder_with_extend() { let mut ctx = array_session().create_execution_ctx(); let array = { - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_null(); builder.append_value("Hello2"); builder.finish() }; - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Hello1"); array.append_to_builder(&mut builder, &mut ctx).unwrap(); @@ -1776,16 +1860,22 @@ mod tests { #[test] fn test_buffer_deduplication() -> VortexResult<()> { let array = { - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("This is a long string that should not be inlined"); builder.append_value("short string"); builder.finish_into_varbinview() }; assert_eq!(array.data_buffers().len(), 1); - let mut builder = - VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_buffer_deduplication( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); let mut ctx = array_session().create_execution_ctx(); @@ -1801,8 +1891,11 @@ mod tests { assert_eq!(builder.completed_block_count(), 1); let array2 = { - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); builder.append_value("This is a long string that should not be inlined"); builder.finish_into_varbinview() }; @@ -1826,8 +1919,11 @@ mod tests { use crate::scalar::Scalar; // Test with Utf8 builder. - let mut utf8_builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); + let mut utf8_builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); // Test appending a valid utf8 value. let utf8_scalar1 = Scalar::utf8("hello", Nullability::Nullable); @@ -1847,8 +1943,11 @@ mod tests { assert_arrays_eq!(&array, &expected, &mut ctx); // Test with Binary builder. - let mut binary_builder = - VarBinViewBuilder::with_capacity(DType::Binary(Nullability::Nullable), 10); + let mut binary_builder = VarBinViewBuilder::with_capacity( + DType::Binary(Nullability::Nullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); let binary_scalar = Scalar::binary(vec![1u8, 2, 3], Nullability::Nullable); binary_builder.append_scalar(&binary_scalar).unwrap(); @@ -1862,8 +1961,11 @@ mod tests { assert_arrays_eq!(&binary_array, &expected, &mut ctx); // Test wrong dtype error. - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 10); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::NonNullable), + 10, + BufferAllocatorRef::statically_allocated(), + ); let wrong_scalar = Scalar::from(42i32); assert!(builder.append_scalar(&wrong_scalar).is_err()); } @@ -1902,6 +2004,7 @@ mod tests { Default::default(), BufferGrowthStrategy::exponential(1024, 4096), 0.0, + BufferAllocatorRef::statically_allocated(), ); // Create a value larger than max_size diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index dab9ccfcfb5..86b25c3e70d 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -36,7 +36,7 @@ use crate::Canonical; use crate::IntoArray; use crate::array::ArrayId; use crate::builders::ArrayBuilder; -use crate::builders::builder_with_capacity_in; +use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::matcher::Matcher; use crate::memory::BufferAllocatorRef; @@ -292,10 +292,10 @@ impl ArrayRef { ExecutionStep::AppendChild(i) => { if current_builder.is_none() { trace_op!(record_builder_start(&array)); - current_builder = Some(builder_with_capacity_in( - ctx.allocator().clone(), + current_builder = Some(builder_with_capacity( array.dtype(), array.len(), + ctx.allocator(), )); } let (parent, child) = unsafe { array.take_slot_unchecked(i) }?; @@ -541,8 +541,7 @@ impl Executable for ArrayRef { ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. trace_op!(record_builder_start(&array)); - let builder = - builder_with_capacity_in(ctx.allocator().clone(), array.dtype(), array.len()); + let builder = builder_with_capacity(array.dtype(), array.len(), ctx.allocator()); let mut builder = execute_into_builder(array, builder, ctx)?; let output = builder.finish(); trace_op!(record_builder_finish(&output)); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 9831a963354..c729184da6a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -709,7 +709,12 @@ fn map_array( let rows = rows.into_iter().collect::>(); let dtype = map_dtype(nullability)?; let map_dtype = dtype.as_map_opt().vortex_expect("map dtype").clone(); - let mut builder = MapBuilder::::with_capacity(map_dtype, nullability, rows.len()); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + nullability, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for row in rows { let scalar = match row { Some(entries) => map_scalar(nullability, entries)?, diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index d5ea8702c53..71d4e78837b 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -354,7 +354,7 @@ fn merge_case_branches( return else_value.cast(output_dtype); } - let builder = builder_with_capacity(&output_dtype, else_value.len()); + let builder = builder_with_capacity(&output_dtype, else_value.len(), ctx.allocator()); let fragmented = spans.len() > else_value.len() / SLICE_CROSSOVER_RUN_LEN; if fragmented { diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index 4285b79388c..248d2ffd1c3 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -22,7 +22,7 @@ use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::ScalarFnArray; use crate::arrays::VariantArray; -use crate::builders::builder_with_capacity_in; +use crate::builders::builder_with_capacity; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; @@ -144,8 +144,7 @@ impl ScalarFnVTable for VariantGet { .map_or(DType::Variant(Nullability::Nullable), DType::as_nullable); if !dtype.is_variant() { - let mut builder = - builder_with_capacity_in(ctx.allocator().clone(), &dtype, input.len()); + let mut builder = builder_with_capacity(&dtype, input.len(), ctx.allocator()); for idx in 0..input.len() { let scalar = input.execute_scalar(idx, ctx)?; let output = variant_get_scalar(&scalar, options, &dtype)?; diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 440a7e99b5a..99922f8d764 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -216,7 +216,7 @@ pub(crate) fn zip_impl( &if_true, &if_false, mask_values.as_ref(), - builder_with_capacity(&return_type, if_true.len()), + builder_with_capacity(&return_type, if_true.len(), ctx.allocator()), ctx, ) } @@ -528,6 +528,7 @@ mod tests { Default::default(), BufferGrowthStrategy::fixed(64 * 1024), 0.0, + vortex_buffer::BufferAllocatorRef::statically_allocated(), ); for _ in 0..100 { builder.append_value("Hello"); @@ -543,6 +544,7 @@ mod tests { Default::default(), BufferGrowthStrategy::fixed(64 * 1024), 0.0, + vortex_buffer::BufferAllocatorRef::statically_allocated(), ); for _ in 0..100 { builder.append_value("Hello2"); diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index e6016ec9705..07ba16e8c91 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -340,8 +340,11 @@ fn append_to_varbin_builder(bencher: Bencher, encoding: StringEncoding) { .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) .bench_values(|(array, mut ctx)| { - let mut builder = - VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + let mut builder = VarBinBuilder::::with_capacity( + array.dtype().clone(), + array.len(), + ctx.allocator(), + ); array.append_to_builder(&mut builder, &mut ctx).unwrap(); builder.finish_into_varbin() }); @@ -356,7 +359,11 @@ fn append_to_view_builder(bencher: Bencher, encoding: StringEncoding) { .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) .bench_values(|(array, mut ctx)| { - let mut builder = VarBinViewBuilder::with_capacity(array.dtype().clone(), array.len()); + let mut builder = VarBinViewBuilder::with_capacity( + array.dtype().clone(), + array.len(), + ctx.allocator().clone(), + ); array.append_to_builder(&mut builder, &mut ctx).unwrap(); builder.finish_into_varbinview() }); diff --git a/vortex-arrow/src/executor/byte.rs b/vortex-arrow/src/executor/byte.rs index 0d653b11d01..15fc5f350a2 100644 --- a/vortex-arrow/src/executor/byte.rs +++ b/vortex-arrow/src/executor/byte.rs @@ -81,7 +81,11 @@ where // The builder's offset type matches the Arrow target, so `varbin_to_byte_array` hands the // offsets buffer straight to Arrow without a cast. - let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + let mut builder = VarBinBuilder::::with_capacity( + array.dtype().clone(), + array.len(), + ctx.allocator(), + ); array.append_to_builder(&mut builder, ctx)?; varbin_to_byte_array::(builder.finish_into_varbin().as_view(), validate_utf8, ctx) } diff --git a/vortex-arrow/src/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs index 0acffa15361..55c9f576ea4 100644 --- a/vortex-arrow/src/executor/decimal.rs +++ b/vortex-arrow/src/executor/decimal.rs @@ -274,7 +274,11 @@ mod tests { #[case] _decimal_type: T, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); @@ -302,7 +306,11 @@ mod tests { use arrow_array::Decimal32Array; let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); @@ -330,7 +338,11 @@ mod tests { use arrow_array::Decimal64Array; let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); @@ -358,7 +370,11 @@ mod tests { #[case] _decimal_type: T, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); + let mut decimal = DecimalBuilder::new::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); decimal.append_value(10); decimal.append_value(11); decimal.append_value(12); diff --git a/vortex-arrow/src/executor/map.rs b/vortex-arrow/src/executor/map.rs index c6b71057a87..f801bd6e392 100644 --- a/vortex-arrow/src/executor/map.rs +++ b/vortex-arrow/src/executor/map.rs @@ -372,7 +372,12 @@ mod tests { false, )?; let dtype = DType::Map(map_dtype.clone(), Nullable); - let mut builder = MapBuilder::::with_capacity(map_dtype, Nullable, 3); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + Nullable, + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); let rows: [&[(i32, &str)]; 3] = [ &[(1, "a"), (2, "b")], &[(3, "c"), (4, "d"), (5, "e")], diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index de0a163abf1..f4abf972476 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -155,13 +155,11 @@ pub async fn parquet_to_vortex_chunks_with_batch_size( fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult { let schema = batch.schema(); let chunk = SESSION.arrow().from_arrow_record_batch(batch, &schema)?; - let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); + let mut ctx = VortexSession::default().create_execution_ctx(); + let mut builder = builder_with_capacity(chunk.dtype(), chunk.len(), ctx.allocator()); // Canonicalize the chunk. - chunk.append_to_builder( - builder.as_mut(), - &mut VortexSession::default().create_execution_ctx(), - )?; + chunk.append_to_builder(builder.as_mut(), &mut ctx)?; Ok(builder.finish()) } diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index 4c7a5f85fa6..2e537769814 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -59,7 +59,11 @@ fn test_dict_compressed() -> VortexResult<()> { #[test] fn test_null_dominated_compressed() -> VortexResult<()> { - let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 100); + let mut builder = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for i in 0..5 { builder.append_value(i as f64); } diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index 2d0a04542cc..c90ba5d1e28 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -77,7 +77,11 @@ fn test_rle_compression() -> VortexResult<()> { #[test] fn test_sparse_compression() -> VortexResult<()> { - let mut array = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 100); + let mut array = PrimitiveBuilder::::with_capacity( + Nullability::Nullable, + 100, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); array.append_value(f32::NAN); array.append_value(-f32::NAN); array.append_value(f32::INFINITY); diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index ec9c4060f64..4d19aa95ed3 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -45,7 +45,11 @@ fn test_strings() -> VortexResult<()> { #[test] fn test_sparse_nulls() -> VortexResult<()> { - let mut strings = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 100); + let mut strings = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::Nullable), + 100, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); strings.append_nulls(99); strings.append_value("one little string"); diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 2523587253c..a65e18a5334 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -41,7 +41,8 @@ impl BufferAllocatorRef { Self(None) } - pub(crate) fn static_ref() -> &'static Self { + /// Return a borrowed reference to the static allocator. + pub fn static_ref() -> &'static Self { &STATIC_ALLOCATOR } diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index 2e5c3ea39b6..e496c34716e 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -6,6 +6,7 @@ use std::ops::Not; use bitvec::view::BitView; use crate::BitBuffer; +use crate::BufferAllocatorRef; use crate::BufferMut; use crate::ByteBufferMut; use crate::bit::collect_bool_words; @@ -119,8 +120,13 @@ impl BitBufferMut { /// Creates a `BitBufferMut` from a [`BitBuffer`] by copying all of the data over. pub fn copy_from(bit_buffer: &BitBuffer) -> Self { + Self::copy_from_in(bit_buffer, bit_buffer.inner().allocator().clone()) + } + + /// Copies a bit buffer with the provided allocator. + pub fn copy_from_in(bit_buffer: &BitBuffer, allocator: BufferAllocatorRef) -> Self { Self { - buffer: ByteBufferMut::copy_from(bit_buffer.inner()), + buffer: ByteBufferMut::copy_from_in(bit_buffer.inner(), allocator), offset: bit_buffer.offset(), len: bit_buffer.len(), } @@ -129,8 +135,14 @@ impl BitBufferMut { /// Create a new empty mutable bit buffer with requested capacity (in bits). #[inline] pub fn with_capacity(capacity: usize) -> Self { + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable bit buffer with the provided allocator. + #[inline] + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { - buffer: BufferMut::with_capacity(capacity.div_ceil(8)), + buffer: BufferMut::with_capacity_in(capacity.div_ceil(8), allocator), offset: 0, len: 0, } @@ -163,6 +175,16 @@ impl BitBufferMut { Self::with_capacity(0) } + /// Create an empty mutable bit buffer with the provided allocator. + pub fn empty_in(allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_in(0, allocator) + } + + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.buffer.allocator() + } + /// Create a new mutable buffer with requested `len` and all bits set to `value`. #[inline] pub fn full(value: bool, len: usize) -> Self { diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index ec14383ce36..6ba14afb6a5 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -760,8 +760,12 @@ fn map_array_from_rows(rows: &[MapRowFixture<'_>], keys_sorted: bool) -> VortexR keys_sorted, )?; let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); - let mut builder = - MapBuilder::::with_capacity(map_dtype, Nullability::Nullable, rows.len()); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + Nullability::Nullable, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for row in rows { let scalar = match row { diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 74a87486e6b..c85808a0f67 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -18,6 +18,7 @@ use vortex::array::dtype::extension::ExtDType; use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; +use vortex::buffer::BufferAllocatorRef; use vortex::buffer::BufferMut; use vortex::dtype::DType; use vortex::dtype::DecimalDType; @@ -107,7 +108,8 @@ fn vector_as_string_blob(vector: &VectorRef, len: usize, dtype: DType) -> ArrayR let data = vector.as_slice_with_len::(len); let validity = vector.validity_ref(len); - let mut builder = VarBinViewBuilder::with_capacity(dtype, len); + let mut builder = + VarBinViewBuilder::with_capacity(dtype, len, BufferAllocatorRef::statically_allocated()); for (i, s) in data.iter().enumerate() { if validity.is_valid(i) { diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 0876be1ca4c..6134ef8501a 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -983,8 +983,11 @@ fn test_geometry() { let mut wkb_binary: Vec = Vec::new(); wkb::writer::write_polygon(&mut wkb_binary, &rect10, &WriteOptions::default()) .expect("serializing WKB"); - let mut geometry = - VarBinBuilder::::with_capacity(DType::Binary(Nullability::NonNullable), 10); + let mut geometry = VarBinBuilder::::with_capacity( + DType::Binary(Nullability::NonNullable), + 10, + vortex::buffer::BufferAllocatorRef::static_ref(), + ); for _ in 0..10 { geometry.append_value(wkb_binary.as_slice()); } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..237f303f506 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1932,8 +1932,12 @@ fn map_array_from_rows(rows: &[MapRowFixture<'_>], keys_sorted: bool) -> VortexR keys_sorted, )?; let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); - let mut builder = - MapBuilder::::with_capacity(map_dtype, Nullability::Nullable, rows.len()); + let mut builder = MapBuilder::::with_capacity( + map_dtype, + Nullability::Nullable, + rows.len(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); for row in rows { let scalar = match row { diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index c3743a1df45..42898d03b1d 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -278,7 +278,11 @@ mod tests { // Create utf8 array (2 buffers: views + data) let (ptr2, eof2) = SequenceId::root().split(); - let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(NonNullable), 5); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(NonNullable), + 5, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); for s in [ "hello world this is long", "another long string", diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index bd68864725e..2d3c4a82b02 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -560,7 +560,7 @@ fn start_encoding( chunk: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let encoder = dict_encoder(chunk, constraints); + let encoder = dict_encoder(chunk, constraints, ctx.allocator().clone()); encode_chunk(encoder, chunk, ctx) } diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index dd0c4d6c13c..c128069148d 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -219,17 +219,29 @@ fn stats_builder_with_capacity( capacity: usize, max_length: usize, ) -> Box { - let values_builder = builder_with_capacity(dtype, capacity); + let values_builder = builder_with_capacity( + dtype, + capacity, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); match stat { Stat::Max => match dtype { DType::Utf8(_) => Box::new(TruncatedMaxBinaryStatsBuilder::::new( values_builder, - BoolBuilder::with_capacity(Nullability::NonNullable, capacity), + BoolBuilder::with_capacity( + Nullability::NonNullable, + capacity, + vortex_buffer::BufferAllocatorRef::static_ref(), + ), max_length, )), DType::Binary(_) => Box::new(TruncatedMaxBinaryStatsBuilder::::new( values_builder, - BoolBuilder::with_capacity(Nullability::NonNullable, capacity), + BoolBuilder::with_capacity( + Nullability::NonNullable, + capacity, + vortex_buffer::BufferAllocatorRef::static_ref(), + ), max_length, )), _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), @@ -237,12 +249,20 @@ fn stats_builder_with_capacity( Stat::Min => match dtype { DType::Utf8(_) => Box::new(TruncatedMinBinaryStatsBuilder::::new( values_builder, - BoolBuilder::with_capacity(Nullability::NonNullable, capacity), + BoolBuilder::with_capacity( + Nullability::NonNullable, + capacity, + vortex_buffer::BufferAllocatorRef::static_ref(), + ), max_length, )), DType::Binary(_) => Box::new(TruncatedMinBinaryStatsBuilder::::new( values_builder, - BoolBuilder::with_capacity(Nullability::NonNullable, capacity), + BoolBuilder::with_capacity( + Nullability::NonNullable, + capacity, + vortex_buffer::BufferAllocatorRef::static_ref(), + ), max_length, )), _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), @@ -523,10 +543,18 @@ mod tests { #[case(DType::Binary(Nullability::NonNullable))] fn truncates_accumulated_stats(#[case] dtype: DType) { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 2); + let mut builder = VarBinViewBuilder::with_capacity( + dtype.clone(), + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Value to be truncated"); builder.append_value("untruncated"); - let mut builder2 = VarBinViewBuilder::with_capacity(dtype, 2); + let mut builder2 = VarBinViewBuilder::with_capacity( + dtype, + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder2.append_value("Another"); builder2.append_value("wait a minute"); let mut acc = @@ -573,7 +601,11 @@ mod tests { #[case(DType::Binary(Nullability::NonNullable))] fn truncated_accumulated_stats_are_inexact(#[case] dtype: DType) { let mut ctx = array_session().create_execution_ctx(); - let mut builder = VarBinViewBuilder::with_capacity(dtype, 2); + let mut builder = VarBinViewBuilder::with_capacity( + dtype, + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Value to be truncated"); builder.append_value("Another truncated value"); let mut acc = StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min], 12); diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 790e5dd6663..bebd5eb1d9b 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -260,8 +260,11 @@ mod tests { let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); - let mut builder = - VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 2); + let mut builder = VarBinViewBuilder::with_capacity( + DType::Utf8(Nullability::NonNullable), + 2, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); builder.append_value("Long value to test that the statistics are actually truncated, it needs a bit of extra padding though"); builder.append_value("Another string that's meant to be smaller than the previous value, though still need extra padding"); let array = builder.finish(); diff --git a/vortex-layout/src/layouts/zoned/builder.rs b/vortex-layout/src/layouts/zoned/builder.rs index 3c86f371fb6..35ce3343859 100644 --- a/vortex-layout/src/layouts/zoned/builder.rs +++ b/vortex-layout/src/layouts/zoned/builder.rs @@ -129,7 +129,11 @@ impl AggregateStatsArrayBuilder { Self { aggregate_fn, dtype: dtype.clone(), - builder: builder_with_capacity(dtype, capacity), + builder: builder_with_capacity( + dtype, + capacity, + vortex_buffer::BufferAllocatorRef::static_ref(), + ), } } diff --git a/vortex-spatial/src/tests/wkb.rs b/vortex-spatial/src/tests/wkb.rs index 40477a464c9..1651c95e7cd 100644 --- a/vortex-spatial/src/tests/wkb.rs +++ b/vortex-spatial/src/tests/wkb.rs @@ -62,8 +62,11 @@ fn wkb_extension_array() -> VortexResult<(Vec, vortex_array::ArrayRef)> { wkb::writer::write_geometry(&mut buf, &test_polygon(), &WriteOptions::default()) .map_err(|e| vortex_err!("writing WKB failed: {e}"))?; - let mut builder = - VarBinBuilder::::with_capacity(DType::Binary(Nullability::NonNullable), 3); + let mut builder = VarBinBuilder::::with_capacity( + DType::Binary(Nullability::NonNullable), + 3, + vortex_buffer::BufferAllocatorRef::static_ref(), + ); builder.append_value(&buf); builder.append_value(&buf); builder.append_value(&buf);