diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index 9292143e406..ba2cac32499 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -17,7 +17,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -406,3 +412,177 @@ fn patched_take_10k_adversarial(bencher: Bencher) { .unwrap() }) } + +// These are the element runs emitted by a sparse FSL gather. The baseline materializes those +// runs inside the measured operation, reproducing the previous BitPacked take path. +fn embedding_runs_fixture(dimension: usize, ptype: PType, bit_width: u8) -> (ArrayRef, ArrayRef) { + const ROWS: usize = 4096; + const SELECTED: usize = 16; + let mut ctx = SESSION.create_execution_ctx(); + let values = fixture(ROWS * dimension, usize::from(bit_width)) + .into_array() + .cast(ptype.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx) + .unwrap() + .into_array(); + let mut rng = StdRng::seed_from_u64(0); + let starts = PrimitiveArray::from_iter( + (0..SELECTED).map(|_| (rng.random_range(0..ROWS) * dimension) as u64), + ) + .into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(dimension as u64, SELECTED).into_array(), + ConstantArray::new(1u64, SELECTED).into_array(), + SELECTED * dimension, + ) + .unwrap() + .into_array(); + (packed, indices) +} + +#[divan::bench(args = [768, 1024, 1536])] +fn embedding_runs(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U32, 16); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [768, 1024, 1536])] +fn embedding_runs_expanded_baseline(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U32, 16); + bench_piecewise_take(bencher, packed, indices, true); +} + +#[divan::bench(args = [512, 1024, 2048, 4096])] +fn embedding_code_runs(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U8, 4); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [512, 1024, 2048, 4096])] +fn embedding_code_runs_expanded_baseline(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U8, 4); + bench_piecewise_take(bencher, packed, indices, true); +} + +// Short runs share chunks; long runs visit distant chunks. Neither crosses the full-decode +// threshold, including the largest case (16 * 1024 selected out of 128 * 1024 values). +fn fragmented_runs_fixture(run_length: usize) -> (ArrayRef, ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + let values = fixture(128 * CHUNK_SIZE, 16).into_array(); + let packed = BitPackedData::encode(&values, 16, &mut ctx) + .unwrap() + .into_array(); + let (num_runs, stride) = if run_length < 64 { + (128, run_length + 1) + } else { + (16, 4 * CHUNK_SIZE) + }; + let starts = + PrimitiveArray::from_iter((0..num_runs).map(|run| (run * stride) as u64)).into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(run_length as u64, num_runs).into_array(), + ConstantArray::new(1u64, num_runs).into_array(), + num_runs * run_length, + ) + .unwrap() + .into_array(); + (packed, indices) +} + +#[divan::bench(args = [1, 2, 7, 8, 64, 512, 1024])] +fn fragmented_runs(bencher: Bencher, run_length: usize) { + let (packed, indices) = fragmented_runs_fixture(run_length); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [1, 2, 7, 8, 64, 512, 1024])] +fn fragmented_runs_expanded_baseline(bencher: Bencher, run_length: usize) { + let (packed, indices) = fragmented_runs_fixture(run_length); + bench_piecewise_take(bencher, packed, indices, true); +} + +// Isolated runs cannot amortize a whole-chunk decode across neighboring runs. Offsets near +// a chunk boundary also exercise sparse fragments of runs above the full-decode threshold. +#[divan::bench( + types = [u8, u16, u32, u64], + args = [ + (8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (160, 0), (161, 0), + (256, 0), (512, 0), (161, 944), (256, 1020), (256, 896), + ], +)] +fn isolated_runs(bencher: Bencher, (run_length, offset): (usize, usize)) { + let bit_width = match T::PTYPE { + PType::U8 => 4, + PType::U16 => 9, + PType::U32 => 16, + PType::U64 => 33, + _ => unreachable!(), + }; + let mut ctx = SESSION.create_execution_ctx(); + let mut rng = StdRng::seed_from_u64(42); + let values = PrimitiveArray::from_iter( + (0..128 * CHUNK_SIZE).map(|_| rng.random_range(0..(1u64 << bit_width))), + ) + .into_array() + .cast(T::PTYPE.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx) + .unwrap() + .into_array(); + let num_runs = 16; + let starts = + PrimitiveArray::from_iter((0..num_runs).map(|run| (run * 4 * CHUNK_SIZE + offset) as u64)) + .into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(run_length as u64, num_runs).into_array(), + ConstantArray::new(1u64, num_runs).into_array(), + num_runs * run_length, + ) + .unwrap() + .into_array(); + bench_piecewise_take(bencher, packed, indices, false); +} + +fn bench_piecewise_take( + bencher: Bencher, + packed: ArrayRef, + indices: ArrayRef, + expanded_baseline: bool, +) { + let mut ctx = SESSION.create_execution_ctx(); + let canonical = packed.clone().execute::(&mut ctx).unwrap(); + assert_arrays_eq!( + packed.take(indices.clone()).unwrap(), + canonical.take(indices.clone()).unwrap(), + &mut ctx + ); + bencher + .counter(ItemsCount::new(indices.len())) + .with_inputs(|| (indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(indices, ctx)| { + let indices = if expanded_baseline { + indices + .clone() + .execute::(ctx) + .unwrap() + .into_array() + } else { + indices.clone() + }; + packed + .take(indices) + .unwrap() + .execute::(ctx) + .unwrap() + }); +}