Faster runend decompress (#2537)

Honestly... not sure why

Author: gatesn

Cargo.lock

@@ -2,10 +2,12 @@
 
 use divan::Bencher;
 use itertools::repeat_n;
+use num_traits::PrimInt;
 use vortex_array::Array;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
+use vortex_dtype::NativePType;
 use vortex_runend::RunEndArray;
 use vortex_runend::compress::runend_encode;
 
@@ -38,13 +40,18 @@ fn compress(bencher: Bencher, (length, run_step): (usize, usize)) {
         .bench_refs(|values| runend_encode(values).unwrap());
 }
 
-#[divan::bench(args = BENCH_ARGS)]
-fn decompress(bencher: Bencher, (length, run_step): (usize, usize)) {
+#[divan::bench(types = [u8, u16, u32, u64], args = BENCH_ARGS)]
+fn decompress<T: NativePType + PrimInt>(bencher: Bencher, (length, run_step): (usize, usize)) {
     let values = PrimitiveArray::new(
         (0..=length)
             .step_by(run_step)
             .enumerate()
-            .flat_map(|(idx, x)| repeat_n(idx as u64, x))
+            .flat_map(|(idx, x)| {
+                repeat_n(
+                    T::from(idx % T::max_value().to_usize().unwrap()).unwrap(),
+                    x,
+                )
+            })
             .collect::<Buffer<_>>(),
         Validity::NonNullable,
     );

encodings/runend/benches/run_end_compress.rs

@@ -25,6 +25,7 @@ arrow-buffer = { workspace = true, optional = true }
 bytes = { workspace = true }
 compio = { workspace = true, optional = true }
 log = { workspace = true, optional = true }
+num-traits = { workspace = true }
 rkyv = { workspace = true, optional = true }
 vortex-error = { workspace = true }
 

vortex-buffer/Cargo.toml

@@ -4,6 +4,7 @@ use std::iter::Iterator;
 
 use arrow_buffer::{ArrowNativeType, MutableBuffer, ScalarBuffer};
 use divan::Bencher;
+use num_traits::PrimInt;
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_error::{VortexExpect, vortex_err};
 
@@ -108,3 +109,19 @@ fn push_arrow_buffer(bencher: Bencher, length: i32) {
             }
         });
 }
+
+#[divan::bench(types = [u8, u16, u32, u64], args = [100, 1_000, 10_000, 100_000, 1_000_000])]
+fn push_n_vortex_buffer<T: PrimInt>(bencher: Bencher, length: usize) {
+    bencher
+        .with_inputs(|| BufferMut::<T>::with_capacity(length))
+        .bench_refs(|buffer| {
+            for _ in 0..100 {
+                unsafe {
+                    buffer.push_n_unchecked(
+                        divan::black_box(T::one()),
+                        divan::black_box(length / 100),
+                    )
+                };
+            }
+        });
+}

vortex-buffer/benches/vortex_buffer.rs

@@ -293,7 +293,8 @@ impl<T> BufferMut<T> {
         let mut dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast();
         // SAFETY: we checked the capacity in the reserve call
         unsafe {
-            for _ in 0..n {
+            let end = dst.add(n);
+            while dst < end {
                 dst.write(item);
                 dst = dst.add(1);
             }