diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 98de7caf95d..c83397b50d8 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -216,6 +216,10 @@ harness = false name = "primitive_zip" harness = false +[[bench]] +name = "listview_zip" +harness = false + [[bench]] name = "take_primitive" harness = false diff --git a/vortex-array/benches/listview_zip.rs b/vortex-array/benches/listview_zip.rs new file mode 100644 index 00000000000..28db66f04a9 --- /dev/null +++ b/vortex-array/benches/listview_zip.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::LEGACY_SESSION; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_mask::Mask; + +fn main() { + divan::main(); +} + +// Smaller than the value-path benches: listview zip cost is dominated by element concatenation and +// per-list canonicalization. A few thousand lists already exercise the select while keeping each +// case well under a few hundred microseconds under CodSpeed's instruction-count simulation, which +// runs ~10x the local walltime. +const LEN: usize = 4_096; + +/// Fragmented (alternating) mask: the worst case for the per-element branch this kernel replaces. +/// The branchless chunked select is mask-shape-independent, so one shape suffices. +fn mask() -> Mask { + Mask::from_iter((0..LEN).map(|i| i.is_multiple_of(2))) +} + +#[divan::bench] +fn nonnull(bencher: Bencher) { + run(bencher, list_view(0, false), list_view(1_000_000, false)); +} + +#[divan::bench] +fn nullable(bencher: Bencher) { + run(bencher, list_view(0, true), list_view(1_000_000, true)); +} + +fn run(bencher: Bencher, if_true: ArrayRef, if_false: ArrayRef) { + let mask = mask(); + bencher + .with_inputs(|| { + ( + if_true.clone(), + if_false.clone(), + mask.clone().into_array(), + LEGACY_SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(t, f, m, ctx)| { + m.zip(t.clone(), f.clone()) + .unwrap() + .execute::(ctx) + .unwrap(); + }); +} + +/// `LEN` single-element lists: `list[i] = [base + i]`. When `nullable`, every 7th list is null +/// (list-level validity backed by a `BoolArray`), exercising the `zip_validity` path. +fn list_view(base: i64, nullable: bool) -> ArrayRef { + let mut elements = BufferMut::::with_capacity(LEN); + elements.extend((0..LEN as i64).map(|i| base + i)); + let offsets: BufferMut = (0..LEN as u64).collect(); + let sizes: BufferMut = std::iter::repeat_n(1u64, LEN).collect(); + + let validity = if nullable { + Validity::Array(BoolArray::from_iter((0..LEN).map(|i| !i.is_multiple_of(7))).into_array()) + } else { + Validity::NonNullable + }; + + ListViewArray::try_new( + elements.freeze().into_array(), + offsets.freeze().into_array(), + sizes.freeze().into_array(), + validity, + ) + .unwrap() + .into_array() +} diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 46537f1fafe..6007113300c 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -1,12 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::ops::BitAnd; use std::ops::BitOr; use std::ops::Not; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -85,24 +87,65 @@ impl ZipKernel for ListView { let mut offsets = BufferMut::::with_capacity(len); let mut sizes = BufferMut::::with_capacity(len); - for ((idx, (out_offsets, out_sizes)), selected) in offsets - .spare_capacity_mut() - .iter_mut() - .zip(sizes.spare_capacity_mut().iter_mut()) - .take(len) - .enumerate() - .zip(mask.iter()) { - if selected { - out_offsets.write(true_offsets[idx]); - out_sizes.write(true_sizes[idx]); - } else { - out_offsets.write(false_offsets[idx] + false_shift); - out_sizes.write(false_sizes[idx]); + let true_offsets = true_offsets.as_slice(); + let true_sizes = true_sizes.as_slice(); + let false_offsets = false_offsets.as_slice(); + let false_sizes = false_sizes.as_slice(); + + let offsets_out = offsets.spare_capacity_mut(); + let sizes_out = sizes.spare_capacity_mut(); + + // We matched `Mask::Values` above, so the bit buffer is materialized. `unaligned_chunks` + // iterates faster than `chunks`: it exposes the byte-aligned body as a plain `&[u64]` + // with no per-word reshifting, isolating any bit misalignment into a leading `prefix` + // and trailing `suffix` word. We blend both sides branchlessly per row so the compiler + // vectorizes the inner select instead of mispredicting a data-dependent branch. + let mask_bits = mask + .values() + .vortex_expect("mask is Mask::Values") + .bit_buffer(); + let unaligned = mask_bits.unaligned_chunks(); + // The prefix word's low `lead` bits are padding; shifting them out aligns row 0 to bit 0, + // after which every chunk and the suffix start cleanly on a row boundary. + let lead = unaligned.lead_padding(); + + let mut select_block = |word: u64, base: usize, n: usize| { + let end = base + n; + // `if_false` views address the second half of the concatenated elements, so shift + // their offsets by `false_shift`; sizes are taken verbatim from the chosen side. + select_column( + word, + &true_offsets[base..end], + &false_offsets[base..end], + false_shift, + &mut offsets_out[base..end], + ); + select_column( + word, + &true_sizes[base..end], + &false_sizes[base..end], + 0, + &mut sizes_out[base..end], + ); + }; + + let mut base = 0; + if let Some(prefix) = unaligned.prefix() { + let n = (64 - lead).min(len); + select_block(prefix >> lead, base, n); + base += n; + } + for &word in unaligned.chunks() { + select_block(word, base, 64); + base += 64; + } + if let Some(suffix) = unaligned.suffix() { + select_block(suffix, base, len - base); } } - // SAFETY: the loop above initialized exactly `len` slots in both buffers. + // SAFETY: `select_column` initialized exactly `len` slots in both buffers. unsafe { offsets.set_len(len); sizes.set_len(len); @@ -122,6 +165,30 @@ impl ZipKernel for ListView { } } +/// Branchlessly select one `u64` column per row from `if_true` or `if_false`. +/// +/// `word` holds the mask bits for this block, bit `j` (LSB-first) selecting row `j`: a set bit keeps +/// `true_vals[j]`, an unset bit keeps `false_vals[j] + false_add`. The bit is expanded to a +/// full-width lane mask and blended, so the inner loop is branch-free and auto-vectorizable. Inputs +/// are sliced to the output length up front so the compiler can elide bounds checks across the block. +#[inline] +fn select_column( + word: u64, + true_vals: &[u64], + false_vals: &[u64], + false_add: u64, + out: &mut [MaybeUninit], +) { + let n = out.len(); + let true_vals = &true_vals[..n]; + let false_vals = &false_vals[..n]; + for j in 0..n { + // 0 for an unset bit, `u64::MAX` for a set bit. + let lane = 0u64.wrapping_sub((word >> j) & 1); + out[j].write((true_vals[j] & lane) | ((false_vals[j] + false_add) & !lane)); + } +} + /// Appends `array`'s element chunks to `chunks`, flattening a top-level [`ChunkedArray`] so the /// concatenated elements never nest chunked arrays. fn push_element_chunks(array: ArrayRef, chunks: &mut Vec) { @@ -164,6 +231,12 @@ fn zip_validity( #[cfg(test)] mod tests { + #![allow( + clippy::cast_possible_truncation, + reason = "test fixtures use small indices that fit the target widths" + )] + + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -200,6 +273,7 @@ mod tests { /// `zip` of two list views selects whole lists per the mask and keeps the list encoding. #[test] fn zip_selects_lists() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[1, 2], [3], [4, 5, 6]] let if_true = list_view( buffer![1i32, 2, 3, 4, 5, 6].into_array(), @@ -216,7 +290,6 @@ mod tests { ); let mask = Mask::from_iter([true, false, true]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -239,6 +312,7 @@ mod tests { /// `zip` selects list-level validity from the chosen side and widens nullability. #[test] fn zip_selects_validity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[1], null, [2]] (list-level nulls) let if_true = list_view( buffer![1i32, 2].into_array(), @@ -256,7 +330,6 @@ mod tests { // true -> if_true, false -> if_false let mask = Mask::from_iter([false, true, true]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -277,6 +350,7 @@ mod tests { /// is nullable. #[test] fn zip_out_of_order_offsets_and_widening() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // [[5, 6], [7], [8, 9]] expressed with out-of-order offsets. let if_true = list_view( buffer![7i32, 8, 9, 5, 6].into_array(), @@ -293,7 +367,6 @@ mod tests { ); let mask = Mask::from_iter([true, true, false]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -311,10 +384,153 @@ mod tests { Ok(()) } + /// Zipping more rows than fit in a single 64-bit mask chunk exercises both the chunked select + /// loop and the trailing remainder, including the `false_shift` applied to `if_false` views. + #[test] + fn zip_spans_multiple_mask_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // 130 single-element lists per side: `if_true[i] = [i]`, `if_false[i] = [1000 + i]`. + let len = 130usize; + let true_elements: Vec = (0..len as i32).collect(); + let false_elements: Vec = (0..len as i32).map(|i| 1000 + i).collect(); + let offsets: Vec = (0..len as u64).collect(); + let sizes: Vec = vec![1; len]; + + let if_true = list_view( + true_elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ); + let if_false = list_view( + false_elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ); + + // A non-trivial pattern that straddles the chunk boundary (index 63/64) and the remainder. + let mask_bits: Vec = (0..len).map(|i| i.is_multiple_of(3) || i == 64).collect(); + let mask = Mask::from_iter(mask_bits.iter().copied()); + + let result = mask + .into_array() + .zip(if_true, if_false)? + .execute::(&mut ctx)?; + assert!(result.is::()); + + // Each row collapses to a single element: `i` when the mask is set, else `1000 + i`. + let expected_elements: Vec = (0..len) + .map(|i| { + if mask_bits[i] { + i as i32 + } else { + 1000 + i as i32 + } + }) + .collect(); + let expected = list_view( + expected_elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + /// A mask whose bit buffer starts at a non-byte-aligned offset (here from slicing a bool array) + /// has non-zero `unaligned_chunks` lead padding, exercising the prefix word alongside the + /// aligned chunk body and the suffix. + #[test] + fn zip_handles_offset_mask() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // 200 single-element lists per side: `if_true[i] = [i]`, `if_false[i] = [1000 + i]`. With a + // 3-bit lead offset the mask spans more than 16 bytes, so `unaligned_chunks` exposes a + // non-empty aligned `chunks` body between the prefix and suffix words. + let len = 200usize; + let true_elements: Vec = (0..len as i32).collect(); + let false_elements: Vec = (0..len as i32).map(|i| 1000 + i).collect(); + let offsets: Vec = (0..len as u64).collect(); + let sizes: Vec = vec![1; len]; + + let single_element_view = |elements: &[i32]| { + list_view( + elements + .iter() + .copied() + .collect::>() + .into_array(), + offsets + .iter() + .copied() + .collect::>() + .into_array(), + sizes.iter().copied().collect::>().into_array(), + Validity::NonNullable, + ) + }; + let if_true = single_element_view(&true_elements); + let if_false = single_element_view(&false_elements); + + // Slice off the first `offset` bits so the mask's bit buffer keeps a sub-byte offset while + // remaining `len` rows long. A non-trivial pattern straddles the prefix/body and chunk + // boundaries within the sliced window. + let offset = 3usize; + let mask_bits: Vec = (0..offset + len) + .map(|i| i.is_multiple_of(3) || i == offset + 64) + .collect(); + let mask = BoolArray::from_iter(mask_bits.iter().copied()) + .into_array() + .slice(offset..offset + len)?; + + let result = mask.zip(if_true, if_false)?.execute::(&mut ctx)?; + assert!(result.is::()); + + // Each row collapses to a single element: `i` when the sliced mask is set, else `1000 + i`. + let expected_elements: Vec = (0..len) + .map(|i| { + if mask_bits[offset + i] { + i as i32 + } else { + 1000 + i as i32 + } + }) + .collect(); + let expected = single_element_view(&expected_elements); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + /// When an input's `elements` is already a [`ChunkedArray`], its chunks are spliced in rather /// than nesting a chunked array inside the concatenated elements. #[test] fn zip_flattens_chunked_elements() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // elements [1, 2, 3] stored as two chunks; lists [[1, 2], [3]]. let chunked_elements = ChunkedArray::try_new( vec![buffer![1i32, 2].into_array(), buffer![3i32].into_array()], @@ -336,7 +552,6 @@ mod tests { ); let mask = Mask::from_iter([true, false]); - let mut ctx = array_session().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)?