diff --git a/Cargo.lock b/Cargo.lock index 11f1514eb3..064405951e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5442,6 +5442,7 @@ dependencies = [ "vortex-buffer", "vortex-dtype", "vortex-error", + "vortex-fastlanes", "vortex-mask", "vortex-scalar", ] diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index bfed1580e3..cd1aeee4ff 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -22,6 +22,7 @@ arrow-buffer = { workspace = true } fastlanes = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } +rand = { workspace = true, optional = true } rkyv = { workspace = true } serde = { workspace = true } vortex-array = { workspace = true } @@ -36,6 +37,10 @@ criterion = { workspace = true } divan = { workspace = true } rand = { workspace = true } vortex-array = { workspace = true, features = ["test-harness"] } +vortex-fastlanes = { path = ".", features = ["test-harness"] } + +[features] +test-harness = ["dep:rand"] [[bench]] name = "bitpacking_take" @@ -47,4 +52,5 @@ harness = false [[bench]] name = "canonicalize_bench" -harness = false \ No newline at end of file +harness = false +required-features = ["test-harness"] diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index 33c63240a5..30d7d5be8e 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -1,116 +1,134 @@ use divan::Bencher; use rand::prelude::StdRng; -use rand::{Rng, SeedableRng}; +use rand::SeedableRng; use vortex_array::array::ChunkedArray; use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder}; -use vortex_array::{Array, IntoArray, IntoArrayVariant, IntoCanonical}; -use vortex_buffer::BufferMut; -use vortex_dtype::NativePType; -use vortex_error::{VortexExpect, VortexUnwrap}; -use vortex_fastlanes::bitpack_to_best_bit_width; +use vortex_array::{IntoArray, IntoCanonical}; +use vortex_error::{VortexExpect as _, VortexUnwrap}; +use vortex_fastlanes::test_harness::make_array; fn main() { divan::main(); } -fn make_array(len: usize) -> Array { +const BENCH_ARGS: [(usize, usize, f64); 10] = [ + (100000, 1, 0.10), + (100000, 1, 0.01), + (100000, 1, 0.00), + (100000, 10, 0.10), + (100000, 10, 0.01), + (100000, 10, 0.00), + (100000, 100, 0.10), + (100000, 100, 0.01), + (100000, 100, 0.00), + (100000, 1000, 0.00), + // Too slow for 1000 samples. Try 10 samples. + // (1000000, 100, 0.00), + // (1000000, 1000, 0.00), + // (10000000, 100, 0.00), +]; + +#[divan::bench(args = BENCH_ARGS)] +fn into_canonical_non_nullable( + bencher: Bencher, + (chunk_len, chunk_count, fraction_patched): (usize, usize, f64), +) { let mut rng = StdRng::seed_from_u64(0); - let values = (0..len) - .map(|_| T::from(rng.gen_range(0..100)).vortex_expect("valid value")) - .collect::>() - .into_array() - .into_primitive() - .vortex_unwrap(); - bitpack_to_best_bit_width(values) - .vortex_unwrap() - .into_array() + let chunks = (0..chunk_count) + .map(|_| { + make_array(&mut rng, chunk_len, fraction_patched, 0.0).vortex_expect("make_array works") + }) + .collect::>(); + let chunked = ChunkedArray::from_iter(chunks).into_array(); + + bencher + .with_inputs(|| chunked.clone()) + .bench_values(|chunked| chunked.into_canonical().vortex_unwrap()); } -#[divan::bench()] -fn test() { - let chunks = (0..10).map(|_| make_array::(100)).collect::>(); - let arr = make_array::(1); - let chunked = ChunkedArray::try_new(chunks, arr.dtype().clone()) - .vortex_unwrap() - .into_array(); +#[divan::bench(args = BENCH_ARGS)] +fn canonical_into_non_nullable( + bencher: Bencher, + (chunk_len, chunk_count, fraction_patched): (usize, usize, f64), +) { + let mut rng = StdRng::seed_from_u64(0); - let into_ca = chunked - .clone() - .into_canonical() - .vortex_unwrap() - .into_primitive() - .vortex_unwrap(); - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(arr.dtype().nullability(), 10 * 100); - chunked - .clone() - .canonicalize_into(&mut primitive_builder) - .vortex_unwrap(); - let ca_into = primitive_builder.finish().vortex_unwrap(); + let chunks = (0..chunk_count) + .map(|_| { + make_array(&mut rng, chunk_len, fraction_patched, 0.0).vortex_expect("make_array works") + }) + .collect::>(); + let chunked = ChunkedArray::from_iter(chunks).into_array(); - assert_eq!( - into_ca.as_slice::(), - ca_into.into_primitive().vortex_unwrap().as_slice::() - ); + bencher + .with_inputs(|| chunked.clone()) + .bench_values(|chunked| { + let mut primitive_builder = PrimitiveBuilder::::with_capacity( + chunked.dtype().nullability(), + chunk_len * chunk_count, + ); + chunked + .canonicalize_into(&mut primitive_builder) + .vortex_unwrap(); + primitive_builder.finish().vortex_unwrap() + }); +} - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(arr.dtype().nullability(), 10 * 100); - primitive_builder.extend_from_array(chunked).vortex_unwrap(); - let ca_into = primitive_builder.finish().vortex_unwrap(); +const NULLABLE_BENCH_ARGS: [(usize, usize, f64); 6] = [ + (100000, 1, 0.10), + (100000, 1, 0.00), + (100000, 10, 0.10), + (100000, 10, 0.00), + (100000, 100, 0.10), + (100000, 100, 0.00), +]; - assert_eq!( - into_ca.as_slice::(), - ca_into.into_primitive().vortex_unwrap().as_slice::() - ); -} +#[divan::bench(args = NULLABLE_BENCH_ARGS)] +fn into_canonical_nullable( + bencher: Bencher, + (chunk_len, chunk_count, fraction_patched): (usize, usize, f64), +) { + let mut rng = StdRng::seed_from_u64(0); -#[divan::bench( - types = [u32], - args = [ - // (1000, 100), - // (100000, 100), - // (1000000, 100), - // (100000, 1000), - (100000, 3), - ] -)] -fn into_canonical(bencher: Bencher, (arr_len, chunk_count): (usize, usize)) { let chunks = (0..chunk_count) - .map(|_| make_array::(arr_len)) + .map(|_| { + make_array(&mut rng, chunk_len, fraction_patched, 0.05) + .vortex_expect("make_array works") + }) .collect::>(); - let arr = make_array::(1); - let chunked = ChunkedArray::try_new(chunks, arr.dtype().clone()).vortex_unwrap(); + let chunked = ChunkedArray::from_iter(chunks).into_array(); - bencher.bench(|| chunked.clone().into_canonical().vortex_unwrap().len()); + bencher + .with_inputs(|| chunked.clone()) + .bench_values(|chunked| chunked.into_canonical().vortex_unwrap()); } -#[divan::bench( - types = [u32], - args = [ - // (1000, 100), - // (100000, 100), - // (1000000, 100), - // (100000, 1000), - (100000, 3), - ] -)] -fn canonical_into(bencher: Bencher, (arr_len, chunk_count): (usize, usize)) { +#[divan::bench(args = NULLABLE_BENCH_ARGS)] +fn canonical_into_nullable( + bencher: Bencher, + (chunk_len, chunk_count, fraction_patched): (usize, usize, f64), +) { + let mut rng = StdRng::seed_from_u64(0); + let chunks = (0..chunk_count) - .map(|_| make_array::(arr_len)) + .map(|_| { + make_array(&mut rng, chunk_len, fraction_patched, 0.05) + .vortex_expect("make_array works") + }) .collect::>(); - let arr = make_array::(1); - let chunked = ChunkedArray::try_new(chunks, arr.dtype().clone()) - .vortex_unwrap() - .into_array(); + let chunked = ChunkedArray::from_iter(chunks).into_array(); - bencher.bench(|| { - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(arr.dtype().nullability(), arr_len * chunk_count); - chunked - .clone() - .canonicalize_into(&mut primitive_builder) - .vortex_unwrap(); - primitive_builder.finish().vortex_unwrap().len() - }); + bencher + .with_inputs(|| chunked.clone()) + .bench_values(|chunked| { + let mut primitive_builder = PrimitiveBuilder::::with_capacity( + chunked.dtype().nullability(), + chunk_len * chunk_count, + ); + chunked + .canonicalize_into(&mut primitive_builder) + .vortex_unwrap(); + primitive_builder.finish().vortex_unwrap() + }); } diff --git a/encodings/fastlanes/src/bitpacking/compress.rs b/encodings/fastlanes/src/bitpacking/compress.rs index d586d349eb..068cc4f714 100644 --- a/encodings/fastlanes/src/bitpacking/compress.rs +++ b/encodings/fastlanes/src/bitpacking/compress.rs @@ -1,13 +1,17 @@ +use std::ops::DerefMut as _; + use arrow_buffer::ArrowNativeType; use fastlanes::BitPacking; use vortex_array::array::PrimitiveArray; +use vortex_array::builders::{ArrayBuilder as _, PrimitiveBuilder}; use vortex_array::patches::Patches; use vortex_array::validity::Validity; use vortex_array::variants::PrimitiveArrayTrait; use vortex_array::IntoArray; -use vortex_buffer::{buffer, Buffer, BufferMut, ByteBuffer}; +use vortex_buffer::{Buffer, BufferMut, ByteBuffer}; use vortex_dtype::{ - match_each_integer_ptype, match_each_unsigned_integer_ptype, NativePType, PType, + match_each_integer_ptype, match_each_integer_ptype_with_unsigned_type, + match_each_unsigned_integer_ptype, NativePType, PType, }; use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult}; use vortex_scalar::Scalar; @@ -196,39 +200,62 @@ pub fn gather_patches( } pub fn unpack(array: BitPackedArray) -> VortexResult { + match_each_integer_ptype_with_unsigned_type!(array.ptype(), |$P, $UnsignedT| { + unpack_primitive::<$P, $UnsignedT>(array) + }) +} + +pub fn unpack_primitive( + array: BitPackedArray, +) -> VortexResult { + let n = array.len(); + let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len()); + assert!(size_of::() == size_of::()); + unpack_into::( + array, + &mut builder, + // SAFETY: UnsignedT is the unsigned verison of T, reinterpreting &[UnsignedT] to + // &[T] is therefore safe. + |x| unsafe { std::mem::transmute(x) }, + // SAFETY: UnsignedT is the unsigned verison of T, reinterpreting &mut [T] to + // &mut [UnsignedT] is therefore safe. + |x| unsafe { std::mem::transmute(x) }, + )?; + assert!(builder.len() == n, "{} {}", builder.len(), n); + builder.finish_into_primitive() +} + +pub(crate) fn unpack_into( + array: BitPackedArray, + // TODO(ngates): do we want to use fastlanes alignment for this buffer? + builder: &mut PrimitiveBuilder, + transmute: F, + transmute_mut: G, +) -> VortexResult<()> +where + F: Fn(&[UnsignedT]) -> &[T], + G: Fn(&mut [T]) -> &mut [UnsignedT], +{ let bit_width = array.bit_width() as usize; let length = array.len(); let offset = array.offset() as usize; - let ptype = array.ptype(); - let mut unpacked = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |$P| { - PrimitiveArray::new( - unpack_primitive::<$P>(array.packed_slice::<$P>(), bit_width, offset, length), - array.validity(), - ) - }); + let last_chunk_length = match (offset + length) % 1024 { + 0 => 1024, + last_chunk_length => last_chunk_length, + }; - // Cast to signed if necessary - if ptype.is_signed_int() { - unpacked = unpacked.reinterpret_cast(ptype); - } + builder.nulls.append_validity(array.validity(), length)?; - if let Some(patches) = array.patches() { - unpacked.patch(patches) - } else { - Ok(unpacked) - } -} - -pub fn unpack_primitive( - packed: &[T], - bit_width: usize, - offset: usize, - length: usize, -) -> Buffer { if bit_width == 0 { - return buffer![T::zero(); length]; + builder.append_zeros(length); + return Ok(()); } + builder.values.reserve(array.len()); + + let builder_current_length = builder.len(); + let packed = array.packed_slice::(); + // How many fastlanes vectors we will process. // Packed array might not start at 0 when the array is sliced. Offset is guaranteed to be < 1024. let num_chunks = (offset + length + 1023) / 1024; @@ -241,42 +268,70 @@ pub fn unpack_primitive( num_chunks * elems_per_chunk ); - // Allocate a result vector. - // TODO(ngates): do we want to use fastlanes alignment for this buffer? - let mut output = BufferMut::with_capacity(num_chunks * 1024 - offset); - - // Handle first chunk if offset is non 0. We have to decode the chunk and skip first offset elements - let first_full_chunk = if offset != 0 { - let chunk: &[T] = &packed[0..elems_per_chunk]; - let mut decoded = [T::zero(); 1024]; + if num_chunks == 1 { + let chunk = &packed[..elems_per_chunk]; + let mut decoded = [UnsignedT::zero(); 1024]; + // SAFETY: + // 1. chunk is elems_per_chunk. + // 2. decoded is exactly 1024. unsafe { BitPacking::unchecked_unpack(bit_width, chunk, &mut decoded) }; - output.extend_from_slice(&decoded[offset..]); - 1 + builder + .values + .extend_from_slice(transmute(&decoded[offset..][..length])); } else { - 0 - }; - - // Loop over all the chunks. - (first_full_chunk..num_chunks).for_each(|i| { - let chunk: &[T] = &packed[i * elems_per_chunk..][0..elems_per_chunk]; - unsafe { - let output_len = output.len(); - output.set_len(output_len + 1024); - BitPacking::unchecked_unpack(bit_width, chunk, &mut output[output_len..][0..1024]); + let first_chunk_is_sliced = offset != 0; + let last_chunk_is_sliced = last_chunk_length != 1024; + let full_chunks_range = + (first_chunk_is_sliced as usize)..(num_chunks - last_chunk_is_sliced as usize); + + if first_chunk_is_sliced { + let chunk = &packed[..elems_per_chunk]; + let mut decoded = [UnsignedT::zero(); 1024]; + // SAFETY: + // 1. chunk is elems_per_chunk. + // 2. decoded is exactly 1024. + unsafe { BitPacking::unchecked_unpack(bit_width, chunk, &mut decoded) }; + builder + .values + .extend_from_slice(transmute(&decoded[offset..])); } - }); + for i in full_chunks_range { + let chunk = &packed[i * elems_per_chunk..][..elems_per_chunk]; + + // SAFETY: + // + // 1. unchecked_unpack only writes into the output and when it is finished, all the outputs + // have been written to. + // + // 2. The output buffer is exactly size 1024. + unsafe { + builder.values.set_len(builder.values.len() + 1024); + BitPacking::unchecked_unpack( + bit_width, + chunk, + &mut transmute_mut(builder.values.deref_mut()) + [builder_current_length + i * 1024 - offset..][..1024], + ); + } + } + if last_chunk_is_sliced { + let chunk = &packed[(num_chunks - 1) * elems_per_chunk..][..elems_per_chunk]; + let mut decoded = [UnsignedT::zero(); 1024]; + // SAFETY: + // 1. chunk is elems_per_chunk. + // 2. decoded is exactly 1024. + unsafe { BitPacking::unchecked_unpack(bit_width, chunk, &mut decoded) }; + builder + .values + .extend_from_slice(transmute(&decoded[..last_chunk_length])); + } + } - // The final chunk may have had padding - output.truncate(length); + if let Some(patches) = array.patches() { + builder.patch(patches, builder_current_length)?; + } - assert_eq!( - output.len(), - length, - "Expected unpacked array to be of length {} but got {}", - length, - output.len() - ); - output.freeze() + Ok(()) } pub fn unpack_single(array: &BitPackedArray, index: usize) -> VortexResult { @@ -381,13 +436,56 @@ pub fn count_exceptions(bit_width: u8, bit_width_freq: &[usize]) -> usize { bit_width_freq[bit_width as usize + 1..].iter().sum() } +#[cfg(feature = "test-harness")] +pub mod test_harness { + use rand::rngs::StdRng; + use rand::Rng as _; + use vortex_array::array::PrimitiveArray; + use vortex_array::validity::Validity; + use vortex_array::{Array, IntoArray, IntoArrayVariant as _}; + use vortex_buffer::BufferMut; + use vortex_error::VortexResult; + + use super::bitpack_encode; + + pub fn make_array( + rng: &mut StdRng, + len: usize, + fraction_patches: f64, + fraction_null: f64, + ) -> VortexResult { + let values = (0..len) + .map(|_| { + let mut v = rng.gen_range(0..100i32); + if rng.gen_bool(fraction_patches) { + v += 1 << 13 + }; + v + }) + .collect::>(); + + let values = if fraction_null == 0.0 { + values.into_array().into_primitive()? + } else { + let validity = Validity::from_iter((0..len).map(|_| !rng.gen_bool(fraction_null))); + PrimitiveArray::new(values, validity) + }; + + bitpack_encode(values, 12).map(IntoArray::into_array) + } +} + #[cfg(test)] #[allow(clippy::cast_possible_truncation)] mod test { - use vortex_array::IntoArrayVariant; + use rand::rngs::StdRng; + use rand::SeedableRng as _; + use vortex_array::array::ChunkedArray; + use vortex_array::{IntoArrayVariant, IntoCanonical as _}; use vortex_error::VortexError; use super::*; + use crate::bitpacking::compress::test_harness::make_array; #[test] fn test_best_bit_width() { @@ -462,4 +560,43 @@ mod test { let err = BitPackedArray::encode(array.as_ref(), 1024u32.ilog2() as u8).unwrap_err(); assert!(matches!(err, VortexError::InvalidArgument(_, _))); } + + #[test] + fn canonicalize_chunked_of_bitpacked() { + let mut rng = StdRng::seed_from_u64(0); + + let chunks = (0..10) + .map(|_| make_array(&mut rng, 100, 0.0, 0.0).unwrap()) + .collect::>(); + let chunked = ChunkedArray::from_iter(chunks).into_array(); + + let into_ca = chunked + .clone() + .into_canonical() + .unwrap() + .into_primitive() + .unwrap(); + let mut primitive_builder = + PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); + chunked + .clone() + .canonicalize_into(&mut primitive_builder) + .unwrap(); + let ca_into = primitive_builder.finish().unwrap(); + + assert_eq!( + into_ca.as_slice::(), + ca_into.into_primitive().unwrap().as_slice::() + ); + + let mut primitive_builder = + PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); + primitive_builder.extend_from_array(chunked).unwrap(); + let ca_into = primitive_builder.finish().unwrap(); + + assert_eq!( + into_ca.as_slice::(), + ca_into.into_primitive().unwrap().as_slice::() + ); + } } diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index 1ed4d2a434..244251e126 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -13,9 +13,9 @@ use vortex_array::vtable::{ CanonicalVTable, StatisticsVTable, ValidateVTable, ValidityVTable, VariantsVTable, VisitorVTable, }; -use vortex_array::{encoding_ids, impl_encoding, Array, Canonical, IntoArray, RkyvMetadata}; +use vortex_array::{encoding_ids, impl_encoding, Array, Canonical, RkyvMetadata}; use vortex_buffer::ByteBuffer; -use vortex_dtype::{DType, NativePType, PType}; +use vortex_dtype::{match_each_integer_ptype_with_unsigned_type, DType, NativePType, PType}; use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult}; use vortex_mask::Mask; @@ -264,8 +264,21 @@ impl CanonicalVTable for BitPackedEncoding { array: BitPackedArray, builder: &mut dyn ArrayBuilder, ) -> VortexResult<()> { - // TODO(joe): add specialised impl - builder.extend_from_array(array.into_array()) + match_each_integer_ptype_with_unsigned_type!(array.ptype(), |$T, $UnsignedT| { + unpack_into::<$T, $UnsignedT, _, _>( + array, + builder + .as_any_mut() + .downcast_mut() + .vortex_expect("bit packed array must canonicalize into a primitive array"), + // SAFETY: UnsignedT is the unsigned verison of T, reinterpreting &[UnsignedT] to + // &[T] is therefore safe. + |x| unsafe { std::mem::transmute(x) }, + // SAFETY: UnsignedT is the unsigned verison of T, reinterpreting &mut [T] to + // &mut [UnsignedT] is therefore safe. + |x| unsafe { std::mem::transmute(x) }, + ) + }) } } diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 067f9c7834..559e3a3de9 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -3,7 +3,6 @@ use std::any::Any; use arrow_buffer::BooleanBufferBuilder; use vortex_dtype::{DType, Nullability}; use vortex_error::{vortex_bail, VortexResult}; -use vortex_mask::AllOr; use crate::array::BoolArray; use crate::builders::lazy_validity_builder::LazyNullBufferBuilder; @@ -82,13 +81,7 @@ impl ArrayBuilder for BoolBuilder { self.inner.append_buffer(&array.boolean_buffer()); - match array.validity_mask()?.boolean_buffer() { - AllOr::All => { - self.nulls.append_n_non_nulls(array.len()); - } - AllOr::None => self.nulls.append_n_nulls(array.len()), - AllOr::Some(validity) => self.nulls.append_buffer(validity.clone()), - } + self.nulls.append_validity_mask(array.validity_mask()?); Ok(()) } diff --git a/vortex-array/src/builders/lazy_validity_builder.rs b/vortex-array/src/builders/lazy_validity_builder.rs index 2cbf5b8735..06fd7221e1 100644 --- a/vortex-array/src/builders/lazy_validity_builder.rs +++ b/vortex-array/src/builders/lazy_validity_builder.rs @@ -2,6 +2,7 @@ use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer}; use vortex_dtype::Nullability; use vortex_dtype::Nullability::{NonNullable, Nullable}; use vortex_error::{vortex_panic, VortexExpect, VortexResult}; +use vortex_mask::Mask; use crate::validity::Validity; @@ -72,12 +73,33 @@ impl LazyNullBufferBuilder { } #[inline] - pub fn append_buffer(&mut self, bool_buffer: BooleanBuffer) { + pub fn append_buffer(&mut self, bool_buffer: &BooleanBuffer) { self.materialize_if_needed(); self.inner .as_mut() .vortex_expect("buffer just materialized") - .append_buffer(&bool_buffer); + .append_buffer(bool_buffer); + } + + pub fn append_validity_mask(&mut self, validity_mask: Mask) { + match validity_mask { + Mask::AllTrue(len) => self.append_n_non_nulls(len), + Mask::AllFalse(len) => self.append_n_nulls(len), + Mask::Values(is_valid) => self.append_buffer(is_valid.boolean_buffer()), + } + } + + pub fn append_validity(&mut self, validity: Validity, length: usize) -> VortexResult<()> { + self.append_validity_mask(validity.to_logical(length)?); + Ok(()) + } + + pub fn set_bit(&mut self, index: usize, v: bool) { + self.materialize_if_needed(); + self.inner + .as_mut() + .vortex_expect("buffer just materialized") + .set_bit(index, v); } pub fn len(&self) -> usize { @@ -85,6 +107,21 @@ impl LazyNullBufferBuilder { self.inner.as_ref().map(|i| i.len()).unwrap_or(self.len) } + pub fn truncate(&mut self, len: usize) { + if let Some(b) = self.inner.as_mut() { + b.truncate(len) + } + self.len = len; + } + + pub fn reserve(&mut self, n: usize) { + self.materialize_if_needed(); + self.inner + .as_mut() + .vortex_expect("buffer just materialized") + .reserve(n); + } + pub fn finish(&mut self) -> Option { self.len = 0; Some(NullBuffer::new(self.inner.take()?.finish())) diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index c0fa2a80f1..ab2a866b65 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use vortex_dtype::Nullability::NonNullable; use vortex_dtype::{DType, NativePType, Nullability}; use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult}; -use vortex_mask::AllOr; use vortex_scalar::{BinaryNumericOperator, ListScalar}; use crate::array::{ConstantArray, ListArray, OffsetPType}; @@ -116,13 +115,7 @@ impl ArrayBuilder for ListBuilder { } fn extend_from_array(&mut self, array: Array) -> VortexResult<()> { - match array.validity_mask()?.boolean_buffer() { - AllOr::All => { - self.nulls.append_n_non_nulls(array.len()); - } - AllOr::None => self.nulls.append_n_nulls(array.len()), - AllOr::Some(validity) => self.nulls.append_buffer(validity.clone()), - } + self.nulls.append_validity_mask(array.validity_mask()?); let list = array.into_canonical()?.into_list()?; diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index ce06ac8ce9..5fd6d0692c 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -1,20 +1,22 @@ use std::any::Any; +use num_traits::AsPrimitive; use vortex_buffer::BufferMut; -use vortex_dtype::{DType, NativePType, Nullability}; +use vortex_dtype::{match_each_unsigned_integer_ptype, DType, NativePType, Nullability}; use vortex_error::{vortex_bail, VortexResult}; -use vortex_mask::{AllOr, Mask}; +use vortex_mask::Mask; use crate::array::{BoolArray, PrimitiveArray}; use crate::builders::lazy_validity_builder::LazyNullBufferBuilder; use crate::builders::ArrayBuilder; +use crate::patches::Patches; use crate::validity::Validity; use crate::variants::PrimitiveArrayTrait; -use crate::{Array, IntoArray, IntoCanonical}; +use crate::{Array, IntoArray, IntoArrayVariant as _, IntoCanonical}; pub struct PrimitiveBuilder { - values: BufferMut, - nulls: LazyNullBufferBuilder, + pub values: BufferMut, + pub nulls: LazyNullBufferBuilder, dtype: DType, } @@ -46,19 +48,81 @@ impl PrimitiveBuilder { } } + pub fn patch(&mut self, patches: Patches, starting_at: usize) -> VortexResult<()> { + let (array_len, indices_offset, indices, values) = patches.into_parts(); + assert!(starting_at + array_len == self.len()); + let indices = indices.into_primitive()?; + let values = values.into_primitive()?; + let validity = values.validity(); + let values = values.as_slice::(); + match_each_unsigned_integer_ptype!(indices.ptype(), |$P| { + self.insert_values_and_validity_at_indices::<$P>(indices, values, validity, starting_at, indices_offset) + }) + } + + fn insert_values_and_validity_at_indices>( + &mut self, + indices: PrimitiveArray, + values: &[T], + validity: Validity, + starting_at: usize, + indices_offset: usize, + ) -> VortexResult<()> { + match validity { + Validity::NonNullable => { + for (compressed_index, decompressed_index) in + indices.as_slice::().iter().enumerate() + { + let decompressed_index = decompressed_index.as_(); + let out_index = starting_at + decompressed_index - indices_offset; + self.values[out_index] = values[compressed_index]; + } + } + _ => { + let validity = validity.to_logical(indices.len())?; + for (compressed_index, decompressed_index) in + indices.as_slice::().iter().enumerate() + { + let decompressed_index = decompressed_index.as_(); + let out_index = starting_at + decompressed_index - indices_offset; + self.values[out_index] = values[compressed_index]; + self.nulls.set_bit(out_index, validity.value(out_index)); + } + } + } + + Ok(()) + } + + pub fn finish_into_primitive(&mut self) -> VortexResult { + let validity = match (self.nulls.finish(), self.dtype().nullability()) { + (None, Nullability::NonNullable) => Validity::NonNullable, + (Some(_), Nullability::NonNullable) => { + vortex_bail!("Non-nullable builder has null values") + } + (None, Nullability::Nullable) => Validity::AllValid, + (Some(nulls), Nullability::Nullable) => { + if nulls.null_count() == nulls.len() { + Validity::AllInvalid + } else { + Validity::Array(BoolArray::from(nulls.into_inner()).into_array()) + } + } + }; + + Ok(PrimitiveArray::new( + std::mem::take(&mut self.values).freeze(), + validity, + )) + } + pub fn extend_with_iterator(&mut self, iter: impl IntoIterator, mask: Mask) { self.values.extend(iter); self.extend_with_validity_mask(mask) } fn extend_with_validity_mask(&mut self, validity_mask: Mask) { - match validity_mask.boolean_buffer() { - AllOr::All => { - self.nulls.append_n_non_nulls(validity_mask.len()); - } - AllOr::None => self.nulls.append_n_nulls(validity_mask.len()), - AllOr::Some(validity) => self.nulls.append_buffer(validity.clone()), - } + self.nulls.append_validity_mask(validity_mask); } } @@ -103,21 +167,6 @@ impl ArrayBuilder for PrimitiveBuilder { } fn finish(&mut self) -> VortexResult { - let validity = match (self.nulls.finish(), self.dtype().nullability()) { - (None, Nullability::NonNullable) => Validity::NonNullable, - (Some(_), Nullability::NonNullable) => { - vortex_bail!("Non-nullable builder has null values") - } - (None, Nullability::Nullable) => Validity::AllValid, - (Some(nulls), Nullability::Nullable) => { - if nulls.null_count() == nulls.len() { - Validity::AllInvalid - } else { - Validity::Array(BoolArray::from(nulls.into_inner()).into_array()) - } - } - }; - - Ok(PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity).into_array()) + self.finish_into_primitive().map(IntoArray::into_array) } } diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 7a5c8e69c0..daf0a2c820 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -4,7 +4,7 @@ use std::cmp::max; use vortex_buffer::{BufferMut, ByteBuffer}; use vortex_dtype::{DType, Nullability}; use vortex_error::{vortex_bail, VortexExpect, VortexResult}; -use vortex_mask::{AllOr, Mask}; +use vortex_mask::Mask; use crate::array::{BinaryView, VarBinViewArray}; use crate::builders::lazy_validity_builder::LazyNullBufferBuilder; @@ -118,14 +118,7 @@ impl VarBinViewBuilder { // Pushes a validity mask into the builder not affecting the views or buffers fn push_only_validity_mask(&mut self, validity_mask: Mask) { - match validity_mask.boolean_buffer() { - AllOr::All => { - self.null_buffer_builder - .append_n_non_nulls(validity_mask.len()); - } - AllOr::None => self.null_buffer_builder.append_n_nulls(validity_mask.len()), - AllOr::Some(validity) => self.null_buffer_builder.append_buffer(validity.clone()), - } + self.null_buffer_builder.append_validity_mask(validity_mask); } } diff --git a/vortex-array/src/builders/varbinview_builder.rs b/vortex-array/src/builders/varbinview_builder.rs index 43f0ee6863..567d7a1d45 100644 --- a/vortex-array/src/builders/varbinview_builder.rs +++ b/vortex-array/src/builders/varbinview_builder.rs @@ -4,7 +4,6 @@ use std::cmp::max; use vortex_buffer::{BufferMut, ByteBuffer, ByteBufferMut}; use vortex_dtype::{DType, Nullability}; use vortex_error::{vortex_bail, VortexExpect, VortexResult}; -use vortex_mask::AllOr; use crate::array::{BinaryView, VarBinViewArray}; use crate::builders::lazy_validity_builder::LazyNullBufferBuilder; @@ -156,13 +155,8 @@ impl ArrayBuilder for VarBinViewBuilder { } })); - match array.validity_mask()?.boolean_buffer() { - AllOr::All => { - self.null_buffer_builder.append_n_non_nulls(array.len()); - } - AllOr::None => self.null_buffer_builder.append_n_nulls(array.len()), - AllOr::Some(validity) => self.null_buffer_builder.append_buffer(validity.clone()), - } + self.null_buffer_builder + .append_validity_mask(array.validity_mask()?); Ok(()) } diff --git a/vortex-dtype/src/ptype.rs b/vortex-dtype/src/ptype.rs index 729ea3a44e..c2dccbf63f 100644 --- a/vortex-dtype/src/ptype.rs +++ b/vortex-dtype/src/ptype.rs @@ -206,6 +206,27 @@ macro_rules! match_each_integer_ptype { }) } +/// Macro to match over each integer PType, binding the corresponding native type (from +/// `NativePType`) and the corresponding unsigned type (also a `NativePType`). +#[macro_export] +macro_rules! match_each_integer_ptype_with_unsigned_type { + ($self:expr, | $_:tt $enc:ident, $_2:tt $unsigned:ident | $($body:tt)*) => ({ + macro_rules! __with__ {( $_ $enc:ident, $_2 $unsigned:ident ) => ( $($body)* )} + use $crate::PType; + match $self { + PType::I8 => __with__! { i8, u8 }, + PType::I16 => __with__! { i16, u16 }, + PType::I32 => __with__! { i32, u32 }, + PType::I64 => __with__! { i64, u64 }, + PType::U8 => __with__! { u8, u8 }, + PType::U16 => __with__! { u16, u16 }, + PType::U32 => __with__! { u32, u32 }, + PType::U64 => __with__! { u64, u64 }, + other => panic!("Unsupported ptype {other}") + } + }) +} + /// Macro to match over each unsigned integer type, binding the corresponding native type (from `NativePType`) #[macro_export] macro_rules! match_each_unsigned_integer_ptype {