From 2ce133097c5980b65bb9401c2c5a5b26428837e0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 27 Aug 2026 21:04:52 +0000 Subject: [PATCH 1/3] Add PatchesV2 with chunk-local patch indices `Patches` addresses outliers with global row indices, which forces a wide index child (u32/u64 for large arrays) and threads `offset` / `offset_within_chunk` saturating adjustments through every chunked lookup and slice. `PatchesV2` stores the same information addressed by chunk-local u16 indices: positions are local to each 1024-value chunk, `chunk_offsets` holds required u32 prefix patch counts rebased on every slice, and an `offset` within the first chunk keeps unaligned slices on the chunk grid. Lookups are an O(1) chunk select plus an in-chunk binary search, and slices are self-contained rather than carrying saturating adjustments. The container addresses patches by position only, so it carries values of any dtype. Tests cover that with randomized model-based cases (search, slice, nested slice, apply, and the global-index round trip against a reference model over fixed seeds) plus extension-typed values. Signed-off-by: Joe Isaacs --- vortex-array/benches/patches_lookup.rs | 116 ++++++ vortex-array/src/lib.rs | 1 + vortex-array/src/patches_v2/mod.rs | 519 +++++++++++++++++++++++++ vortex-array/src/patches_v2/tests.rs | 473 ++++++++++++++++++++++ 4 files changed, 1109 insertions(+) create mode 100644 vortex-array/src/patches_v2/mod.rs create mode 100644 vortex-array/src/patches_v2/tests.rs 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/mod.rs b/vortex-array/src/patches_v2/mod.rs new file mode 100644 index 00000000000..a3a5273a776 --- /dev/null +++ b/vortex-array/src/patches_v2/mod.rs @@ -0,0 +1,519 @@ +// 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 std::sync::LazyLock; +use std::sync::atomic::AtomicBool; +#[cfg(any(test, feature = "_test-harness"))] +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +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; + +static PATCHES_V2_SCATTER: AtomicBool = AtomicBool::new(false); + +/// Returns whether decompression scatters chunked patches through [`PatchesV2`]. +/// +/// Converting a global-index patch set per decompression costs a pass and allocations over the +/// patch set, so the chunk-local scatter stays opt-in until the stored layout is chunk-local. +/// Enabled by [`force_patches_v2_scatter`] or `VORTEX_PATCHES_V2_SCATTER=1`. +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); +} + +#[cfg(any(test, feature = "_test-harness"))] +static SCATTERS: AtomicU64 = AtomicU64::new(0); + +/// The number of [`PatchesV2::apply_into`] scatters performed, for tests that need to prove a +/// read path took the chunk-local form rather than silently falling back. +#[cfg(any(test, feature = "_test-harness"))] +pub fn patches_v2_scatter_count() -> u64 { + SCATTERS.load(Ordering::Relaxed) +} + +/// 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(()) + } + + /// Scatter the patch values over `out`, which holds the decoded base values. + /// + /// This is the decompression scatter for canonical primitive output: each patched position + /// is overwritten with its patch value. + pub fn apply_into( + &self, + out: &mut [T], + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + vortex_ensure!( + out.len() == self.array_len, + "PatchesV2 apply_into expects {} elements, got {}", + self.array_len, + out.len() + ); + #[cfg(any(test, feature = "_test-harness"))] + SCATTERS.fetch_add(1, Ordering::Relaxed); + let values = self.values.clone().execute::(ctx)?; + let values = values.as_slice::(); + self.apply_each(ctx, |logical, ordinal| out[logical] = values[ordinal]) + } + + /// 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; diff --git a/vortex-array/src/patches_v2/tests.rs b/vortex-array/src/patches_v2/tests.rs new file mode 100644 index 00000000000..eac66983f59 --- /dev/null +++ b/vortex-array/src/patches_v2/tests.rs @@ -0,0 +1,473 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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 apply_into_scatters_patch_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // Patches at rows 1 and 3 of a 4-row array. + let indices = PrimitiveArray::new(buffer![1u64, 3], Validity::NonNullable); + let values = PrimitiveArray::new(buffer![10u64, 20], Validity::NonNullable); + let global = Patches::new(4, 0, indices.into_array(), values.into_array(), None)?; + let patches = PatchesV2::from_patches(&global, &mut ctx)?; + + let mut out = [7u64; 4]; + patches.apply_into(&mut out, &mut ctx)?; + assert_eq!(out, [7, 10, 7, 20]); + + // The scatter is not integer-only: float output works the same way. + let values = PrimitiveArray::new(buffer![1.5f64, -2.5], Validity::NonNullable); + let global = Patches::new( + 4, + 0, + PrimitiveArray::new(buffer![1u64, 3], Validity::NonNullable).into_array(), + values.into_array(), + None, + )?; + let patches = PatchesV2::from_patches(&global, &mut ctx)?; + let mut out = [0.0f64; 4]; + patches.apply_into(&mut out, &mut ctx)?; + assert_eq!(out, [0.0, 1.5, 0.0, -2.5]); + + // A length mismatch is an error rather than a panic. + let mut wrong = [0.0f64; 3]; + assert!(patches.apply_into(&mut wrong, &mut ctx).is_err()); + 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(()) +} + +/// Randomized model-based tests. +/// +/// Each case builds a [`Model`] describing which logical rows carry a patch, constructs the +/// equivalent [`PatchesV2`], and asserts every operation agrees with the model. Seeds are fixed +/// so failures reproduce. +mod property { + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + + use super::*; + + /// The reference semantics a [`PatchesV2`] must implement: a sorted set of patched logical + /// rows, each carrying one value. + #[derive(Clone, Debug)] + struct Model { + array_len: usize, + /// Strictly increasing logical row positions carrying a patch. + positions: Vec, + values: Vec, + } + + impl Model { + /// Draw an array length spanning several chunks and a random subset of patched rows. + fn generate(rng: &mut StdRng) -> Self { + let array_len = rng.random_range(1usize..4096); + let num_patches = rng.random_range(1usize..=array_len.min(64)); + let mut positions: Vec = Vec::with_capacity(num_patches); + while positions.len() < num_patches { + let candidate = rng.random_range(0..array_len); + if let Err(idx) = positions.binary_search(&candidate) { + positions.insert(idx, candidate); + } + } + let values = (0..positions.len()).map(|_| rng.random::()).collect(); + Self { + array_len, + positions, + values, + } + } + + fn search(&self, index: usize) -> SearchResult { + if index >= self.array_len { + return SearchResult::NotFound(self.positions.len()); + } + match self.positions.binary_search(&index) { + Ok(ordinal) => SearchResult::Found(ordinal), + Err(ordinal) => SearchResult::NotFound(ordinal), + } + } + + /// The model restricted to `range`, with positions rebased onto the slice. + fn slice(&self, range: Range) -> Option { + let positions: Vec = self + .positions + .iter() + .filter(|&&position| range.contains(&position)) + .map(|&position| position - range.start) + .collect(); + if positions.is_empty() { + return None; + } + let values = self + .positions + .iter() + .zip(&self.values) + .filter(|&(&position, _)| range.contains(&position)) + .map(|(_, &value)| value) + .collect(); + Some(Self { + array_len: range.len(), + positions, + values, + }) + } + + /// Build the equivalent patch set by way of a global-index [`Patches`]. + fn build(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let indices: Vec = self.positions.iter().map(|&p| p as u64).collect(); + let global = Patches::new( + self.array_len, + 0, + PrimitiveArray::new(Buffer::from(indices), Validity::NonNullable).into_array(), + PrimitiveArray::new(Buffer::from(self.values.clone()), Validity::NonNullable) + .into_array(), + None, + )?; + PatchesV2::from_patches(&global, ctx) + } + + /// Assert every [`PatchesV2`] operation agrees with this model. + fn assert_matches(&self, patches: &PatchesV2, ctx: &mut ExecutionCtx) -> VortexResult<()> { + assert_eq!(patches.array_len(), self.array_len); + assert_eq!(patches.num_patches(), self.positions.len()); + + // Constructed components pass validation. + PatchesV2::try_new( + patches.array_len(), + patches.offset(), + patches.indices().clone(), + patches.values().clone(), + patches.chunk_offsets().clone(), + ctx, + )?; + + let view = patches.view().expect("canonical children"); + for index in 0..self.array_len { + let expected = self.search(index); + assert_eq!( + patches.search_index(index, ctx)?, + expected, + "search_index({index}) on {self:?}" + ); + assert_eq!(view.search_index(index), expected, "view search({index})"); + + let expected_value = expected + .to_found() + .map(|ordinal| Scalar::primitive(self.values[ordinal], NonNullable)); + assert_eq!( + patches.get_patched(index, ctx)?, + expected_value, + "get_patched({index})" + ); + } + // Out-of-bounds lookups report the end of the patch list rather than panicking. + assert_eq!( + patches.search_index(self.array_len, ctx)?, + SearchResult::NotFound(self.positions.len()) + ); + + let mut visited = Vec::new(); + patches.apply_each(ctx, |logical, ordinal| visited.push((logical, ordinal)))?; + let expected: Vec<(usize, usize)> = self + .positions + .iter() + .enumerate() + .map(|(ordinal, &position)| (position, ordinal)) + .collect(); + assert_eq!(visited, expected, "apply_each on {self:?}"); + + // Replace-mode scatter writes exactly the patched rows. + let mut out = vec![0u64; self.array_len]; + patches.apply_into(&mut out, ctx)?; + let mut expected_out = vec![0u64; self.array_len]; + for (&position, &value) in self.positions.iter().zip(&self.values) { + expected_out[position] = value; + } + assert_eq!(out, expected_out, "apply_into on {self:?}"); + Ok(()) + } + } + + #[test] + fn model_matches_across_seeds() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + for seed in 0..64u64 { + let mut rng = StdRng::seed_from_u64(seed); + let model = Model::generate(&mut rng); + let patches = model.build(&mut ctx)?; + model.assert_matches(&patches, &mut ctx)?; + } + Ok(()) + } + + #[test] + fn slices_match_the_model() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + for seed in 0..64u64 { + let mut rng = StdRng::seed_from_u64(seed ^ 0x5115); + let model = Model::generate(&mut rng); + let patches = model.build(&mut ctx)?; + + for _ in 0..8 { + let start = rng.random_range(0..model.array_len); + let end = rng.random_range(start..=model.array_len); + let sliced = patches.slice(start..end, &mut ctx)?; + match model.slice(start..end) { + // An empty range, or one holding no patches, slices away entirely. + None => assert!( + sliced.is_none(), + "expected no patches in {start}..{end} of {model:?}" + ), + Some(expected) => { + let sliced = + sliced.unwrap_or_else(|| panic!("patches remain in {start}..{end}")); + expected.assert_matches(&sliced, &mut ctx)?; + + // Slicing a slice rebases onto the inner range. + let inner_start = rng.random_range(0..expected.array_len); + let inner_end = rng.random_range(inner_start..=expected.array_len); + let inner = sliced.slice(inner_start..inner_end, &mut ctx)?; + match expected.slice(inner_start..inner_end) { + None => assert!(inner.is_none()), + Some(expected_inner) => { + let inner = inner.expect("patches remain in inner slice"); + expected_inner.assert_matches(&inner, &mut ctx)?; + } + } + } + } + } + } + Ok(()) + } + + #[test] + fn global_index_roundtrip() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + for seed in 0..64u64 { + let mut rng = StdRng::seed_from_u64(seed ^ 0x9001); + let model = Model::generate(&mut rng); + let patches = model.build(&mut ctx)?; + + let back = patches.to_patches(&mut ctx)?; + let globals = back.indices().clone().execute::(&mut ctx)?; + let expected: Vec = model.positions.iter().map(|&p| p as u64).collect(); + assert_eq!(globals.as_slice::(), expected.as_slice()); + assert_eq!(back.array_len(), model.array_len); + + // Round-tripping back through the chunk-local form is a fixed point. + let again = PatchesV2::from_patches(&back, &mut ctx)?; + model.assert_matches(&again, &mut ctx)?; + } + Ok(()) + } +} + +/// The container addresses patches by position, so it carries values of any dtype — including +/// extension types, which [`Sparse`] itself cannot yet canonicalize. +/// +/// [`Sparse`]: https://docs.rs/vortex-sparse +#[test] +fn carries_extension_typed_values() -> VortexResult<()> { + use crate::arrays::ExtensionArray; + use crate::arrays::FixedSizeListArray; + use crate::dtype::extension::ExtDType; + use crate::extension::uuid::Uuid; + use crate::extension::uuid::UuidMetadata; + + let mut ctx = array_session().create_execution_ctx(); + let storage = FixedSizeListArray::try_new( + PrimitiveArray::new(Buffer::from(vec![7u8; 32]), Validity::NonNullable).into_array(), + 16, + Validity::NonNullable, + 2, + )? + .into_array(); + let ext_dtype = + ExtDType::try_with_vtable(Uuid, UuidMetadata::default(), storage.dtype().clone())?.erased(); + let values = ExtensionArray::new(ext_dtype, storage).into_array(); + let expected_dtype = values.dtype().clone(); + + let global = Patches::new( + 2048, + 0, + PrimitiveArray::new(buffer![10u64, 1500], Validity::NonNullable).into_array(), + values, + None, + )?; + let patches = PatchesV2::from_patches(&global, &mut ctx)?; + + assert_eq!(patches.dtype(), &expected_dtype); + assert_eq!(patches.num_patches(), 2); + assert_eq!(patches.search_index(10, &mut ctx)?, SearchResult::Found(0)); + assert_eq!( + patches.search_index(1500, &mut ctx)?, + SearchResult::Found(1) + ); + assert!(patches.get_patched(10, &mut ctx)?.is_some()); + assert_eq!(patches.get_patched(11, &mut ctx)?, None); + + // Slicing keeps the value dtype intact. + let sliced = patches.slice(1000..2000, &mut ctx)?.expect("patch remains"); + assert_eq!(sliced.dtype(), &expected_dtype); + assert_eq!(sliced.num_patches(), 1); + assert_eq!(sliced.search_index(500, &mut ctx)?, SearchResult::Found(0)); + + // And the round trip back to global indices is dtype-agnostic too. + assert_eq!(patches.to_patches(&mut ctx)?.dtype(), &expected_dtype); + Ok(()) +} From c814da61809244a89cb0337c39b76813eb5671ac Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 27 Aug 2026 21:04:52 +0000 Subject: [PATCH 2/3] Scatter chunked patches through PatchesV2 on read paths Wires the chunk-local form into two real decompression paths, behind `VORTEX_PATCHES_V2_SCATTER=1` / `force_patches_v2_scatter`: - BitPacked and FoR decompression scatter chunked patch sets through `PatchesV2::apply_each` instead of materializing global indices. - Sparse canonicalization of a constant fill scatters through `PatchesV2::apply_into` when the types match, i.e. when the result is a non-nullable primitive and there is no validity to patch, so the whole operation is a value scatter over the fill. Converting a global-index patch set per decompression costs a pass and allocations over the patch set, so both stay opt-in until the stored layout is chunk-local; the default paths are unchanged and the branch is kept out of line in the bitpacked scatter loop. Tests assert both paths agree with their defaults on real patched arrays, and that a `Sparse` array rebuilt through `PatchesV2` is identical to the original across every dtype `Sparse` supports: null, bool, the primitive widths, decimal, utf8 (VarBin and VarBinView), binary, list, fixed-size list and struct. Signed-off-by: Joe Isaacs --- .github/workflows/sql-bench-matrix.yml | 1 + .../bitpacking/array/bitpack_decompress.rs | 29 +- .../fastlanes/src/bitpacking/array/mod.rs | 31 +- encodings/sparse/src/canonical.rs | 14 + encodings/sparse/tests/patches_v2_stand_in.rs | 314 ++++++++++++++++++ 5 files changed, 378 insertions(+), 11 deletions(-) create mode 100644 encodings/sparse/tests/patches_v2_stand_in.rs 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..416bceb3409 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -16,6 +16,8 @@ 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::patches_v2::use_patches_v2_scatter; use vortex_array::scalar::Scalar; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -142,11 +144,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 +169,22 @@ pub(crate) fn apply_patches_to_uninit_range T>( + dst: &mut UninitRange, + patches: &Patches, + values: &[S], + ctx: &mut ExecutionCtx, + f: F, +) -> VortexResult<()> { + let v2 = PatchesV2::from_patches(patches, ctx)?; + v2.apply_each(ctx, |logical, ordinal| { + dst.set_value(logical, f(values[ordinal])); + }) +} + 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..805cbea228a 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -331,6 +331,7 @@ mod test { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::patches_v2::force_patches_v2_scatter; use vortex_buffer::Buffer; use vortex_session::VortexSession; @@ -385,15 +386,25 @@ mod test { let packed_with_patches = BitPackedData::encode(&parray, 9, &mut ctx).unwrap(); assert!(packed_with_patches.patches().is_some()); - let packed_primitive = packed_with_patches - .as_array() - .clone() - .execute::(&mut ctx) - .unwrap(); - assert_arrays_eq!( - packed_primitive, - PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable), - &mut ctx - ); + + // Both scatter paths must decompress to the original values. Toggling the chunk-local + // scatter around the decode pins down that the two agree on a real patched array. + for chunk_local in [false, true] { + force_patches_v2_scatter(chunk_local); + let packed_primitive = packed_with_patches + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_arrays_eq!( + packed_primitive, + PrimitiveArray::new( + values.clone(), + vortex_array::validity::Validity::NonNullable + ), + &mut ctx + ); + } + force_patches_v2_scatter(false); } } diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 727c094cd84..e89f01b0d25 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -53,6 +53,8 @@ use vortex_array::match_each_native_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::match_smallest_list_offset_type; use vortex_array::patches::Patches; +use vortex_array::patches_v2::PatchesV2; +use vortex_array::patches_v2::use_patches_v2_scatter; use vortex_array::scalar::DecimalScalar; use vortex_array::scalar::ListScalar; use vortex_array::scalar::Scalar; @@ -60,6 +62,7 @@ use vortex_array::scalar::StructScalar; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; @@ -705,6 +708,17 @@ fn execute_sparse_primitives TryFrom<&'a Scalar, Error ) }; + // A non-nullable result has no validity to patch, so the whole canonicalization is a value + // scatter over a constant fill -- exactly what PatchesV2 addresses. Chunked patch sets take + // that path when the chunk-local scatter is enabled; everything else falls back to `patch`, + // which also handles patching validity. + if use_patches_v2_scatter() && matches!(validity, Validity::NonNullable) { + let mut out = BufferMut::full(primitive_fill, patches.array_len()); + let v2 = PatchesV2::from_patches(patches, ctx)?; + v2.apply_into(out.as_mut_slice(), ctx)?; + return Ok(PrimitiveArray::new(out.freeze(), Validity::NonNullable).into_array()); + } + let parray = PrimitiveArray::new(buffer![primitive_fill; patches.array_len()], validity); Ok(parray.patch(patches, ctx)?.into_array()) diff --git a/encodings/sparse/tests/patches_v2_stand_in.rs b/encodings/sparse/tests/patches_v2_stand_in.rs new file mode 100644 index 00000000000..12aafabccb5 --- /dev/null +++ b/encodings/sparse/tests/patches_v2_stand_in.rs @@ -0,0 +1,314 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! [`PatchesV2`] must be able to stand in for the patch set inside a constant-fill +//! [`Sparse`] array, for every dtype `Sparse` supports. +//! +//! Each case builds a `Sparse` array from patch values of one dtype, converts those patches +//! through [`PatchesV2`] and back, and asserts the rebuilt array is identical. The container +//! addresses patches by position only, so this pins down that nothing in the chunk-local +//! addressing is dtype-specific. + +#![allow(clippy::cast_possible_truncation, clippy::tests_outside_test_module)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::NullArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability::Nullable; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::patches::Patches; +use vortex_array::patches_v2::PatchesV2; +use vortex_array::patches_v2::force_patches_v2_scatter; +use vortex_array::patches_v2::patches_v2_scatter_count; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; +use vortex_sparse::Sparse; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_sparse::initialize(&session); + session +}); + +/// Rows carrying a patch, spread across chunk boundaries of a 3000-row array. +const POSITIONS: [u64; 4] = [3, 1023, 1024, 2050]; +const ARRAY_LEN: usize = 3000; + +fn primitive_values(values: [T; 4]) -> ArrayRef { + PrimitiveArray::from_option_iter(values.map(Some)).into_array() +} + +/// Four patch values of the named dtype, always nullable so `Scalar::null` is a legal fill. +fn values_for(case: &str) -> VortexResult { + Ok(match case { + "null" => NullArray::new(4).into_array(), + "bool" => BoolArray::from_iter([Some(true), None, Some(false), Some(true)]).into_array(), + "u8" => primitive_values([1u8, 2, 3, 4]), + "i16" => primitive_values([-1i16, 2, -3, 4]), + "u32" => primitive_values([1u32, 2, 3, 4]), + "i64" => primitive_values([-1i64, 2, -3, 4]), + "f32" => primitive_values([1.5f32, -2.5, 3.5, f32::NAN]), + "f64" => primitive_values([1.5f64, -2.5, 3.5, f64::INFINITY]), + "decimal" => DecimalArray::new( + buffer![100i128, 200, 300, 4000], + DecimalDType::new(3, 2), + Validity::from_iter([true, true, true, false]), + ) + .into_array(), + "utf8_varbin" => VarBinArray::from_iter( + [ + Some("a"), + None, + Some("ccc"), + Some("a string too long to inline"), + ], + DType::Utf8(Nullable), + ) + .into_array(), + "utf8_varbinview" => VarBinViewArray::from_iter( + [ + Some("a"), + None, + Some("ccc"), + Some("a string too long to inline"), + ], + DType::Utf8(Nullable), + ) + .into_array(), + "binary" => VarBinArray::from_iter( + [ + Some(vec![1u8, 2]), + None, + Some(vec![3u8]), + Some(vec![4u8; 40]), + ], + DType::Binary(Nullable), + ) + .into_array(), + "list" => ListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + buffer![0u32, 1, 3, 3, 6].into_array(), + Validity::AllValid, + )? + .into_array(), + "fixed_size_list" => FixedSizeListArray::try_new( + buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12].into_array(), + 3, + Validity::AllValid, + 4, + )? + .into_array(), + "struct" => StructArray::try_new_with_dtype( + vec![ + primitive_values([1i32, 2, 3, 4]), + VarBinViewArray::from_iter( + [Some("x"), Some("y"), None, Some("z")], + DType::Utf8(Nullable), + ) + .into_array(), + ], + StructFields::new( + FieldNames::from_iter(["a", "b"]), + vec![ + DType::Primitive(PType::I32, Nullable), + DType::Utf8(Nullable), + ], + ), + 4, + Validity::AllValid, + )? + .into_array(), + other => vortex_bail!("unknown case {other}"), + }) +} + +/// Build the equivalent chunk-local patch set for `values` at [`POSITIONS`]. +fn build(values: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<(Patches, PatchesV2)> { + let indices = buffer![POSITIONS[0], POSITIONS[1], POSITIONS[2], POSITIONS[3]].into_array(); + let patches = Patches::new(ARRAY_LEN, 0, indices, values.clone(), None)?; + let v2 = PatchesV2::from_patches(&patches, ctx)?; + Ok((patches, v2)) +} + +/// A `Sparse` array rebuilt through `PatchesV2` is identical to the original, for every dtype. +#[rstest] +#[case::null("null")] +#[case::bool("bool")] +#[case::u8("u8")] +#[case::i16("i16")] +#[case::u32("u32")] +#[case::i64("i64")] +#[case::f32("f32")] +#[case::f64("f64")] +#[case::decimal("decimal")] +#[case::utf8_varbin("utf8_varbin")] +#[case::utf8_varbinview("utf8_varbinview")] +#[case::binary("binary")] +#[case::list("list")] +#[case::fixed_size_list("fixed_size_list")] +#[case::struct_("struct")] +fn stands_in_for_constant_fill_sparse(#[case] case: &str) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = values_for(case)?; + let fill = Scalar::null(values.dtype().clone()); + let (patches, v2) = build(&values, &mut ctx)?; + + assert_eq!( + v2.dtype(), + values.dtype(), + "{case}: dtype is carried through" + ); + assert_eq!(v2.num_patches(), POSITIONS.len()); + + // Positions survive the round trip through chunk-local addressing. + let original = Sparse::try_new_from_patches(patches, fill.clone())?.into_array(); + let rebuilt = + Sparse::try_new_from_patches(v2.to_patches(&mut ctx)?, fill.clone())?.into_array(); + assert_arrays_eq!(original, rebuilt, &mut ctx); + + // Scalar lookups match the canonicalized array at every patched and unpatched row. + let canonical = original.execute::(&mut ctx)?.into_array(); + for (ordinal, &position) in POSITIONS.iter().enumerate() { + let position = position as usize; + let patched = v2 + .get_patched(position, &mut ctx)? + .unwrap_or_else(|| panic!("{case}: row {position} is patched")); + assert_eq!( + patched, + canonical.execute_scalar(position, &mut ctx)?, + "{case}: patch value at row {position}" + ); + assert_eq!(patched, values.execute_scalar(ordinal, &mut ctx)?); + // The next row, when it is not itself patched, falls back to the fill value. + let next = position + 1; + if !POSITIONS.contains(&(next as u64)) { + assert_eq!(v2.get_patched(next, &mut ctx)?, None, "{case}: row {next}"); + assert_eq!(canonical.execute_scalar(next, &mut ctx)?, fill); + } + } + Ok(()) +} + +/// Slicing the patch set matches slicing the equivalent `Sparse` array, for every dtype. +#[rstest] +#[case::bool("bool")] +#[case::i64("i64")] +#[case::f64("f64")] +#[case::decimal("decimal")] +#[case::utf8_varbinview("utf8_varbinview")] +#[case::binary("binary")] +#[case::list("list")] +#[case::fixed_size_list("fixed_size_list")] +#[case::struct_("struct")] +fn slices_like_sparse(#[case] case: &str) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = values_for(case)?; + let fill = Scalar::null(values.dtype().clone()); + let (patches, v2) = build(&values, &mut ctx)?; + + // 1023..2051 drops the patch at row 3, keeps the other three, and starts unaligned in the + // middle of the first chunk. + let range = 1023..2051; + let sliced = v2 + .slice(range.clone(), &mut ctx)? + .expect("patches remain in slice"); + assert_eq!(sliced.num_patches(), 3); + assert_eq!(sliced.offset(), 1023); + assert_eq!(sliced.array_len(), range.len()); + + let original = Sparse::try_new_from_patches(patches, fill.clone())? + .into_array() + .slice(range)?; + let rebuilt = Sparse::try_new_from_patches(sliced.to_patches(&mut ctx)?, fill)?.into_array(); + assert_arrays_eq!(original, rebuilt, &mut ctx); + Ok(()) +} + +/// Canonicalizing a constant-fill `Sparse` array through the chunk-local scatter produces the +/// same array as the default patch path. +/// +/// The scatter only applies where the types match: a non-nullable primitive result, where there +/// is no validity to patch and canonicalization is purely a value scatter over the fill. +/// +/// The dtypes are covered in one test rather than as `rstest` cases because the scatter switch is +/// process-global, so cases toggling it in parallel would race. +#[test] +fn scatter_matches_default_sparse_execution() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // Non-nullable values and a non-null fill are what put this on the chunk-local path. + let cases: Vec<(&str, ArrayRef)> = vec![ + ( + "u8", + PrimitiveArray::new(buffer![1u8, 2, 3, 4], Validity::NonNullable).into_array(), + ), + ( + "i16", + PrimitiveArray::new(buffer![-1i16, 2, -3, 4], Validity::NonNullable).into_array(), + ), + ( + "u32", + PrimitiveArray::new(buffer![1u32, 2, 3, 4], Validity::NonNullable).into_array(), + ), + ( + "i64", + PrimitiveArray::new(buffer![-1i64, 2, -3, 4], Validity::NonNullable).into_array(), + ), + ( + "f32", + PrimitiveArray::new(buffer![1.5f32, -2.5, 3.5, 4.5], Validity::NonNullable) + .into_array(), + ), + ( + "f64", + PrimitiveArray::new(buffer![1.5f64, -2.5, 3.5, 4.5], Validity::NonNullable) + .into_array(), + ), + ]; + + for (case, values) in cases { + let fill = values.execute_scalar(0, &mut ctx)?; + let indices = buffer![POSITIONS[0], POSITIONS[1], POSITIONS[2], POSITIONS[3]].into_array(); + let sparse = Sparse::try_new(indices, values, ARRAY_LEN, fill)?.into_array(); + + force_patches_v2_scatter(false); + let default_path = sparse.clone().execute::(&mut ctx)?.into_array(); + + // Guard against this going vacuous: the scatter count must move, or the flag silently + // did nothing and both sides would be the same code path. + let before = patches_v2_scatter_count(); + force_patches_v2_scatter(true); + let chunk_local_path = sparse.execute::(&mut ctx)?.into_array(); + force_patches_v2_scatter(false); + assert!( + patches_v2_scatter_count() > before, + "{case}: expected canonicalization to take the chunk-local scatter" + ); + + assert_arrays_eq!(default_path, chunk_local_path, &mut ctx); + } + Ok(()) +} From 582719679f632d8a60d1eb006b8b5500920fefdb Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 27 Aug 2026 21:04:52 +0000 Subject: [PATCH 3/3] Bitpack patch index children in btrblocks compress_patches Patch indices are sorted and bounded by the array length, so they bitpack well. Compressing the index child narrows patch sets whose indices would otherwise be stored at full width. This stays opt-in behind `force_patch_index_bitpack` while the golden corpus is regenerated, so default output is unchanged. Signed-off-by: Joe Isaacs --- vortex-btrblocks/src/lib.rs | 1 + vortex-btrblocks/src/schemes/patches.rs | 71 ++++++++++++++++++++++--- 2 files changed, 64 insertions(+), 8 deletions(-) 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()) +}