diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index cae3630d6b8..a31dbaebb8b 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -102,6 +102,7 @@ jobs: timeout-minutes: 120 env: VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" + VORTEX_PATCHES_V2_SCATTER: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" # Makes python output nicer COLUMNS: 120 diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 692e7dcdd7f..7a1f498b950 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -2,6 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::mem::MaybeUninit; +use std::sync::LazyLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use fastlanes::BitPacking; use itertools::Itertools; @@ -16,6 +20,7 @@ use vortex_array::dtype::NativePType; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::patches::Patches; +use vortex_array::patches_v2::PatchesV2; use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -142,11 +147,20 @@ pub(crate) fn apply_patches_to_uninit_range VortexResult<()> { assert_eq!(patches.array_len(), dst.len()); - let indices = patches.indices().clone().execute::(ctx)?; let values = patches.values().clone().execute::(ctx)?; assert!(values.all_valid(ctx)?, "Patch values must be all valid"); let values = values.as_slice::(); + // When enabled, chunked patch sets scatter through the chunk-local PatchesV2 form to + // exercise it on the real decompression path. Converting per decompression costs a pass and + // allocations over the patch set, so this stays opt-in until the stored layout is + // chunk-local; the default path below is unchanged and the branch stays out of line to keep + // it out of the hot scatter loop's codegen. + if use_patches_v2_scatter() && patches.chunk_offsets().is_some() { + return apply_patches_v2(dst, patches, values, ctx, f); + } + + let indices = patches.indices().clone().execute::(ctx)?; match_each_unsigned_integer_ptype!(indices.ptype(), |P| { for (index, &value) in indices.as_slice::

().iter().zip_eq(values) { dst.set_value( @@ -158,6 +172,52 @@ pub(crate) fn apply_patches_to_uninit_range T>( + dst: &mut UninitRange, + patches: &Patches, + values: &[S], + ctx: &mut ExecutionCtx, + f: F, +) -> VortexResult<()> { + static ANNOUNCE: std::sync::Once = std::sync::Once::new(); + ANNOUNCE.call_once(|| eprintln!("vortex: PatchesV2 decompression scatter active")); + let v2 = PatchesV2::from_patches(patches, ctx)?; + v2.apply_each(ctx, |logical, ordinal| { + dst.set_value(logical, f(values[ordinal])); + })?; + PATCHES_V2_APPLIES.fetch_add(1, Ordering::Relaxed); + Ok(()) +} + +static PATCHES_V2_APPLIES: AtomicU64 = AtomicU64::new(0); +static PATCHES_V2_SCATTER: AtomicBool = AtomicBool::new(false); + +/// Returns whether decompression scatters chunked patches through [`PatchesV2`]. +/// +/// Enabled by [`force_patches_v2_scatter`] or the `VORTEX_PATCHES_V2_SCATTER=1` environment +/// variable. +pub fn use_patches_v2_scatter() -> bool { + static FROM_ENV: LazyLock = + LazyLock::new(|| std::env::var("VORTEX_PATCHES_V2_SCATTER").is_ok_and(|v| v == "1")); + PATCHES_V2_SCATTER.load(Ordering::Relaxed) || *FROM_ENV +} + +/// Force the chunk-local patch scatter on or off for this process. +pub fn force_patches_v2_scatter(enabled: bool) { + PATCHES_V2_SCATTER.store(enabled, Ordering::Relaxed); +} + +/// The number of decompressions that scattered patches through [`PatchesV2`]. +/// +/// This instruments the chunk-local patch path so integration tests can assert real read paths +/// exercise it. +pub fn patches_v2_apply_count() -> u64 { + PATCHES_V2_APPLIES.load(Ordering::Relaxed) +} + pub fn unpack_single(array: ArrayView<'_, BitPacked>, index: usize) -> Scalar { let bit_width = array.bit_width() as usize; let ptype = array.dtype().as_ptype(); diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 03fa3ed7f4c..cac4b7d82d4 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -385,11 +385,17 @@ mod test { let packed_with_patches = BitPackedData::encode(&parray, 9, &mut ctx).unwrap(); assert!(packed_with_patches.patches().is_some()); + crate::bitpack_decompress::force_patches_v2_scatter(true); + let applies_before = crate::bitpack_decompress::patches_v2_apply_count(); let packed_primitive = packed_with_patches .as_array() .clone() .execute::(&mut ctx) .unwrap(); + assert!( + crate::bitpack_decompress::patches_v2_apply_count() > applies_before, + "expected the chunk-local patch path" + ); assert_arrays_eq!( packed_primitive, PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable), diff --git a/vortex-array/benches/patches_lookup.rs b/vortex-array/benches/patches_lookup.rs index 262e3144495..19a2d1aaf08 100644 --- a/vortex-array/benches/patches_lookup.rs +++ b/vortex-array/benches/patches_lookup.rs @@ -9,8 +9,11 @@ use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::patches::PATCH_CHUNK_SIZE; use vortex_array::patches::Patches; +use vortex_array::patches_v2::PatchesV2; use vortex_buffer::Buffer; fn main() { @@ -165,3 +168,116 @@ fn search_index_full_range_random(bencher: Bencher) { fn search_index_full_range_random_chunked(bencher: Bencher) { bench_search_index(bencher, full_range_patches(true), queries_full_range()); } + +fn patches_v2_from(patches: &Patches) -> PatchesV2 { + let mut ctx = array_session().create_execution_ctx(); + PatchesV2::from_patches(patches, &mut ctx).unwrap() +} + +fn bench_search_index_v2(bencher: Bencher, patches: PatchesV2, queries: Vec) { + let mut ctx = array_session().create_execution_ctx(); + bencher + .with_inputs(|| (&patches, &queries)) + .bench_local_refs(|(patches, queries)| { + for &q in queries.iter() { + divan::black_box(patches.search_index(q, &mut ctx).unwrap()); + } + }); +} + +#[divan::bench] +fn search_index_below_min_v2(bencher: Bencher) { + bench_search_index_v2( + bencher, + patches_v2_from(&narrow_band_patches(false)), + queries_below_min(), + ); +} + +#[divan::bench] +fn search_index_above_max_v2(bencher: Bencher) { + bench_search_index_v2( + bencher, + patches_v2_from(&narrow_band_patches(false)), + queries_above_max(), + ); +} + +#[divan::bench] +fn search_index_mixed_out_of_range_v2(bencher: Bencher) { + bench_search_index_v2( + bencher, + patches_v2_from(&narrow_band_patches(false)), + queries_mixed_out_of_range(), + ); +} + +#[divan::bench] +fn search_index_in_range_v2(bencher: Bencher) { + bench_search_index_v2( + bencher, + patches_v2_from(&narrow_band_patches(false)), + queries_in_range(), + ); +} + +#[divan::bench] +fn search_index_full_range_random_v2(bencher: Bencher) { + bench_search_index_v2( + bencher, + patches_v2_from(&full_range_patches(false)), + queries_full_range(), + ); +} + +fn bench_apply_v1(bencher: Bencher, patches: Patches) { + let mut ctx = array_session().create_execution_ctx(); + let indices = patches + .indices() + .clone() + .execute::(&mut ctx) + .unwrap(); + bencher + .with_inputs(|| vec![0i64; ARRAY_LEN]) + .bench_local_values(|mut dst| { + for &index in indices.as_slice::() { + dst[index as usize] = 1; + } + divan::black_box(dst); + }); +} + +fn bench_apply_v2(bencher: Bencher, patches: PatchesV2) { + let mut ctx = array_session().create_execution_ctx(); + bencher + .with_inputs(|| vec![0i64; ARRAY_LEN]) + .bench_local_values(|mut dst| { + patches + .apply_each(&mut ctx, |logical, _ordinal| dst[logical] = 1) + .unwrap(); + divan::black_box(dst); + }); +} + +#[divan::bench] +fn apply_full_range(bencher: Bencher) { + bench_apply_v1(bencher, full_range_patches(false)); +} + +#[divan::bench] +fn apply_full_range_v2(bencher: Bencher) { + bench_apply_v2(bencher, patches_v2_from(&full_range_patches(false))); +} + +#[divan::bench] +fn slice_unaligned(bencher: Bencher) { + let patches = full_range_patches(true); + bencher.bench(|| divan::black_box(patches.slice(1_000..900_000).unwrap())); +} + +#[divan::bench] +fn slice_unaligned_v2(bencher: Bencher) { + let patches = patches_v2_from(&full_range_patches(true)); + let mut ctx = array_session().create_execution_ctx(); + bencher.bench_local(|| divan::black_box(patches.slice(1_000..900_000, &mut ctx).unwrap())); +} diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 9439dc7dd1b..299f4176199 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -133,6 +133,7 @@ pub mod normalize; pub mod optimizer; mod partial_ord; pub mod patches; +pub mod patches_v2; pub mod scalar; pub mod scalar_fn; pub mod search_sorted; diff --git a/vortex-array/src/patches_v2.rs b/vortex-array/src/patches_v2.rs new file mode 100644 index 00000000000..b8c248f4d0c --- /dev/null +++ b/vortex-array/src/patches_v2.rs @@ -0,0 +1,620 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A patch set addressed by chunk-local indices. +//! +//! [`PatchesV2`] stores the same information as [`Patches`]: sparse exception values for an +//! array. It differs in how patch positions are addressed: +//! +//! - `indices` holds `u16` positions **local to each 1024-value chunk** instead of global row +//! indices, so the index child stays two bytes per patch at any array length. +//! - `chunk_offsets` is required, holds `u32` prefix patch counts with a leading zero, and is +//! rebased on every slice, so chunk lookups never need the saturating-adjustment bookkeeping +//! that global offsets force onto [`Patches`]. +//! +//! An `offset` in `0..PATCH_CHUNK_SIZE` places logical element zero inside the first chunk, so +//! slices at unaligned positions keep constant-time chunk addressing: logical index `i` lives at +//! grid position `offset + i`, in chunk `(offset + i) / 1024` at local position +//! `(offset + i) % 1024`. +//! +//! [`Patches`]: crate::patches::Patches + +use std::ops::Range; + +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ArrayView; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::PType; +use crate::patches::PATCH_CHUNK_SIZE; +use crate::patches::Patches; +use crate::scalar::Scalar; +use crate::search_sorted::SearchResult; +use crate::validity::Validity; + +/// Sparse patch values addressed by chunk-local `u16` indices. +#[derive(Debug, Clone)] +pub struct PatchesV2 { + array_len: usize, + /// Grid position of logical element zero, in `0..PATCH_CHUNK_SIZE`. + offset: usize, + /// Chunk-local `u16` patch positions, sorted within each chunk. + indices: ArrayRef, + /// One patch value per index. + values: ArrayRef, + /// `u32` prefix patch counts per chunk, with a leading zero. + chunk_offsets: ArrayRef, +} + +impl PatchesV2 { + /// Construct and validate a new patch set. + /// + /// Validation canonicalizes the index and chunk-offset children, so callers on a hot path + /// with already-validated components should prefer [`Self::new_unchecked`]. + pub fn try_new( + array_len: usize, + offset: usize, + indices: ArrayRef, + values: ArrayRef, + chunk_offsets: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + vortex_ensure!( + offset < PATCH_CHUNK_SIZE, + "PatchesV2 offset must be within the first chunk" + ); + vortex_ensure!( + indices.len() == values.len(), + "PatchesV2 indices and values must have the same length" + ); + vortex_ensure!(!indices.is_empty(), "PatchesV2 must not be empty"); + vortex_ensure!( + indices.len() <= array_len, + "PatchesV2 cannot have more patches than rows" + ); + vortex_ensure!( + indices.dtype() == &DType::Primitive(PType::U16, NonNullable), + "PatchesV2 indices must be non-nullable u16, got {}", + indices.dtype() + ); + vortex_ensure!( + chunk_offsets.dtype() == &DType::Primitive(PType::U32, NonNullable), + "PatchesV2 chunk offsets must be non-nullable u32, got {}", + chunk_offsets.dtype() + ); + let chunk_count = (offset + array_len).div_ceil(PATCH_CHUNK_SIZE); + vortex_ensure!( + chunk_offsets.len() == chunk_count + 1, + "PatchesV2 expects {} chunk offsets, got {}", + chunk_count + 1, + chunk_offsets.len() + ); + + let local_indices = indices.clone().execute::(ctx)?; + let local_indices = local_indices.as_slice::(); + let offsets = chunk_offsets.clone().execute::(ctx)?; + let offsets = offsets.as_slice::(); + vortex_ensure!( + offsets.first() == Some(&0), + "PatchesV2 chunk offsets must start at zero" + ); + vortex_ensure!( + usize::try_from(offsets[chunk_count])? == indices.len(), + "PatchesV2 chunk offsets must end at the patch count" + ); + for chunk_idx in 0..chunk_count { + let chunk = + usize::try_from(offsets[chunk_idx])?..usize::try_from(offsets[chunk_idx + 1])?; + vortex_ensure!( + chunk.start <= chunk.end, + "PatchesV2 chunk offsets must not decrease" + ); + let chunk_grid_len = grid_range(offset, array_len, chunk_idx, chunk_count); + let locals = &local_indices[chunk]; + vortex_ensure!( + locals.windows(2).all(|pair| pair[0] < pair[1]), + "PatchesV2 indices must be strictly sorted within each chunk" + ); + vortex_ensure!( + locals + .iter() + .all(|&local| chunk_grid_len.contains(&usize::from(local))), + "PatchesV2 chunk {chunk_idx} contains out-of-range indices" + ); + } + + Ok(unsafe { Self::new_unchecked(array_len, offset, indices, values, chunk_offsets) }) + } + + /// Construct a patch set without validating the components. + /// + /// # Safety + /// + /// Callers must uphold every invariant checked by [`Self::try_new`]: matching child lengths, + /// non-nullable `u16` indices strictly sorted within each chunk and inside the sliced grid + /// range, and non-decreasing `u32` chunk offsets starting at zero and ending at the patch + /// count, with one entry per chunk plus one. + pub unsafe fn new_unchecked( + array_len: usize, + offset: usize, + indices: ArrayRef, + values: ArrayRef, + chunk_offsets: ArrayRef, + ) -> Self { + Self { + array_len, + offset, + indices, + values, + chunk_offsets, + } + } + + /// Convert a global-index [`Patches`] into chunk-local form. + pub fn from_patches(patches: &Patches, ctx: &mut ExecutionCtx) -> VortexResult { + let array_len = patches.array_len(); + let offset = patches.offset() % PATCH_CHUNK_SIZE; + let chunk_count = (offset + array_len).div_ceil(PATCH_CHUNK_SIZE); + let global = patches.indices().clone().execute::(ctx)?; + let mut locals = Vec::with_capacity(global.len()); + let mut chunk_offsets = vec![0u32; chunk_count + 1]; + let patches_offset = patches.offset(); + crate::match_each_unsigned_integer_ptype!(global.ptype(), |P| { + for &index in global.as_slice::

() { + // Rebase from the source offset onto this grid, which starts at `offset`. + let index: usize = index.as_(); + let grid = index - patches_offset + offset; + locals.push(u16::try_from(grid % PATCH_CHUNK_SIZE)?); + chunk_offsets[grid / PATCH_CHUNK_SIZE + 1] += 1; + } + }); + for chunk_idx in 0..chunk_count { + chunk_offsets[chunk_idx + 1] += chunk_offsets[chunk_idx]; + } + Ok(unsafe { + Self::new_unchecked( + array_len, + offset, + PrimitiveArray::new(Buffer::from(locals), Validity::NonNullable).into_array(), + patches.values().clone(), + PrimitiveArray::new(Buffer::from(chunk_offsets), Validity::NonNullable) + .into_array(), + ) + }) + } + + /// Convert back into a global-index [`Patches`]. + pub fn to_patches(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let (locals, offsets) = self.canonical_parts(ctx)?; + let mut globals = Vec::with_capacity(locals.len()); + for chunk_idx in 0..offsets.len() - 1 { + let chunk = + usize::try_from(offsets[chunk_idx])?..usize::try_from(offsets[chunk_idx + 1])?; + for &local in &locals[chunk] { + globals.push(u64::try_from( + chunk_idx * PATCH_CHUNK_SIZE + usize::from(local) - self.offset, + )?); + } + } + Patches::new( + self.array_len, + 0, + PrimitiveArray::new(Buffer::from(globals), Validity::NonNullable).into_array(), + self.values.clone(), + None, + ) + } + + /// Returns the length of the patched array. + pub fn array_len(&self) -> usize { + self.array_len + } + + /// Returns the number of patches. + pub fn num_patches(&self) -> usize { + self.indices.len() + } + + /// Returns the dtype of the patch values. + pub fn dtype(&self) -> &DType { + self.values.dtype() + } + + /// Returns the chunk-local patch indices. + pub fn indices(&self) -> &ArrayRef { + &self.indices + } + + /// Returns the patch values. + pub fn values(&self) -> &ArrayRef { + &self.values + } + + /// Returns the per-chunk patch count prefix sums. + pub fn chunk_offsets(&self) -> &ArrayRef { + &self.chunk_offsets + } + + /// Returns the grid position of logical element zero. + pub fn offset(&self) -> usize { + self.offset + } + + /// Search for a patch at logical `index`. + /// + /// Returns [`SearchResult::Found`] with the patch ordinal, or [`SearchResult::NotFound`] + /// with the insertion point. + pub fn search_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(view) = self.view() { + return Ok(view.search_index(index)); + } + if index >= self.array_len { + return Ok(SearchResult::NotFound(self.num_patches())); + } + let (locals, offsets) = self.canonical_parts(ctx)?; + Ok(search_local(&locals, &offsets, self.offset + index)) + } + + /// Borrow a resolved view over canonical index and chunk-offset children. + /// + /// Returns `None` when either child is not a canonical primitive array. Hot loops should + /// resolve the view once and query it repeatedly; each call performs the downcasts. + pub fn view(&self) -> Option> { + let locals = self.indices.as_opt::()?; + let offsets = self.chunk_offsets.as_opt::()?; + Some(PatchesV2View { + locals, + offsets, + offset: self.offset, + array_len: self.array_len, + }) + } + + /// Visit every patch as `(logical_index, patch_ordinal)`, in patch order. + /// + /// This is the decompression primitive: callers scatter the canonicalized patch values over + /// a decoded buffer without materializing global indices. + pub fn apply_each( + &self, + ctx: &mut ExecutionCtx, + mut apply: impl FnMut(usize, usize), + ) -> VortexResult<()> { + if let Some(view) = self.view() { + apply_each_parts( + view.locals.as_slice::(), + view.offsets.as_slice::(), + self.offset, + &mut apply, + ); + return Ok(()); + } + let (locals, offsets) = self.canonical_parts(ctx)?; + apply_each_parts(&locals, &offsets, self.offset, &mut apply); + Ok(()) + } + + /// Return the patch value at logical `index`, if one exists. + pub fn get_patched( + &self, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.search_index(index, ctx)? + .to_found() + .map(|patch_idx| self.values.execute_scalar(patch_idx, ctx)) + .transpose() + } + + /// Slice the patch set to `range`, returning `None` when no patches remain. + /// + /// The chunk offsets are rebased so the result is self-contained: no saturating adjustments + /// are carried forward, unlike [`Patches::slice`]. + pub fn slice(&self, range: Range, ctx: &mut ExecutionCtx) -> VortexResult> { + vortex_ensure!( + range.end <= self.array_len, + "PatchesV2 slice is out of bounds" + ); + if range.is_empty() { + return Ok(None); + } + let (locals, offsets) = self.canonical_parts(ctx)?; + let grid_start = self.offset + range.start; + let grid_end = self.offset + range.end; + let patch_start = search_local(&locals, &offsets, grid_start).to_index(); + let patch_end = search_local(&locals, &offsets, grid_end).to_index(); + if patch_start == patch_end { + return Ok(None); + } + + let chunk_start = grid_start / PATCH_CHUNK_SIZE; + let chunk_end = grid_end.div_ceil(PATCH_CHUNK_SIZE); + let rebased: Vec = (chunk_start..=chunk_end) + .map(|chunk_idx| { + let offset = usize::try_from(offsets[chunk_idx])?.clamp(patch_start, patch_end) + - patch_start; + Ok(u32::try_from(offset)?) + }) + .collect::>()?; + Ok(Some(unsafe { + Self::new_unchecked( + range.len(), + grid_start % PATCH_CHUNK_SIZE, + self.indices.slice(patch_start..patch_end)?, + self.values.slice(patch_start..patch_end)?, + PrimitiveArray::new(Buffer::from(rebased), Validity::NonNullable).into_array(), + ) + })) + } + + /// Execute the index and chunk-offset children into typed buffers. + /// + /// This is the slow path for encoded children; canonical children are read in place by the + /// callers' downcast fast paths. + fn canonical_parts(&self, ctx: &mut ExecutionCtx) -> VortexResult<(Buffer, Buffer)> { + let locals = self + .indices + .clone() + .execute::(ctx)? + .into_buffer::(); + let offsets = self + .chunk_offsets + .clone() + .execute::(ctx)? + .into_buffer::(); + Ok((locals, offsets)) + } +} + +/// A resolved, borrowed view over a [`PatchesV2`] with canonical children. +/// +/// Constructed via [`PatchesV2::view`]; queries are plain slice reads with no dispatch, +/// allocation, or error paths, so this is the form hot loops should hold. +#[derive(Clone, Debug)] +pub struct PatchesV2View<'a> { + locals: ArrayView<'a, Primitive>, + offsets: ArrayView<'a, Primitive>, + offset: usize, + array_len: usize, +} + +impl PatchesV2View<'_> { + /// Search for a patch at logical `index`. + pub fn search_index(&self, index: usize) -> SearchResult { + if index >= self.array_len { + return SearchResult::NotFound(self.locals.len()); + } + search_local( + self.locals.as_slice::(), + self.offsets.as_slice::(), + self.offset + index, + ) + } + + /// Returns the patch ordinal at logical `index`, if one exists. + pub fn patch_ordinal(&self, index: usize) -> Option { + self.search_index(index).to_found() + } +} + +/// Walk every patch as `(logical_index, patch_ordinal)` from resolved parts. +fn apply_each_parts( + locals: &[u16], + offsets: &[u32], + offset: usize, + apply: &mut impl FnMut(usize, usize), +) { + // Walk patches with a chunk cursor so sparse patch sets skip empty chunks cheaply. + let mut chunk_idx = 0usize; + for (ordinal, &local) in locals.iter().enumerate() { + while offsets[chunk_idx + 1] as usize <= ordinal { + chunk_idx += 1; + } + apply( + chunk_idx * PATCH_CHUNK_SIZE + usize::from(local) - offset, + ordinal, + ); + } +} + +/// The grid-local index range a chunk may address, honoring first- and last-chunk trims. +fn grid_range( + offset: usize, + array_len: usize, + chunk_idx: usize, + chunk_count: usize, +) -> Range { + let start = if chunk_idx == 0 { offset } else { 0 }; + let stop = if chunk_idx == chunk_count - 1 { + (offset + array_len) - chunk_idx * PATCH_CHUNK_SIZE + } else { + PATCH_CHUNK_SIZE + }; + start..stop +} + +/// Search the flat local-index buffer for grid position `grid`. +/// +/// Chunk selection is constant time via the offsets; the in-chunk search is a binary search +/// over at most [`PATCH_CHUNK_SIZE`] `u16` values. +fn search_local(locals: &[u16], offsets: &[u32], grid: usize) -> SearchResult { + let chunk_idx = grid / PATCH_CHUNK_SIZE; + if chunk_idx >= offsets.len() - 1 { + return SearchResult::NotFound(locals.len()); + } + let chunk = offsets[chunk_idx] as usize..offsets[chunk_idx + 1] as usize; + let local = + u16::try_from(grid % PATCH_CHUNK_SIZE).vortex_expect("chunk-local index fits in u16"); + match locals[chunk.clone()].binary_search(&local) { + Ok(idx) => SearchResult::Found(chunk.start + idx), + Err(idx) => SearchResult::NotFound(chunk.start + idx), + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::patches::Patches; + + fn test_patches(ctx: &mut ExecutionCtx) -> VortexResult { + // Patches at global rows 5, 100, 1023, 1024, 2050 in a 3000-row array. + let indices = + PrimitiveArray::new(buffer![5u64, 100, 1023, 1024, 2050], Validity::NonNullable); + let values = PrimitiveArray::new(buffer![50u64, 51, 52, 53, 54], Validity::NonNullable); + let global = Patches::new(3000, 0, indices.into_array(), values.into_array(), None)?; + PatchesV2::from_patches(&global, ctx) + } + + #[test] + fn from_global_roundtrip() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let patches = test_patches(&mut ctx)?; + assert_eq!(patches.num_patches(), 5); + assert_eq!(patches.offset(), 0); + + let back = patches.to_patches(&mut ctx)?; + let globals = back.indices().clone().execute::(&mut ctx)?; + assert_eq!(globals.as_slice::(), &[5, 100, 1023, 1024, 2050]); + Ok(()) + } + + #[test] + fn validates_components() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let patches = test_patches(&mut ctx)?; + PatchesV2::try_new( + patches.array_len(), + patches.offset(), + patches.indices().clone(), + patches.values().clone(), + patches.chunk_offsets().clone(), + &mut ctx, + )?; + + // Unsorted local indices within one chunk are rejected. + let unsorted = PatchesV2::try_new( + 3000, + 0, + PrimitiveArray::new(buffer![100u16, 5], Validity::NonNullable).into_array(), + PrimitiveArray::new(buffer![1u64, 2], Validity::NonNullable).into_array(), + PrimitiveArray::new(buffer![0u32, 2, 2, 2], Validity::NonNullable).into_array(), + &mut ctx, + ); + assert!(unsorted.is_err()); + Ok(()) + } + + #[test] + fn search_across_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let patches = test_patches(&mut ctx)?; + assert_eq!( + patches.search_index(1023, &mut ctx)?, + SearchResult::Found(2) + ); + assert_eq!( + patches.search_index(1024, &mut ctx)?, + SearchResult::Found(3) + ); + assert_eq!( + patches.search_index(1500, &mut ctx)?, + SearchResult::NotFound(4) + ); + assert_eq!( + patches.search_index(2999, &mut ctx)?, + SearchResult::NotFound(5) + ); + + let value = patches.get_patched(2050, &mut ctx)?; + assert_eq!(value, Some(Scalar::primitive(54u64, NonNullable))); + assert_eq!(patches.get_patched(2051, &mut ctx)?, None); + Ok(()) + } + + #[test] + fn apply_each_visits_all_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let patches = test_patches(&mut ctx)?; + let mut visited = Vec::new(); + patches.apply_each(&mut ctx, |logical, ordinal| { + visited.push((logical, ordinal)) + })?; + assert_eq!( + visited, + vec![(5, 0), (100, 1), (1023, 2), (1024, 3), (2050, 4)] + ); + + // A sliced patch set reports logical indices relative to the slice. + let sliced = patches + .slice(100..2050, &mut ctx)? + .expect("patches remain in slice"); + let mut visited = Vec::new(); + sliced.apply_each(&mut ctx, |logical, ordinal| { + visited.push((logical, ordinal)) + })?; + assert_eq!(visited, vec![(0, 0), (923, 1), (924, 2)]); + Ok(()) + } + + #[test] + fn view_matches_generic_search() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let patches = test_patches(&mut ctx)?; + let view = patches.view().expect("canonical children"); + for index in [0, 5, 100, 1023, 1024, 1500, 2050, 2999, 5000] { + assert_eq!( + view.search_index(index), + patches.search_index(index, &mut ctx)? + ); + } + Ok(()) + } + + #[test] + fn slice_unaligned() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let patches = test_patches(&mut ctx)?; + + // Slice 100..2050 keeps rows 100, 1023, 1024 and drops 5 and 2050. + let sliced = patches + .slice(100..2050, &mut ctx)? + .expect("patches remain in slice"); + assert_eq!(sliced.array_len(), 1950); + assert_eq!(sliced.num_patches(), 3); + assert_eq!(sliced.offset(), 100); + assert_eq!(sliced.search_index(0, &mut ctx)?, SearchResult::Found(0)); + assert_eq!(sliced.search_index(923, &mut ctx)?, SearchResult::Found(1)); + assert_eq!(sliced.search_index(924, &mut ctx)?, SearchResult::Found(2)); + assert_eq!( + sliced.get_patched(924, &mut ctx)?, + Some(Scalar::primitive(53u64, NonNullable)) + ); + assert_eq!(sliced.get_patched(925, &mut ctx)?, None); + + // Slicing a slice rebases again. + let inner = sliced + .slice(900..1000, &mut ctx)? + .expect("patches remain in inner slice"); + assert_eq!(inner.num_patches(), 2); + assert_eq!(inner.search_index(23, &mut ctx)?, SearchResult::Found(0)); + assert_eq!(inner.search_index(24, &mut ctx)?, SearchResult::Found(1)); + + // A gap with no patches slices to None. + assert!(patches.slice(1100..2000, &mut ctx)?.is_none()); + Ok(()) + } +} diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 1ca05c86b4e..348d990f2c7 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -80,6 +80,7 @@ pub use builder::ALL_SCHEMES; pub use builder::BtrBlocksCompressorBuilder; pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; +pub use schemes::patches::force_patch_index_bitpack; pub use vortex_compressor::CascadingCompressor; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; diff --git a/vortex-btrblocks/src/schemes/patches.rs b/vortex-btrblocks/src/schemes/patches.rs index 69ca8450f12..75a7d55e79f 100644 --- a/vortex-btrblocks/src/schemes/patches.rs +++ b/vortex-btrblocks/src/schemes/patches.rs @@ -1,25 +1,50 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::patches::PATCH_CHUNK_SIZE; use vortex_array::patches::Patches; -use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_fastlanes::bitpack_compress::bitpack_encode; + +static PATCH_INDEX_BITPACK: AtomicBool = AtomicBool::new(false); + +/// Toggles bitpacking of patch index children. +/// +/// Off by default: on TPC-H shaped data the per-array patch sets are too small for FastLanes +/// packing to pay for itself, while the extra array node makes serialized trees larger and cold +/// file opens measurably slower. `VORTEX_PATCH_INDEX_BITPACK=1` or this toggle turns it on for +/// dense-patch workloads and for size measurements. +pub fn force_patch_index_bitpack(enabled: bool) { + PATCH_INDEX_BITPACK.store(enabled, Ordering::Relaxed); +} + +fn patch_index_bitpack() -> bool { + static FROM_ENV: std::sync::LazyLock = std::sync::LazyLock::new(|| { + std::env::var("VORTEX_PATCH_INDEX_BITPACK").is_ok_and(|v| v == "1") + }); + PATCH_INDEX_BITPACK.load(Ordering::Relaxed) || *FROM_ENV +} -/// Compresses the given patches by downscaling integers and checking for constant values. +/// Compresses the given patches by downscaling and bitpacking integers and checking for constant +/// values. pub fn compress_patches(patches: Patches, ctx: &mut ExecutionCtx) -> VortexResult { - // Downscale the patch indices. + // Downscale and bitpack the patch indices. let indices = patches .indices() .clone() .execute::(ctx)? - .narrow(ctx)? - .into_array(); + .narrow(ctx)?; + let indices = bitpack_index_child(indices, ctx)?; // Check if the values are constant. let values = patches.values(); @@ -39,9 +64,8 @@ pub fn compress_patches(patches: Patches, ctx: &mut ExecutionCtx) -> VortexResul let offsets_primitive = offsets .clone() .execute::(ctx)? - .narrow(ctx)? - .into_array(); - Ok::(offsets_primitive) + .narrow(ctx)?; + bitpack_index_child(offsets_primitive, ctx) }) .transpose()?; @@ -53,3 +77,34 @@ pub fn compress_patches(patches: Patches, ctx: &mut ExecutionCtx) -> VortexResul chunk_offsets, ) } + +/// Bitpacks a non-nullable index child (patch indices or chunk offsets) at the exact bit width of +/// its maximum, so the packed form never needs patches of its own. +/// +/// FastLanes packs in 1024-value chunks and pads the tail, so short children stay unpacked — +/// padding would outweigh the width saving. +fn bitpack_index_child(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult { + if !patch_index_bitpack() + || array.len() < PATCH_CHUNK_SIZE + || array.dtype().is_nullable() + || !array.ptype().is_unsigned_int() + { + return Ok(array.into_array()); + } + let bit_width: u32 = match_each_unsigned_integer_ptype!(array.ptype(), |P| { + let Some(max) = array.statistics().compute_max::

(ctx) else { + return Ok(array.into_array()); + }; + if max == 0 { + return Ok(array.into_array()); + } + max.ilog2() + 1 + }); + let Ok(bit_width) = u8::try_from(bit_width) else { + return Ok(array.into_array()); + }; + if usize::from(bit_width) >= array.ptype().bit_width() { + return Ok(array.into_array()); + } + Ok(bitpack_encode(&array, bit_width, None, ctx)?.into_array()) +} diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 07069f6309a..20386fb6196 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -468,3 +468,331 @@ fn trace_scan_take_on_compressed_table() -> VortexResult<()> { Ok(()) } + +/// Compress `batches` lineitem batches at `scale_factor` and decompress them, returning how many +/// scatters went through the chunk-local `PatchesV2` path. +fn count_patches_v2_applies(scale_factor: f64, batch_size: usize) -> VortexResult { + use tpchgen::generators::OrderGenerator; + use tpchgen::generators::PartSuppGenerator; + use tpchgen_arrow::OrderArrow; + use tpchgen_arrow::PartSuppArrow; + use vortex_fastlanes::bitpack_decompress::force_patches_v2_scatter; + use vortex_fastlanes::bitpack_decompress::patches_v2_apply_count; + + force_patches_v2_scatter(true); + let before = patches_v2_apply_count(); + let mut ctx = execution_ctx(); + let session = trace_session(); + let mut roundtrip = |batch: RecordBatch| -> VortexResult<()> { + let schema = batch.schema(); + let array = session.arrow().from_arrow_record_batch(batch, &schema)?; + let compressed = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + // A canonical struct keeps compressed children, so decode each column explicitly. + let columns: Vec = compressed + .as_::() + .iter_unmasked_fields() + .cloned() + .collect(); + for column in columns { + let decoded = column.execute::(&mut ctx)?.into_array(); + assert_eq!(decoded.len(), array.len()); + } + Ok(()) + }; + + let lineitem = LineItemGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + for batch in LineItemArrow::new(lineitem).with_batch_size(batch_size) { + roundtrip(batch)?; + } + let orders = OrderGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + for batch in OrderArrow::new(orders).with_batch_size(batch_size) { + roundtrip(batch)?; + } + let partsupp = PartSuppGenerator::new_with_text_pool(scale_factor, 1, 1, &TEXT_POOL); + for batch in PartSuppArrow::new(partsupp).with_batch_size(batch_size) { + roundtrip(batch)?; + } + Ok(patches_v2_apply_count() - before) +} + +/// Lineitem compression produces bit-packed columns with patches, and decompressing them goes +/// through the chunk-local `PatchesV2` scatter. +#[test] +fn lineitem_decompress_uses_patches_v2() -> VortexResult<()> { + let applies = count_patches_v2_applies(0.01, 1 << 14)?; + assert!( + applies > 0, + "expected decompression to scatter patches through PatchesV2" + ); + Ok(()) +} + +/// The full TPC-H scale-factor-1 lineitem validation. Run manually: +/// `TPCH_SF=1 cargo test --release -p vortex-btrblocks lineitem_decompress -- --ignored --nocapture` +#[test] +#[ignore = "generates and compresses six million lineitem rows"] +fn lineitem_sf1_decompress_uses_patches_v2() -> VortexResult<()> { + let scale_factor = std::env::var("TPCH_SF") + .ok() + .and_then(|sf| sf.parse().ok()) + .unwrap_or(1.0); + let applies = count_patches_v2_applies(scale_factor, 1 << 16)?; + println!("PatchesV2 scatters at SF {scale_factor}: {applies}"); + assert!( + applies > 0, + "expected decompression to scatter patches through PatchesV2" + ); + Ok(()) +} + +/// Per-table TPC-H report: compress every table at the given scale factor, decode every column, +/// and report how many patched arrays each table produced and how many scattered through +/// `PatchesV2`. Run manually: +/// `cargo test --release -p vortex-btrblocks tpch_per_table -- --ignored --nocapture` +#[test] +#[ignore = "generates and compresses all TPC-H tables at scale factor one"] +fn tpch_per_table_patches_v2_report() -> VortexResult<()> { + use tpchgen::generators::CustomerGenerator; + use tpchgen::generators::NationGenerator; + use tpchgen::generators::OrderGenerator; + use tpchgen::generators::PartGenerator; + use tpchgen::generators::PartSuppGenerator; + use tpchgen::generators::RegionGenerator; + use tpchgen::generators::SupplierGenerator; + use tpchgen_arrow::CustomerArrow; + use tpchgen_arrow::NationArrow; + use tpchgen_arrow::OrderArrow; + use tpchgen_arrow::PartArrow; + use tpchgen_arrow::PartSuppArrow; + use tpchgen_arrow::RegionArrow; + use tpchgen_arrow::SupplierArrow; + use vortex_fastlanes::bitpack_decompress::force_patches_v2_scatter; + use vortex_fastlanes::bitpack_decompress::patches_v2_apply_count; + + force_patches_v2_scatter(true); + let scale_factor: f64 = std::env::var("TPCH_SF") + .ok() + .and_then(|sf| sf.parse().ok()) + .unwrap_or(1.0); + let batch_size = 1 << 16; + let mut ctx = execution_ctx(); + let session = trace_session(); + + let mut report = |table: &str, + batches: &mut dyn Iterator| + -> VortexResult<()> { + let before = patches_v2_apply_count(); + let mut rows = 0usize; + let mut patched_arrays = 0usize; + let mut bitpacked_arrays = 0usize; + for batch in batches { + let schema = batch.schema(); + let array = session.arrow().from_arrow_record_batch(batch, &schema)?; + rows += array.len(); + let compressed = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + let tree = compressed.display_tree().to_string(); + patched_arrays += tree.matches("patch_indices").count(); + bitpacked_arrays += tree.matches("fastlanes.bitpacked").count(); + let columns: Vec<(String, ArrayRef)> = compressed + .as_::() + .iter_unmasked_fields() + .cloned() + .enumerate() + .map(|(idx, column)| (format!("column{idx}"), column)) + .collect(); + for (name, column) in columns { + let column_tree = column.display_tree().to_string(); + let column_patched = column_tree.matches("patch_indices").count(); + let column_before = patches_v2_apply_count(); + column.execute::(&mut ctx)?; + let column_applies = patches_v2_apply_count() - column_before; + if column_patched > 0 && std::env::var("TPCH_COLUMN_DEBUG").is_ok() { + let root = column_tree.lines().next().unwrap_or("").trim().to_string(); + println!( + " {table} {name}: patched={column_patched} scatters={column_applies} root={root}" + ); + if column_applies == 0 && std::env::var("TPCH_TREE_DEBUG").is_ok() { + println!("{column_tree}"); + } + } + } + } + let applies = patches_v2_apply_count() - before; + println!( + "{table:<10} rows={rows:<8} bitpacked_arrays={bitpacked_arrays:<4} \ + patched_arrays={patched_arrays:<4} patches_v2_scatters={applies}" + ); + Ok(()) + }; + + let lineitem = LineItemGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "lineitem", + &mut LineItemArrow::new(lineitem).with_batch_size(batch_size), + )?; + let orders = OrderGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "orders", + &mut OrderArrow::new(orders).with_batch_size(batch_size), + )?; + let partsupp = PartSuppGenerator::new_with_text_pool(scale_factor, 1, 1, &TEXT_POOL); + report( + "partsupp", + &mut PartSuppArrow::new(partsupp).with_batch_size(batch_size), + )?; + let customer = CustomerGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "customer", + &mut CustomerArrow::new(customer).with_batch_size(batch_size), + )?; + let part = PartGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "part", + &mut PartArrow::new(part).with_batch_size(batch_size), + )?; + let supplier = SupplierGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "supplier", + &mut SupplierArrow::new(supplier).with_batch_size(batch_size), + )?; + report( + "nation", + &mut NationArrow::new(NationGenerator::default()).with_batch_size(batch_size), + )?; + report( + "region", + &mut RegionArrow::new(RegionGenerator::default()).with_batch_size(batch_size), + )?; + Ok(()) +} + +/// Reports compressed sizes with and without bitpacked patch index children, per TPC-H table. +/// +/// Run with `TPCH_SF=1` in release for the report the PR quotes: +/// `TPCH_SF=1 cargo test -p vortex-btrblocks --release tpch_patch_index_compression_report -- --ignored --nocapture` +#[test] +#[ignore = "slow: generates TPC-H tables; run explicitly with --ignored"] +fn tpch_patch_index_compression_report() -> VortexResult<()> { + use tpchgen::generators::OrderGenerator; + use tpchgen::generators::PartSuppGenerator; + use tpchgen_arrow::OrderArrow; + use tpchgen_arrow::PartSuppArrow; + + use crate::force_patch_index_bitpack; + + let scale_factor: f64 = std::env::var("TPCH_SF") + .ok() + .and_then(|sf| sf.parse().ok()) + .unwrap_or(0.1); + let batch_size = 1 << 16; + let mut ctx = execution_ctx(); + let session = trace_session(); + + let mut report = + |table: &str, batches: &mut dyn Iterator| -> VortexResult<()> { + let mut raw = 0u64; + let mut packed = 0u64; + let mut verified = false; + for batch in batches { + let schema = batch.schema(); + let array = session.arrow().from_arrow_record_batch(batch, &schema)?; + + force_patch_index_bitpack(false); + let plain = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + force_patch_index_bitpack(true); + let bitpacked = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + raw += plain.nbytes(); + packed += bitpacked.nbytes(); + + // Prove the bitpacked patch indices decode losslessly, once per table. + if !verified { + let expected = plain.clone().execute::(&mut ctx)?.into_array(); + let actual = bitpacked + .clone() + .execute::(&mut ctx)? + .into_array(); + assert_arrays_eq!(expected, actual, &mut ctx); + verified = true; + } + } + let delta = raw as i64 - packed as i64; + println!( + "{table:<10} plain={raw:<12} bitpacked_patch_indices={packed:<12} saved={delta} \ + ({:.3}%)", + delta as f64 * 100.0 / raw as f64 + ); + Ok(()) + }; + + let lineitem = LineItemGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "lineitem", + &mut LineItemArrow::new(lineitem).with_batch_size(batch_size), + )?; + let orders = OrderGenerator::new_with_distributions_and_text_pool( + scale_factor, + 1, + 1, + Distributions::static_default(), + &TEXT_POOL, + ); + report( + "orders", + &mut OrderArrow::new(orders).with_batch_size(batch_size), + )?; + let partsupp = PartSuppGenerator::new_with_text_pool(scale_factor, 1, 1, &TEXT_POOL); + report( + "partsupp", + &mut PartSuppArrow::new(partsupp).with_batch_size(batch_size), + )?; + force_patch_index_bitpack(true); + Ok(()) +}